From f86e897e7ee139998bcab20b1d74d7af4c77af85 Mon Sep 17 00:00:00 2001 From: allegroai <> Date: Mon, 9 Oct 2023 13:48:46 +0300 Subject: [PATCH] Add --upload-files allowing to easily pass local files/folder into a remote session --- README.md | 27 +++++++++------ clearml_session/__main__.py | 37 +++++++++++++++------ clearml_session/interactive_session_task.py | 30 +++++++++++++++++ 3 files changed, 74 insertions(+), 20 deletions(-) diff --git a/README.md b/README.md index a1f9eed..cf88e77 100644 --- a/README.md +++ b/README.md @@ -139,7 +139,7 @@ VSCode server available at http://localhost:8898/ Connection is up and running 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) 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 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) or "Shutdown" to shut down remote interactive session @@ -219,7 +219,7 @@ clearml-session --help ``` console clearml-session - CLI for launching JupyterLab / VSCode on a remote machine usage: clearml-session [-h] [--version] [--attach [ATTACH]] - [--shutdown [SHUTDOWN]] [--interactive] + [--shutdown [SHUTDOWN]] [--shell] [--debugging-session DEBUGGING_SESSION] [--queue QUEUE] [--docker DOCKER] [--docker-args DOCKER_ARGS] [--public-ip [true/false]] @@ -228,6 +228,7 @@ usage: clearml-session [-h] [--version] [--attach [ATTACH]] [--vscode-version VSCODE_VERSION] [--vscode-extensions VSCODE_EXTENSIONS] [--jupyter-lab [true/false]] + [--upload-files UPLOAD_FILES] [--git-credentials [true/false]] [--user-folder USER_FOLDER] [--packages [PACKAGES [PACKAGES ...]]] @@ -239,9 +240,9 @@ usage: clearml-session [-h] [--version] [--attach [ATTACH]] [--keepalive [true/false]] [--queue-excluded-tag [QUEUE_EXCLUDED_TAG [QUEUE_EXCLUDED_TAG ...]]] [--queue-include-tag [QUEUE_INCLUDE_TAG [QUEUE_INCLUDE_TAG ...]]] - [--skip-docker-network] [--password PASSWORD] - [--username USERNAME] [--force_dropbear [true/false]] - [--verbose] [--yes] + [--skip-docker-network [true/false]] + [--password PASSWORD] [--username USERNAME] + [--force_dropbear [true/false]] [--verbose] [--yes] 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: previous session) --shutdown [SHUTDOWN], -S [SHUTDOWN] - Shut down an active session (default: previous session) - --interactive, -I open the SSH session directly, notice quiting the SSH session will - Not shutdown the remote session + Shut down an active session (default: previous + session) + --shell Open the SSH shell session directly, notice quiting + the SSH session will Not shutdown the remote session --debugging-session DEBUGGING_SESSION Pass existing Task id (experiment), create a copy of the experiment on a remote machine, and launch @@ -290,6 +292,11 @@ optional arguments: --jupyter-lab [true/false] Install Jupyter-Lab on interactive session (default: 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] If true, local .git-credentials file is sent to the interactive session. (default: false) @@ -332,7 +339,7 @@ optional arguments: --queue-include-tag [QUEUE_INCLUDE_TAG [QUEUE_INCLUDE_TAG ...]] Advanced: Only include queues with this specific tag from the selection - --skip-docker-network + --skip-docker-network [true/false] Advanced: If set, `--network host` is **not** passed to docker (assumes k8s network ingestion) (default: false) diff --git a/clearml_session/__main__.py b/clearml_session/__main__.py index 84d5418..ecd5461 100644 --- a/clearml_session/__main__.py +++ b/clearml_session/__main__.py @@ -5,6 +5,7 @@ import logging import os import subprocess import sys +from pathlib import Path from argparse import ArgumentParser, FileType from functools import reduce 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]) + # 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 # noinspection PyProtectedMember task._edit(script=task_script) @@ -482,7 +487,8 @@ def load_state(state_file): # never reload --verbose and --yes states state.pop('verbose', None) state.pop('yes', None) - state.pop('interactive', None) + state.pop('shell', None) + state.pop('upload_files', None) return state @@ -861,12 +867,12 @@ def monitor_ssh_tunnel(state, task): connect_message = ( '\nConnection is up and running\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' 'or \"Shutdown\" to shut down remote interactive session' ) 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(): ssh_process, ssh_password = start_ssh_tunnel( @@ -903,9 +909,9 @@ def monitor_ssh_tunnel(state, task): connect_state['reconnect'] = False # if interactive start with SSH interactive - if state.pop('interactive', None): + if state.pop('shell', None): 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 # wait for user input @@ -928,7 +934,7 @@ def monitor_ssh_tunnel(state, task): if not user_input: print(short_console_msg) continue - elif user_input.lower() in ('i', 'interactive',): + elif user_input.lower() in ('s', 'shell',): interactive_ssh(ssh_process) continue elif user_input.lower() == 'shutdown': @@ -972,8 +978,8 @@ def setup_parser(parser): help='Attach to running interactive session (default: previous session)') parser.add_argument("--shutdown", "-S", default=None, const="", nargs="?", 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 " + parser.add_argument("--shell", action='store_true', default=None, + help="Open the SSH shell session directly, notice quiting the SSH session " "will Not shutdown the remote session") 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, ' @@ -1005,6 +1011,10 @@ def setup_parser(parser): parser.add_argument('--jupyter-lab', default=True, nargs='?', const='true', metavar='true/false', type=lambda x: (str(x).strip().lower() in ('true', 'yes')), 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', type=lambda x: (str(x).strip().lower() in ('true', 'yes')), 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') parser.add_argument('--queue-include-tag', default=None, nargs='*', 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 ' '(assumes k8s network ingestion) (default: false)') parser.add_argument('--password', type=str, default=None, @@ -1098,7 +1109,7 @@ def cli(): if args.verbose: state['verbose'] = args.verbose - state['interactive'] = bool(args.interactive) + state['shell'] = bool(args.shell) client = APIClient() @@ -1114,6 +1125,12 @@ def cli(): print("Session #{} shut down, goodbye!".format(task.id)) 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 task = _get_previous_session(client, args, state, task_id=args.attach) diff --git a/clearml_session/interactive_session_task.py b/clearml_session/interactive_session_task.py index 2a72708..5a5306d 100644 --- a/clearml_session/interactive_session_task.py +++ b/clearml_session/interactive_session_task.py @@ -852,6 +852,36 @@ def setup_user_env(param, task): except Exception: 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