Revert "Fix pending kernels again" (#853)

Author: blink1073

jupyter_client/manager.py

@@ -57,37 +57,34 @@ class _ShutdownStatus(Enum):
 F = t.TypeVar('F', bound=t.Callable[..., t.Any])
 
 
-def in_pending_state(prefix: str = '') -> t.Callable[[F], F]:
-    def decorator(method: F) -> F:
-        """Sets the kernel to a pending state by
-        creating a fresh Future for the KernelManager's `ready`
-        attribute. Once the method is finished, set the Future's results.
-        """
-
-        @t.no_type_check
-        @functools.wraps(method)
-        async def wrapper(self, *args, **kwargs):
-            # Create a future for the decorated method
-            name = f"{prefix}_ready"
-            future = getattr(self, name)
-            if not future or future.done():
-                future = self._future_factory()
-                setattr(self, name, future)
-            try:
-                # call wrapped method, await, and set the result or exception.
-                out = await method(self, *args, **kwargs)
-                # Add a small sleep to ensure tests can capture the state before done
-                await asyncio.sleep(0.01)
-                future.set_result(None)
-                return out
-            except Exception as e:
-                future.set_exception(e)
-                self.log.exception(future.exception())
-                raise e
+def in_pending_state(method: F) -> F:
+    """Sets the kernel to a pending state by
+    creating a fresh Future for the KernelManager's `ready`
+    attribute. Once the method is finished, set the Future's results.
+    """
 
-        return t.cast(F, wrapper)
+    @t.no_type_check
+    @functools.wraps(method)
+    async def wrapper(self, *args, **kwargs):
+        # Create a future for the decorated method
+        try:
+            self._ready = Future()
+        except RuntimeError:
+            # No event loop running, use concurrent future
+            self._ready = CFuture()
+        try:
+            # call wrapped method, await, and set the result or exception.
+            out = await method(self, *args, **kwargs)
+            # Add a small sleep to ensure tests can capture the state before done
+            await asyncio.sleep(0.01)
+            self._ready.set_result(None)
+            return out
+        except Exception as e:
+            self._ready.set_exception(e)
+            self.log.exception(self._ready.exception())
+            raise e
 
-    return decorator
+    return t.cast(F, wrapper)
 
 
 class KernelManager(ConnectionFileMixin):
@@ -117,14 +114,18 @@ def _emit(self, *, action: str) -> None:
             data={"action": action, "kernel_id": self.kernel_id, "caller": "kernel_manager"},
         )
 
-    _ready: t.Optional[CFuture]
-    _shutdown_ready: t.Optional[CFuture]
+    _ready: t.Union[Future, CFuture]
 
     def __init__(self, *args, **kwargs):
         super().__init__(**kwargs)
         self._shutdown_status = _ShutdownStatus.Unset
-        self._ready = None
-        self._shutdown_ready = None
+        # Create a place holder future.
+        try:
+            asyncio.get_running_loop()
+            self._ready = Future()
+        except RuntimeError:
+            # No event loop running, use concurrent future
+            self._ready = CFuture()
 
     _created_context: Bool = Bool(False)
 
@@ -142,8 +143,6 @@ def _context_default(self) -> zmq.Context:
     )
     client_factory: Type = Type(klass="jupyter_client.KernelClient")
 
-    _future_factory: t.Type[CFuture] = CFuture
-
     @default("client_factory")  # type:ignore[misc]
     def _client_factory_default(self) -> Type:
         return import_item(self.client_class)
@@ -209,21 +208,10 @@ def _default_cache_ports(self) -> bool:
         return self.transport == "tcp"
 
     @property
-    def ready(self) -> CFuture:
-        """A future that resolves when the kernel process has started."""
-        if not self._ready:
-            self._ready = self._future_factory()
-        assert self._ready is not None
+    def ready(self) -> t.Union[CFuture, Future]:
+        """A future that resolves when the kernel process has started for the first time"""
         return self._ready
 
-    @property
-    def shutdown_ready(self) -> CFuture:
-        """A future that resolves when the kernel process has shut down."""
-        if not self._shutdown_ready:
-            self._shutdown_ready = self._future_factory()
-        assert self._shutdown_ready is not None
-        return self._shutdown_ready
-
     @property
     def ipykernel(self) -> bool:
         return self.kernel_name in {"python", "python2", "python3"}
@@ -407,7 +395,7 @@ async def _async_post_start_kernel(self, **kw: t.Any) -> None:
 
     post_start_kernel = run_sync(_async_post_start_kernel)
 
-    @in_pending_state()
+    @in_pending_state
     async def _async_start_kernel(self, **kw: t.Any) -> None:
         """Starts a kernel on this host in a separate process.
 
@@ -503,7 +491,7 @@ async def _async_cleanup_resources(self, restart: bool = False) -> None:
 
     cleanup_resources = run_sync(_async_cleanup_resources)
 
-    @in_pending_state('_shutdown')
+    @in_pending_state
     async def _async_shutdown_kernel(self, now: bool = False, restart: bool = False) -> None:
         """Attempts to stop the kernel process cleanly.
 
@@ -522,8 +510,6 @@ async def _async_shutdown_kernel(self, now: bool = False, restart: bool = False)
             Will this kernel be restarted after it is shutdown. When this
             is True, connection files will not be cleaned up.
         """
-        # Reset the start ready future.
-        self._ready = self._future_factory()
         self._emit(action="shutdown_started")
         self.shutting_down = True  # Used by restarter to prevent race condition
         # Stop monitoring for restarting while we shutdown.
@@ -696,8 +682,6 @@ class AsyncKernelManager(KernelManager):
     # The PyZMQ Context to use for communication with the kernel.
     context: Instance = Instance(zmq.asyncio.Context)
 
-    _future_factory: t.Type[Future] = Future  # type:ignore[assignment]
-
     @default("context")  # type:ignore[misc]
     def _context_default(self) -> zmq.asyncio.Context:
         self._created_context = True

jupyter_client/multikernelmanager.py

@@ -221,7 +221,9 @@ async def _async_start_kernel(
             self._kernels[kernel_id] = km
         else:
             await task
-            await asyncio.wrap_future(km.ready)
+            # raise an exception if one occurred during kernel startup.
+            if km.ready.exception():
+                raise km.ready.exception()  # type: ignore
 
         return kernel_id
 
@@ -251,17 +253,15 @@ async def _async_shutdown_kernel(
             try:
                 await task
                 km = self.get_kernel(kernel_id)
-                await asyncio.wrap_future(km.ready)
+                await t.cast(asyncio.Future, km.ready)
             except asyncio.CancelledError:
                 pass
             except Exception:
                 self.remove_kernel(kernel_id)
                 return
         km = self.get_kernel(kernel_id)
         # If a pending kernel raised an exception, remove it.
-        try:
-            await asyncio.wrap_future(km.ready)
-        except Exception:
+        if not km.ready.cancelled() and km.ready.exception():
             self.remove_kernel(kernel_id)
             return
         stopper = ensure_async(km.shutdown_kernel(now, restart))
@@ -270,7 +270,9 @@ async def _async_shutdown_kernel(
         # Await the kernel if not using pending kernels.
         if not self._using_pending_kernels():
             await fut
-            await asyncio.wrap_future(km.shutdown_ready)
+            # raise an exception if one occurred during kernel shutdown.
+            if km.ready.exception():
+                raise km.ready.exception()  # type: ignore
 
     shutdown_kernel = run_sync(_async_shutdown_kernel)
 
@@ -313,7 +315,7 @@ async def _async_shutdown_all(self, now: bool = False) -> None:
         if self._using_pending_kernels():
             for km in kms:
                 try:
-                    await km.shutdown_ready
+                    await km.ready
                 except asyncio.CancelledError:
                     self._pending_kernels[km.kernel_id].cancel()
                 except Exception:

tests/test_kernelmanager.py

@@ -135,13 +135,13 @@ def _(*expected_list):
 
 
 @pytest.fixture(params=[AsyncKernelManager, AsyncKMSubclass])
-async def async_km(request, config, jp_event_logger):
+def async_km(request, config, jp_event_logger):
     km = request.param(config=config, event_logger=jp_event_logger)
     return km
 
 
 @pytest.fixture
-async def async_km_subclass(config, jp_event_logger):
+def async_km_subclass(config, jp_event_logger):
     km = AsyncKMSubclass(config=config, event_logger=jp_event_logger)
     return km
 
@@ -489,11 +489,11 @@ async def test_lifecycle(
         await async_km.start_kernel(stdout=PIPE, stderr=PIPE)
         is_alive = await async_km.is_alive()
         assert is_alive
-        await async_km.ready
+        is_ready = async_km.ready.done()
+        assert is_ready
         await async_km.restart_kernel(now=True)
         is_alive = await async_km.is_alive()
         assert is_alive
-        await async_km.ready
         await async_km.interrupt_kernel()
         assert isinstance(async_km, AsyncKernelManager)
         await async_km.shutdown_kernel(now=True)

tests/test_multikernelmanager.py

@@ -429,7 +429,7 @@ async def test_use_pending_kernels(self):
         assert isinstance(k, AsyncKernelManager)
         await ensure_future(km.shutdown_kernel(kid, now=True))
         # Wait for the kernel to shutdown
-        await kernel.shutdown_ready
+        await kernel.ready
         assert kid not in km, f"{kid} not in {km}"
 
     @gen_test
@@ -443,7 +443,7 @@ async def test_use_pending_kernels_early_restart(self):
         await kernel.ready
         await ensure_future(km.shutdown_kernel(kid, now=True))
         # Wait for the kernel to shutdown
-        await kernel.shutdown_ready
+        await kernel.ready
         assert kid not in km, f"{kid} not in {km}"
 
     @gen_test
@@ -455,7 +455,7 @@ async def test_use_pending_kernels_early_shutdown(self):
         # Try shutting down while the kernel is pending
         await ensure_future(km.shutdown_kernel(kid, now=True))
         # Wait for the kernel to shutdown
-        await kernel.shutdown_ready
+        await kernel.ready
         assert kid not in km, f"{kid} not in {km}"
 
     @gen_test
@@ -470,7 +470,7 @@ async def test_use_pending_kernels_early_interrupt(self):
         await kernel.ready
         await ensure_future(km.shutdown_kernel(kid, now=True))
         # Wait for the kernel to shutdown
-        await kernel.shutdown_ready
+        await kernel.ready
         assert kid not in km, f"{kid} not in {km}"
 
     @gen_test