ClearML/clearml-session
1
0
Fork 0
mirror of https://github.com/clearml/clearml-session synced 2025-05-05 20:44:57 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add --interactive ('i') for direct SSH console connection

Browse Source
This commit is contained in:
allegroai 2023-09-03 13:08:58 +03:00
parent f6b7e9f4fa
commit 12a5e81f1c
3 changed files with 60 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
README.md
View File

@ -140,7 +140,7 @@ VSCode server available at http://localhost:8898/
Connection is up and running Connection is up and running
Enter "r" (or "reconnect") to reconnect the session (for example after suspend) Enter "r" (or "reconnect") to reconnect the session (for example after suspend)
Ctrl-C (or "quit") to abort (remote session remains active) Ctrl-C (or "quit") to abort (remote session remains active)
or "Shutdown" to shutdown remote interactive session or "Shutdown" to shut down remote interactive session
``` ```
Click on the JupyterLab link (http://localhost:8878/?token=xyz) Click on the JupyterLab link (http://localhost:8878/?token=xyz)
@ -171,7 +171,7 @@ It will shut down the remote session, free the resource and close the CLI
``` console ``` console
Enter "r" (or "reconnect") to reconnect the session (for example after suspend) Enter "r" (or "reconnect") to reconnect the session (for example after suspend)
Ctrl-C (or "quit") to abort (remote session remains active) Ctrl-C (or "quit") to abort (remote session remains active)
or "Shutdown" to shutdown remote interactive session or "Shutdown" to shut down remote interactive session
shutdown shutdown
@ -217,7 +217,7 @@ clearml-session --help
``` console ``` console
clearml-session - CLI for launching JupyterLab / VSCode on a remote machine clearml-session - CLI for launching JupyterLab / VSCode on a remote machine
usage: clearml-session [-h] [--version] [--attach [ATTACH]] usage: clearml-session [-h] [--version] [--attach [ATTACH]]
[--shutdown [SHUTDOWN]] [--shutdown [SHUTDOWN]] [--interactive]
[--debugging-session DEBUGGING_SESSION] [--queue QUEUE] [--debugging-session DEBUGGING_SESSION] [--queue QUEUE]
[--docker DOCKER] [--docker-args DOCKER_ARGS] [--docker DOCKER] [--docker-args DOCKER_ARGS]
[--public-ip [true/false]] [--public-ip [true/false]]
@ -249,8 +249,9 @@ optional arguments:
--attach [ATTACH] Attach to running interactive session (default: --attach [ATTACH] Attach to running interactive session (default:
previous session) previous session)
--shutdown [SHUTDOWN], -S [SHUTDOWN] --shutdown [SHUTDOWN], -S [SHUTDOWN]
Shut down an active session (default: previous Shut down an active session (default: previous session)
session) --interactive, -I open the SSH session directly, notice quiting the SSH session will
Not shutdown the remote session
--debugging-session DEBUGGING_SESSION --debugging-session DEBUGGING_SESSION
Pass existing Task id (experiment), create a copy of Pass existing Task id (experiment), create a copy of
the experiment on a remote machine, and launch the experiment on a remote machine, and launch

54
clearml_session/__main__.py
View File

@ -481,6 +481,7 @@ def load_state(state_file):
# never reload --verbose and --yes states # never reload --verbose and --yes states
state.pop('verbose', None) state.pop('verbose', None)
state.pop('yes', None) state.pop('yes', None)
state.pop('interactive', None)
return state return state
@ -753,10 +754,26 @@ def start_ssh_tunnel(username, remote_address, ssh_port, ssh_password, local_rem
child.terminate(force=True) child.terminate(force=True)
child = None child = None
print('\n') print('\n')
child.logfile = None
return child, ssh_password return child, ssh_password
def monitor_ssh_tunnel(state, task): def monitor_ssh_tunnel(state, task):
def interactive_ssh(p):
import struct, fcntl, termios, signal, sys # noqa
def sigwinch_passthrough(sig, data):
s = struct.pack("HHHH", 0, 0, 0, 0)
a = struct.unpack('hhhh', fcntl.ioctl(sys.stdout.fileno(), termios.TIOCGWINSZ, s))
if not p.closed:
p.setwinsize(a[0], a[1])
print("Switching to active SSH session, press ``Ctrl - ]`` to leave")
# Note this 'p' is global and used in sigwinch_passthrough.
signal.signal(signal.SIGWINCH, sigwinch_passthrough)
p.interact()
print("\nSSH session running in background\n")
print('Setting up connection to remote session') print('Setting up connection to remote session')
local_jupyter_port, local_jupyter_port_, local_ssh_port, local_ssh_port_, local_vscode_port, local_vscode_port_ = \ local_jupyter_port, local_jupyter_port_, local_ssh_port, local_ssh_port_, local_vscode_port, local_vscode_port_ = \
_get_available_ports([8878, 8878+1, 8022, 8022+1, 8898, 8898+1]) _get_available_ports([8878, 8878+1, 8022, 8022+1, 8898, 8898+1])
@ -780,9 +797,18 @@ def monitor_ssh_tunnel(state, task):
default_section = _get_config_section_name()[0] default_section = _get_config_section_name()[0]
local_remote_pair_list = [] local_remote_pair_list = []
shutdown = False
try: try:
while task.get_status() == 'in_progress': while task.get_status() == 'in_progress':
if not all([ssh_port, jupyter_token, jupyter_port, internal_ssh_port, ssh_password, remote_address]): if not all([
ssh_port,
not state.get('jupyter_lab') or jupyter_token,
not state.get('jupyter_lab') or jupyter_port,
not state.get('vscode_server') or vscode_port,
internal_ssh_port,
ssh_password,
remote_address
]):
task.reload() task.reload()
task_parameters = task.get_parameters() task_parameters = task.get_parameters()
if Session.check_min_api_version("2.13"): if Session.check_min_api_version("2.13"):
@ -859,8 +885,9 @@ def monitor_ssh_tunnel(state, task):
print('\nConnection is up and running\n' print('\nConnection is up and running\n'
'Enter \"r\" (or \"reconnect\") to reconnect the session (for example after suspend)\n' 'Enter \"r\" (or \"reconnect\") to reconnect the session (for example after suspend)\n'
'Ctrl-C (or "quit") to abort (remote session remains active)\n' '`i` (or "interactive") to connect to the SSH session\n'
'or \"Shutdown\" to shutdown remote interactive session') '`Ctrl-C` (or "quit") to abort (remote session remains active)\n'
'or \"Shutdown\" to shut down remote interactive session')
else: else:
logging.getLogger().warning('SSH tunneling failed, retrying in {} seconds'.format(3)) logging.getLogger().warning('SSH tunneling failed, retrying in {} seconds'.format(3))
sleep(3.) sleep(3.)
@ -868,6 +895,12 @@ def monitor_ssh_tunnel(state, task):
connect_state['reconnect'] = False connect_state['reconnect'] = False
# if interactive start with SSH interactive
if state.pop('interactive', None):
interactive_ssh(ssh_process)
# if we are in --interactive, when we leave the session we should leave the process
break
# wait for user input # wait for user input
user_input = _read_std_input(timeout=sleep_period) user_input = _read_std_input(timeout=sleep_period)
if user_input is None: if user_input is None:
@ -885,9 +918,13 @@ def monitor_ssh_tunnel(state, task):
pass pass
continue continue
if user_input.lower() == 'shutdown': if user_input.lower() in ('i', 'interactive',):
interactive_ssh(ssh_process)
continue
elif user_input.lower() == 'shutdown':
print('Shutting down interactive session') print('Shutting down interactive session')
task.mark_stopped() task.mark_stopped()
shutdown = True
break break
elif user_input.lower() in ('r', 'reconnect', ): elif user_input.lower() in ('r', 'reconnect', ):
print('Reconnecting to interactive session') print('Reconnecting to interactive session')
@ -901,7 +938,7 @@ def monitor_ssh_tunnel(state, task):
else: else:
print('unknown command: \'{}\''.format(user_input)) print('unknown command: \'{}\''.format(user_input))
print('Interactive session ended') print("Remote session shutdown" if shutdown else "Remote session still running!")
except KeyboardInterrupt: except KeyboardInterrupt:
print('\nUser aborted') print('\nUser aborted')
@ -925,6 +962,9 @@ def setup_parser(parser):
help='Attach to running interactive session (default: previous session)') help='Attach to running interactive session (default: previous session)')
parser.add_argument("--shutdown", "-S", default=None, const="", nargs="?", parser.add_argument("--shutdown", "-S", default=None, const="", nargs="?",
help="Shut down an active session (default: previous session)") help="Shut down an active session (default: previous session)")
parser.add_argument("--interactive", "-I", action='store_true', default=None,
help="open the SSH session directly, notice quiting the SSH session "
"will Not shutdown the remote session")
parser.add_argument('--debugging-session', type=str, default=None, parser.add_argument('--debugging-session', type=str, default=None,
help='Pass existing Task id (experiment), create a copy of the experiment on a remote machine, ' help='Pass existing Task id (experiment), create a copy of the experiment on a remote machine, '
'and launch jupyter/ssh for interactive access. Example --debugging-session <task_id>') 'and launch jupyter/ssh for interactive access. Example --debugging-session <task_id>')
@ -1048,6 +1088,8 @@ def cli():
if args.verbose: if args.verbose:
state['verbose'] = args.verbose state['verbose'] = args.verbose
state['interactive'] = bool(args.interactive)
client = APIClient() client = APIClient()
if args.shutdown is not None: if args.shutdown is not None:
@ -1107,7 +1149,7 @@ def cli():
monitor_ssh_tunnel(state, task) monitor_ssh_tunnel(state, task)
# we are done # we are done
print('Leaving interactive session') print('Goodbye')
def _get_previous_session( def _get_previous_session(

12
clearml_session/interactive_session_task.py
View File

@ -243,7 +243,7 @@ def monitor_jupyter_server(fd, local_filename, process, task, jupyter_port, host
pass pass
def start_vscode_server(hostname, hostnames, param, task, env): def start_vscode_server(hostname, hostnames, param, task, env, bind_ip="127.0.0.1"):
if not param.get("vscode_server"): if not param.get("vscode_server"):
return return
@ -368,7 +368,7 @@ def start_vscode_server(hostname, hostnames, param, task, env):
"--auth", "--auth",
"none", "none",
"--bind-addr", "--bind-addr",
"127.0.0.1:{}".format(port), "{}:{}".format(bind_ip, port),
"--user-data-dir", user_folder, "--user-data-dir", user_folder,
"--extensions-dir", exts_folder, "--extensions-dir", exts_folder,
] + vscode_extensions_cmd, ] + vscode_extensions_cmd,
@ -401,8 +401,8 @@ def start_vscode_server(hostname, hostnames, param, task, env):
proc = subprocess.Popen( proc = subprocess.Popen(
['bash', '-c', ['bash', '-c',
'{} --auth none --bind-addr 127.0.0.1:{} --disable-update-check {} {}'.format( '{} --auth none --bind-addr {}:{} --disable-update-check {} {}'.format(
vscode_path, port, vscode_path, bind_ip, port,
'--user-data-dir \"{}\"'.format(user_folder) if user_folder else '', '--user-data-dir \"{}\"'.format(user_folder) if user_folder else '',
'--extensions-dir \"{}\"'.format(exts_folder) if exts_folder else '')], '--extensions-dir \"{}\"'.format(exts_folder) if exts_folder else '')],
env=env, env=env,
@ -425,7 +425,7 @@ def start_vscode_server(hostname, hostnames, param, task, env):
task.set_parameter(name='properties/vscode_port', value=str(port)) task.set_parameter(name='properties/vscode_port', value=str(port))
def start_jupyter_server(hostname, hostnames, param, task, env): def start_jupyter_server(hostname, hostnames, param, task, env, bind_ip="127.0.0.1"):
if not param.get('jupyterlab', True): if not param.get('jupyterlab', True):
print('no jupyterlab to monitor - going to sleep') print('no jupyterlab to monitor - going to sleep')
while True: while True:
@ -467,7 +467,7 @@ def start_jupyter_server(hostname, hostnames, param, task, env):
"--no-browser", "--no-browser",
"--allow-root", "--allow-root",
"--ip", "--ip",
"127.0.0.1", bind_ip,
"--port", "--port",
str(port), str(port),
], ],