diff --git a/jupyter_server/services/kernels/kernelmanager.py b/jupyter_server/services/kernels/kernelmanager.py index 65582844cb..461b8b4351 100644 --- a/jupyter_server/services/kernels/kernelmanager.py +++ b/jupyter_server/services/kernels/kernelmanager.py @@ -188,6 +188,11 @@ def cwd_for_path(self, path): os_path = os.path.dirname(os_path) return os_path + async def _remove_kernel_when_ready(self, kernel_id, kernel_awaitable): + await super()._remove_kernel_when_ready(kernel_id, kernel_awaitable) + self._kernel_connections.pop(kernel_id, None) + self._kernel_ports.pop(kernel_id, None) + async def start_kernel(self, kernel_id=None, path=None, **kwargs): """Start a kernel for a session and return its kernel_id. @@ -384,19 +389,12 @@ def shutdown_kernel(self, kernel_id, now=False, restart=False): self._check_kernel_id(kernel_id) self.stop_watching_activity(kernel_id) self.stop_buffering(kernel_id) - self._kernel_connections.pop(kernel_id, None) # Decrease the metric of number of kernels # running for the relevant kernel type by 1 KERNEL_CURRENTLY_RUNNING_TOTAL.labels(type=self._kernels[kernel_id].kernel_name).dec() self.pinned_superclass.shutdown_kernel(self, kernel_id, now=now, restart=restart) - # Unlike its async sibling method in AsyncMappingKernelManager, removing the kernel_id - # from the connections dictionary isn't as problematic before the shutdown since the - # method is synchronous. However, we'll keep the relative call orders the same from - # a maintenance perspective. - self._kernel_connections.pop(kernel_id, None) - self._kernel_ports.pop(kernel_id, None) async def restart_kernel(self, kernel_id, now=False): """Restart a kernel by kernel_id""" @@ -655,6 +653,4 @@ async def shutdown_kernel(self, kernel_id, now=False, restart=False): ret = await self.pinned_superclass.shutdown_kernel( self, kernel_id, now=now, restart=restart ) - self._kernel_connections.pop(kernel_id, None) - self._kernel_ports.pop(kernel_id, None) return ret