Backport PR #4412: Enable restart_kernel for async usage

minrk · minrk · commit ab6445233ef7 · 2019-03-05T15:15:41.000+01:00
Converted `MappingKernelManager.restart_kernel` to a coroutine so that projects that take advantage of async kernel startup can also realize appropriate behavior relative to restarts. To take full advantage of async kernel startup (with restarts), this PR along with its "[sibling PR](jupyter/jupyter_client#425)" will be required. Please note, however, that each of these PRs can be independently installed w/o affecting today's notebook applications (as far as I can determine via testing the possible combinations). That said, BOTH of these PRs will be required for usage by Enterprise Gateway - which incurs very long kernel startup times - or other projects that require concurrent kernel starts. It would be ideal to have this PR back-ported to python 2.7-relative release since EG has an immediate need. Signed-off-by: Min RK <benjaminrk@gmail.com>
diff --git a/notebook/services/kernels/kernelmanager.py b/notebook/services/kernels/kernelmanager.py
@@ -280,10 +280,11 @@ def shutdown_kernel(self, kernel_id, now=False):
         self.last_kernel_activity = utcnow()
         return super(MappingKernelManager, self).shutdown_kernel(kernel_id, now=now)
 
+    @gen.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)
+        yield gen.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()
@@ -319,7 +320,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)
-        return future
+        raise gen.Return(future)
 
     def notify_connect(self, kernel_id):
         """Notice a new connection to a kernel"""
