Add --upload-files allowing to easily pass local files/folder into a remote session

This commit is contained in:
allegroai 2023-10-09 13:48:46 +03:00
parent b5e4c5db76
commit f86e897e7e
3 changed files with 74 additions and 20 deletions

View File

@ -139,7 +139,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)
`i` (or "interactive") to connect to the SSH session `s` (or "shell") to connect to the SSH session
`Ctrl-C` (or "quit") to abort (remote session remains active) `Ctrl-C` (or "quit") to abort (remote session remains active)
or "Shutdown" to shut down remote interactive session or "Shutdown" to shut down remote interactive session
``` ```
@ -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)
`i` (or "interactive") to connect to the SSH session `s` (or "shell") to connect to the SSH session
`Ctrl-C` (or "quit") to abort (remote session remains active) `Ctrl-C` (or "quit") to abort (remote session remains active)
or "Shutdown" to shut down remote interactive session or "Shutdown" to shut down remote interactive session
@ -219,7 +219,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]] [--interactive] [--shutdown [SHUTDOWN]] [--shell]
[--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]]
@ -228,6 +228,7 @@ usage: clearml-session [-h] [--version] [--attach [ATTACH]]
[--vscode-version VSCODE_VERSION] [--vscode-version VSCODE_VERSION]
[--vscode-extensions VSCODE_EXTENSIONS] [--vscode-extensions VSCODE_EXTENSIONS]
[--jupyter-lab [true/false]] [--jupyter-lab [true/false]]
[--upload-files UPLOAD_FILES]
[--git-credentials [true/false]] [--git-credentials [true/false]]
[--user-folder USER_FOLDER] [--user-folder USER_FOLDER]
[--packages [PACKAGES [PACKAGES ...]]] [--packages [PACKAGES [PACKAGES ...]]]
@ -239,9 +240,9 @@ usage: clearml-session [-h] [--version] [--attach [ATTACH]]
[--keepalive [true/false]] [--keepalive [true/false]]
[--queue-excluded-tag [QUEUE_EXCLUDED_TAG [QUEUE_EXCLUDED_TAG ...]]] [--queue-excluded-tag [QUEUE_EXCLUDED_TAG [QUEUE_EXCLUDED_TAG ...]]]
[--queue-include-tag [QUEUE_INCLUDE_TAG [QUEUE_INCLUDE_TAG ...]]] [--queue-include-tag [QUEUE_INCLUDE_TAG [QUEUE_INCLUDE_TAG ...]]]
[--skip-docker-network] [--password PASSWORD] [--skip-docker-network [true/false]]
[--username USERNAME] [--force_dropbear [true/false]] [--password PASSWORD] [--username USERNAME]
[--verbose] [--yes] [--force_dropbear [true/false]] [--verbose] [--yes]
clearml-session - CLI for launching JupyterLab / VSCode on a remote machine clearml-session - CLI for launching JupyterLab / VSCode on a remote machine
@ -251,9 +252,10 @@ 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 session) Shut down an active session (default: previous
--interactive, -I open the SSH session directly, notice quiting the SSH session will session)
Not shutdown the remote session --shell Open the SSH shell 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
@ -290,6 +292,11 @@ optional arguments:
--jupyter-lab [true/false] --jupyter-lab [true/false]
Install Jupyter-Lab on interactive session (default: Install Jupyter-Lab on interactive session (default:
true) true)
--upload-files UPLOAD_FILES
Advanced: Upload local files/folders to the remote
session. Example: `/my/local/data/` will upload the
local folder and extract it into the container in
~/session-files/
--git-credentials [true/false] --git-credentials [true/false]
If true, local .git-credentials file is sent to the If true, local .git-credentials file is sent to the
interactive session. (default: false) interactive session. (default: false)
@ -332,7 +339,7 @@ optional arguments:
--queue-include-tag [QUEUE_INCLUDE_TAG [QUEUE_INCLUDE_TAG ...]] --queue-include-tag [QUEUE_INCLUDE_TAG [QUEUE_INCLUDE_TAG ...]]
Advanced: Only include queues with this specific tag Advanced: Only include queues with this specific tag
from the selection from the selection
--skip-docker-network --skip-docker-network [true/false]
Advanced: If set, `--network host` is **not** passed Advanced: If set, `--network host` is **not** passed
to docker (assumes k8s network ingestion) (default: to docker (assumes k8s network ingestion) (default:
false) false)

View File

@ -5,6 +5,7 @@ import logging
import os import os
import subprocess import subprocess
import sys import sys
from pathlib import Path
from argparse import ArgumentParser, FileType from argparse import ArgumentParser, FileType
from functools import reduce from functools import reduce
from getpass import getpass from getpass import getpass
@ -201,6 +202,10 @@ def create_base_task(state, project_name=None, task_name=None):
task.set_system_tags([system_tag]) task.set_system_tags([system_tag])
# if we need to upload data now is the time
if state.get("upload_files"):
task.upload_artifact(name="session-files", artifact_object=Path(state.get("upload_files")).expanduser())
# only update the data at the end, so reload requests are smaller # only update the data at the end, so reload requests are smaller
# noinspection PyProtectedMember # noinspection PyProtectedMember
task._edit(script=task_script) task._edit(script=task_script)
@ -482,7 +487,8 @@ 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) state.pop('shell', None)
state.pop('upload_files', None)
return state return state
@ -861,12 +867,12 @@ def monitor_ssh_tunnel(state, task):
connect_message = ( connect_message = (
'\nConnection is up and running\n' '\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'
'`i` (or "interactive") to connect to the SSH session\n' '`s` (or "shell") to connect to the SSH session\n'
'`Ctrl-C` (or "quit") to abort (remote session remains active)\n' '`Ctrl-C` (or "quit") to abort (remote session remains active)\n'
'or \"Shutdown\" to shut down remote interactive session' 'or \"Shutdown\" to shut down remote interactive session'
) )
short_console_msg = \ short_console_msg = \
"Enter \"r\" (\"reconnect\"), `i` (\"interactive\"), `Ctrl-C` (\"quit\") or \"Shutdown\"" "Enter \"r\" (\"reconnect\"), `s` (\"shell\"), `Ctrl-C` (\"quit\") or \"Shutdown\""
if not ssh_process or not ssh_process.isalive(): if not ssh_process or not ssh_process.isalive():
ssh_process, ssh_password = start_ssh_tunnel( ssh_process, ssh_password = start_ssh_tunnel(
@ -903,9 +909,9 @@ def monitor_ssh_tunnel(state, task):
connect_state['reconnect'] = False connect_state['reconnect'] = False
# if interactive start with SSH interactive # if interactive start with SSH interactive
if state.pop('interactive', None): if state.pop('shell', None):
interactive_ssh(ssh_process) interactive_ssh(ssh_process)
# if we are in --interactive, when we leave the session we should leave the process # if we are in --shell, when we leave the session we should leave the process
break break
# wait for user input # wait for user input
@ -928,7 +934,7 @@ def monitor_ssh_tunnel(state, task):
if not user_input: if not user_input:
print(short_console_msg) print(short_console_msg)
continue continue
elif user_input.lower() in ('i', 'interactive',): elif user_input.lower() in ('s', 'shell',):
interactive_ssh(ssh_process) interactive_ssh(ssh_process)
continue continue
elif user_input.lower() == 'shutdown': elif user_input.lower() == 'shutdown':
@ -972,8 +978,8 @@ 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, parser.add_argument("--shell", action='store_true', default=None,
help="open the SSH session directly, notice quiting the SSH session " help="Open the SSH shell session directly, notice quiting the SSH session "
"will Not shutdown the remote 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, '
@ -1005,6 +1011,10 @@ def setup_parser(parser):
parser.add_argument('--jupyter-lab', default=True, nargs='?', const='true', metavar='true/false', parser.add_argument('--jupyter-lab', default=True, nargs='?', const='true', metavar='true/false',
type=lambda x: (str(x).strip().lower() in ('true', 'yes')), type=lambda x: (str(x).strip().lower() in ('true', 'yes')),
help='Install Jupyter-Lab on interactive session (default: true)') help='Install Jupyter-Lab on interactive session (default: true)')
parser.add_argument('--upload-files', type=str, default=None,
help='Advanced: Upload local files/folders to the remote session. '
'Example: `/my/local/data/` will upload the local folder and extract it '
'into the container in ~/session-files/')
parser.add_argument('--git-credentials', default=False, nargs='?', const='true', metavar='true/false', parser.add_argument('--git-credentials', default=False, nargs='?', const='true', metavar='true/false',
type=lambda x: (str(x).strip().lower() in ('true', 'yes')), type=lambda x: (str(x).strip().lower() in ('true', 'yes')),
help='If true, local .git-credentials file is sent to the interactive session. ' help='If true, local .git-credentials file is sent to the interactive session. '
@ -1042,7 +1052,8 @@ def setup_parser(parser):
help='Advanced: Excluded queues with this specific tag from the selection') help='Advanced: Excluded queues with this specific tag from the selection')
parser.add_argument('--queue-include-tag', default=None, nargs='*', parser.add_argument('--queue-include-tag', default=None, nargs='*',
help='Advanced: Only include queues with this specific tag from the selection') help='Advanced: Only include queues with this specific tag from the selection')
parser.add_argument('--skip-docker-network', action='store_true', default=None, parser.add_argument('--skip-docker-network', default=None, nargs='?', const='true', metavar='true/false',
type=lambda x: (str(x).strip().lower() in ('true', 'yes')),
help='Advanced: If set, `--network host` is **not** passed to docker ' help='Advanced: If set, `--network host` is **not** passed to docker '
'(assumes k8s network ingestion) (default: false)') '(assumes k8s network ingestion) (default: false)')
parser.add_argument('--password', type=str, default=None, parser.add_argument('--password', type=str, default=None,
@ -1098,7 +1109,7 @@ def cli():
if args.verbose: if args.verbose:
state['verbose'] = args.verbose state['verbose'] = args.verbose
state['interactive'] = bool(args.interactive) state['shell'] = bool(args.shell)
client = APIClient() client = APIClient()
@ -1114,6 +1125,12 @@ def cli():
print("Session #{} shut down, goodbye!".format(task.id)) print("Session #{} shut down, goodbye!".format(task.id))
return 0 return 0
# check if upload folder/files exist
if args.upload_files:
if not Path(args.upload_files).expanduser().exists():
print("Requested file/folder `{}` does not exist, exiting".format(args.upload_files))
return 1
# get previous session, if it is running # get previous session, if it is running
task = _get_previous_session(client, args, state, task_id=args.attach) task = _get_previous_session(client, args, state, task_id=args.attach)

View File

@ -852,6 +852,36 @@ def setup_user_env(param, task):
except Exception: except Exception:
print('Could not write {} file'.format(git_config_file)) print('Could not write {} file'.format(git_config_file))
# check if we need to retrieve remote files for the session
if "session-files" in task.artifacts:
try:
target_dir = os.path.expanduser("~/session-files/")
cached_files_folder = task.artifacts["session-files"].get_local_copy(
extract_archive=True, force_download=True, raise_on_error=True)
# noinspection PyBroadException
try:
# first try a simple, move, if we fail, copy and delete
os.replace(cached_files_folder, target_dir)
except Exception:
import shutil
Path(target_dir).mkdir(parents=True, exist_ok=True)
if Path(cached_files_folder).is_dir():
shutil.copytree(
src=cached_files_folder,
dst=target_dir,
symlinks=True,
ignore_dangling_symlinks=True,
dirs_exist_ok=True)
shutil.rmtree(cached_files_folder)
else:
target_file = Path(cached_files_folder).name
# we need to remove the taskid prefix from the cache folder
target_file = (Path(target_dir) / (".".join(target_file.split(".")[1:]))).as_posix()
shutil.copy(cached_files_folder, target_file, follow_symlinks=False)
os.unlink(cached_files_folder)
except Exception as ex:
print("\nWARNING: Failed downloading remote session files! {}\n".format(ex))
return env return env