def escape_queue(s):
"""Escapes the path to a queue, e.g. preserves ~ at the begining.
"""
if isinstance(s, PosixPath):
s = unicode_(s)
elif isinstance(s, bytes):
s = s.decode('utf-8')
if s.startswith('~/'):
return '~/' + shell_escape(s[2:])
else:
return shell_escape(s)
def escape_queue(s):
"""Escapes the path to a queue, e.g. preserves ~ at the begining.
"""
if isinstance(s, PosixPath):
s = unicode_(s)
elif isinstance(s, bytes):
s = s.decode('utf-8')
if s.startswith('~/'):
return '~/' + shell_escape(s[2:])
else:
return shell_escape(s)
def submit(self, job_id, directory, script=None):
"""Submits a job to the queue.
If the runtime is not there, it will be installed. If it is a broken
chain of links, error.
"""
if job_id is None:
job_id = '%s_%s_%s' % (Path(directory).unicodename,
self.destination['username'],
make_unique_name())
else:
check_jobid(job_id)
queue = self._get_queue()
if queue is None:
queue = self._setup()
if script is None:
script = 'start.sh'
# Create directory
ret, target = self._call('%s %s' % (
shell_escape(queue / 'commands/new_job'),
job_id),
True)
if ret == 4:
raise JobAlreadyExists
elif ret != 0:
raise JobNotFound("Couldn't create job")
target = PosixPath(target)
logger.debug("Server created directory %s", target)
# Upload to directory
try:
scp_client = self.get_scp_client()
scp_client.put(str(Path(directory)),
str(target),
recursive=True)
except BaseException as e:
try:
self.delete(job_id)
except BaseException:
raise e
raise
logger.debug("Files uploaded")
# Submit job
self.check_call('%s %s %s %s' % (
shell_escape(queue / 'commands/submit'),
job_id, shell_escape(target),
shell_escape(script)))
logger.info("Submitted job %s", job_id)
return job_id
def submit(self, job_id, directory, script=None):
"""Submits a job to the queue.
If the runtime is not there, it will be installed. If it is a broken
chain of links, error.
"""
if job_id is None:
job_id = '%s_%s_%s' % (Path(directory).unicodename,
self.destination['username'],
make_unique_name())
else:
check_jobid(job_id)
queue = self._get_queue()
if queue is None:
queue = self._setup()
if script is None:
script = 'start.sh'
# Create directory
ret, target = self._call(
'%s %s' % (shell_escape(queue / 'commands/new_job'), job_id), True)
if ret == 4:
raise JobAlreadyExists
elif ret != 0:
raise JobNotFound("Couldn't create job")
target = PosixPath(target)
logger.debug("Server created directory %s", target)
# Upload to directory
try:
scp_client = self.get_scp_client()
scp_client.put(str(Path(directory)), str(target), recursive=True)
except BaseException as e:
try:
self.delete(job_id)
except BaseException:
raise e
raise
logger.debug("Files uploaded")
# Submit job
self.check_call('%s %s %s %s' %
(shell_escape(queue / 'commands/submit'), job_id,
shell_escape(target), shell_escape(script)))
logger.info("Submitted job %s", job_id)
return job_id
def _call(self, cmd, get_output):
"""Calls a command through the SSH connection.
Remote stderr gets printed to this program's stderr. Output is captured
and may be returned.
"""
server_err = self.server_logger()
chan = self.get_client().get_transport().open_session()
try:
logger.debug("Invoking %r%s", cmd,
" (stdout)" if get_output else "")
chan.exec_command('/bin/sh -c %s' % shell_escape(cmd))
output = b''
while True:
r, w, e = select.select([chan], [], [])
if chan not in r:
continue # pragma: no cover
recvd = False
while chan.recv_stderr_ready():
data = chan.recv_stderr(1024)
server_err.append(data)
recvd = True
while chan.recv_ready():
data = chan.recv(1024)
if get_output:
output += data
recvd = True
if not recvd and chan.exit_status_ready():
break
output = output.rstrip(b'\r\n')
return chan.recv_exit_status(), output
finally:
server_err.done()
chan.close()
def status(self, job_id):
"""Gets the status of a previously-submitted job.
"""
check_jobid(job_id)
queue = self._get_queue()
if queue is None:
raise QueueDoesntExist
ret, output = self._call('%s %s' % (
shell_escape(queue / 'commands/status'),
job_id),
True)
if ret == 0:
directory, result = output.splitlines()
result = result.decode('utf-8')
return RemoteQueue.JOB_DONE, PosixPath(directory), result
elif ret == 2:
directory = output.splitlines()[0]
return RemoteQueue.JOB_RUNNING, PosixPath(directory), None
elif ret == 3:
raise JobNotFound
else:
raise RemoteCommandFailure(command="commands/status",
ret=ret)
def _call(self, cmd, get_output):
"""Calls a command through the SSH connection.
Remote stderr gets printed to this program's stderr. Output is captured
and may be returned.
"""
server_err = self.server_logger()
chan = self.get_client().get_transport().open_session()
try:
logger.debug("Invoking %r%s",
cmd, " (stdout)" if get_output else "")
chan.exec_command('/bin/sh -c %s' % shell_escape(cmd))
output = b''
while True:
r, w, e = select.select([chan], [], [])
if chan not in r:
continue # pragma: no cover
recvd = False
while chan.recv_stderr_ready():
data = chan.recv_stderr(1024)
server_err.append(data)
recvd = True
while chan.recv_ready():
data = chan.recv(1024)
if get_output:
output += data
recvd = True
if not recvd and chan.exit_status_ready():
break
output = output.rstrip(b'\r\n')
return chan.recv_exit_status(), output
finally:
server_err.done()
chan.close()
def cleanup(self, kill=False):
queue = self._get_queue()
if queue is not None:
# Kill jobs
for job_id, info in self.list():
if info['status'] == 'running':
if not kill:
raise JobStillRunning("Can't cleanup, some jobs are "
"still running")
else:
logger.info("Killing running job %s", job_id)
self.kill(job_id)
# Remove queue
logger.info("Removing queue at %s", queue)
self.check_call('rm -rf -- %s' % shell_escape(queue))
# Remove links
for link in self._links:
self.check_call('rm -rf -- %s' % shell_escape(link))
return True
def _setup(self):
"""Actually installs the runtime.
"""
# Expands ~user in queue
if self.queue.path[0:1] == b'/':
queue = self.queue
else:
if self.queue.path[0:1] == b'~':
output = self.check_output('echo %s' %
escape_queue(self.queue))
queue = PosixPath(output.rstrip(b'\r\n'))
else:
output = self.check_output('pwd')
queue = PosixPath(output.rstrip(b'\r\n')) / self.queue
logger.debug("Resolved to %s", queue)
# Select runtime
if not self.setup_runtime:
# Autoselect
if self._call('which qsub', False)[0] == 0:
logger.debug("qsub is available, using runtime 'pbs'")
runtime = 'pbs'
else:
logger.debug("qsub not found, using runtime 'default'")
runtime = 'default'
else:
runtime = self.setup_runtime
if self.need_runtime is not None and runtime not in self.need_runtime:
raise ValueError("About to setup runtime %s but that wouldn't "
"match explicitely allowed runtimes" % runtime)
logger.info("Installing runtime %s%s at %s",
runtime,
"" if self.setup_runtime else " (auto)",
self.queue)
# Uploads runtime
scp_client = self.get_scp_client()
filename = pkg_resources.resource_filename('tej',
'remotes/%s' % runtime)
scp_client.put(filename, str(queue), recursive=True)
logger.debug("Files uploaded")
# Runs post-setup script
self.check_call('/bin/sh %s' % shell_escape(queue / 'commands/setup'))
logger.debug("Post-setup script done")
self._queue = queue
return queue
def cleanup(self, kill=False):
queue = self._get_queue()
if queue is not None:
# Kill jobs
for job_id, info in self.list():
if info['status'] == 'running':
if not kill:
raise JobStillRunning("Can't cleanup, some jobs are "
"still running")
else:
logger.info("Killing running job %s", job_id)
self.kill(job_id)
# Remove queue
logger.info("Removing queue at %s", queue)
self.check_call('rm -rf -- %s' % shell_escape(queue))
# Remove links
for link in self._links:
self.check_call('rm -rf -- %s' % shell_escape(link))
return True
def kill(self, job_id):
"""Kills a job on the server.
"""
check_jobid(job_id)
queue = self._get_queue()
if queue is None:
raise QueueDoesntExist
ret, output = self._call(
'%s %s' % (shell_escape(queue / 'commands/kill'), job_id), False)
if ret == 3:
raise JobNotFound
elif ret != 0:
raise RemoteCommandFailure(command='commands/kill', ret=ret)
def _setup(self):
"""Actually installs the runtime.
"""
# Expands ~user in queue
if self.queue.path[0:1] == b'/':
queue = self.queue
else:
if self.queue.path[0:1] == b'~':
output = self.check_output('echo %s' %
escape_queue(self.queue))
queue = PosixPath(output.rstrip(b'\r\n'))
else:
output = self.check_output('pwd')
queue = PosixPath(output.rstrip(b'\r\n')) / self.queue
logger.debug("Resolved to %s", queue)
# Select runtime
if not self.setup_runtime:
# Autoselect
if self._call('which qsub', False)[0] == 0:
logger.debug("qsub is available, using runtime 'pbs'")
runtime = 'pbs'
else:
logger.debug("qsub not found, using runtime 'default'")
runtime = 'default'
else:
runtime = self.setup_runtime
if self.need_runtime is not None and runtime not in self.need_runtime:
raise ValueError("About to setup runtime %s but that wouldn't "
"match explicitely allowed runtimes" % runtime)
logger.info("Installing runtime %s%s at %s", runtime,
"" if self.setup_runtime else " (auto)", self.queue)
# Uploads runtime
scp_client = self.get_scp_client()
filename = pkg_resources.resource_filename('tej',
'remotes/%s' % runtime)
scp_client.put(filename, str(queue), recursive=True)
logger.debug("Files uploaded")
# Runs post-setup script
self.check_call('/bin/sh %s' % shell_escape(queue / 'commands/setup'))
logger.debug("Post-setup script done")
self._queue = queue
return queue
def delete(self, job_id):
"""Deletes a job from the server.
"""
check_jobid(job_id)
queue = self._get_queue()
if queue is None:
raise QueueDoesntExist
ret, output = self._call(
'%s %s' % (shell_escape(queue / 'commands/delete'), job_id), False)
if ret == 3:
raise JobNotFound
elif ret == 2:
raise JobStillRunning
elif ret != 0:
raise RemoteCommandFailure(command='commands/delete', ret=ret)
def kill(self, job_id):
"""Kills a job on the server.
"""
check_jobid(job_id)
queue = self._get_queue()
if queue is None:
raise QueueDoesntExist
ret, output = self._call('%s %s' % (
shell_escape(queue / 'commands/kill'),
job_id),
False)
if ret == 3:
raise JobNotFound
elif ret != 0:
raise RemoteCommandFailure(command='commands/kill',
ret=ret)
def delete(self, job_id):
"""Deletes a job from the server.
"""
check_jobid(job_id)
queue = self._get_queue()
if queue is None:
raise QueueDoesntExist
ret, output = self._call('%s %s' % (
shell_escape(queue / 'commands/delete'),
job_id),
False)
if ret == 3:
raise JobNotFound
elif ret == 2:
raise JobStillRunning
elif ret != 0:
raise RemoteCommandFailure(command='commands/delete',
ret=ret)
def status(self, job_id):
"""Gets the status of a previously-submitted job.
"""
check_jobid(job_id)
queue = self._get_queue()
if queue is None:
raise QueueDoesntExist
ret, output = self._call(
'%s %s' % (shell_escape(queue / 'commands/status'), job_id), True)
if ret == 0:
directory, result = output.splitlines()
result = result.decode('utf-8')
return RemoteQueue.JOB_DONE, PosixPath(directory), result
elif ret == 2:
directory = output.splitlines()[0]
return RemoteQueue.JOB_RUNNING, PosixPath(directory), None
elif ret == 3:
raise JobNotFound
else:
raise RemoteCommandFailure(command="commands/status", ret=ret)
def list(self):
"""Lists the jobs on the server.
"""
queue = self._get_queue()
if queue is None:
raise QueueDoesntExist
output = self.check_output('%s' %
shell_escape(queue / 'commands/list'))
job_id, info = None, None
for line in output.splitlines():
line = line.decode('utf-8')
if line.startswith(' '):
key, value = line[4:].split(': ', 1)
info[key] = value
else:
if job_id is not None:
yield job_id, info
job_id = line
info = {}
if job_id is not None:
yield job_id, info
def list(self):
"""Lists the jobs on the server.
"""
queue = self._get_queue()
if queue is None:
raise QueueDoesntExist
output = self.check_output('%s' %
shell_escape(queue / 'commands/list'))
job_id, info = None, None
for line in output.splitlines():
line = line.decode('utf-8')
if line.startswith(' '):
key, value = line[4:].split(': ', 1)
info[key] = value
else:
if job_id is not None:
yield job_id, info
job_id = line
info = {}
if job_id is not None:
yield job_id, info
def test_shell_escape(self):
self.assertEqual(shell_escape("test"), "test")
self.assertEqual(shell_escape("hello world"), '"hello world"')
self.assertEqual(shell_escape('some"thing'), '"some\\"thing"')