Merge pull request #34 from MSeal/zombieProcFix

Zombie Process Fix

Author: MSeal

.travis.yml

@@ -2,8 +2,6 @@ language: python
 sudo: false
 matrix:
   include:
-    - python: 3.5
-      env: TOXENV=py35
     - python: 3.6
       env: TOXENV=py36
     - python: 3.7

README.md

@@ -1,7 +1,6 @@
 [![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/jupyter/nbclient/master?filepath=binder%2Frun_nbclient.ipynb)
 [![Travis Build Status](https://travis-ci.org/jupyter/nbclient.svg?branch=master)](https://travis-ci.org/jupyter/nbclient)
 [![image](https://codecov.io/github/jupyter/nbclient/coverage.svg?branch=master)](https://codecov.io/github/jupyter/nbclient?branch=master)
-[![Python 3.5](https://img.shields.io/badge/python-3.5-blue.svg)](https://www.python.org/downloads/release/python-350/)
 [![Python 3.6](https://img.shields.io/badge/python-3.6-blue.svg)](https://www.python.org/downloads/release/python-360/)
 [![Python 3.7](https://img.shields.io/badge/python-3.7-blue.svg)](https://www.python.org/downloads/release/python-370/)
 [![Python 3.8](https://img.shields.io/badge/python-3.8-blue.svg)](https://www.python.org/downloads/release/python-380/)
@@ -26,7 +25,7 @@ This library used to be part of [nbconvert](https://nbconvert.readthedocs.io/en/
 
 ## Python Version Support
 
-This library currently supports python 3.5+ versions. As minor python
+This library currently supports python 3.6+ versions. As minor python
 versions are officially sunset by the python org nbclient will similarly
 drop support in the future.
 

docs/index.rst

@@ -25,7 +25,7 @@ applications.
 Python Version Support
 ----------------------
 
-This library currently supports python 3.5+ verisons. As minor python
+This library currently supports python 3.6+ verisons. As minor python
 versions are officially sunset by the python org nbclient will similarly
 drop support in the future.
 

nbclient/client.py

@@ -2,9 +2,8 @@
 import base64
 from textwrap import dedent
 
-# 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 async_generator import asynccontextmanager
+from contextlib import contextmanager
 
 from time import monotonic
 from queue import Empty
@@ -15,8 +14,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, ensure_async
 
 
 def timestamp():
@@ -324,7 +329,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 ensure_async(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 ensure_async(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 ensure_async(self.km.cleanup())
+            await ensure_async(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 +372,43 @@ 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 ensure_async(self.km.start_kernel(extra_arguments=self.extra_arguments, **kwargs))
 
         self.kc = self.km.client()
-        self.kc.start_channels()
+        await ensure_async(self.kc.start_channels())
         try:
-            await self.kc.wait_for_ready(timeout=self.startup_timeout)
+            await ensure_async(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 +421,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
+            yield
         finally:
-            self.kc.stop_channels()
-            self.kc = None
+            await self._async_cleanup_kernel()
 
     async def async_execute(self, **kwargs):
         """
@@ -392,15 +438,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 ensure_async(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 +497,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 ensure_async(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 +521,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 ensure_async(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 +545,46 @@ 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 ensure_async(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 ensure_async(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 ensure_async(
+                    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,8 +650,12 @@ 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(
-            cell.source, store_history=store_history, stop_on_error=not self.allow_errors
+        parent_msg_id = await ensure_async(
+            self.kc.execute(
+                cell.source,
+                store_history=store_history,
+                stop_on_error=not self.allow_errors
+            )
         )
         # We launched a code cell to execute
         self.code_cells_executed += 1
@@ -607,11 +665,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