diff --git a/notebook/services/kernels/handlers.py b/notebook/services/kernels/handlers.py index ef90bf9122..7fef1963d2 100644 --- a/notebook/services/kernels/handlers.py +++ b/notebook/services/kernels/handlers.py @@ -29,7 +29,11 @@ class MainKernelHandler(APIHandler): @gen.coroutine def get(self): km = self.kernel_manager - kernels = yield gen.maybe_future(km.list_kernels()) + res = km.list_kernels() + if not isinstance(res, list): + kernels = res + else: + kernels = yield res self.finish(json.dumps(kernels, default=date_default)) @web.authenticated @@ -82,7 +86,9 @@ def post(self, kernel_id, action): if action == 'restart': try: - yield gen.maybe_future(km.restart_kernel(kernel_id)) + res = km.restart_kernel(kernel_id) + if res is not None: + yield from res except Exception as e: self.log.error("Exception restarting kernel", exc_info=True) self.set_status(500) diff --git a/notebook/services/kernels/kernelmanager.py b/notebook/services/kernels/kernelmanager.py index 6d2158ef8a..366bce7b2c 100644 --- a/notebook/services/kernels/kernelmanager.py +++ b/notebook/services/kernels/kernelmanager.py @@ -28,6 +28,8 @@ from notebook.prometheus.metrics import KERNEL_CURRENTLY_RUNNING_TOTAL +import asyncio + class MappingKernelManager(MultiKernelManager): """A KernelManager that handles notebook mapping and HTTP error handling""" @@ -158,9 +160,21 @@ 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 gen.maybe_future( - super(MappingKernelManager, self).start_kernel(**kwargs) - ) + + sup = super(MappingKernelManager, self) + async_sk = getattr(sup, 'start_kernel_async', None) + if async_sk is not None: + self.log.debug('dispathcing start to async super') + res = super().start_kernel_async(**kwargs) + else: + self.log.debug('dispathcing start to sync super') + res = super().start_kernel(**kwargs) + + if isinstance(res, str): + kernel_id = res + else: + kernel_id = yield res + self._kernel_connections[kernel_id] = 0 self.start_watching_activity(kernel_id) self.log.info("Kernel started: %s" % kernel_id) @@ -296,10 +310,19 @@ def shutdown_kernel(self, kernel_id, now=False): return super(MappingKernelManager, self).shutdown_kernel(kernel_id, now=now) + @asyncio.coroutine def restart_kernel(self, kernel_id): """Restart a kernel by kernel_id""" self._check_kernel_id(kernel_id) - super(MappingKernelManager, self).restart_kernel(kernel_id) + sup = super(MappingKernelManager, self) + + if hasattr(sup, 'restart_kernel_async'): + self.log.debug('dispatching restart_kernel to async super') + yield from sup.restart_kernel_async(kernel_id) + else: + self.log.debug('dispatching restart_kernel to sync super') + sup.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()