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

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]
+                       [--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
 
@@ -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)
+                        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
                         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)

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,
+    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")
     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)
 

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