handle shutdown and restart for pending kernels in the multikernelmanager

Author: Zsailer

jupyter_client/multikernelmanager.py

@@ -214,6 +214,14 @@ async def _async_shutdown_kernel(
         restart : bool
             Will the kernel be restarted?
         """
+        kernel = self.get_kernel(kernel_id)
+        if getattr(self, 'use_pending_kernels', False):
+            # Make sure the previous kernel started before trying to shutdown.
+            if not kernel.ready.done():
+                raise RuntimeError("Kernel is in a pending state. Cannot shutdown.")
+            # If the kernel didn't start properly, no need to shutdown.
+            elif kernel.ready.exception():
+                return
         self.log.info("Kernel shutdown: %s" % kernel_id)
         if kernel_id in self._starting_kernels:
             try:
@@ -291,8 +299,7 @@ def signal_kernel(self, kernel_id: str, signum: int) -> None:
         """
         self.log.info("Signaled Kernel %s with %s" % (kernel_id, signum))
 
-    @kernel_method
-    def restart_kernel(self, kernel_id: str, now: bool = False) -> None:
+    async def _async_restart_kernel(self, kernel_id: str, now: bool = False) -> None:
         """Restart a kernel by its uuid, keeping the same ports.
 
         Parameters
@@ -307,7 +314,15 @@ def restart_kernel(self, kernel_id: str, now: bool = False) -> None:
             In all cases the kernel is restarted, the only difference is whether
             it is given a chance to perform a clean shutdown or not.
         """
+        kernel = self.get_kernel(kernel_id)
+        if getattr(self, 'use_pending_kernels', False):
+            if not kernel.ready.done():
+                raise RuntimeError("Kernel is in a pending state. Cannot restart.")
+        out = await ensure_async(kernel.restart_kernel(now=now))
         self.log.info("Kernel restarted: %s" % kernel_id)
+        return out
+
+    restart_kernel = run_sync(_async_restart_kernel)
 
     @kernel_method
     def is_alive(self, kernel_id: str) -> bool:
@@ -475,5 +490,6 @@ class AsyncMultiKernelManager(MultiKernelManager):
     ).tag(config=True)
 
     start_kernel = MultiKernelManager._async_start_kernel
+    restart_kernel = MultiKernelManager._async_restart_kernel
     shutdown_kernel = MultiKernelManager._async_shutdown_kernel
     shutdown_all = MultiKernelManager._async_shutdown_all

jupyter_client/tests/test_multikernelmanager.py

@@ -29,6 +29,14 @@
 TIMEOUT = 30
 
 
+async def now(awaitable):
+    """Use this function ensure that this awaitable
+    happens before other awaitables defined after it.
+    """
+    (out,) = await asyncio.gather(awaitable)
+    return out
+
+
 class TestKernelManager(TestCase):
     def setUp(self):
         self.env_patch = test_env()
@@ -357,32 +365,32 @@ async def test_shutdown_all_while_starting(self):
     @gen_test
     async def test_use_pending_kernels(self):
         km = self._get_pending_kernels_km()
-        kid = await km.start_kernel(stdout=PIPE, stderr=PIPE)
+        kid = await now(km.start_kernel(stdout=PIPE, stderr=PIPE))
         kernel = km.get_kernel(kid)
         assert not kernel.ready.done()
         assert kid in km
         assert kid in km.list_kernel_ids()
         assert len(km) == 1, f"{len(km)} != {1}"
-        await kernel.ready
-        await km.restart_kernel(kid, now=True)
+        await now(kernel.ready)
+        await now(km.restart_kernel(kid, now=True))
         assert await km.is_alive(kid)
         assert kid in km.list_kernel_ids()
-        await km.interrupt_kernel(kid)
+        await now(km.interrupt_kernel(kid))
         k = km.get_kernel(kid)
         assert isinstance(k, AsyncKernelManager)
-        await km.shutdown_kernel(kid, now=True)
+        await now(km.shutdown_kernel(kid, now=True))
         assert kid not in km, f"{kid} not in {km}"
 
     @gen_test
     async def test_use_pending_kernels_early_restart(self):
         km = self._get_pending_kernels_km()
-        kid = await km.start_kernel(stdout=PIPE, stderr=PIPE)
+        kid = await now(km.start_kernel(stdout=PIPE, stderr=PIPE))
         kernel = km.get_kernel(kid)
         assert not kernel.ready.done()
         with pytest.raises(RuntimeError):
             await km.restart_kernel(kid, now=True)
-        await kernel.ready
-        await km.shutdown_kernel(kid, now=True)
+        await now(kernel.ready)
+        await now(km.shutdown_kernel(kid, now=True))
         assert kid not in km, f"{kid} not in {km}"
 
     @gen_test