Async cell execution

.github/workflows/main.yml

@@ -18,7 +18,7 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu-latest, macos-latest]
-        python-version: [3.5, 3.6, 3.7]
+        python-version: [3.6, 3.7, 3.8]
 
     steps:
     - uses: actions/checkout@v1
@@ -34,25 +34,25 @@ jobs:
     - name: Update conda
       run: |
         conda update -y -n base conda setuptools
-    
+
     - name: Init conda
       run: |
         conda init bash
         conda info -a
 
     - name: Create the conda environment
-      run: conda create -q -y -n voila-tests -c conda-forge python=$PYTHON_VERSION pip jupyterlab_pygments==0.1.0 nbconvert=5.5 pytest-cov nodejs flake8 ipywidgets matplotlib xeus-cling
+      run: conda create -q -y -n voila-tests -c conda-forge python=$PYTHON_VERSION pip jupyterlab_pygments==0.1.0 pytest-cov nodejs flake8 ipywidgets matplotlib xeus-cling
       env:
         PYTHON_VERSION: ${{ matrix.python-version }}
 
-    - name: Install dependencies 
+    - name: Install dependencies
       run: |
         source "$CONDA/etc/profile.d/conda.sh"
         conda activate voila-tests
         whereis python
         python --version
-        
-        python -m pip install ".[test]"
+
+        python -m pip install --ignore-installed ".[test]"
         cd tests/test_template; pip install .; cd ../../;
 
     - name: Flake8

.travis.yml

@@ -4,9 +4,9 @@ os:
  - osx
 env:
  matrix:
-  - PYTHON_VERSION=3.5
   - PYTHON_VERSION=3.6
   - PYTHON_VERSION=3.7
+  - PYTHON_VERSION=3.8
 before_install:
   - if [[ $TRAVIS_OS_NAME == osx ]]; then ulimit -n 2048; fi
   - if [[ $TRAVIS_OS_NAME == linux ]]; then sudo apt-get update; fi
@@ -18,10 +18,10 @@ before_install:
   - conda config --set always_yes yes --set changeps1 no
   - conda update -q conda
   - conda info -a
-  - conda create -q -n test-environment -c conda-forge python=$PYTHON_VERSION jupyterlab_pygments==0.1.0 nbconvert=5.5 pytest-cov nodejs flake8 ipywidgets matplotlib xeus-cling
+  - conda create -q -n test-environment -c conda-forge python=$PYTHON_VERSION jupyterlab_pygments==0.1.0 pytest-cov nodejs flake8 ipywidgets matplotlib xeus-cling
   - source activate test-environment
 install:
-  - pip install ".[test]"
+  - pip install --ignore-installed ".[test]"
   - cd tests/test_template; pip install .; cd ../../;
 script:
   - VOILA_TEST_DEBUG=1 VOILA_TEST_XEUS_CLING=1 py.test tests/ --async-test-timeout=240

setup.py

@@ -377,15 +377,17 @@ def get_data_files():
     'install_requires': [
         'async_generator',
         'jupyter_server>=0.1.0,<0.2.0',
-        'nbconvert>=5.5.0,<6',
+        'jupyter_client>=6.0.0',
+        'nbclient @ git+https://github.com/jupyter/nbclient',
+        'nbconvert @ git+https://github.com/jupyter/nbconvert',
         'jupyterlab_pygments>=0.1.0,<0.2',
         'pygments>=2.4.1,<3'  # Explicitly requiring pygments which is a second-order dependency.
                               # An older versions is generally installed already and is otherwise not updated by pip.
     ],
     'extras_require': {
         'test': [
             'mock',
-            'pytest<4',
+            'pytest',
             'pytest-tornado',
             'matplotlib',
             'ipywidgets'

share/jupyter/voila/templates/default/nbconvert_templates/lab.tpl

@@ -1,4 +1,4 @@
-{%- extends 'display_priority.tpl' -%}
+{%- extends 'display_priority.j2' -%}
 
 {% block codecell %}
 {%- if not cell.outputs -%}

share/jupyter/voila/templates/default/nbconvert_templates/voila.tpl

@@ -62,7 +62,7 @@ a.anchor-link {
 
 {%- block body_loop -%}
   {# from this point on, the kernel is started #}
-  {%- with kernel_id = kernel_start() -%}
+  {%- with kernel_id = kernel_start(nb) -%}
     <script id="jupyter-config-data" type="application/json">
     {
         "baseUrl": "{{resources.base_url}}",

voila/app.py

@@ -59,7 +59,7 @@
 from ._version import __version__
 from .static_file_handler import MultiStaticFileHandler, WhiteListFileHandler
 from .configuration import VoilaConfiguration
-from .execute import VoilaExecutePreprocessor
+from .execute import VoilaExecutor
 from .exporter import VoilaExporter
 from .csspreprocessor import VoilaCSSPreprocessor
 
@@ -129,7 +129,7 @@ class Voila(Application):
     }
     classes = [
         VoilaConfiguration,
-        VoilaExecutePreprocessor,
+        VoilaExecutor,
         VoilaExporter,
         VoilaCSSPreprocessor
     ]

voila/execute.py

@@ -11,7 +11,8 @@
 from time import monotonic
 
 from nbconvert.preprocessors import ClearOutputPreprocessor
-from nbconvert.preprocessors.execute import CellExecutionError, ExecutePreprocessor
+from nbclient.exceptions import CellExecutionError
+from nbclient import NotebookClient
 from nbformat.v4 import output_from_msg
 import zmq
 
@@ -113,7 +114,7 @@ def set_state(self, state):
                 self.msg_id = msg_id
 
 
-class VoilaExecutePreprocessor(ExecutePreprocessor):
+class VoilaExecutor(NotebookClient):
     """Execute, but respect the output widget behaviour"""
     cell_error_instruction = Unicode(
         'Please run Voila with --debug to see the error message.',
@@ -124,21 +125,21 @@ class VoilaExecutePreprocessor(ExecutePreprocessor):
     )
 
     cell_timeout_instruction = Unicode(
-        'Please run Voila with --VoilaExecutePreprocessor.interrupt_on_timeout=True to continue executing the rest of the notebook.',
+        'Please run Voila with --VoilaExecutor.interrupt_on_timeout=True to continue executing the rest of the notebook.',
         config=True,
         help=(
             'instruction given to user to continue execution on timeout'
         )
     )
 
-    def __init__(self, **kwargs):
-        super(VoilaExecutePreprocessor, self).__init__(**kwargs)
+    def __init__(self, nb, km=None, **kwargs):
+        super(VoilaExecutor, self).__init__(nb, km=km, **kwargs)
         self.output_hook_stack = collections.defaultdict(list)  # maps to list of hooks, where the last is used
         self.output_objects = {}
 
-    def preprocess(self, nb, resources, km=None):
+    def execute(self, nb, resources, km=None):
         try:
-            result = super(VoilaExecutePreprocessor, self).preprocess(nb, resources=resources, km=km)
+            result = super(VoilaExecutor, self).execute()
         except CellExecutionError as e:
             self.log.error(e)
             result = (nb, resources)
@@ -149,11 +150,11 @@ def preprocess(self, nb, resources, km=None):
 
         return result
 
-    def preprocess_cell(self, cell, resources, cell_index, store_history=True):
+    async def execute_cell(self, cell, resources, cell_index, store_history=True):
         try:
             # TODO: pass store_history as a 5th argument when we can require nbconver >=5.6.1
-            # result = super(VoilaExecutePreprocessor, self).preprocess_cell(cell, resources, cell_index, store_history)
-            result = super(VoilaExecutePreprocessor, self).preprocess_cell(cell, resources, cell_index)
+            # result = super(VoilaExecutor, self).execute_cell(cell, resources, cell_index, store_history)
+            result = await super(VoilaExecutor, self).async_execute_cell(cell, cell_index)
         except TimeoutError as e:
             self.log.error(e)
             self.show_code_cell_timeout(cell)
@@ -186,10 +187,10 @@ def output(self, outs, msg, display_id, cell_index):
             hook = self.output_hook_stack[parent_msg_id][-1]
             hook.output(outs, msg, display_id, cell_index)
             return
-        super(VoilaExecutePreprocessor, self).output(outs, msg, display_id, cell_index)
+        super(VoilaExecutor, self).output(outs, msg, display_id, cell_index)
 
     def handle_comm_msg(self, outs, msg, cell_index):
-        super(VoilaExecutePreprocessor, self).handle_comm_msg(outs, msg, cell_index)
+        super(VoilaExecutor, self).handle_comm_msg(outs, msg, cell_index)
         self.log.debug('comm msg: %r', msg)
         if msg['msg_type'] == 'comm_open' and msg['content'].get('target_name') == 'jupyter.widget':
             content = msg['content']
@@ -213,7 +214,7 @@ def clear_output(self, outs, msg, cell_index):
             hook = self.output_hook_stack[parent_msg_id][-1]
             hook.clear_output(outs, msg, cell_index)
             return
-        super(VoilaExecutePreprocessor, self).clear_output(outs, msg, cell_index)
+        super(VoilaExecutor, self).clear_output(outs, msg, cell_index)
 
     def strip_notebook_errors(self, nb):
         """Strip error messages and traceback from a Notebook."""
@@ -344,5 +345,5 @@ def executenb(nb, cwd=None, km=None, **kwargs):
         resources['metadata'] = {'path': cwd}  # pragma: no cover
     # Clear any stale output, in case of exception
     nb, resources = ClearOutputPreprocessor().preprocess(nb, resources)
-    ep = VoilaExecutePreprocessor(**kwargs)
-    return ep.preprocess(nb, resources, km=km)[0]
+    executor = VoilaExecutor(nb, km=km, **kwargs)
+    return executor.execute(nb, resources, km=km)

voila/exporter.py

@@ -19,12 +19,6 @@
 from nbconvert.exporters.templateexporter import TemplateExporter
 from nbconvert.filters.highlight import Highlight2HTML
 
-from .threading import async_generator_to_thread
-
-# As long as we support Python35, we use this library to get as async
-# generators: https://pypi.org/project/async_generator/
-from async_generator import async_generator, yield_
-
 
 class VoilaMarkdownRenderer(IPythonRenderer):
     """Custom markdown renderer that inlines images"""
@@ -83,7 +77,6 @@ def _default_preprocessors(self):
     def default_template_file(self):
         return 'voila.tpl'
 
-    @async_generator
     async def generate_from_notebook_node(self, nb, resources=None, extra_context={}, **kw):
         # this replaces from_notebook_node, but calls template.generate instead of template.render
         langinfo = nb.metadata.get('language_info', {})
@@ -106,24 +99,12 @@ async def generate_from_notebook_node(self, nb, resources=None, extra_context={}
                 'no_prompt': self.exclude_input_prompt and self.exclude_output_prompt,
                 }
 
-        # Jinja with Python3.5 does not support async (generators), which
-        # means that it's not async all the way down. Which means that we
-        # cannot use coroutines for the cell_generator, and that they will
-        # block the IO loop. In that case we will run the iterator in a
-        # thread instead.
-
-        @async_generator_to_thread
-        @async_generator
-        async def async_jinja_generator():
-            # Top level variables are passed to the template_exporter here.
-            for output in self.template.generate(nb=nb_copy, resources=resources, **extra_context):
-                await yield_((output, resources))
-
-        async for output, resources in async_jinja_generator():
-            await yield_((output, resources))
+        async for output in self.template.generate_async(nb=nb_copy, resources=resources, **extra_context):
+            yield (output, resources)
 
     @property
     def environment(self):
+        self.enable_async = True
         env = super(type(self), self).environment
         if 'jinja2.ext.do' not in env.extensions:
             env.add_extension('jinja2.ext.do')

voila/handler.py

@@ -18,7 +18,7 @@
 
 from nbconvert.preprocessors import ClearOutputPreprocessor
 
-from .execute import executenb, VoilaExecutePreprocessor
+from .execute import executenb, VoilaExecutor
 from .exporter import VoilaExporter
 
 
@@ -96,16 +96,11 @@ async def get(self, path=None):
         # Template should first call kernel_start, and then decide to use notebook_execute
         # or cell_generator to implement progressive cell rendering
         extra_context = {
-            # NOTE: we can remove the lambda is we use jinja's async feature, which will automatically await the future
-            'kernel_start': lambda: self._jinja_kernel_start().result(),  # pass the result (not the future) to the template
+            'kernel_start': self._jinja_kernel_start,
             'cell_generator': self._jinja_cell_generator,
             'notebook_execute': self._jinja_notebook_execute,
         }
 
-        # Currenly _jinja_kernel_start is executed from a different thread, which causes the websocket connection from
-        # the frontend to fail. Instead, we start it beforehand, and just return the kernel_id in _jinja_kernel_start
-        self.kernel_id = await tornado.gen.maybe_future(self.kernel_manager.start_kernel(kernel_name=self.notebook.metadata.kernelspec.name, path=self.cwd))
-
         # Compose reply
         self.set_header('Content-Type', 'text/html')
         # render notebook in snippets, and flush them out to the browser can render progresssively
@@ -118,12 +113,23 @@ async def get(self, path=None):
     def redirect_to_file(self, path):
         self.redirect(url_path_join(self.base_url, 'voila', 'files', path))
 
-    @tornado.gen.coroutine
-    def _jinja_kernel_start(self):
+    async def _jinja_kernel_start(self, nb):
         assert not self.kernel_started, "kernel was already started"
-        # See command above aboout not being able to start the kernel from a different thread
+        kernel_id = await self.kernel_manager.start_kernel(kernel_name=self.notebook.metadata.kernelspec.name, path=self.cwd)
+        km = self.kernel_manager.get_kernel(kernel_id)
+        km.client_class = 'jupyter_client.asynchronous.AsyncKernelClient'
+        self.executor = VoilaExecutor(nb, km=km, config=self.traitlet_config)
+        self.executor.kc = km.client()
+        self.executor.kc.start_channels()
+        try:
+            await self.executor.kc.wait_for_ready(timeout=self.executor.startup_timeout)
+        except RuntimeError:
+            self.executor.kc.stop_channels()
+            self.executor.km.shutdown_kernel()
+            raise
+        self.executor.kc.allow_stdin = False
         self.kernel_started = True
-        return self.kernel_id
+        return kernel_id
 
     def _jinja_notebook_execute(self, nb, kernel_id):
         km = self.kernel_manager.get_kernel(kernel_id)
@@ -133,12 +139,10 @@ def _jinja_notebook_execute(self, nb, kernel_id):
         # see the updated variable (it seems to be local to our block)
         nb.cells = result.cells
 
-    def _jinja_cell_generator(self, nb, kernel_id):
+    async def _jinja_cell_generator(self, nb, kernel_id):
         """Generator that will execute a single notebook cell at a time"""
-        km = self.kernel_manager.get_kernel(kernel_id)
-
         nb, resources = ClearOutputPreprocessor().preprocess(nb, {'metadata': {'path': self.cwd}})
-        ep = VoilaExecutePreprocessor(config=self.traitlet_config)
+        ep = VoilaExecutor(config=self.traitlet_config)
 
         stop_execution = False
         with ep.setup_preprocessor(nb, resources, km=km):
@@ -151,7 +155,9 @@ def _jinja_cell_generator(self, nb, kernel_id):
                     res = (cell, resources)
                     stop_execution = True
 
-                yield res[0]
+        for cell_idx, cell in enumerate(nb.cells):
+            res = await self.executor.execute_cell(cell, None, cell_idx, store_history=False)
+            yield res
 
     # @tornado.gen.coroutine
     async def load_notebook(self, path):