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 10, 2020
1 parent 7abd4be commit 5bada42
Show file tree
Hide file tree
Showing 3 changed files with 17 additions and 16 deletions.
11 changes: 8 additions & 3 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 @@ -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

0 comments on commit 5bada42

Please sign in to comment.