From 811ac3e4998e1c6e569477138c17baa75dcecb58 Mon Sep 17 00:00:00 2001 From: David Brochart Date: Fri, 13 Mar 2020 12:10:30 +0100 Subject: [PATCH 1/2] Use AsyncKernelManager from jupyter_client --- nbclient/client.py | 30 ++++++++++++++-------------- nbclient/tests/fake_kernelmanager.py | 8 ++++---- nbclient/tests/test_client.py | 11 +++++++--- requirements.txt | 2 +- 4 files changed, 28 insertions(+), 23 deletions(-) diff --git a/nbclient/client.py b/nbclient/client.py index d196a339..11551c50 100644 --- a/nbclient/client.py +++ b/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,7 +332,7 @@ 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() @@ -340,7 +340,7 @@ async def start_new_kernel_client(self, **kwargs): 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 diff --git a/nbclient/tests/fake_kernelmanager.py b/nbclient/tests/fake_kernelmanager.py index fb76a689..893f9176 100644 --- a/nbclient/tests/fake_kernelmanager.py +++ b/nbclient/tests/fake_kernelmanager.py @@ -1,7 +1,7 @@ -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): @@ -9,10 +9,10 @@ def __init__(self, *args, **kwargs): 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') diff --git a/nbclient/tests/test_client.py b/nbclient/tests/test_client.py index 988f8539..7cb1a7f6 100644 --- a/nbclient/tests/test_client.py +++ b/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): diff --git a/requirements.txt b/requirements.txt index 0a9e2ad2..6f129963 100644 --- a/requirements.txt +++ b/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 From f1e79487f44ad48ce10bdbb89987b6d347b1a270 Mon Sep 17 00:00:00 2001 From: David Brochart Date: Fri, 20 Mar 2020 08:58:13 +0100 Subject: [PATCH 2/2] Require jupyter_client>=6.1.0 --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 6f129963..19441cd4 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ traitlets>=4.2 -jupyter_client @ git+https://github.com/jupyter/jupyter_client@master +jupyter_client>=6.1.0 nbformat>=5.0 async_generator nest_asyncio