Skip to content

Commit

Permalink
Convert to async/await and apply touch ups
Browse files Browse the repository at this point in the history
  • Loading branch information
kevin-bates committed Mar 17, 2020
1 parent 7abd4be commit 4300468
Show file tree
Hide file tree
Showing 6 changed files with 53 additions and 19 deletions.
1 change: 1 addition & 0 deletions appveyor.yml
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ install:
# currently a bug with the version that the anaconda installs, so we will just install it with pip
- cmd: pip install nbconvert
- cmd: python setup.py build
- cmd: pip uninstall -y jupyter_client
- cmd: pip install .[test]

test_script:
Expand Down
13 changes: 9 additions & 4 deletions notebook/notebookapp.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
from __future__ import absolute_import, print_function

import notebook
import asyncio
import binascii
import datetime
import errno
Expand Down Expand Up @@ -583,7 +584,7 @@ class NotebookApp(JupyterApp):
flags = flags

classes = [
KernelManager, Session, MappingKernelManager, KernelSpecManager,
KernelManager, Session, MappingKernelManager, AsyncMappingKernelManager, KernelSpecManager,
ContentsManager, FileContentsManager, NotebookNotary,
GatewayKernelManager, GatewayKernelSpecManager, GatewaySessionManager, GatewayClient,
]
Expand Down Expand Up @@ -1392,10 +1393,10 @@ def init_configurables(self):
# Ensure the appropriate jupyter_client is in place.
if isinstance(self.kernel_manager, AsyncMappingKernelManager):
if not async_kernel_mgmt_available:
raise ValueError("You're using `AsyncMappingKernelManager` without an appropriate "
raise ValueError("You are using `AsyncMappingKernelManager` without an appropriate "
"jupyter_client installed! Upgrade jupyter_client or change kernel managers.")
else:
self.log.info("Asynchronous kernel management has been configured via '{}'.".
self.log.info("Asynchronous kernel management has been configured to use '{}'.".
format(self.kernel_manager.__class__.__name__))

self.contents_manager = self.contents_manager_class(
Expand Down Expand Up @@ -1800,7 +1801,11 @@ def cleanup_kernels(self):
n_kernels = len(self.kernel_manager.list_kernel_ids())
kernel_msg = trans.ngettext('Shutting down %d kernel', 'Shutting down %d kernels', n_kernels)
self.log.info(kernel_msg % n_kernels)
self.kernel_manager.shutdown_all()
# If we're using async kernel management, we need to invoke the async method via the event loop.
if isinstance(self.kernel_manager, AsyncMappingKernelManager):
asyncio.get_event_loop().run_until_complete(self.kernel_manager.shutdown_all())
else:
self.kernel_manager.shutdown_all()

def notebook_info(self, kernel_count=True):
"Return the current working directory and the server url information"
Expand Down
2 changes: 1 addition & 1 deletion notebook/services/kernels/handlers.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,7 @@ class KernelActionHandler(APIHandler):
def post(self, kernel_id, action):
km = self.kernel_manager
if action == 'interrupt':
km.interrupt_kernel(kernel_id)
yield maybe_future(km.interrupt_kernel(kernel_id))
self.set_status(204)
if action == 'restart':

Expand Down
20 changes: 8 additions & 12 deletions notebook/services/kernels/kernelmanager.py
Original file line number Diff line number Diff line change
Expand Up @@ -401,8 +401,7 @@ def _handle_kernel_died(self, kernel_id):
self.log.warning("Kernel %s died, removing from map.", kernel_id)
self.remove_kernel(kernel_id)

@gen.coroutine
def start_kernel(self, kernel_id=None, path=None, **kwargs):
async def start_kernel(self, kernel_id=None, path=None, **kwargs):
"""Start a kernel for a session and return its kernel_id.
Parameters
Expand All @@ -421,7 +420,7 @@ def start_kernel(self, kernel_id=None, path=None, **kwargs):
if kernel_id is None:
if path is not None:
kwargs['cwd'] = self.cwd_for_path(path)
kernel_id = yield super(AsyncMappingKernelManager, self).start_kernel(**kwargs)
kernel_id = await super(AsyncMappingKernelManager, self).start_kernel(**kwargs)

self._kernel_connections[kernel_id] = 0
self.start_watching_activity(kernel_id)
Expand All @@ -443,11 +442,9 @@ def start_kernel(self, kernel_id=None, path=None, **kwargs):
self._check_kernel_id(kernel_id)
self.log.info("Using existing kernel: %s" % kernel_id)

# py2-compat
raise gen.Return(kernel_id)
return kernel_id

@gen.coroutine
def shutdown_kernel(self, kernel_id, now=False, restart=False):
async def shutdown_kernel(self, kernel_id, now=False, restart=False):
"""Shutdown a kernel by kernel_id"""
self._check_kernel_id(kernel_id)
kernel = self._kernels[kernel_id]
Expand All @@ -464,13 +461,12 @@ def shutdown_kernel(self, kernel_id, now=False, restart=False):
type=self._kernels[kernel_id].kernel_name
).dec()

yield super(AsyncMappingKernelManager, self).shutdown_kernel(kernel_id, now=now, restart=restart)
await super(AsyncMappingKernelManager, self).shutdown_kernel(kernel_id, now=now, restart=restart)

@gen.coroutine
def restart_kernel(self, kernel_id, now=False):
async def restart_kernel(self, kernel_id, now=False):
"""Restart a kernel by kernel_id"""
self._check_kernel_id(kernel_id)
yield super(AsyncMappingKernelManager, self).restart_kernel(kernel_id, now=now)
await super(AsyncMappingKernelManager, self).restart_kernel(kernel_id, now=now)
kernel = self.get_kernel(kernel_id)
# return a Future that will resolve when the kernel has successfully restarted
channel = kernel.connect_shell()
Expand Down Expand Up @@ -506,7 +502,7 @@ def on_restart_failed():
channel.on_recv(on_reply)
loop = IOLoop.current()
timeout = loop.add_timeout(loop.time() + self.kernel_info_timeout, on_timeout)
raise gen.Return(future)
return future

def kernel_model(self, kernel_id):
"""Return a JSON-safe dict representing a kernel
Expand Down
33 changes: 32 additions & 1 deletion notebook/services/sessions/tests/test_sessions_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,13 +9,22 @@
import shutil
import time

pjoin = os.path.join
from unittest import SkipTest

from notebook.utils import url_path_join
from notebook.tests.launchnotebook import NotebookTestBase, assert_http_error
from nbformat.v4 import new_notebook
from nbformat import write

try:
from jupyter_client import AsyncMultiKernelManager
async_testing_enabled = True
except ImportError:
async_testing_enabled = False

pjoin = os.path.join


class SessionAPI(object):
"""Wrapper for notebook API calls."""
def __init__(self, request):
Expand Down Expand Up @@ -77,6 +86,7 @@ def modify_kernel_id(self, id, kernel_id):
def delete(self, id):
return self._req('DELETE', id)


class SessionAPITest(NotebookTestBase):
"""Test the sessions web service API"""
def setUp(self):
Expand Down Expand Up @@ -254,3 +264,24 @@ def test_modify_kernel_id(self):
kernel.pop('last_activity')
[ k.pop('last_activity') for k in kernel_list ]
self.assertEqual(kernel_list, [kernel])


class AsyncSessionAPITest(SessionAPITest):
"""Test the sessions web service API using the AsyncMappingKernelManager"""

@classmethod
def get_argv(cls):
argv = super(AsyncSessionAPITest, cls).get_argv()

# before we extend the argv with the class, ensure that appropriate jupyter_client is available.
# if not available, don't set kernel_manager_class, resulting in the repeat of sync-based tests.
if async_testing_enabled:
argv.extend(['--NotebookApp.kernel_manager_class='
'notebook.services.kernels.kernelmanager.AsyncMappingKernelManager'])
return argv

def setUp(self):
if not async_testing_enabled:
raise SkipTest("AsyncSessionAPITest.{test_method} skipped due to down-level jupyter_client!".
format(test_method=self._testMethodName))
super(AsyncSessionAPITest, self).setUp()
3 changes: 2 additions & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -105,7 +105,8 @@
'ipython_genutils',
'traitlets>=4.2.1',
'jupyter_core>=4.6.1',
'jupyter_client>=5.3.4',
# 'jupyter_client>=5.3.4',
'jupyter_client @ git+https://github.com/jupyter/jupyter_client',
'nbformat',
'nbconvert',
'ipykernel', # bless IPython kernel for now
Expand Down

0 comments on commit 4300468

Please sign in to comment.