From d68dbe8a89aa34206027d18a1fa3d29bd9b8befd Mon Sep 17 00:00:00 2001 From: Timothy Guo Date: Mon, 14 Aug 2023 17:09:40 -0400 Subject: [PATCH 1/2] Added restart_in_place flag --- jupyter_server/services/kernels/handlers.py | 6 +++--- .../services/kernels/kernelmanager.py | 21 ++++++++++++++----- 2 files changed, 19 insertions(+), 8 deletions(-) diff --git a/jupyter_server/services/kernels/handlers.py b/jupyter_server/services/kernels/handlers.py index 0ebebf4a51..f312a84768 100644 --- a/jupyter_server/services/kernels/handlers.py +++ b/jupyter_server/services/kernels/handlers.py @@ -94,9 +94,9 @@ async def post(self, kernel_id, action): if action == "interrupt": await ensure_async(km.interrupt_kernel(kernel_id)) self.set_status(204) - if action == "restart": + if action == "restart" or action == "restart-in-place": try: - await km.restart_kernel(kernel_id) + await km.restart_kernel(kernel_id, restart_in_place=action == "restart-in-place") except Exception as e: message = "Exception restarting kernel" self.log.error(message, exc_info=True) @@ -113,7 +113,7 @@ async def post(self, kernel_id, action): # URL to handler mappings # ----------------------------------------------------------------------------- _kernel_id_regex = r"(?P\w+-\w+-\w+-\w+-\w+)" -_kernel_action_regex = r"(?Prestart|interrupt)" +_kernel_action_regex = r"(?Prestart|interrupt|restart-in-place)" default_handlers = [ (r"/api/kernels", MainKernelHandler), diff --git a/jupyter_server/services/kernels/kernelmanager.py b/jupyter_server/services/kernels/kernelmanager.py index f45fdff193..777df2f348 100644 --- a/jupyter_server/services/kernels/kernelmanager.py +++ b/jupyter_server/services/kernels/kernelmanager.py @@ -17,7 +17,10 @@ from typing import Optional from jupyter_client.ioloop.manager import AsyncIOLoopKernelManager -from jupyter_client.multikernelmanager import AsyncMultiKernelManager, MultiKernelManager +from jupyter_client.multikernelmanager import ( + AsyncMultiKernelManager, + MultiKernelManager, +) from jupyter_client.session import Session from jupyter_core.paths import exists from jupyter_core.utils import ensure_async @@ -207,7 +210,11 @@ async def _remove_kernel_when_ready(self, kernel_id, kernel_awaitable): # TODO DEC 2022: Revise the type-ignore once the signatures have been changed upstream # https://github.com/jupyter/jupyter_client/pull/905 async def _async_start_kernel( # type:ignore[override] - self, *, kernel_id: Optional[str] = None, path: Optional[ApiPath] = None, **kwargs: str + self, + *, + kernel_id: Optional[str] = None, + path: Optional[ApiPath] = None, + **kwargs: str, ) -> str: """Start a kernel for a session and return its kernel_id. @@ -434,10 +441,12 @@ async def _async_shutdown_kernel(self, kernel_id, now=False, restart=False): shutdown_kernel = _async_shutdown_kernel - async def _async_restart_kernel(self, kernel_id, now=False): + async def _async_restart_kernel(self, kernel_id, now=False, restart_in_place=False): """Restart a kernel by kernel_id""" self._check_kernel_id(kernel_id) - await self.pinned_superclass._async_restart_kernel(self, kernel_id, now=now) + await self.pinned_superclass._async_restart_kernel( + self, kernel_id, now=now, restart_in_place=restart_in_place + ) kernel = self.get_kernel(kernel_id) # return a Future that will resolve when the kernel has successfully restarted channel = kernel.connect_shell() @@ -743,7 +752,9 @@ async def wrapped_method(self, *args, **kwargs): "action": action, "status": "success", "msg": success_msg.format( - kernel_id=self.kernel_id, kernel_name=self.kernel_name, action=action + kernel_id=self.kernel_id, + kernel_name=self.kernel_name, + action=action, ), } if self.kernel_id: From 5424408d36ad0c132f9ac7641f07f9365479bc21 Mon Sep 17 00:00:00 2001 From: Timothy Guo Date: Mon, 14 Aug 2023 17:42:13 -0400 Subject: [PATCH 2/2] Revert formatting changes --- jupyter_server/services/kernels/kernelmanager.py | 15 +++------------ 1 file changed, 3 insertions(+), 12 deletions(-) diff --git a/jupyter_server/services/kernels/kernelmanager.py b/jupyter_server/services/kernels/kernelmanager.py index 777df2f348..1fd5a22998 100644 --- a/jupyter_server/services/kernels/kernelmanager.py +++ b/jupyter_server/services/kernels/kernelmanager.py @@ -17,10 +17,7 @@ from typing import Optional from jupyter_client.ioloop.manager import AsyncIOLoopKernelManager -from jupyter_client.multikernelmanager import ( - AsyncMultiKernelManager, - MultiKernelManager, -) +from jupyter_client.multikernelmanager import AsyncMultiKernelManager, MultiKernelManager from jupyter_client.session import Session from jupyter_core.paths import exists from jupyter_core.utils import ensure_async @@ -210,11 +207,7 @@ async def _remove_kernel_when_ready(self, kernel_id, kernel_awaitable): # TODO DEC 2022: Revise the type-ignore once the signatures have been changed upstream # https://github.com/jupyter/jupyter_client/pull/905 async def _async_start_kernel( # type:ignore[override] - self, - *, - kernel_id: Optional[str] = None, - path: Optional[ApiPath] = None, - **kwargs: str, + self, *, kernel_id: Optional[str] = None, path: Optional[ApiPath] = None, **kwargs: str ) -> str: """Start a kernel for a session and return its kernel_id. @@ -752,9 +745,7 @@ async def wrapped_method(self, *args, **kwargs): "action": action, "status": "success", "msg": success_msg.format( - kernel_id=self.kernel_id, - kernel_name=self.kernel_name, - action=action, + kernel_id=self.kernel_id, kernel_name=self.kernel_name, action=action ), } if self.kernel_id: