fix(kernel): stop restarts waiting 2.5 s to kill the old kernel

deepnote_core/kernel.py

@@ -0,0 +1,59 @@
+"""Kernel class for Deepnote kernels, set with ``c.IPKernelApp.kernel_class``.
+
+It works around ipykernel 6 bugs that make restarts wait 2.5 s for
+jupyter_client to kill the old kernel, and that fail control requests sent
+while the kernel starts. This module is imported before the shell exists, so it
+must not import deepnote_toolkit.
+"""
+
+import asyncio
+import atexit
+from typing import TYPE_CHECKING, Any, List
+
+from ipykernel.ipkernel import IPythonKernel
+
+if TYPE_CHECKING:
+    import psutil
+
+
+def _is_running(process: "psutil.Process") -> bool:
+    """Return whether the process exists and isn't a zombie."""
+    # Imported here because psutil is optional since ipykernel 7, which lists no
+    # children without it
+    import psutil
+
+    try:
+        return process.status() != psutil.STATUS_ZOMBIE
+    except psutil.NoSuchProcess:
+        return False
+
+
+class DeepnoteKernel(IPythonKernel):
+    """IPython kernel with fixes for ipykernel 6 start and shutdown bugs."""
+
+    def __init__(self, **kwargs: Any) -> None:
+        """Create the kernel and register an exit hook for the control thread."""
+        super().__init__(**kwargs)
+        # Exit hooks run last-in, first-out and ipykernel registers its own first,
+        # so this stops the control thread while IOPub can still serve its flushes
+        atexit.register(self._stop_control_thread)
+
+    def start(self) -> None:
+        """Start handling messages, with the control lock created first."""
+        # ipykernel 6 creates it only after control requests start arriving, so an
+        # early one fails. Fixed upstream in 7.4 (ipython/ipykernel#1568).
+        self._control_lock = asyncio.Lock()
+        super().start()
+
+    def _stop_control_thread(self) -> None:
+        """Stop the control thread and wait up to 2 s for it to finish."""
+        thread = self.control_thread
+        if thread is not None and thread.is_alive():
+            thread.stop()
+            thread.join(timeout=2)
+
+    def _process_children(self) -> List["psutil.Process"]:
+        """List child processes, leaving out exited ones nobody has reaped."""
+        # ipykernel's shutdown signals listed children until none are left, which
+        # never happens for an exited one
+        return [child for child in super()._process_children() if _is_running(child)]

deepnote_core/resources/ipython/profile_default/ipython_kernel_config.py

@@ -334,7 +334,7 @@
 #  This should allow easy re-use of the IPKernelApp entry point to configure and
 #  launch kernels other than IPython's own.
 #  Default: 'ipykernel.ipkernel.IPythonKernel'
-# c.IPKernelApp.kernel_class = 'ipykernel.ipkernel.IPythonKernel'
+c.IPKernelApp.kernel_class = "deepnote_core.kernel.DeepnoteKernel"  # noqa: F821
 
 ## The date format used by logging formatters for %(asctime)s
 #  See also: Application.log_datefmt

tests/unit/test_kernel.py

@@ -0,0 +1,157 @@
+"""Tests for the Deepnote kernel class."""
+
+import asyncio
+import json
+import os
+import shutil
+import subprocess
+import sys
+import time
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Iterator, List, Tuple
+
+import psutil
+import pytest
+from ipykernel.control import ControlThread
+from ipykernel.ipkernel import IPythonKernel
+from jupyter_client.blocking.client import BlockingKernelClient
+from jupyter_client.kernelspec import KernelSpecManager
+from jupyter_client.manager import KernelManager
+
+from deepnote_core.config.resources import get_resources_source_path
+from deepnote_core.kernel import DeepnoteKernel, _is_running
+
+
+@pytest.fixture
+def kernel(tmp_path: Path) -> Iterator[Tuple[KernelManager, BlockingKernelClient]]:
+    """A kernel started with the toolkit's IPython profile, without runtime init."""
+    ipython_dir = tmp_path / "ipython"
+    shutil.copytree(get_resources_source_path() / "ipython", ipython_dir)
+
+    # A kernelspec of our own, so the kernel runs on this interpreter
+    spec_dir = tmp_path / "kernels" / "toolkit-test"
+    spec_dir.mkdir(parents=True)
+    argv = [sys.executable, "-m", "ipykernel_launcher", "-f", "{connection_file}"]
+    (spec_dir / "kernel.json").write_text(
+        json.dumps({"argv": argv, "display_name": "Test", "language": "python"})
+    )
+
+    manager = KernelManager(
+        kernel_name="toolkit-test",
+        kernel_spec_manager=KernelSpecManager(kernel_dirs=[str(tmp_path / "kernels")]),
+    )
+    env = {
+        **os.environ,
+        "IPYTHONDIR": str(ipython_dir),
+        "DISABLE_KERNEL_INITIALIZATION_CUSTOM_HOOKS": "1",
+    }
+    manager.start_kernel(env=env)
+    client = manager.client()
+    client.start_channels()
+    try:
+        client.wait_for_ready(timeout=60)
+        yield manager, client
+    finally:
+        client.stop_channels()
+        if manager.is_alive():
+            manager.shutdown_kernel(now=True)
+
+
+class TestKernelWithProfile:
+    """Kernels started with the toolkit's IPython profile."""
+
+    def test_runs_deepnote_kernel(
+        self, kernel: Tuple[KernelManager, BlockingKernelClient]
+    ) -> None:
+        """The profile makes the kernel app create a DeepnoteKernel."""
+        _, client = kernel
+        reply = client.execute_interactive(
+            "",
+            user_expressions={"kernel_class": "type(get_ipython().kernel).__name__"},
+            timeout=30,
+        )
+
+        result = reply["content"]["user_expressions"]["kernel_class"]
+        assert result["data"]["text/plain"] == "'DeepnoteKernel'"
+
+    @pytest.mark.skipif(
+        sys.platform != "linux",
+        reason="Only Linux lists exited children in the kernel's process group",
+    )
+    def test_exited_child_does_not_delay_shutdown(
+        self, kernel: Tuple[KernelManager, BlockingKernelClient]
+    ) -> None:
+        """A kernel with an exited, unreaped child shuts down without SIGTERM."""
+        manager, client = kernel
+        code = (
+            "import subprocess, time, psutil\n"
+            "child = subprocess.Popen(['true'])\n"
+            "while psutil.Process(child.pid).status() != psutil.STATUS_ZOMBIE:\n"
+            "    time.sleep(0.01)\n"
+        )
+        reply = client.execute_interactive(code, timeout=30)
+        assert reply["content"]["status"] == "ok"
+
+        process = manager.provisioner.process
+        manager.shutdown_kernel()
+
+        # A kernel that jupyter_client had to SIGTERM ends with -15
+        assert process.returncode == 0
+
+
+class TestStart:
+    """Starting the kernel's message handlers."""
+
+    def test_creates_control_lock_before_ipykernel_starts(
+        self, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """The control lock exists before ipykernel starts handling control requests."""
+        locks: List[asyncio.Lock] = []
+        monkeypatch.setattr(
+            IPythonKernel, "start", lambda kernel: locks.append(kernel._control_lock)
+        )
+        kernel = DeepnoteKernel.__new__(DeepnoteKernel)
+
+        kernel.start()
+
+        assert isinstance(locks[0], asyncio.Lock)
+
+
+class TestIsRunning:
+    """Telling running child processes from exited ones."""
+
+    def test_only_running_process_counts(self) -> None:
+        """A running process counts, an unreaped or reaped one doesn't."""
+        running = subprocess.Popen(["sleep", "30"])
+        exited = subprocess.Popen(["true"])
+        try:
+            exited_process = psutil.Process(exited.pid)
+            while exited_process.status() != psutil.STATUS_ZOMBIE:
+                time.sleep(0.01)
+
+            assert _is_running(psutil.Process(running.pid))
+            assert not _is_running(exited_process)
+
+            exited.wait()
+            assert not _is_running(exited_process)
+        finally:
+            running.kill()
+            running.wait()
+
+
+class TestStopControlThread:
+    """The exit hook, called on a stand-in because a real kernel needs sockets."""
+
+    def test_stops_running_thread(self) -> None:
+        """The control thread is stopped and joined."""
+        thread = ControlThread(daemon=True)
+        thread.start()
+
+        DeepnoteKernel._stop_control_thread(SimpleNamespace(control_thread=thread))
+
+        assert not thread.is_alive()
+
+    def test_does_nothing_without_thread(self) -> None:
+        """A kernel without a control thread is left alone."""
+        DeepnoteKernel._stop_control_thread(SimpleNamespace(control_thread=None))