Update for jupyter_kernel_mgmt APIs

Author: takluyver

jupyter_ssh_kernels/__main__.py

@@ -1,21 +1,20 @@
 from jupyter_client.client2 import BlockingKernelClient2
-from jupyter_client.manager2 import shutdown
 from .provider import SSHKernelProvider
 
 # Try starting a remote kernel and connecting to it.
-km = SSHKernelProvider().launch('mydesktop')
+conn_info, km = SSHKernelProvider().launch('mydesktop')
 print("Started remote kernel")
 print()
-print(km.get_connection_info())
+print(conn_info)
 print()
 
-kc = BlockingKernelClient2(km.get_connection_info(), km)
+kc = BlockingKernelClient2(conn_info, km)
 print("Getting kernel info...")
 print(kc.kernel_info(reply=True)['content'])
 print()
 
 import time
 time.sleep(5)
 print("Shutting down...")
-shutdown(kc, km)
+kc.shutdown_or_terminate()
 print("Shutdown complete")

jupyter_ssh_kernels/launcher.py

@@ -0,0 +1,39 @@
+import json
+from jupyter_client.manager2 import KernelManager2ABC
+from paramiko import SSHClient
+from pathlib import Path
+import re
+
+def remote_launch_py(kinfo):
+    rl_path = Path(__file__).parent / 'remote_launch.py'
+    with rl_path.open() as f:
+        return f.read().format(**kinfo)
+
+CONN_FILE_RE = r"!!!Started, connection_file: (.*) !!!"
+
+def launch(server, remote_code):
+    client = SSHClient()
+    client.load_system_host_keys()
+    client.connect(server)
+    stdin, stdout, stderr = client.exec_command('python3')
+    stdin.write(remote_code)
+    stdin.close()
+    stdin.channel.shutdown_write()  # Sends EOF
+
+    for line in stdout:
+        m = re.match(CONN_FILE_RE, line)
+        if m:
+            remote_conn_file = m.group(1)
+            break
+    else:
+        print('Exit status', stdin.channel.exit_status)
+        print("Stderr:")
+        print(stderr.read().decode())
+        print("Stdout:")
+        print(stdout.read())
+        raise RuntimeError("Remote kernel failed to start")
+
+    print("Remote connection file:", remote_conn_file)
+    with client.open_sftp() as sftp:
+        with sftp.open(remote_conn_file) as f:
+            return json.load(f)

jupyter_ssh_kernels/manager.py

@@ -3,7 +3,7 @@
 from pathlib import Path
 import pytoml
 
-from .manager import SSHKernelManager, remote_launch_py
+from .manager import launch, remote_launch_py
 
 class SSHKernelProvider(KernelProviderBase):
     id = "simple-ssh"
@@ -29,7 +29,7 @@ def find_kernels(self):
     def launch(self, name, cwd=None):
         for candidate_name, kinfo in self.find_kernels():
             if candidate_name == name:
-                return SSHKernelManager(kinfo['address'],
-                                        remote_launch_py(kinfo))
+                conn_info = launch(kinfo['address'], remote_launch_py(kinfo))
+                return (conn_info, None)
 
         raise KeyError(name)

jupyter_ssh_kernels/provider.py

@@ -1,6 +1,9 @@
 """Sent to Python on the remote system to launch a kernel."""
 
-from jupyter_client.manager2 import KernelManager2
-km = KernelManager2(kernel_cmd={argv!r}, cwd={cwd!r}, ip={address!r})
-print("!!!Started, connection_file:", km.connection_file, '!!!', flush=True)
+from jupyter_kernel_mgmt.subproc import SubprocessKernelLauncher
+from jupyter_kernel_mgmt.nanny import KernelNanny
+l = SubprocessKernelLauncher(kernel_cmd={argv!r}, cwd={cwd!r}, ip={address!r})
+conn_info, mgr = l.launch()
+nanny = KernelNanny(conn_info, mgr)
+print("!!!Started, connection_file:", nanny.connection_file, '!!!', flush=True)
 km.wait(None)

jupyter_ssh_kernels/remote_launch.py

@@ -10,5 +10,5 @@ home-page = "https://github.com/takluyver/jupyter_ssh_kernels"
 description-file = "README.rst"
 classifiers = ["License :: OSI Approved :: MIT License"]
 
-[tool.flit.entrypoints."jupyter_client.kernel_providers"]
+[tool.flit.entrypoints."jupyter_kernel_mgmt.kernel_type_providers"]
 simple-ssh = "jupyter_ssh_kernels:SSHKernelProvider"