Skip to content

Commit

Permalink
Update all coroutine/yield methods to async/await, share another method
Browse files Browse the repository at this point in the history
  • Loading branch information
kevin-bates committed Mar 20, 2020
1 parent 5709862 commit 788bb0e
Showing 1 changed file with 15 additions and 25 deletions.
40 changes: 15 additions & 25 deletions notebook/services/kernels/kernelmanager.py
Original file line number Diff line number Diff line change
Expand Up @@ -139,6 +139,11 @@ def __init__(self, **kwargs):
self.kernel_culler = KernelCuller(parent=self)
self.activity_monitor = ActivityMonitor(parent=self)

def _handle_kernel_died(self, kernel_id):
"""notice that a kernel died"""
self.log.warning("Kernel %s died, removing from map.", kernel_id)
self.remove_kernel(kernel_id)

def kernel_model(self, kernel_id):
"""Return a JSON-safe dict representing a kernel
Expand Down Expand Up @@ -269,13 +274,7 @@ def __init__(self, **kwargs):
# Methods for managing kernels and sessions
# -------------------------------------------------------------------------

def _handle_kernel_died(self, kernel_id):
"""notice that a kernel died"""
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 @@ -294,7 +293,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 maybe_future(
kernel_id = await maybe_future(
super(MappingKernelManager, self).start_kernel(**kwargs)
)
self._kernel_connections[kernel_id] = 0
Expand All @@ -317,8 +316,7 @@ 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

def shutdown_kernel(self, kernel_id, now=False):
"""Shutdown a kernel by kernel_id"""
Expand All @@ -338,11 +336,10 @@ def shutdown_kernel(self, kernel_id, now=False):

return super(MappingKernelManager, self).shutdown_kernel(kernel_id, now=now)

@gen.coroutine
def restart_kernel(self, kernel_id):
async def restart_kernel(self, kernel_id):
"""Restart a kernel by kernel_id"""
self._check_kernel_id(kernel_id)
yield maybe_future(super(MappingKernelManager, self).restart_kernel(kernel_id))
await maybe_future(super(MappingKernelManager, self).restart_kernel(kernel_id))
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 @@ -379,7 +376,7 @@ def on_restart_failed():
loop = IOLoop.current()
timeout = loop.add_timeout(loop.time() + self.kernel_info_timeout, on_timeout)
# wait for restart to complete
yield future
await future


class AsyncMappingKernelManager(MappingKernelManagerBase, AsyncMultiKernelManager):
Expand All @@ -396,11 +393,6 @@ def __init__(self, **kwargs):
# Methods for managing kernels and sessions
# -------------------------------------------------------------------------

def _handle_kernel_died(self, kernel_id):
"""notice that a kernel died"""
self.log.warning("Kernel %s died, removing from map.", kernel_id)
self.remove_kernel(kernel_id)

async def start_kernel(self, kernel_id=None, path=None, **kwargs):
"""Start a kernel for a session and return its kernel_id.
Expand Down Expand Up @@ -695,16 +687,14 @@ def _log_info(self):
log_msg.append(".")
self.log.info(''.join(log_msg))

@gen.coroutine
def cull_kernels(self):
async def cull_kernels(self):
self.log.debug("Polling every %s seconds for kernels %s > %s seconds...",
self.parent.cull_interval, self.cull_state, self.parent.cull_idle_timeout)
# Get a separate list of kernels to avoid conflicting updates while iterating
for kernel_id in self.parent.list_kernel_ids():
yield self.parent.cull_kernel_if_idle(kernel_id)
await self.parent.cull_kernel_if_idle(kernel_id)

@gen.coroutine
def cull_kernel_if_idle(self, kernel_id):
async def cull_kernel_if_idle(self, kernel_id):

# Get the kernel model and use that to determine cullability...
try:
Expand All @@ -729,7 +719,7 @@ def cull_kernel_if_idle(self, kernel_id):
self.log.warning(
"Culling '%s' kernel '%s' (%s) with %d connections due to %s seconds of inactivity.",
model['execution_state'], model['name'], kernel_id, connections, idle_duration)
yield maybe_future(self.parent.shutdown_kernel(kernel_id))
await maybe_future(self.parent.shutdown_kernel(kernel_id))
except KeyError:
pass # KeyErrors are somewhat expected since the kernel can be shutdown as the culling check is made.
except Exception as e: # other errors are not as expected, so we'll make some noise, but continue.
Expand Down

0 comments on commit 788bb0e

Please sign in to comment.