Fix support for jupyterlab v3 and add venv agent support

clearml_session/interactive_session_task.py

@@ -3,6 +3,8 @@ import os
 import socket
 import subprocess
 import sys
+from time import sleep
+
 import requests
 from copy import deepcopy
 from tempfile import mkstemp
@@ -81,7 +83,7 @@ __allocated_ports = []
 def get_free_port(range_min, range_max):
     global __allocated_ports
     used_ports = [i.laddr.port for i in psutil.net_connections()]
-    port = [i for i in range(range_min, range_max) if i not in used_ports and i not in __allocated_ports][0]
+    port = next(i for i in range(range_min, range_max) if i not in used_ports and i not in __allocated_ports)
     __allocated_ports.append(port)
     return port
 
@@ -172,7 +174,7 @@ def monitor_jupyter_server(fd, local_filename, process, task, jupyter_port, host
         for line in new_lines:
             if "http://" not in line and "https://" not in line:
                 continue
-            parts = line.split('/?token=', 1)
+            parts = line.split('?token=', 1)
             if len(parts) != 2:
                 continue
             token = parts[1]
@@ -222,6 +224,7 @@ def start_vscode_server(hostname, hostnames, param, task, env):
     # find a free tcp port
     port = get_free_port(9000, 9100)
 
+    if os.geteuid() == 0:
         # installing VSCODE:
         try:
             python_ext = StorageManager.get_local_copy(
@@ -234,6 +237,18 @@ def start_vscode_server(hostname, hostnames, param, task, env):
         except Exception as ex:
             print("Failed installing vscode server: {}".format(ex))
             return
+        vscode_path = 'code-server'
+    else:
+        python_ext = None
+        # check if code-server exists
+        # noinspection PyBroadException
+        try:
+            vscode_path = subprocess.check_output('which code-server', shell=True).decode().strip()
+            assert vscode_path
+        except Exception:
+            print('Error: Cannot install code-server (not root) and could not find code-server executable, skipping.')
+            task.set_parameter(name='properties/vscode_port', value=str(-1))
+            return
 
     cwd = (
         os.path.expandvars(os.path.expanduser(param["user_base_directory"]))
@@ -255,17 +270,16 @@ def start_vscode_server(hostname, hostnames, param, task, env):
         fd, local_filename = mkstemp()
         subprocess.Popen(
             [
-                "code-server",
+                vscode_path,
                 "--auth",
                 "none",
                 "--bind-addr",
                 "127.0.0.1:{}".format(port),
                 "--user-data-dir", user_folder,
                 "--extensions-dir", exts_folder,
-                "--install-extension", python_ext,
                 "--install-extension", "ms-toolsai.jupyter",
                 # "--install-extension", "donjayamanne.python-extension-pack"
-            ],
+            ] + ["--install-extension", python_ext] if python_ext else [],
             env=env,
             stdout=fd,
             stderr=fd,
@@ -292,8 +306,8 @@ def start_vscode_server(hostname, hostnames, param, task, env):
             pass
         proc = subprocess.Popen(
             ['bash', '-c',
-             'code-server --auth none --bind-addr 127.0.0.1:{} --disable-update-check '
+             '{} --auth none --bind-addr 127.0.0.1:{} --disable-update-check '
-             '--user-data-dir {} --extensions-dir {}'.format(port, user_folder, exts_folder)],
+             '--user-data-dir {} --extensions-dir {}'.format(vscode_path, port, user_folder, exts_folder)],
             env=env,
             stdout=fd,
             stderr=fd,
@@ -313,6 +327,12 @@ def start_vscode_server(hostname, hostnames, param, task, env):
 
 
 def start_jupyter_server(hostname, hostnames, param, task, env):
+    if not param.get('jupyterlab', True):
+        print('no jupyterlab to monitor - going to sleep')
+        while True:
+            sleep(10.)
+        return
+
     # execute jupyter notebook
     fd, local_filename = mkstemp()
     cwd = (
@@ -324,6 +344,10 @@ def start_jupyter_server(hostname, hostnames, param, task, env):
     # find a free tcp port
     port = get_free_port(8888, 9000)
 
+    # if we are not running as root, make sure the sys executable is in the PATH
+    env = dict(**env)
+    env['PATH'] = '{}:{}'.format(Path(sys.executable).parent.as_posix(), env.get('PATH', ''))
+
     # make sure we have the needed cwd
     # noinspection PyBroadException
     try:
@@ -342,7 +366,7 @@ def start_jupyter_server(hostname, hostnames, param, task, env):
             "--no-browser",
             "--allow-root",
             "--ip",
-            "0.0.0.0",
+            "127.0.0.1",
             "--port",
             str(port),
         ],
@@ -365,6 +389,8 @@ def setup_ssh_server(hostname, hostnames, param, task):
         port = get_free_port(10022, 15000)
         proxy_port = get_free_port(10022, 15000)
 
+        # if we are root, install open-ssh
+        if os.geteuid() == 0:
             # noinspection SpellCheckingInspection
             os.system(
                 "export PYTHONPATH=\"\" && "
@@ -388,11 +414,44 @@ def setup_ssh_server(hostname, hostnames, param, task):
                     trains_config_file=os.environ.get("CLEARML_CONFIG_FILE") or os.environ.get("TRAINS_CONFIG_FILE"),
                 )
             )
+            sshd_path = '/usr/sbin/sshd'
+            ssh_config_path = '/etc/ssh/'
+            custom_ssh_conf = None
+        else:
+            # check if sshd exists
+            # noinspection PyBroadException
+            try:
+                sshd_path = subprocess.check_output('which sshd', shell=True).decode().strip()
+                ssh_config_path = os.path.join(os.getcwd(), '.clearml_session_sshd')
+                Path(ssh_config_path).mkdir(parents=True, exist_ok=True)
+                custom_ssh_conf = os.path.join(ssh_config_path, 'sshd_config')
+                with open(custom_ssh_conf, 'wt') as f:
+                    conf = \
+                        "PermitRootLogin yes" + "\n"\
+                        "ClientAliveInterval 10" + "\n"\
+                        "ClientAliveCountMax 20" + "\n"\
+                        "AllowTcpForwarding yes" + "\n"\
+                        "UsePAM yes" + "\n"\
+                        "AuthorizedKeysFile {}".format(os.path.join(ssh_config_path, 'authorized_keys')) + "\n"\
+                        "PidFile {}".format(os.path.join(ssh_config_path, 'sshd.pid')) + "\n"\
+                        "AcceptEnv TRAINS_API_ACCESS_KEY TRAINS_API_SECRET_KEY "\
+                        "CLEARML_API_ACCESS_KEY CLEARML_API_SECRET_KEY"+"\n"
+                    for k in default_ssh_fingerprint:
+                        filename = os.path.join(ssh_config_path, '{}'.format(k.replace('__pub', '.pub')))
+                        conf += "HostKey {}\n".format(filename)
+
+                    f.write(conf)
+            except Exception:
+                print('Error: Cannot install sshd (not root) and could not find sshd executable, leaving!')
+                return
+            # clear the ssh password, we cannot change it
+            ssh_password = None
+            task.set_parameter('{}/ssh_password'.format(config_section_name), '')
 
         # create fingerprint files
-        Path('/etc/ssh/').mkdir(parents=True, exist_ok=True)
+        Path(ssh_config_path).mkdir(parents=True, exist_ok=True)
         for k, v in default_ssh_fingerprint.items():
-            filename = '/etc/ssh/{}'.format(k.replace('__pub', '.pub'))
+            filename = os.path.join(ssh_config_path, '{}'.format(k.replace('__pub', '.pub')))
             try:
                 os.unlink(filename)
             except Exception:  # noqa
@@ -400,12 +459,20 @@ def setup_ssh_server(hostname, hostnames, param, task):
             if v:
                 with open(filename, 'wt') as f:
                     f.write(v + (' root@{}'.format(hostname) if filename.endswith('.pub') else ''))
-                os.chmod(filename, 0o644 if filename.endswith('.pub') else 0o600)
+                os.chmod(filename, 0o600 if filename.endswith('.pub') else 0o600)
 
-        # run server
+        # run server in foreground so it gets killed with us
-        result = os.system("/usr/sbin/sshd -p {port}".format(port=port))
+        proc_args = [sshd_path, "-D", "-p", str(port)] + (["-f", custom_ssh_conf] if custom_ssh_conf else [])
+        proc = subprocess.Popen(args=proc_args)
+        # noinspection PyBroadException
+        try:
+            result = proc.wait(timeout=1)
+        except Exception:
+            result = 0
+
+        if result != 0:
+            raise ValueError("Failed launching sshd: ", proc_args)
 
-        if result == 0:
         # noinspection PyBroadException
         try:
             TcpProxy(listen_port=proxy_port, target_port=port, proxy_state={}, verbose=False,  # noqa
@@ -424,10 +491,9 @@ def setup_ssh_server(hostname, hostnames, param, task):
                 hostname, hostnames, port, ssh_password
             )
         )
-        else:
+
-            raise ValueError()
     except Exception as ex:
-        print("{}\n\n#\n# Error: SSH server could not be launched\n#\n".format(ex))
+        print("Error: {}\n\n#\n# Error: SSH server could not be launched\n#\n".format(ex))
 
 
 def setup_user_env(param, task):
@@ -593,6 +659,7 @@ def main():
         "user_key": None,
         "user_secret": None,
         "vscode_server": True,
+        "jupyterlab": True,
         "public_ip": False,
     }
     task = init_task(param, default_ssh_fingerprint)