Async jupyter client

.travis.yml

@@ -10,22 +10,22 @@ matrix:
       env: TOXENV=py37
     - python: 3.8
       env: TOXENV=py38
-    - python: 3.6
+    - python: 3.8
       env: TOXENV=flake8
-    - python: 3.6
+    - python: 3.8
       env: TOXENV=dist
-    - python: 3.6
+    - python: 3.8
       env: TOXENV=docs
-    - python: 3.6
+    - python: 3.8
       env: TOXENV=manifest
 install:
   - pip install tox coverage codecov
 script:
   - tox -e $TOXENV
 after_success:
   - test $TRAVIS_BRANCH = "master" &&
-    test $TOXENV = "py36" &&
+    test $TOXENV = "py38" &&
     coverage xml -i
   - test $TRAVIS_BRANCH = "master" &&
-    test $TOXENV = "py36" &&
+    test $TOXENV = "py38" &&
     codecov

nbclient/execute.py

@@ -1,8 +1,13 @@
 import base64
 from textwrap import dedent
-from contextlib import contextmanager
+
+# 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 time import monotonic
 from queue import Empty
+import asyncio
 
 from traitlets.config.configurable import LoggingConfigurable
 from traitlets import List, Unicode, Bool, Enum, Any, Type, Dict, Integer, default
@@ -286,9 +291,10 @@ def start_kernel_manager(self):
             self.km = self.kernel_manager_class(config=self.config)
         else:
             self.km = self.kernel_manager_class(kernel_name=self.kernel_name, config=self.config)
+        self.km.client_class = 'jupyter_client.asynchronous.AsyncKernelClient'
         return self.km
 
-    def start_new_kernel_client(self, **kwargs):
+    async def start_new_kernel_client(self, **kwargs):
         """Creates a new kernel client.
 
         Parameters
@@ -315,16 +321,17 @@ def start_new_kernel_client(self, **kwargs):
         self.kc = self.km.client()
         self.kc.start_channels()
         try:
-            self.kc.wait_for_ready(timeout=self.startup_timeout)
+            await self.kc.wait_for_ready(timeout=self.startup_timeout)
         except RuntimeError:
             self.kc.stop_channels()
             self.km.shutdown_kernel()
             raise
         self.kc.allow_stdin = False
         return self.kc
 
-    @contextmanager
-    def setup_kernel(self, **kwargs):
+    @asynccontextmanager
+    @async_generator  # needed for python 3.5 compatibility
+    async def setup_kernel(self, **kwargs):
         """
         Context manager for setting up the kernel to execute a notebook.
 
@@ -337,17 +344,30 @@ def setup_kernel(self, **kwargs):
             self.start_kernel_manager()
 
         if not self.km.has_kernel:
-            self.start_new_kernel_client(**kwargs)
+            await self.start_new_kernel_client(**kwargs)
         try:
-            yield
+            await yield_(None)  # would just yield in python >3.5
         finally:
             self.kc.stop_channels()
             self.kc = None
 
     # TODO: Remove non-kwarg arguments
     def execute(self, **kwargs):
         """
-        Executes each code cell.
+        Executes each code cell (blocking).
+
+        Returns
+        -------
+        nb : NotebookNode
+            The executed notebook.
+        """
+        loop = get_loop()
+        return loop.run_until_complete(self.async_execute(**kwargs))
+
+    # TODO: Remove non-kwarg arguments
+    async def async_execute(self, **kwargs):
+        """
+        Executes each code cell asynchronously.
 
         Returns
         -------
@@ -356,11 +376,11 @@ def execute(self, **kwargs):
         """
         self.reset_execution_trackers()
 
-        with self.setup_kernel(**kwargs):
+        async with self.setup_kernel(**kwargs):
             self.log.info("Executing notebook with kernel: %s" % self.kernel_name)
             for index, cell in enumerate(self.nb.cells):
-                self.execute_cell(cell, index)
-            info_msg = self._wait_for_reply(self.kc.kernel_info())
+                await self.async_execute_cell(cell, index)
+            info_msg = await self._wait_for_reply(self.kc.kernel_info())
             self.nb.metadata['language_info'] = info_msg['content']['language_info']
             self.set_widgets_metadata()
 
@@ -388,14 +408,23 @@ def set_widgets_metadata(self):
 
     def execute_cell(self, cell, cell_index, store_history=True):
         """
-        Executes a single code cell.
+        Executes a single code cell (blocking).
+
+        To execute all cells see :meth:`execute`.
+        """
+        loop = get_loop()
+        return loop.run_until_complete(self.async_execute_cell(cell, cell_index, store_history))
+
+    async def async_execute_cell(self, cell, cell_index, store_history=True):
+        """
+        Executes a single code cell asynchronously.
 
         To execute all cells see :meth:`execute`.
         """
         if cell.cell_type != 'code' or not cell.source.strip():
             return cell
 
-        reply, outputs = self.run_cell(cell, cell_index, store_history)
+        reply, outputs = await self.async_run_cell(cell, cell_index, store_history)
         # Backwards compatibility for processes that wrap run_cell
         cell.outputs = outputs
 
@@ -432,16 +461,22 @@ def _update_display_id(self, display_id, msg):
                 outputs[output_idx]['data'] = out['data']
                 outputs[output_idx]['metadata'] = out['metadata']
 
-    def _poll_for_reply(self, msg_id, cell=None, timeout=None):
-        try:
-            # check with timeout if kernel is still alive
-            msg = self.kc.shell_channel.get_msg(timeout=timeout)
-            if msg['parent_header'].get('msg_id') == msg_id:
-                return msg
-        except Empty:
-            # received no message, check if kernel is still alive
-            self._check_alive()
-            # kernel still alive, wait for a message
+    async def _poll_for_reply(self, msg_id, cell=None, timeout=None):
+        if timeout is not None:
+            deadline = monotonic() + timeout
+        while True:
+            try:
+                # check with timeout if kernel is still alive
+                msg = await self.kc.shell_channel.get_msg(timeout=timeout)
+                if msg['parent_header'].get('msg_id') == msg_id:
+                    return msg
+                else:
+                    if timeout is not None:
+                        timeout = max(0, deadline - monotonic())
+            except Empty:
+                # received no message, check if kernel is still alive
+                self._check_alive()
+                self._handle_timeout(timeout, cell)
 
     def _get_timeout(self, cell):
         if self.timeout_func is not None and cell is not None:
@@ -469,14 +504,14 @@ def _check_alive(self):
             self.log.error("Kernel died while waiting for execute reply.")
             raise DeadKernelError("Kernel died")
 
-    def _wait_for_reply(self, msg_id, cell=None):
+    async def _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 = self.kc.shell_channel.get_msg(timeout=self.shell_timeout_interval)
+                msg = await self.kc.shell_channel.get_msg(timeout=self.shell_timeout_interval)
             except Empty:
                 self._check_alive()
                 cummulative_time += self.shell_timeout_interval
@@ -502,14 +537,41 @@ def _passed_deadline(self, deadline):
         return False
 
     def run_cell(self, cell, cell_index=0, store_history=False):
+        loop = get_loop()
+        return loop.run_until_complete(self.async_run_cell(cell, cell_index, store_history))
+
+    async def poll_output_msg(self, parent_msg_id, cell, cell_index, timeout=None):
+        if timeout is not None:
+            deadline = monotonic() + timeout
+        while True:
+            try:
+                msg = await self.kc.iopub_channel.get_msg(timeout=timeout)
+            except Empty:
+                # timeout
+                if self.raise_on_iopub_timeout:
+                    raise CellTimeoutError.error_from_timeout_and_cell(
+                        "Timeout waiting for IOPub output", self.iopub_timeout, cell
+                    )
+                else:
+                    self.log.warning("Timeout waiting for IOPub output")
+                    return
+            if msg['parent_header'].get('msg_id') != parent_msg_id:
+                # not an output from our execution
+                pass
+            else:
+                try:
+                    # Will raise CellExecutionComplete when completed
+                    self.process_message(msg, cell, cell_index)
+                except CellExecutionComplete:
+                    return
+            if timeout is not None:
+                timeout = max(0, deadline - monotonic())
+
+    async def async_run_cell(self, cell, cell_index=0, store_history=False):
         parent_msg_id = self.kc.execute(
             cell.source, store_history=store_history, stop_on_error=not self.allow_errors
         )
         self.log.debug("Executing cell:\n%s", cell.source)
-        exec_timeout = self._get_timeout(cell)
-        deadline = None
-        if exec_timeout is not None:
-            deadline = monotonic() + exec_timeout
 
         cell.outputs = []
         self.clear_before_next_output = False
@@ -520,57 +582,14 @@ def run_cell(self, cell, cell_index=0, store_history=False):
         # after exec_reply was obtained from shell_channel, leading to the
         # aforementioned dropped data.
 
-        # These two variables are used to track what still needs polling:
-        # more_output=true => continue to poll the iopub_channel
-        more_output = True
-        # polling_exec_reply=true => continue to poll the shell_channel
-        polling_exec_reply = True
-
-        while more_output or polling_exec_reply:
-            if polling_exec_reply:
-                if self._passed_deadline(deadline):
-                    self._handle_timeout(exec_timeout, cell)
-                    polling_exec_reply = False
-                    continue
-
-                # Avoid exceeding the execution timeout (deadline), but stop
-                # after at most 1s so we can poll output from iopub_channel.
-                timeout = self._timeout_with_deadline(1, deadline)
-                exec_reply = self._poll_for_reply(parent_msg_id, cell, timeout)
-                if exec_reply is not None:
-                    polling_exec_reply = False
-
-            if more_output:
-                try:
-                    timeout = self.iopub_timeout
-                    if polling_exec_reply:
-                        # Avoid exceeding the execution timeout (deadline) while
-                        # polling for output.
-                        timeout = self._timeout_with_deadline(timeout, deadline)
-                    msg = self.kc.iopub_channel.get_msg(timeout=timeout)
-                except Empty:
-                    if polling_exec_reply:
-                        # Still waiting for execution to finish so we expect that
-                        # output may not always be produced yet.
-                        continue
-
-                    if self.raise_on_iopub_timeout:
-                        raise CellTimeoutError.error_from_timeout_and_cell(
-                            "Timeout waiting for IOPub output", self.iopub_timeout, cell
-                        )
-                    else:
-                        self.log.warning("Timeout waiting for IOPub output")
-                        more_output = False
-                        continue
-            if msg['parent_header'].get('msg_id') != parent_msg_id:
-                # not an output from our execution
-                continue
-
-            try:
-                # Will raise CellExecutionComplete when completed
-                self.process_message(msg, cell, cell_index)
-            except CellExecutionComplete:
-                more_output = False
+        exec_timeout = self._get_timeout(cell)
+        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)
+        if not task_poll_output_msg.done():
+            task_poll_output_msg.cancel()
+            await self.poll_output_msg(parent_msg_id, cell, cell_index, self.iopub_timeout)
 
         # Return cell.outputs still for backwards compatibility
         return exec_reply, cell.outputs
@@ -721,3 +740,12 @@ def executenb(nb, cwd=None, km=None, **kwargs):
     if cwd is not None:
         resources['metadata'] = {'path': cwd}
     return Executor(nb=nb, resources=resources, km=km, **kwargs).execute()
+
+
+def get_loop():
+    try:
+        loop = asyncio.get_event_loop()
+    except RuntimeError:
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+    return loop