fix: switch to kernel_manager_class

myst_nb/core/execute/__init__.py

@@ -10,7 +10,7 @@
 
 if TYPE_CHECKING:
     from nbformat import NotebookNode
-    from jupyter_client.manager import KernelManager
+    from jupyter_client import KernelManager
 
     from myst_nb.core.config import NbParserConfig
     from myst_nb.core.loggers import LoggerType
@@ -23,7 +23,7 @@ def create_client(
     logger: LoggerType,
     read_fmt: None | dict = None,
     *,
-    kernel_manager: KernelManager | None = None,
+    kernel_manager_class: type[KernelManager] | None = None,
 ) -> NotebookClientBase:
     """Create a notebook execution client, to update its outputs.
 
@@ -63,19 +63,19 @@ def create_client(
 
     if nb_config.execution_mode in ("auto", "force"):
         return NotebookClientDirect(
-            notebook, path, nb_config, logger, kernel_manager=kernel_manager
+            notebook, path, nb_config, logger, kernel_manager_class=kernel_manager_class
         )
 
     if nb_config.execution_mode == "cache":
         return NotebookClientCache(
             *(notebook, path, nb_config, logger),
             read_fmt=read_fmt,
-            kernel_manager=kernel_manager,
+            kernel_manager_class=kernel_manager_class,
         )
 
     if nb_config.execution_mode == "inline":
         return NotebookClientInline(
-            notebook, path, nb_config, logger, kernel_manager=kernel_manager
+            notebook, path, nb_config, logger, kernel_manager_class=kernel_manager_class
         )
 
     return NotebookClientBase(notebook, path, nb_config, logger)

myst_nb/core/execute/base.py

@@ -58,15 +58,15 @@ def __init__(
         nb_config: NbParserConfig,
         logger: LoggerType,
         *,
-        kernel_manager: KernelManager | None = None,
+        kernel_manager_class: type[KernelManager] | None = None,
         **kwargs: Any,
     ):
         """Initialize the client."""
         self._notebook = notebook
         self._path = path
         self._nb_config = nb_config
         self._logger = logger
-        self._kernel_manager = kernel_manager
+        self._kernel_manager_class = kernel_manager_class
         self._kwargs = kwargs
 
         self._glue_data: dict[str, NotebookNode] = {}

myst_nb/core/execute/cache.py

@@ -76,7 +76,11 @@ def start_client(self):
                 allow_errors=self.nb_config.execution_allow_errors,
                 timeout=self.nb_config.execution_timeout,
                 meta_override=True,  # TODO still support this?
-                km=self._kernel_manager,
+                **(
+                    dict(kernel_manager_class=self._kernel_manager_class)
+                    if self._kernel_manager_class
+                    else {}
+                ),
             )
 
         # handle success / failure cases

myst_nb/core/execute/direct.py

@@ -44,7 +44,11 @@ def start_client(self):
                 allow_errors=self.nb_config.execution_allow_errors,
                 timeout=self.nb_config.execution_timeout,
                 meta_override=True,  # TODO still support this?
-                km=self._kernel_manager,
+                **(
+                    dict(kernel_manager_class=self._kernel_manager_class)
+                    if self._kernel_manager_class
+                    else {}
+                ),
             )
 
         if result.err is not None:

myst_nb/core/execute/inline.py

@@ -9,6 +9,7 @@
 from tempfile import mkdtemp
 import time
 import traceback
+from typing import Any, cast
 
 from nbclient.client import (
     CellControlSignal,
@@ -56,7 +57,12 @@ def start_client(self):
             resources=resources,
             allow_errors=self.nb_config.execution_allow_errors,
             timeout=self.nb_config.execution_timeout,
-            km=self._kernel_manager,
+            **cast(
+                "dict[str, Any]",
+                dict(kernel_manager_class=self._kernel_manager_class)
+                if self._kernel_manager_class
+                else {},
+            ),
         )
         self._client.reset_execution_trackers()
         if self._client.km is None:

myst_nb/docutils_.py

@@ -87,9 +87,14 @@ class Parser(MystParser):
 
     config_section = "myst-nb parser"
 
-    def __init__(self, *args, kernel_manager: KernelManager | None = None, **kwargs):
+    def __init__(
+        self,
+        *args,
+        kernel_manager_class: type[KernelManager] | None = None,
+        **kwargs,
+    ):
         super().__init__(*args, **kwargs)
-        self.kernel_manager = kernel_manager
+        self.kernel_manager_class = kernel_manager_class
 
     def parse(self, inputstring: str, document: nodes.document) -> None:
         # register/unregister special directives and roles
@@ -197,7 +202,7 @@ def _parse(self, inputstring: str, document: nodes.document) -> None:
         # this may execute the notebook immediately or during the page render
         with create_client(
             *(notebook, document_source, nb_config, logger),
-            kernel_manager=self.kernel_manager,
+            kernel_manager_class=self.kernel_manager_class,
         ) as nb_client:
             mdit_parser.options["nb_client"] = nb_client
             # convert to docutils AST, which is added to the document

myst_nb/sphinx_.py

@@ -68,9 +68,14 @@ class Parser(MystParser):
 
     env: SphinxEnvType
 
-    def __init__(self, *args, kernel_manager: KernelManager | None = None, **kwargs):
+    def __init__(
+        self,
+        *args,
+        kernel_manager_class: type[KernelManager] | None = None,
+        **kwargs,
+    ):
         super().__init__(*args, **kwargs)
-        self.kernel_manager = kernel_manager
+        self.kernel_manager_class = kernel_manager_class
 
     def parse(self, inputstring: str, document: nodes.document) -> None:
         """Parse source text.
@@ -163,7 +168,7 @@ def parse(self, inputstring: str, document: nodes.document) -> None:
         # this may execute the notebook immediately or during the page render
         with create_client(
             *(notebook, document_path, nb_config, logger, nb_reader.read_fmt),
-            kernel_manager=self.kernel_manager,
+            kernel_manager_class=self.kernel_manager_class,
         ) as nb_client:
             mdit_parser.options["nb_client"] = nb_client
             # convert to docutils AST, which is added to the document

tests/test_execute_kernel_manager.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import gc
+import warnings
+from typing import TYPE_CHECKING
+
+import nbformat
+import pytest
+from jupyter_client import AsyncKernelManager
+
+from myst_nb.core.config import NbParserConfig
+from myst_nb.core.execute import create_client
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+
+class _FakeLogger:
+    def info(self, *args, **kwargs):
+        pass
+
+    def warning(self, *args, **kwargs):
+        pass
+
+    def debug(self, *args, **kwargs):
+        pass
+
+
+def _notebook():
+    return nbformat.v4.new_notebook(
+        metadata=nbformat.NotebookNode(
+            kernelspec=dict(name="python3", display_name="", language="python")
+        ),
+        cells=[nbformat.v4.new_code_cell("1 + 1")],
+    )
+
+
+@pytest.mark.parametrize("execution_mode", ["force", "cache", "inline"])
+def test_external_kernel_manager_client_is_closed(tmp_path: Path, execution_mode):
+    nb_path = tmp_path / "nb.ipynb"
+    nbformat.write(_notebook(), nb_path)
+    nb_config = NbParserConfig(
+        execution_mode=execution_mode,
+        execution_in_temp=True,
+        execution_cache_path=str(tmp_path / ".jupyter_cache"),
+    )
+
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        with create_client(
+            _notebook(),
+            str(nb_path),
+            nb_config,
+            _FakeLogger(),
+            kernel_manager_class=AsyncKernelManager,
+        ) as client:
+            if execution_mode == "inline":
+                # inline mode executes lazily, on request for a cell's outputs
+                client.code_cell_outputs(0)
+        assert client.exec_metadata["succeeded"]
+        gc.collect()
+
+    unclosed = [
+        w
+        for w in caught
+        if issubclass(w.category, ResourceWarning) and "Unclosed" in str(w.message)
+    ]
+    assert not unclosed, [str(w.message) for w in unclosed]