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))