Add --remote-ssh-port --vscode-version

Upgrade vscode version to 3.9.2 and vscode py ext 2021.3.658691958
Add SSH compression by default

clearml_session/__main__.py

@@ -422,6 +422,8 @@ def clone_task(state, project_id):
     task_params["{}/jupyterlab".format(section)] = bool(state.get('jupyter_lab'))
     task_params["{}/vscode_server".format(section)] = bool(state.get('vscode_server'))
     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 ''
     if state.get('user_folder'):
         task_params['{}/user_base_directory'.format(section)] = state.get('user_folder')
     docker = state.get('docker') or task.data.execution.docker_cmd
@@ -543,6 +545,7 @@ def start_ssh_tunnel(username, remote_address, ssh_port, ssh_password, local_rem
     args = ['-N', '-C',
             '{}@{}'.format(username, remote_address), '-p', '{}'.format(ssh_port),
             '-o', 'UserKnownHostsFile=/dev/null',
+            '-o', 'Compression=yes',
             '-o', 'StrictHostKeyChecking=no',
             '-o', 'ServerAliveInterval=10',
             '-o', 'ServerAliveCountMax=10', ]
@@ -767,9 +770,14 @@ def setup_parser(parser):
                         type=lambda x: (str(x).strip().lower() in ('true', 'yes')),
                         help='If True register the public IP of the remote machine. Set if running on the cloud. '
                              'Default: false (use for local / on-premises)')
+    parser.add_argument('--remote-ssh-port', type=str, default=None,
+                        help='Set the remote ssh server port, running on the agent`s machine. (default: 10022)')
     parser.add_argument('--vscode-server', default=True, nargs='?', const='true', metavar='true/false',
                         type=lambda x: (str(x).strip().lower() in ('true', 'yes')),
                         help='Install vscode server (code-server) on interactive session (default: true)')
+    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('--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)')

clearml_session/interactive_session_task.py

@@ -217,6 +217,16 @@ def start_vscode_server(hostname, hostnames, param, task, env):
     if not param.get("vscode_server"):
         return
 
+    # get vscode version and python extension version
+    # they are extremely flaky, this combination works, most do not.
+    vscode_version = '3.9.2'
+    python_ext_version = '2021.3.658691958'
+    if param.get("vscode_version"):
+        vscode_version_parts = param.get("vscode_version").split(':')
+        vscode_version = vscode_version_parts[0]
+        if len(vscode_version_parts) > 1:
+            python_ext_version = vscode_version_parts[1]
+
     # make a copy of env and remove the pythonpath from it.
     env = dict(**env)
     env.pop('PYTHONPATH', None)
@@ -228,10 +238,12 @@ def start_vscode_server(hostname, hostnames, param, task, env):
         # installing VSCODE:
         try:
             python_ext = StorageManager.get_local_copy(
-                'https://github.com/microsoft/vscode-python/releases/download/2020.10.332292344/ms-python-release.vsix',
+                'https://github.com/microsoft/vscode-python/releases/download/{}/ms-python-release.vsix'.format(
+                    python_ext_version),
                 extract_archive=False)
             code_server_deb = StorageManager.get_local_copy(
-                'https://github.com/cdr/code-server/releases/download/v3.7.4/code-server_3.7.4_amd64.deb',
+                'https://github.com/cdr/code-server/releases/download/'
+                'v{version}/code-server_{version}_amd64.deb'.format(version=vscode_version),
                 extract_archive=False)
             os.system("dpkg -i {}".format(code_server_deb))
         except Exception as ex:
@@ -386,8 +398,11 @@ def setup_ssh_server(hostname, hostnames, param, task):
     ssh_password = param.get("ssh_password", "training")
     # noinspection PyBroadException
     try:
-        port = get_free_port(10022, 15000)
+        ssh_port = param.get("ssh_ports") or "10022:15000"
-        proxy_port = get_free_port(10022, 15000)
+        min_port = int(ssh_port.split(":")[0])
+        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)
 
         # if we are root, install open-ssh
         if os.geteuid() == 0:
@@ -666,8 +681,10 @@ def main():
         "user_key": None,
         "user_secret": None,
         "vscode_server": True,
+        "vscode_version": '',
         "jupyterlab": True,
         "public_ip": False,
+        "ssh_ports": None,
     }
     task = init_task(param, default_ssh_fingerprint)