Add --vscode-extensions enabling to specify bscode extensions to be installed (e.g. "ms-python.pylint,ms-python.flake8")

Add --force_dropbear to force using dropbear instead of SSHd

clearml_session/__main__.py

@@ -537,6 +537,8 @@ def clone_task(state, project_id=None):
     task_params["{}/public_ip".format(section)] = bool(state.get('public_ip'))
     task_params["{}/ssh_ports".format(section)] = state.get('remote_ssh_port') or ''
     task_params["{}/vscode_version".format(section)] = state.get('vscode_version') or ''
+    task_params["{}/vscode_extensions".format(section)] = state.get('vscode_extensions') or ''
+    task_params["{}/force_dropbear".format(section)] = bool(state.get('force_dropbear'))
     if state.get('user_folder'):
         task_params['{}/user_base_directory'.format(section)] = state.get('user_folder')
     docker = state.get('docker') or task.get_base_docker()
@@ -946,6 +948,10 @@ def setup_parser(parser):
     parser.add_argument('--vscode-version', type=str, default=None,
                         help='Set vscode server (code-server) version, as well as vscode python extension version '
                              '<vscode:python-ext> (example: "3.7.4:2020.10.332292344")')
+    parser.add_argument('--vscode-extensions', type=str, default=None,
+                        help='Install additional vscode extensions, as well as vscode python extension '
+                             '(example: "ms-python.python,ms-python.black-formatter,'
+                             'ms-python.pylint,ms-python.flake8")')
     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)')
@@ -995,13 +1001,16 @@ def setup_parser(parser):
     parser.add_argument('--username', type=str, default=None,
                         help='Advanced: Select ssh username for the interactive session '
                              '(default: `root` or previously used one)')
+    parser.add_argument('--force_dropbear', default=None, nargs='?', const='true', metavar='true/false',
+                        type=lambda x: (str(x).strip().lower() in ('true', 'yes')),
+                        help='Force using `dropbear` instead of SSHd')
     parser.add_argument('--verbose', action='store_true', default=None,
                         help='Advanced: If set, print verbose progress information, '
                              'e.g. the remote machine setup process log')
-    parser.add_argument("--yes", "-y",
+    parser.add_argument('--yes', '-y',
-                        action="store_true", default=False,
+                        action='store_true', default=False,
-                        help="Automatic yes to prompts; assume \"yes\" as answer "
+                        help='Automatic yes to prompts; assume \"yes\" as answer '
-                             "to all prompts and run non-interactively",)
+                             'to all prompts and run non-interactively',)
 
 
 def get_version():
@@ -1142,11 +1151,11 @@ def _get_previous_session(
             print("{} active session {}".format(verb, task_id))
         else:
             choice = input("{} active session id={} [Y]/n? ".format(question_verb, task_id))
-            if str(choice).strip().lower() in ("y", "yes"):
+            if str(choice).strip().lower() in ("", "y", "yes"):
                 return Task.get_task(task_id=task_id)
 
     # multiple sessions running
-    print('Active sessions:')
+    print("{} active session:".format(verb))
     try:
         prev_task_id = state.get('task_id')
         default_i = next(i for i, (tid, _, _) in enumerate(running_task_ids_created) if prev_task_id == tid)

clearml_session/interactive_session_task.py

@@ -249,8 +249,8 @@ def start_vscode_server(hostname, hostnames, param, task, env):
 
     # get vscode version and python extension version
     # they are extremely flaky, this combination works, most do not.
-    vscode_version = '4.10.0'
+    vscode_version = '4.14.1'
-    python_ext_version = '2021.10.1365161279'
+    python_ext_version = '2023.12.0'
     if param.get("vscode_version"):
         vscode_version_parts = param.get("vscode_version").split(':')
         vscode_version = vscode_version_parts[0]
@@ -338,6 +338,30 @@ def start_vscode_server(hostname, hostnames, param, task, env):
             else:
                 exts_folder = os.path.expanduser("~/.local/share/code-server/extensions/")
         else:
+            vscode_extensions = param.get("vscode_extensions") or ""
+            vscode_extensions_cmd = []
+            jupyter_ext_version = True
+            for ext in vscode_extensions.split(","):
+                ext = ext.strip()
+                if not ext:
+                    continue
+
+                if ext.startswith("ms-python.python"):
+                    python_ext_version = python_ext = None
+                elif ext.startswith("ms-toolsai.jupyter"):
+                    jupyter_ext_version = None
+
+                vscode_extensions_cmd += ["--install-extension", ext]
+
+            if python_ext:
+                vscode_extensions_cmd += ["--install-extension", "{}".format(python_ext)]
+            elif python_ext_version:
+                vscode_extensions_cmd += ["--install-extension", "ms-python.python@{}".format(python_ext_version)]
+
+            if jupyter_ext_version:
+                vscode_extensions_cmd += ["--install-extension", "ms-toolsai.jupyter"]
+
+            print("VScode extensions: {}".format(vscode_extensions_cmd))
             subprocess.Popen(
                 [
                     vscode_path,
@@ -347,9 +371,7 @@ def start_vscode_server(hostname, hostnames, param, task, env):
                     "127.0.0.1:{}".format(port),
                     "--user-data-dir", user_folder,
                     "--extensions-dir", exts_folder,
-                    "--install-extension", "ms-toolsai.jupyter",
+                ] + vscode_extensions_cmd,
-                    # "--install-extension", "donjayamanne.python-extension-pack"
-                ] + ["--install-extension", "ms-python.python@{}".format(python_ext_version)] if python_ext else [],
                 env=env,
                 stdout=fd,
                 stderr=fd,
@@ -397,6 +419,7 @@ def start_vscode_server(hostname, hostnames, param, task, env):
 
     except Exception as ex:
         print('Failed running vscode server: {}'.format(ex))
+        task.set_parameter(name='properties/vscode_port', value=str(-1))
         return
 
     task.set_parameter(name='properties/vscode_port', value=str(port))
@@ -475,10 +498,10 @@ def setup_ssh_server(hostname, hostnames, param, task, env):
         max_port = max(min_port+32, int(ssh_port.split(":")[-1]))
         port = get_free_port(min_port, max_port)
         proxy_port = get_free_port(min_port, max_port)
-        use_dropbear = False
+        use_dropbear = bool(param.get("force_dropbear", False))
 
         # if we are root, install open-ssh
-        if os.geteuid() == 0:
+        if not use_dropbear and os.geteuid() == 0:
             # noinspection SpellCheckingInspection
             os.system(
                 "export PYTHONPATH=\"\" && "
@@ -535,7 +558,7 @@ def setup_ssh_server(hostname, hostnames, param, task, env):
                     use_dropbear = True
 
                 except Exception:
-                    print('Error: failed locating SSHd and dailed fetching `dropbear`, leaving!')
+                    print('Error: failed locating SSHd and failed fetching `dropbear`, leaving!')
                     return
 
             # noinspection PyBroadException
@@ -588,7 +611,7 @@ def setup_ssh_server(hostname, hostnames, param, task, env):
                     f.write(v + (' {}@{}'.format(
                         getpass.getuser() or "root", hostname) if filename.endswith('.pub') else ''))
                 os.chmod(filename, 0o600 if filename.endswith('.pub') else 0o600)
-            keys_filename[k] = filename
+                keys_filename[k] = filename
 
         # run server in foreground so it gets killed with us
         if use_dropbear:
@@ -596,7 +619,7 @@ def setup_ssh_server(hostname, hostnames, param, task, env):
             dropbear_key_files = []
             for k, ssh_key_file in keys_filename.items():
                 # skip over the public keys, there is no need for them
-                if ssh_key_file.endswith(".pub"):
+                if not ssh_key_file or ssh_key_file.endswith(".pub"):
                     continue
                 drop_key_file = ssh_key_file + ".dropbear"
                 try:
@@ -607,8 +630,9 @@ def setup_ssh_server(hostname, hostnames, param, task, env):
                 except Exception:
                     pass
             proc_args = [sshd_path, "dropbear", "-e", "-K", "30", "-I", "0", "-F", "-p", str(port)] + dropbear_key_files
-            # this is a copy oof `env` so there is nothing to worry about
+            # this is a copy of `env` so there is nothing to worry about
-            env["DROPBEAR_CLEARML_FIXED_PASSWORD"] = ssh_password
+            if ssh_password:
+                env["DROPBEAR_CLEARML_FIXED_PASSWORD"] = ssh_password
         else:
             proc_args = [sshd_path, "-D", "-p", str(port)] + (["-f", custom_ssh_conf] if custom_ssh_conf else [])
 
@@ -711,7 +735,7 @@ def setup_user_env(param, task):
     except Exception as e:
         print("Error: Exception while trying to create symlink. The Application will continue...")
         print(e)
-        pass
+
     # set default user credentials
     if param.get("user_key") and param.get("user_secret"):
         os.system("echo 'export CLEARML_API_ACCESS_KEY=\"{}\"' >> ~/.bashrc".format(
@@ -889,9 +913,11 @@ def main():
         "user_secret": None,
         "vscode_server": True,
         "vscode_version": '',
+        "vscode_extensions": '',
         "jupyterlab": True,
         "public_ip": False,
         "ssh_ports": None,
+        "force_dropbear": False,
     }
     task = init_task(param, default_ssh_fingerprint)