Files
dj/model-bridge/ssh_task.py
T

33 lines
1.2 KiB
Python

"""Run an explicitly supplied task script on the user's model server."""
import argparse
import os
from pathlib import Path
import paramiko
p = argparse.ArgumentParser()
p.add_argument('script')
p.add_argument('--upload', nargs=2, action='append', default=[])
p.add_argument('--download', nargs=2, action='append', default=[])
a = p.parse_args()
c = paramiko.SSHClient()
c.load_system_host_keys()
c.set_missing_host_key_policy(paramiko.WarningPolicy())
c.connect('192.168.110.111', username='ps', password=os.environ['HUOBAO_SSH_PASSWORD'], look_for_keys=False, allow_agent=False, timeout=15)
try:
with c.open_sftp() as sftp:
for source, destination in a.upload:
sftp.put(source, destination)
stdin, stdout, stderr = c.exec_command('python3 -', timeout=60)
stdin.write(Path(a.script).read_text(encoding='utf-8'))
stdin.channel.shutdown_write()
print(stdout.read().decode('utf-8', errors='replace'))
print(stderr.read().decode('utf-8', errors='replace'))
code = stdout.channel.recv_exit_status()
if code == 0 and a.download:
with c.open_sftp() as sftp:
for source, destination in a.download:
sftp.get(source, destination)
raise SystemExit(code)
finally:
c.close()