def eval(self,
command: str,
*,
log_dir: Optional[Path],
is_silent=True) -> EvalResult:
with monit.section(f'Exec: {command}', is_silent=is_silent):
stdin, stdout, stderr = self.client.exec_command(command)
out = stdout.read()
err = stderr.read()
exit_code = stdout.channel.recv_exit_status()
if log_dir:
with open(str(log_dir / 'stdout.log'), 'wb') as f:
f.write(out)
with open(str(log_dir / 'stderr.log'), 'wb') as f:
f.write(err)
with open(str(log_dir / 'exit_code'), 'w') as f:
f.write(str(exit_code))
if exit_code != 0:
monit.fail()
return EvalResult(exit_code,
out.decode('utf-8').strip(),
err.decode('utf-8').strip())
def stream(self, command: str, *, log_dir: Path, ui_mode: UIMode = UIMode.dots, is_silent=True):
with monit.section(f'Exec: {command}', is_silent=is_silent):
thread = ExecutorThread(command,
log_dir=log_dir, ui_mode=ui_mode)
thread.start()
thread.join()
if thread.exit_code != 0:
monit.fail()
return thread.exit_code
def execute(client: SSHClient, command: str, *, is_silent=True):
with monit.section(f'Exec: {command}', is_silent=is_silent):
stdin, stdout, stderr = client.exec_command(command)
out = stdout.read().decode('utf-8').strip()
err = stderr.read().decode('utf-8').strip()
exit_code = stdout.channel.recv_exit_status()
if err != '':
logger.log("Errors:", Text.warning)
print(err)
if exit_code != 0:
monit.fail()
return exit_code, out
def run(command: List[str]):
client = connect()
_, home_path = execute(client, 'pwd')
with monit.section("Setup server"):
logger.log()
if setup_server(client, home_path) != 0:
monit.fail()
fail("Failed to setup server")
return
logger.log()
with monit.section("RSync"):
logger.log()
if rsync_project() != 0:
monit.fail()
fail("Failed to run rsync")
return
logger.log()
with monit.section("Update python packages"):
logger.log()
if update_packages(client, home_path) != 0:
monit.fail()
fail("Failed to update packages")
return
logger.log('\n\n' + '-' * 40 + '\n\n')
with monit.section("Run command"):
logger.log()
if run_command(client, home_path, command) != 0:
monit.fail()
fail("Failed to run command")
return
def autocomplete():
prompt = request.json['prompt']
if not prompt:
return jsonify({'success': False})
with monit.section('Predict') as s:
acquired = lock.acquire(blocking=False)
if acquired:
res = predictor.get_token(prompt, token_chars=TOKEN_CHARS)
lock.release()
s.message = f'{json.dumps(prompt[-5:])} -> {json.dumps(res)}'
return jsonify({'success': True, 'prediction': res})
else:
monit.fail()
return jsonify({'success': False})
def create(self):
self._create()
with monit.section(f'Create {self.service_path}'):
logger.log([
'Start the monitoring service with:\n',
('systemctl --user start labml.service', Text.value),
'\n, stop with:\n',
('systemctl --user stop labml.service', Text.value),
'\n and check status with:\n',
('systemctl --user status labml.service', Text.value),
])
with monit.section('Starting service'):
ret = os.system('systemctl --user start labml.service')
if ret != 0:
monit.fail()
def _create(self):
if self.service_path.exists():
logger.log(f'{self.service_path} exists')
return
if not self.service_path.parent.exists():
self.service_path.parent.mkdir(parents=True)
with open(str(self.service_path), 'w') as f:
f.write(self._get_service_file_content())
with monit.section('Reload systemd daemon'):
ret = os.system('systemctl --user daemon-reload')
if ret != 0:
monit.fail()
with monit.section('Enable service'):
ret = os.system('systemctl --user enable labml.service')
if ret != 0:
monit.fail()
def execute_stream(client: SSHClient, command: str, *, is_silent=True):
with monit.section(f'Exec: {command}', is_silent=is_silent):
channel = client.get_transport().open_session()
channel.exec_command(command)
while True:
while channel.recv_ready():
sys.stdout.write(channel.recv(1024).decode('utf-8'))
while channel.recv_stderr_ready():
sys.stderr.write(channel.recv_stderr(1024).decode('utf-8'))
if channel.exit_status_ready():
break
while channel.recv_ready():
sys.stdout.write(channel.recv(1024).decode('utf-8'))
while channel.recv_stderr_ready():
sys.stderr.write(channel.recv_stderr(1024).decode('utf-8'))
exit_code = channel.recv_exit_status()
if exit_code != 0:
monit.fail()
return exit_code
def autocomplete():
prefix = request.json['prompt']
if not prefix:
return jsonify({'success': False})
with monit.section('Predict') as s:
acquired = lock.acquire(blocking=False)
if acquired:
stripped, prompt = predictor.rstrip(prefix)
rest = prefix[len(stripped):]
prediction_complete = NextWordPredictionComplete(rest, 5)
prompt = torch.tensor(prompt, dtype=torch.long).unsqueeze(-1)
predictions = predictor.get_next_word(prompt, None, rest, [1.], prediction_complete, 5)
predictions.sort(key=lambda x: -x.prob)
results = [pred.text[len(rest):] for pred in predictions]
probs = [pred.prob for pred in predictions]
lock.release()
s.message = f'{json.dumps(prefix[-5:])} -> {json.dumps(results)}'
return jsonify({'success': True, 'prediction': results, 'probs': probs})
else:
monit.fail()
return jsonify({'success': False})
