Use AsyncKernelManager from jupyter_client

Author: davidbrochart

nbclient/client.py

@@ -218,9 +218,9 @@ class NotebookClient(LoggingConfigurable):
     @default('kernel_manager_class')
     def _kernel_manager_class_default(self):
         """Use a dynamic default to avoid importing jupyter_client at startup"""
-        from jupyter_client import KernelManager
+        from jupyter_client import AsyncKernelManager
 
-        return KernelManager
+        return AsyncKernelManager
 
     _display_id_map = Dict(
         help=dedent(
@@ -317,8 +317,8 @@ async def start_new_kernel_client(self, **kwargs):
         ----------
         kwargs :
             Any options for `self.kernel_manager_class.start_kernel()`. Because
-            that defaults to KernelManager, this will likely include options
-            accepted by `KernelManager.start_kernel()``, which includes `cwd`.
+            that defaults to AsyncKernelManager, this will likely include options
+            accepted by `AsyncKernelManager.start_kernel()``, which includes `cwd`.
 
         Returns
         -------
@@ -332,15 +332,15 @@ async def start_new_kernel_client(self, **kwargs):
         if self.km.ipykernel and self.ipython_hist_file:
             self.extra_arguments += ['--HistoryManager.hist_file={}'.format(self.ipython_hist_file)]
 
-        self.km.start_kernel(extra_arguments=self.extra_arguments, **kwargs)
+        await self.km.start_kernel(extra_arguments=self.extra_arguments, **kwargs)
 
         self.kc = self.km.client()
         self.kc.start_channels()
         try:
             await self.kc.wait_for_ready(timeout=self.startup_timeout)
         except RuntimeError:
             self.kc.stop_channels()
-            self.km.shutdown_kernel()
+            await self.km.shutdown_kernel()
             raise
         self.kc.allow_stdin = False
         return self.kc
@@ -470,8 +470,8 @@ async def _poll_for_reply(self, msg_id, cell, timeout, task_poll_output_msg):
                         timeout = max(0, deadline - monotonic())
             except Empty:
                 # received no message, check if kernel is still alive
-                self._check_alive()
-                self._handle_timeout(timeout, cell)
+                await self._check_alive()
+                await self._handle_timeout(timeout, cell)
 
     async def _poll_output_msg(self, parent_msg_id, cell, cell_index):
         while True:
@@ -494,18 +494,18 @@ def _get_timeout(self, cell):
 
         return timeout
 
-    def _handle_timeout(self, timeout, cell=None):
+    async def _handle_timeout(self, timeout, cell=None):
         self.log.error("Timeout waiting for execute reply (%is)." % timeout)
         if self.interrupt_on_timeout:
             self.log.error("Interrupting kernel")
-            self.km.interrupt_kernel()
+            await self.km.interrupt_kernel()
         else:
             raise CellTimeoutError.error_from_timeout_and_cell(
                 "Cell execution timed out", timeout, cell
             )
 
-    def _check_alive(self):
-        if not self.kc.is_alive():
+    async def _check_alive(self):
+        if not await self.kc.is_alive():
             self.log.error("Kernel died while waiting for execute reply.")
             raise DeadKernelError("Kernel died")
 
@@ -518,10 +518,10 @@ async def _wait_for_reply(self, msg_id, cell=None):
             try:
                 msg = await self.kc.shell_channel.get_msg(timeout=self.shell_timeout_interval)
             except Empty:
-                self._check_alive()
+                await self._check_alive()
                 cummulative_time += self.shell_timeout_interval
                 if timeout and cummulative_time > timeout:
-                    self._handle_timeout(timeout, cell)
+                    await self._handle_timeout(timeout, cell)
                     break
             else:
                 if msg['parent_header'].get('msg_id') == msg_id:
@@ -800,7 +800,7 @@ def execute(nb, cwd=None, km=None, **kwargs):
       The notebook object to be executed
     cwd : str, optional
       If supplied, the kernel will run in this directory
-    km : KernelManager, optional
+    km : AsyncKernelManager, optional
       If supplied, the specified kernel manager will be used for code execution.
     kwargs :
       Any other options for ExecutePreprocessor, e.g. timeout, kernel_name

nbclient/tests/fake_kernelmanager.py

@@ -1,18 +1,18 @@
-from jupyter_client.manager import KernelManager
+from jupyter_client.manager import AsyncKernelManager
 
 
-class FakeCustomKernelManager(KernelManager):
+class FakeCustomKernelManager(AsyncKernelManager):
     expected_methods = {'__init__': 0, 'client': 0, 'start_kernel': 0}
 
     def __init__(self, *args, **kwargs):
         self.log.info('FakeCustomKernelManager initialized')
         self.expected_methods['__init__'] += 1
         super(FakeCustomKernelManager, self).__init__(*args, **kwargs)
 
-    def start_kernel(self, *args, **kwargs):
+    async def start_kernel(self, *args, **kwargs):
         self.log.info('FakeCustomKernelManager started a kernel')
         self.expected_methods['start_kernel'] += 1
-        return super(FakeCustomKernelManager, self).start_kernel(*args, **kwargs)
+        return await super(FakeCustomKernelManager, self).start_kernel(*args, **kwargs)
 
     def client(self, *args, **kwargs):
         self.log.info('FakeCustomKernelManager created a client')

nbclient/tests/test_client.py

@@ -144,7 +144,7 @@ def prepared_wrapper(func):
         def test_mock_wrapper(self):
             """
             This inner function wrapper populates the executor object with
-            the fake kernel client. This client has it's iopub and shell
+            the fake kernel client. This client has its iopub and shell
             channels mocked so as to fake the setup handshake and return
             the messages passed into prepare_cell_mocks as the execute_cell loop
             processes them.
@@ -161,6 +161,7 @@ def test_mock_wrapper(self):
                 iopub_channel=MagicMock(get_msg=message_mock),
                 shell_channel=MagicMock(get_msg=shell_channel_message_mock()),
                 execute=MagicMock(return_value=parent_id),
+                is_alive=MagicMock(return_value=make_async(True))
             )
             executor.parent_id = parent_id
             return func(self, executor, cell_mock, message_mock)
@@ -491,7 +492,7 @@ def test_kernel_death(self):
         km = executor.start_kernel_manager()
 
         with patch.object(km, "is_alive") as alive_mock:
-            alive_mock.return_value = False
+            alive_mock.return_value = make_async(False)
             # Will be a RuntimeError or subclass DeadKernelError depending
             # on if jupyter_client or nbconvert catches the dead client first
             with pytest.raises(RuntimeError):
@@ -672,7 +673,11 @@ def test_busy_message(self, executor, cell_mock, message_mock):
     )
     def test_deadline_exec_reply(self, executor, cell_mock, message_mock):
         # exec_reply is never received, so we expect to hit the timeout.
-        executor.kc.shell_channel.get_msg = MagicMock(side_effect=Empty())
+        async def get_msg(timeout):
+            await asyncio.sleep(timeout)
+            raise Empty
+
+        executor.kc.shell_channel.get_msg = get_msg
         executor.timeout = 1
 
         with pytest.raises(TimeoutError):

requirements.txt

@@ -1,5 +1,5 @@
 traitlets>=4.2
-jupyter_client>=6.0.0
+jupyter_client @ git+https://github.com/jupyter/jupyter_client@master
 nbformat>=5.0
 async_generator
 nest_asyncio