Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[WIP] Refactor Remote Client Kernel Management #23720

Draft
wants to merge 4 commits into
base: master
Choose a base branch
from
+343 −592
Draft

[WIP] Refactor Remote Client Kernel Management #23720

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
ab18cf7
feat: refactor jupyter base api to incorportate kernel manager logic
hlouzada Feb 13, 2025
9991d9c
feat: remove deprecated old kernel manager api
hlouzada Feb 13, 2025
1f6419c
fix: temporary workaround to manage kernels with new api
hlouzada Feb 13, 2025
336fc59
fix: use new base remote client exception
hlouzada Feb 13, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
149 changes: 7 additions & 142 deletions spyder/plugins/remoteclient/api/manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,8 +32,6 @@
from spyder.plugins.remoteclient.api.protocol import (
ConnectionInfo,
ConnectionStatus,
KernelConnectionInfo,
KernelInfo,
RemoteClientLog,
SSHClientOptions,
)
Expand Down Expand Up @@ -75,7 +73,9 @@ class SpyderRemoteAPIManager:

REGISTERED_MODULE_APIS: typing.ClassVar[
dict[str, type[SpyderBaseJupyterAPIType]]
] = {}
] = {
JupyterAPI.__qualname__: JupyterAPI,
}

JUPYTER_SERVER_TIMEOUT = 5 # seconds

Expand Down Expand Up @@ -714,9 +714,7 @@ async def stop_remote_server(self):
f"{self._server_info['pid']}"
)
try:
async with JupyterAPI(
self.server_url, api_token=self.api_token
) as jupyter:
async with self.get_jupyter_api() as jupyter:
await jupyter.shutdown_server()
except Exception as err:
self.logger.exception("Error stopping remote server", exc_info=err)
Expand Down Expand Up @@ -778,139 +776,6 @@ def get_api(
raise ValueError(f"API {api} is not registered")

return partial(api_class, manager=self)

# ---- Kernel Management
async def start_new_kernel_ensure_server(
self, _retries=5
) -> KernelConnectionInfo:
"""Launch a new kernel ensuring the remote server is running.

Parameters
----------
options : SSHClientOptions
The options to use for the SSH connection.

Returns
-------
KernelConnectionInfo
The kernel connection information.
"""
if not await self.ensure_connection_and_server():
self.logger.error(
"Cannot launch kernel, remote server is not running"
)
return {}

# This is necessary to avoid an error when the server has not started
# before
await asyncio.sleep(1)
kernel_id = await self.start_new_kernel()

retries = 0
while not kernel_id and retries < _retries:
await asyncio.sleep(1)
kernel_id = await self.start_new_kernel()
self.logger.debug(
f"Server might not be ready yet, retrying kernel launch "
f"({retries + 1}/{_retries})"
)
retries += 1

return kernel_id

async def get_kernel_info_ensure_server(
self, kernel_id, _retries=5
) -> KernelConnectionInfo:
"""Launch a new kernel ensuring the remote server is running.

Parameters
----------
options : SSHClientOptions
The options to use for the SSH connection.

Returns
-------
KernelConnectionInfo
The kernel connection information.
"""
if not await self.ensure_connection_and_server():
self.logger.error(
"Cannot launch kernel, remote server is not running"
)
return {}

# This is necessary to avoid an error when the server has not started
# before
await asyncio.sleep(1)
kernel_info = await self.get_kernel_info(kernel_id)

retries = 0
while not kernel_info and retries < _retries:
await asyncio.sleep(1)
kernel_info = await self.get_kernel_info(kernel_id)
self.logger.debug(
f"Server might not be ready yet, retrying kernel launch "
f"({retries + 1}/{_retries})"
)
retries += 1

return kernel_info

async def start_new_kernel(self, kernel_spec=None) -> KernelInfo:
"""Start new kernel."""
async with JupyterAPI(
self.server_url, api_token=self.api_token
) as jupyter:
response = await jupyter.create_kernel(kernel_spec=kernel_spec)
self.logger.info(f"Kernel started with ID {response['id']}")
return response

async def list_kernels(self) -> list[KernelInfo]:
"""List kernels."""
async with JupyterAPI(
self.server_url, api_token=self.api_token
) as jupyter:
response = await jupyter.list_kernels()

self.logger.info(f"Kernels listed for {self.peer_host}")
return response

async def get_kernel_info(self, kernel_id) -> KernelInfo:
"""Get kernel info."""
async with JupyterAPI(
self.server_url, api_token=self.api_token
) as jupyter:
response = await jupyter.get_kernel(kernel_id=kernel_id)

self.logger.info(f"Kernel info retrieved for ID {kernel_id}")
return response

async def terminate_kernel(self, kernel_id) -> bool:
"""Terminate kernel."""
async with JupyterAPI(
self.server_url, api_token=self.api_token
) as jupyter:
response = await jupyter.delete_kernel(kernel_id=kernel_id)

self.logger.info(f"Kernel terminated for ID {kernel_id}")
return response

async def interrupt_kernel(self, kernel_id) -> bool:
"""Interrupt kernel."""
async with JupyterAPI(
self.server_url, api_token=self.api_token
) as jupyter:
response = await jupyter.interrupt_kernel(kernel_id=kernel_id)

self.logger.info(f"Kernel interrupted for ID {kernel_id}")
return response

async def restart_kernel(self, kernel_id) -> bool:
"""Restart kernel."""
async with JupyterAPI(
self.server_url, api_token=self.api_token
) as jupyter:
response = await jupyter.restart_kernel(kernel_id=kernel_id)

self.logger.info(f"Kernel restarted for ID {kernel_id}")
return response

def get_jupyter_api(self) -> JupyterAPI:
return JupyterAPI(manager=self)
Loading
Loading