Fixed a number of async issues and enhanced process cleanup

nbclient/client.py

@@ -5,6 +5,7 @@
 # For python 3.5 compatibility we import asynccontextmanager from async_generator instead of
 # contextlib, and we `await yield_()` instead of just `yield`
 from async_generator import asynccontextmanager, async_generator, yield_
+from contextlib import contextmanager
 
 from time import monotonic
 from queue import Empty
@@ -15,8 +16,14 @@
 
 from nbformat.v4 import output_from_msg
 
-from .exceptions import CellTimeoutError, DeadKernelError, CellExecutionComplete, CellExecutionError
-from .util import run_sync
+from .exceptions import (
+ CellControlSignal,
+ CellTimeoutError,
+ DeadKernelError,
+ CellExecutionComplete,
+ CellExecutionError
+)
+from .util import run_sync, await_or_block
 
 
 def timestamp():
@@ -324,7 +331,28 @@ def start_kernel_manager(self):
  self.km.client_class = 'jupyter_client.asynchronous.AsyncKernelClient'
  return self.km
 
- async def start_new_kernel_client(self, **kwargs):
+ async def _async_cleanup_kernel(self):
+ try:
+ # Send a polite shutdown request
+ await await_or_block(self.kc.shutdown)
+ try:
+ # Queue the manager to kill the process, sometimes the built-in and above
+ # shutdowns have not been successful or called yet, so give a direct kill
+ # call here and recover gracefully if it's already dead.
+ await await_or_block(self.km.shutdown_kernel, now=True)
+ except RuntimeError as e:
+ # The error isn't specialized, so we have to check the message
+ if 'No kernel is running!' not in str(e):
+ raise
+ finally:
+ # Remove any state left over even if we failed to stop the kernel
+ await await_or_block(self.km.cleanup)
+ await await_or_block(self.kc.stop_channels)
+ self.kc = None
+
+ _cleanup_kernel = run_sync(_async_cleanup_kernel)
+
+ async def async_start_new_kernel_client(self, **kwargs):
  """Creates a new kernel client.
 
  Parameters
@@ -346,22 +374,44 @@ 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)]
 
- await self.km.start_kernel(extra_arguments=self.extra_arguments, **kwargs)
+ await await_or_block(self.km.start_kernel, extra_arguments=self.extra_arguments, **kwargs)
 
  self.kc = self.km.client()
- self.kc.start_channels()
+ await await_or_block(self.kc.start_channels)
  try:
- await self.kc.wait_for_ready(timeout=self.startup_timeout)
+ await await_or_block(self.kc.wait_for_ready, timeout=self.startup_timeout)
  except RuntimeError:
- self.kc.stop_channels()
- await self.km.shutdown_kernel()
+ await self._async_cleanup_kernel()
  raise
  self.kc.allow_stdin = False
  return self.kc
 
+ start_new_kernel_client = run_sync(async_start_new_kernel_client)
+
+ @contextmanager
+ def setup_kernel(self, **kwargs):
+ """
+ Context manager for setting up the kernel to execute a notebook.
+
+ The assigns the Kernel Manager (`self.km`) if missing and Kernel Client(`self.kc`).
+
+ When control returns from the yield it stops the client's zmq channels, and shuts
+ down the kernel.
+ """
+ # Can't use run_until_complete on an asynccontextmanager function :(
+ if self.km is None:
+ self.start_kernel_manager()
+
+ if not self.km.has_kernel:
+ self.start_new_kernel_client(**kwargs)
+ try:
+ yield
+ finally:
+ self._cleanup_kernel()
+
  @asynccontextmanager
  @async_generator # needed for python 3.5 compatibility
- async def setup_kernel(self, **kwargs):
+ async def async_setup_kernel(self, **kwargs):
  """
  Context manager for setting up the kernel to execute a notebook.
 
@@ -374,12 +424,11 @@ async def setup_kernel(self, **kwargs):
  self.start_kernel_manager()
 
  if not self.km.has_kernel:
- await self.start_new_kernel_client(**kwargs)
+ await self.async_start_new_kernel_client(**kwargs)
  try:
  await yield_(None) # would just yield in python >3.5
  finally:
- self.kc.stop_channels()
- self.kc = None
+ await self._async_cleanup_kernel()
 
  async def async_execute(self, **kwargs):
  """
@@ -392,15 +441,16 @@ async def async_execute(self, **kwargs):
  """
  self.reset_execution_trackers()
 
- async with self.setup_kernel(**kwargs):
+ async with self.async_setup_kernel(**kwargs):
  self.log.info("Executing notebook with kernel: %s" % self.kernel_name)
  for index, cell in enumerate(self.nb.cells):
  # Ignore `'execution_count' in content` as it's always 1
  # when store_history is False
  await self.async_execute_cell(
  cell, index, execution_count=self.code_cells_executed + 1
  )
- info_msg = await self._wait_for_reply(self.kc.kernel_info())
+ msg_id = await await_or_block(self.kc.kernel_info)
+ info_msg = await self.async_wait_for_reply(msg_id)
  self.nb.metadata['language_info'] = info_msg['content']['language_info']
  self.set_widgets_metadata()
 
@@ -450,12 +500,12 @@ def _update_display_id(self, display_id, msg):
  outputs[output_idx]['data'] = out['data']
  outputs[output_idx]['metadata'] = out['metadata']
 
- async def _poll_for_reply(self, msg_id, cell, timeout, task_poll_output_msg):
+ async def _async_poll_for_reply(self, msg_id, cell, timeout, task_poll_output_msg):
  if timeout is not None:
  deadline = monotonic() + timeout
  while True:
  try:
- msg = await self.kc.shell_channel.get_msg(timeout=timeout)
+ msg = await await_or_block(self.kc.shell_channel.get_msg, timeout=timeout)
  if msg['parent_header'].get('msg_id') == msg_id:
  if self.record_timing:
  cell['metadata']['execution']['shell.execute_reply'] = timestamp()
@@ -474,12 +524,12 @@ 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
- await self._check_alive()
- await self._handle_timeout(timeout, cell)
+ await self._async_check_alive()
+ await self._async_handle_timeout(timeout, cell)
 
- async def _poll_output_msg(self, parent_msg_id, cell, cell_index):
+ async def _async_poll_output_msg(self, parent_msg_id, cell, cell_index):
  while True:
- msg = await self.kc.iopub_channel.get_msg(timeout=None)
+ msg = await await_or_block(self.kc.iopub_channel.get_msg, timeout=None)
  if msg['parent_header'].get('msg_id') == parent_msg_id:
  try:
  # Will raise CellExecutionComplete when completed
@@ -498,39 +548,42 @@ def _get_timeout(self, cell):
 
  return timeout
 
- async def _handle_timeout(self, timeout, cell=None):
+ async def _async_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")
- await self.km.interrupt_kernel()
+ await await_or_block(self.km.interrupt_kernel)
  else:
  raise CellTimeoutError.error_from_timeout_and_cell(
  "Cell execution timed out", timeout, cell
  )
 
- async def _check_alive(self):
- if not await self.kc.is_alive():
+ async def _async_check_alive(self):
+ if not await await_or_block(self.kc.is_alive):
  self.log.error("Kernel died while waiting for execute reply.")
  raise DeadKernelError("Kernel died")
 
- async def _wait_for_reply(self, msg_id, cell=None):
+ async def async_wait_for_reply(self, msg_id, cell=None):
  # wait for finish, with timeout
  timeout = self._get_timeout(cell)
  cummulative_time = 0
- self.shell_timeout_interval = 5
  while True:
  try:
- msg = await self.kc.shell_channel.get_msg(timeout=self.shell_timeout_interval)
+ msg = await await_or_block(self.kc.shell_channel.get_msg, timeout=self.shell_timeout_interval)
  except Empty:
- await self._check_alive()
+ await self._async_check_alive()
  cummulative_time += self.shell_timeout_interval
  if timeout and cummulative_time > timeout:
- await self._handle_timeout(timeout, cell)
+ await self._async_async_handle_timeout(timeout, cell)
  break
  else:
  if msg['parent_header'].get('msg_id') == msg_id:
  return msg
 
+ wait_for_reply = run_sync(async_wait_for_reply)
+ # Backwards compatability naming for papermill
+ _wait_for_reply = wait_for_reply
+
  def _timeout_with_deadline(self, timeout, deadline):
  if deadline is not None and deadline - monotonic() < timeout:
  timeout = deadline - monotonic()
@@ -596,7 +649,7 @@ async def async_execute_cell(self, cell, cell_index, execution_count=None, store
  cell['metadata']['execution'] = {}
 
  self.log.debug("Executing cell:\n%s", cell.source)
- parent_msg_id = self.kc.execute(
+ parent_msg_id = await await_or_block(self.kc.execute,
  cell.source, store_history=store_history, stop_on_error=not self.allow_errors
  )
  # We launched a code cell to execute
@@ -607,11 +660,20 @@ async def async_execute_cell(self, cell, cell_index, execution_count=None, store
  self.clear_before_next_output = False
 
  task_poll_output_msg = asyncio.ensure_future(
- self._poll_output_msg(parent_msg_id, cell, cell_index)
- )
- exec_reply = await self._poll_for_reply(
- parent_msg_id, cell, exec_timeout, task_poll_output_msg
+ self._async_poll_output_msg(parent_msg_id, cell, cell_index)
  )
+ try:
+ exec_reply = await self._async_poll_for_reply(
+ parent_msg_id, cell, exec_timeout, task_poll_output_msg
+ )
+ except Exception as e:
+ # Best effort to cancel request if it hasn't been resolved
+ try:
+ # Check if the task_poll_output is doing the raising for us
+ if not isinstance(e, CellControlSignal):
+ task_poll_output_msg.cancel()
+ finally:
+ raise
 
  if execution_count:
  cell['execution_count'] = execution_count

nbclient/exceptions.py

@@ -1,4 +1,13 @@
-class CellTimeoutError(TimeoutError):
+class CellControlSignal(Exception):
+ """
+ A custom exception used to indicate that the exception is used for cell
+ control actions (not the best model, but it's needed to cover existing
+ behavior without major refactors).
+ """
+ pass
+
+
+class CellTimeoutError(TimeoutError, CellControlSignal):
  """
  A custom exception to capture when a cell has timed out during execution.
  """
@@ -21,7 +30,7 @@ class DeadKernelError(RuntimeError):
  pass
 
 
-class CellExecutionComplete(Exception):
+class CellExecutionComplete(CellControlSignal):
  """
  Used as a control signal for cell execution across execute_cell and
  process_message function calls. Raised when all execution requests
@@ -32,7 +41,7 @@ class CellExecutionComplete(Exception):
  pass
 
 
-class CellExecutionError(Exception):
+class CellExecutionError(CellControlSignal):
  """
  Custom exception to propagate exceptions that are raised during
  notebook execution to the caller. This is mostly useful when

nbclient/tests/test_client.py

@@ -36,6 +36,10 @@
 IPY_MAJOR = IPython.version_info[0]
 
 
+class AsyncMock(Mock):
+ pass
+
+
 def make_async(mock_value):
  async def _():
  return mock_value
@@ -116,7 +120,7 @@ def prepare_cell_mocks(*messages, reply_msg=None):
  def shell_channel_message_mock():
  # Return the message generator for
  # self.kc.shell_channel.get_msg => {'parent_header': {'msg_id': parent_id}}
- return MagicMock(
+ return AsyncMock(
  return_value=make_async(NBClientTestsBase.merge_dicts(
  {
  'parent_header': {'msg_id': parent_id},
@@ -129,7 +133,7 @@ def shell_channel_message_mock():
  def iopub_messages_mock():
  # Return the message generator for
  # self.kc.iopub_channel.get_msg => messages[i]
- return Mock(
+ return AsyncMock(
  side_effect=[
  # Default the parent_header so mocks don't need to include this
  make_async(
@@ -386,6 +390,16 @@ def get_time_from_str(s):
  assert status_idle - cell_end < delta
 
 
+def test_synchronous_setup_kernel():
+ nb = nbformat.v4.new_notebook()
+ executor = NotebookClient(nb)
+ with executor.setup_kernel():
+ # Prove it initalized client
+ assert executor.kc is not None
+ # Prove it removed the client (and hopefully cleaned up)
+ assert executor.kc is None
+
+
 class TestExecute(NBClientTestsBase):
  """Contains test functions for execute.py"""
 

nbclient/util.py

@@ -5,6 +5,8 @@
 
 import asyncio
 
+from typing import Coroutine
+
 
 def run_sync(coro):
  """Runs a coroutine and blocks until it has executed.
@@ -45,3 +47,17 @@ def wrapped(self, *args, **kwargs):
  return result
  wrapped.__doc__ = coro.__doc__
  return wrapped
+
+
+async def await_or_block(func, *args, **kwargs):
+ """Awaits the function if it's an asynchronous function. Otherwise block
+ on execution.
+ """
+ if asyncio.iscoroutinefunction(func):
+ return await func(*args, **kwargs)
+ else:
+ result = func(*args, **kwargs)
+ # Mocks mask that the function is a coroutine :/
+ if isinstance(result, Coroutine):
+ return await result
+ return result