def tmux_run_sync(tmux_window, cmd, check_interval=0.2, max_wait_sec=600):
"""Uses tmux send-keys command, adds file locking to block until command
finishes executing."""
global tmux_counter
if not os.path.exists('/tmp/tmux'):
_ossystem('mkdir -p /tmp/tmux')
ts = str(u.now_micros())
cmd_fn_in = '/tmp/tmux/' + str(tmux_counter) + '.' + ts + '.in'
cmd_fn_out = '/tmp/tmux/' + str(tmux_counter) + '.' + ts + '.out'
open(cmd_fn_in, 'w').write(cmd + '\n')
modified_cmd = '%s && echo $? > %s' % (cmd, cmd_fn_out)
start_time = time.time()
_ossystem("tmux send-keys -t {} '{}' Enter".format(tmux_window,
modified_cmd))
while True:
if time.time() - start_time > max_wait_sec:
assert False, "Timeout %s exceeded for %s" % (max_wait_sec, cmd)
if not os.path.exists(cmd_fn_out):
time.sleep(check_interval)
continue
contents = open(cmd_fn_out).read()
# if empty wait a bit to allow for race condition
if len(contents) == 0:
time.sleep(check_interval)
contents = open(cmd_fn_out).read()
contents = contents.strip()
assert contents == '0', "Command %s returned status %s" % (cmd,
contents)
break
def run_and_capture_output(self, cmd, sync=True, ignore_errors=False):
assert '|' not in cmd, "don't support piping (since we append piping here)"
ts = str(u.now_micros())
cmd_stdout_fn = self.remote_scratch + '/' + str(
self._run_counter) + '.' + ts + '.out'
cmd = f'{cmd} | tee {cmd_stdout_fn}'
self.run(cmd, sync, ignore_errors)
return self.file_read(cmd_stdout_fn)
def run(self,
cmd,
max_wait_sec=600,
check_interval=1,
ignore_errors=False,
wait_to_finish=True):
"""Runs command in tmux session. No need for multiple tmux sessions per
task, so assume tmux session/window is always called tmux:0"""
self.cmd_idx += 1
self.log("tmux> %s", cmd)
if cmd.startswith("upload "):
_, fname = cmd.split()
fname = fname.replace("~", os.environ["HOME"])
self.upload(fname)
# locking to wait for command to finish
ts = str(u.now_micros())
cmd_fn_out = '/tmp/tmux/' + str(self.cmd_idx) + '.' + ts + '.out'
modified_cmd = '%s; echo $? > %s' % (cmd, cmd_fn_out)
tmux_window = 'tmux:0'
tmux_cmd = "tmux send-keys -t {} '{}' Enter".format(
tmux_window, modified_cmd)
self.run_sync(tmux_cmd)
if not wait_to_finish:
return
start_time = time.time()
while True:
if time.time() - start_time > max_wait_sec:
assert False, "Timeout %s exceeded for %s" % (max_wait_sec,
cmd)
if not self.file_exists(cmd_fn_out):
self.log("waiting %s for %s" % (check_interval, cmd))
time.sleep(check_interval)
continue
contents = self.file_read(cmd_fn_out)
# if empty wait a bit to allow for race condition
if len(contents) == 0:
time.sleep(check_interval)
contents = task.file_read(cmd_fn_out)
contents = contents.strip()
if contents != '0':
if not ignore_errors:
assert False, "Command %s returned status %s" % (cmd,
contents)
else:
self.log("Warning: command %s returned status %s" %
(cmd, contents))
break
def __init__(self, tmux_window, job, task_id):
self.tmux_window = tmux_window
self.job = job
self.ip = '127.0.0.1' # hostname/ip address
self.id = task_id
self.port = portpicker.pick_unused_port()
print("Assigning %s:%s to port %s" %
(self.job.name, self.id, self.port))
self.connect_instructions = 'tmux a -t ' + self.tmux_window
self.last_stdout = '<unavailable>' # compatiblity with aws.py:Task
self.last_stderr = '<unavailable>'
self.scratch = SCRATCH_PREFIX
self.taskdir = "{}/{}.{}/{}".format(TASKDIR_PREFIX, job.name,
u.now_micros(), self.id)
self.run('mkdir -p ' + self.taskdir)
self.run('cd ' + self.taskdir)
def run(self,
cmd,
sync=True,
ignore_errors=False,
max_wait_sec=600,
check_interval=0.5):
"""Runs command in tmux session. No need for multiple tmux sessions per
task, so assume tmux session/window is always called tmux:0"""
assert self._run_command_available, "Have you done wait_until_ready?"
cmd = cmd.strip()
self._run_counter += 1
self.log("tmux> %s", cmd)
# todo: match logic in tmux_session (upload magic handling done
# in init instead of run)
if cmd.startswith('%upload'):
self._upload_handler(cmd)
return
# locking to wait for command to finish
ts = str(u.now_micros())
cmd_fn_out = self.remote_scratch + '/' + str(
self._run_counter) + '.' + ts + '.out'
cmd = _strip_comment(cmd)
assert not '&' in cmd, "cmd '%s' contains &, that breaks things" % (
cmd, )
modified_cmd = '%s; echo $? > %s' % (cmd, cmd_fn_out)
tmux_window = self._tmux_session_name + ':0'
tmux_cmd = "tmux send-keys -t {} {} Enter".format(
tmux_window, shlex.quote(modified_cmd))
self._run_ssh(tmux_cmd)
if not sync:
return
start_time = time.time()
while True:
if time.time() - start_time > max_wait_sec:
assert False, "Timeout %s exceeded for %s" % (max_wait_sec,
cmd)
if not self.file_exists(cmd_fn_out):
self.log("waiting for %s" % (cmd, ))
time.sleep(check_interval)
continue
contents = self.file_read(cmd_fn_out)
# if empty wait a bit to allow for race condition
if len(contents) == 0:
time.sleep(check_interval)
contents = task.file_read(cmd_fn_out)
contents = contents.strip()
if contents != '0':
if not ignore_errors:
assert False, "Command %s returned status %s" % (cmd,
contents)
else:
self.log("Warning: command %s returned status %s" %
(cmd, contents))
break
def file_read(self, remote_fn):
# self.log("file_read")
tmp_fn = self.scratch + '/' + str(u.now_micros())
self.download(remote_fn, tmp_fn)
return open(tmp_fn).read()
def file_write(self, remote_fn, contents):
tmp_fn = self.scratch + '/' + str(u.now_micros())
open(tmp_fn, 'w').write(contents)
self.upload(tmp_fn, remote_fn)
def _make_temp_fn(self): """Returns temporary filename for this task.""" return self.scratch+'/file_write.'+str(u.now_micros())
def file_read(self, remote_fn):
# TODO: create tasklogdir for everything for given job
tmp_fn = '/tmp/tmux/' + str(u.now_micros())
self.download(remote_fn, tmp_fn)
return open(tmp_fn).read()
def file_write(self, remote_fn, contents):
# TODO: create tasklogdir for everything for given job
tmp_fn = '/tmp/tmux/' + str(u.now_micros())
open(tmp_fn, 'w').write(contents)
self.upload(tmp_fn, remote_fn)
def file_write(self, contents, fn):
local_fn = '/tmp/' + str(u.now_micros())
with open(local_fn, 'w') as f:
f.write(contents)
self.upload(local_fn, os.path.basename(fn))
