Make sure kernel is cleaned up in case an error occurred while starting kernel client (#234)

* Make sure kernel is cleaned up in case an error ocurred while starting kernel client

* Add test

* Run pre-commit

* Include kernel id in log

* Fix [no-untyped-def] warning

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Pin traitlets>=5.2.2

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: David Brochart <[email protected]>

Author: 3 people

nbclient/cli.py

@@ -35,12 +35,12 @@ class NbClientApp(JupyterApp):
     An application used to execute notebook files (``*.ipynb``)
     """
 
-    version = __version__
+    version = Unicode(__version__)
     name = 'jupyter-execute'
     aliases = nbclient_aliases
     flags = nbclient_flags
 
-    description = Unicode("An application used to execute notebook files (*.ipynb)")
+    description = "An application used to execute notebook files (*.ipynb)"
     notebooks = List([], help="Path of notebooks to convert").tag(config=True)
     timeout: int = Integer(
         None,

nbclient/client.py

@@ -39,7 +39,7 @@
 from .util import ensure_async, run_hook, run_sync
 
 
-def timestamp(msg: t.Optional[Dict] = None) -> str:
+def timestamp(msg: t.Optional[t.Dict] = None) -> str:
     if msg and 'header' in msg:  # The test mocks don't provide a header, so tolerate that
         msg_header = msg['header']
         if 'date' in msg_header and isinstance(msg_header['date'], datetime.datetime):
@@ -288,7 +288,9 @@ class NotebookClient(LoggingConfigurable):
         """,
     ).tag(config=True)
 
-    kernel_manager_class: KernelManager = Type(config=True, help='The kernel manager class to use.')
+    kernel_manager_class = Type(
+        config=True, klass=KernelManager, help='The kernel manager class to use.'
+    )
 
     on_notebook_start: t.Optional[t.Callable] = Callable(
         default_value=None,
@@ -390,7 +392,7 @@ def _kernel_manager_class_default(self) -> t.Type[KernelManager]:
 
         return AsyncKernelManager
 
-    _display_id_map: t.Dict[str, t.Dict] = Dict(
+    _display_id_map = Dict(
         help=dedent(
             """
               mapping of locations of outputs with a given display_id
@@ -423,7 +425,7 @@ def _kernel_manager_class_default(self) -> t.Type[KernelManager]:
             """,
     ).tag(config=True)
 
-    resources: t.Dict = Dict(
+    resources = Dict(
         help=dedent(
             """
             Additional resources used in the conversion process. For example,
@@ -557,11 +559,16 @@ async def async_start_new_kernel_client(self) -> KernelClient:
             Kernel client as created by the kernel manager ``km``.
         """
         assert self.km is not None
-        self.kc = self.km.client()
-        await ensure_async(self.kc.start_channels())  # type:ignore[func-returns-value]
         try:
-            await ensure_async(self.kc.wait_for_ready(timeout=self.startup_timeout))
-        except RuntimeError:
+            self.kc = self.km.client()
+            await ensure_async(self.kc.start_channels())  # type:ignore[func-returns-value]
+            await ensure_async(self.kc.wait_for_ready(timeout=self.startup_timeout))  # type:ignore
+        except Exception as e:
+            self.log.error(
+                "Error occurred while starting new kernel client for kernel {}: {}".format(
+                    self.km.kernel_id, str(e)
+                )
+            )
             await self._async_cleanup_kernel()
             raise
         self.kc.allow_stdin = False
@@ -846,7 +853,7 @@ async def _async_handle_timeout(
 
     async def _async_check_alive(self) -> None:
         assert self.kc is not None
-        if not await ensure_async(self.kc.is_alive()):
+        if not await ensure_async(self.kc.is_alive()):  # type:ignore
             self.log.error("Kernel died while waiting for execute reply.")
             raise DeadKernelError("Kernel died")
 

nbclient/tests/test_client.py

@@ -15,7 +15,7 @@
 import nbformat
 import pytest
 import xmltodict
-from jupyter_client import KernelManager
+from jupyter_client import KernelClient, KernelManager
 from jupyter_client.kernelspec import KernelSpecManager
 from nbconvert.filters import strip_ansi
 from nbformat import NotebookNode
@@ -211,7 +211,11 @@ def test_mock_wrapper(self):
             cell_mock = NotebookNode(
                 source='"foo" = "bar"', metadata={}, cell_type='code', outputs=[]
             )
-            executor = NotebookClient({})  # type:ignore
+
+            class NotebookClientWithParentID(NotebookClient):
+                parent_id: str
+
+            executor = NotebookClientWithParentID({})  # type:ignore
             executor.nb = {'cells': [cell_mock]}  # type:ignore
 
             # self.kc.iopub_channel.get_msg => message_mock.side_effect[i]
@@ -508,6 +512,38 @@ def test_start_new_kernel_history_file_setting():
     kc.stop_channels()
 
 
+def test_start_new_kernel_client_cleans_up_kernel_on_failure():
+    class FakeClient(KernelClient):
+        def start_channels(
+            self,
+            shell: bool = True,
+            iopub: bool = True,
+            stdin: bool = True,
+            hb: bool = True,
+            control: bool = True,
+        ) -> None:
+            raise Exception("Any error")
+
+        def stop_channels(self) -> None:
+            pass
+
+    nb = nbformat.v4.new_notebook()
+    km = KernelManager()
+    km.client_factory = FakeClient
+    executor = NotebookClient(nb, km=km)
+    executor.start_new_kernel()
+    assert km.has_kernel
+    assert executor.km is not None
+
+    with pytest.raises(Exception) as err:
+        executor.start_new_kernel_client()
+
+    assert str(err.value.args[0]) == "Any error"
+    assert executor.kc is None
+    assert executor.km is None
+    assert not km.has_kernel
+
+
 class TestExecute(NBClientTestsBase):
     """Contains test functions for execute.py"""
 

requirements.txt

@@ -1,4 +1,4 @@
 jupyter_client>=6.1.5
 nbformat>=5.0
 nest_asyncio
-traitlets>=5.0.0
+traitlets>=5.2.2