diff --git a/clearml_session/__main__.py b/clearml_session/__main__.py index aef73a7..28d586a 100644 --- a/clearml_session/__main__.py +++ b/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 ' ' (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", - action="store_true", default=False, - help="Automatic yes to prompts; assume \"yes\" as answer " - "to all prompts and run non-interactively",) + parser.add_argument('--yes', '-y', + action='store_true', default=False, + help='Automatic yes to prompts; assume \"yes\" as answer ' + '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) diff --git a/clearml_session/interactive_session_task.py b/clearml_session/interactive_session_task.py index 6d9233a..f3de9a7 100644 --- a/clearml_session/interactive_session_task.py +++ b/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' - python_ext_version = '2021.10.1365161279' + vscode_version = '4.14.1' + 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", - # "--install-extension", "donjayamanne.python-extension-pack" - ] + ["--install-extension", "ms-python.python@{}".format(python_ext_version)] if python_ext else [], + ] + vscode_extensions_cmd, 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 - env["DROPBEAR_CLEARML_FIXED_PASSWORD"] = ssh_password + # this is a copy of `env` so there is nothing to worry about + 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)