Support external kernels (#961)

* Support external kernels

* Ignore linter rule

* Check if connection directory exists

* Use load_connection_info()

Author: davidbrochart

jupyter_client/manager.py

@@ -85,7 +85,8 @@ async def wrapper(self, *args, **kwargs):
             out = await method(self, *args, **kwargs)
             # Add a small sleep to ensure tests can capture the state before done
             await asyncio.sleep(0.01)
-            self._ready.set_result(None)
+            if self.owns_kernel:
+                self._ready.set_result(None)
             return out
         except Exception as e:
             self._ready.set_exception(e)
@@ -105,6 +106,7 @@ class KernelManager(ConnectionFileMixin):
 
     def __init__(self, *args, **kwargs):
         """Initialize a kernel manager."""
+        self._owns_kernel = kwargs.pop("owns_kernel", True)
         super().__init__(**kwargs)
         self._shutdown_status = _ShutdownStatus.Unset
         self._attempted_start = False
@@ -495,6 +497,9 @@ async def _async_shutdown_kernel(self, now: bool = False, restart: bool = False)
             Will this kernel be restarted after it is shutdown. When this
             is True, connection files will not be cleaned up.
         """
+        if not self.owns_kernel:
+            return
+
         self.shutting_down = True  # Used by restarter to prevent race condition
         # Stop monitoring for restarting while we shutdown.
         self.stop_restarter()
@@ -558,6 +563,10 @@ async def _async_restart_kernel(
 
     restart_kernel = run_sync(_async_restart_kernel)
 
+    @property
+    def owns_kernel(self) -> bool:
+        return self._owns_kernel
+
     @property
     def has_kernel(self) -> bool:
         """Has a kernel process been started that we are actively managing."""
@@ -646,6 +655,9 @@ async def _async_signal_kernel(self, signum: int) -> None:
 
     async def _async_is_alive(self) -> bool:
         """Is the kernel process still running?"""
+        if not self.owns_kernel:
+            return True
+
         if self.has_kernel:
             assert self.provisioner is not None
             ret = await self.provisioner.poll()

jupyter_client/multikernelmanager.py

@@ -2,20 +2,23 @@
 # Copyright (c) Jupyter Development Team.
 # Distributed under the terms of the Modified BSD License.
 import asyncio
+import json
 import os
 import socket
 import typing as t
 import uuid
 from functools import wraps
+from pathlib import Path
 
 import zmq
 from traitlets import Any, Bool, Dict, DottedObjectName, Instance, Unicode, default, observe
 from traitlets.config.configurable import LoggingConfigurable
 from traitlets.utils.importstring import import_item
 
+from .connect import KernelConnectionInfo
 from .kernelspec import NATIVE_KERNEL_NAME, KernelSpecManager
 from .manager import KernelManager
-from .utils import ensure_async, run_sync
+from .utils import ensure_async, run_sync, utcnow
 
 
 class DuplicateKernelError(Exception):
@@ -105,9 +108,14 @@ def _context_default(self) -> zmq.Context:
         return zmq.Context()
 
     connection_dir = Unicode("")
+    external_connection_dir = Unicode(None, allow_none=True)
 
     _kernels = Dict()
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.kernel_id_to_connection_file = {}
+
     def __del__(self):
         """Handle garbage collection.  Destroy context if applicable."""
         if self._created_context and self.context and not self.context.closed:
@@ -123,6 +131,51 @@ def __del__(self):
 
     def list_kernel_ids(self) -> t.List[str]:
         """Return a list of the kernel ids of the active kernels."""
+        if self.external_connection_dir is not None:
+            external_connection_dir = Path(self.external_connection_dir)
+            if external_connection_dir.is_dir():
+                connection_files = [p for p in external_connection_dir.iterdir() if p.is_file()]
+
+                # remove kernels (whose connection file has disappeared) from our list
+                k = list(self.kernel_id_to_connection_file.keys())
+                v = list(self.kernel_id_to_connection_file.values())
+                for connection_file in list(self.kernel_id_to_connection_file.values()):
+                    if connection_file not in connection_files:
+                        kernel_id = k[v.index(connection_file)]
+                        del self.kernel_id_to_connection_file[kernel_id]
+                        del self._kernels[kernel_id]
+
+                # add kernels (whose connection file appeared) to our list
+                for connection_file in connection_files:
+                    if connection_file in self.kernel_id_to_connection_file.values():
+                        continue
+                    try:
+                        connection_info: KernelConnectionInfo = json.loads(
+                            connection_file.read_text()
+                        )
+                    except Exception:  # noqa: S112
+                        continue
+                    self.log.debug("Loading connection file %s", connection_file)
+                    if not ("kernel_name" in connection_info and "key" in connection_info):
+                        continue
+                    # it looks like a connection file
+                    kernel_id = self.new_kernel_id()
+                    self.kernel_id_to_connection_file[kernel_id] = connection_file
+                    km = self.kernel_manager_factory(
+                        parent=self,
+                        log=self.log,
+                        owns_kernel=False,
+                    )
+                    km.load_connection_info(connection_info)
+                    km.last_activity = utcnow()
+                    km.execution_state = "idle"
+                    km.connections = 1
+                    km.kernel_id = kernel_id
+                    km.kernel_name = connection_info["kernel_name"]
+                    km.ready.set_result(None)
+
+                    self._kernels[kernel_id] = km
+
         # Create a copy so we can iterate over kernels in operations
         # that delete keys.
         return list(self._kernels.keys())

jupyter_client/utils.py

@@ -4,6 +4,7 @@
 - vendor functions from ipython_genutils that should be retired at some point.
 """
 import os
+from datetime import datetime, timedelta, tzinfo
 
 from jupyter_core.utils import ensure_async, run_sync  # noqa: F401  # noqa: F401
 
@@ -83,3 +84,35 @@ def _expand_path(s):
     if os.name == "nt":
         s = s.replace("IPYTHON_TEMP", "$\\")
     return s
+
+
+# constant for zero offset
+ZERO = timedelta(0)
+
+
+class tzUTC(tzinfo):  # noqa
+    """tzinfo object for UTC (zero offset)"""
+
+    def utcoffset(self, d):
+        """Compute utcoffset."""
+        return ZERO
+
+    def dst(self, d):
+        """Compute dst."""
+        return ZERO
+
+
+UTC = tzUTC()  # type:ignore
+
+
+def utc_aware(unaware):
+    """decorator for adding UTC tzinfo to datetime's utcfoo methods"""
+
+    def utc_method(*args, **kwargs):
+        dt = unaware(*args, **kwargs)
+        return dt.replace(tzinfo=UTC)
+
+    return utc_method
+
+
+utcnow = utc_aware(datetime.utcnow)