From 8a4beb0d153d31dab487d6ec655c7d2c53b6a5b2 Mon Sep 17 00:00:00 2001 From: Kevin Bates Date: Tue, 10 Mar 2020 16:40:39 -0700 Subject: [PATCH] Convert to async/await and apply touch ups --- notebook/notebookapp.py | 13 +++++--- notebook/services/kernels/handlers.py | 2 +- notebook/services/kernels/kernelmanager.py | 20 +++++------ .../sessions/tests/test_sessions_api.py | 33 ++++++++++++++++++- 4 files changed, 50 insertions(+), 18 deletions(-) diff --git a/notebook/notebookapp.py b/notebook/notebookapp.py index 39552dc85..0c4244926 100755 --- a/notebook/notebookapp.py +++ b/notebook/notebookapp.py @@ -7,6 +7,7 @@ from __future__ import absolute_import, print_function import notebook +import asyncio import binascii import datetime import errno @@ -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, ] @@ -1392,10 +1393,10 @@ class NotebookApp(JupyterApp): # 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( @@ -1800,7 +1801,11 @@ class NotebookApp(JupyterApp): 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" diff --git a/notebook/services/kernels/handlers.py b/notebook/services/kernels/handlers.py index 8484dccdf..bca99ce1b 100644 --- a/notebook/services/kernels/handlers.py +++ b/notebook/services/kernels/handlers.py @@ -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': diff --git a/notebook/services/kernels/kernelmanager.py b/notebook/services/kernels/kernelmanager.py index 6c0bfb31c..ff6fe721f 100644 --- a/notebook/services/kernels/kernelmanager.py +++ b/notebook/services/kernels/kernelmanager.py @@ -401,8 +401,7 @@ class AsyncMappingKernelManager(AsyncMultiKernelManager, MappingKernelManagerBas 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 @@ -421,7 +420,7 @@ class AsyncMappingKernelManager(AsyncMultiKernelManager, MappingKernelManagerBas 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) @@ -443,11 +442,9 @@ class AsyncMappingKernelManager(AsyncMultiKernelManager, MappingKernelManagerBas 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] @@ -464,13 +461,12 @@ class AsyncMappingKernelManager(AsyncMultiKernelManager, MappingKernelManagerBas 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() @@ -506,7 +502,7 @@ class AsyncMappingKernelManager(AsyncMultiKernelManager, MappingKernelManagerBas 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 diff --git a/notebook/services/sessions/tests/test_sessions_api.py b/notebook/services/sessions/tests/test_sessions_api.py index 9c551fc79..9323b5966 100644 --- a/notebook/services/sessions/tests/test_sessions_api.py +++ b/notebook/services/sessions/tests/test_sessions_api.py @@ -9,13 +9,22 @@ import requests 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): @@ -77,6 +86,7 @@ class SessionAPI(object): def delete(self, id): return self._req('DELETE', id) + class SessionAPITest(NotebookTestBase): """Test the sessions web service API""" def setUp(self): @@ -254,3 +264,24 @@ class SessionAPITest(NotebookTestBase): 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()