diff --git a/deepnote_core/kernel.py b/deepnote_core/kernel.py new file mode 100644 index 00000000..27a35e04 --- /dev/null +++ b/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)] diff --git a/deepnote_core/resources/ipython/profile_default/ipython_kernel_config.py b/deepnote_core/resources/ipython/profile_default/ipython_kernel_config.py index 656d9158..0bf81790 100644 --- a/deepnote_core/resources/ipython/profile_default/ipython_kernel_config.py +++ b/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 diff --git a/tests/unit/test_kernel.py b/tests/unit/test_kernel.py new file mode 100644 index 00000000..0d44309d --- /dev/null +++ b/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))