diff --git a/myst_nb/core/execute/__init__.py b/myst_nb/core/execute/__init__.py index b3458714..8a633929 100644 --- a/myst_nb/core/execute/__init__.py +++ b/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) diff --git a/myst_nb/core/execute/base.py b/myst_nb/core/execute/base.py index 8fc12bee..8b6bd0b8 100644 --- a/myst_nb/core/execute/base.py +++ b/myst_nb/core/execute/base.py @@ -58,7 +58,7 @@ def __init__( nb_config: NbParserConfig, logger: LoggerType, *, - kernel_manager: KernelManager | None = None, + kernel_manager_class: type[KernelManager] | None = None, **kwargs: Any, ): """Initialize the client.""" @@ -66,7 +66,7 @@ def __init__( 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] = {} diff --git a/myst_nb/core/execute/cache.py b/myst_nb/core/execute/cache.py index 1918bb94..e4e9dd64 100644 --- a/myst_nb/core/execute/cache.py +++ b/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 diff --git a/myst_nb/core/execute/direct.py b/myst_nb/core/execute/direct.py index b971ceb5..92337861 100644 --- a/myst_nb/core/execute/direct.py +++ b/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: diff --git a/myst_nb/core/execute/inline.py b/myst_nb/core/execute/inline.py index c86a329d..473f759e 100644 --- a/myst_nb/core/execute/inline.py +++ b/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: diff --git a/myst_nb/docutils_.py b/myst_nb/docutils_.py index a141a080..3abece41 100644 --- a/myst_nb/docutils_.py +++ b/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 diff --git a/myst_nb/sphinx_.py b/myst_nb/sphinx_.py index 5948e55d..c6c03d3f 100644 --- a/myst_nb/sphinx_.py +++ b/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 diff --git a/tests/test_execute_kernel_manager.py b/tests/test_execute_kernel_manager.py new file mode 100644 index 00000000..c8a0c8e3 --- /dev/null +++ b/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]