shared shell registry
This commit is contained in:
@@ -1,7 +1,11 @@
|
||||
from collections.abc import AsyncIterator
|
||||
from datetime import UTC, datetime
|
||||
from time import monotonic
|
||||
from uuid import uuid4
|
||||
|
||||
from ..utils import generate_good_names
|
||||
from ..models import AppInfo, ShellInfo
|
||||
from ..registry import ShellRegistry
|
||||
from ..utils import generate_call_id
|
||||
from .shell import JupyterShell
|
||||
from .transport import InputRequest, KernelMessage
|
||||
|
||||
@@ -10,36 +14,61 @@ class JupyterApp:
|
||||
"""Coordinate named persistent shells backed by Jupyter kernels."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.shells: dict[str, JupyterShell] = {}
|
||||
self.last_shell: str | None = None
|
||||
self.app_id = uuid4().hex
|
||||
self.created_at = datetime.now(UTC).isoformat()
|
||||
self._started_at = monotonic()
|
||||
self._total_calls = 0
|
||||
self._shell_registry = ShellRegistry(JupyterShell)
|
||||
self._calls: dict[str, JupyterShell] = {}
|
||||
self._issued_call_ids: set[str] = set()
|
||||
|
||||
def _new_shell(self) -> tuple[str, JupyterShell]:
|
||||
name = generate_good_names()
|
||||
while name in self.shells:
|
||||
name = generate_good_names()
|
||||
shell = JupyterShell(name)
|
||||
self.shells[name] = shell
|
||||
self.last_shell = name
|
||||
return name, shell
|
||||
@property
|
||||
def shells(self) -> dict[str, JupyterShell]:
|
||||
"""Expose the managed shells for compatibility with the sync app."""
|
||||
return self._shell_registry.shells
|
||||
|
||||
@property
|
||||
def last_shell(self) -> str | None:
|
||||
"""Return the name of the most recently selected shell."""
|
||||
return self._shell_registry.last_shell
|
||||
|
||||
def _select_shell(self, shell_name: str) -> tuple[str, JupyterShell]:
|
||||
if shell_name == "new":
|
||||
return self._new_shell()
|
||||
if shell_name == "last":
|
||||
if self.last_shell is None:
|
||||
return self._new_shell()
|
||||
return self.last_shell, self.shells[self.last_shell]
|
||||
if shell_name not in self.shells:
|
||||
self.shells[shell_name] = JupyterShell(shell_name)
|
||||
self.last_shell = shell_name
|
||||
return shell_name, self.shells[shell_name]
|
||||
return self._shell_registry.select(shell_name)
|
||||
|
||||
def get_shell(self, shell_name: str) -> JupyterShell:
|
||||
"""Return a named Jupyter shell."""
|
||||
if shell_name not in self.shells:
|
||||
raise KeyError(f"Shell {shell_name} does not exist")
|
||||
return self.shells[shell_name]
|
||||
def list_shells(self) -> list[ShellInfo]:
|
||||
"""Return metadata for every known Jupyter shell."""
|
||||
return self._shell_registry.list_shells()
|
||||
|
||||
def create_shell(self) -> ShellInfo:
|
||||
"""Create a fresh Jupyter shell and return its metadata."""
|
||||
return self._shell_registry.create_shell()
|
||||
|
||||
def get_shell(self, shell_name: str) -> ShellInfo:
|
||||
"""Return frontend-safe metadata for one Jupyter shell."""
|
||||
return self._shell_registry.get_shell(shell_name)
|
||||
|
||||
def get_last_shell(self) -> ShellInfo:
|
||||
"""Return metadata for the most recently selected Jupyter shell."""
|
||||
return self._shell_registry.get_last_shell()
|
||||
|
||||
def info(self) -> AppInfo:
|
||||
"""Return runtime metadata without starting any kernels."""
|
||||
return AppInfo(
|
||||
app_id=self.app_id,
|
||||
created_at=self.created_at,
|
||||
uptime_seconds=monotonic() - self._started_at,
|
||||
shell_count=len(self.shells),
|
||||
last_shell=self.last_shell,
|
||||
total_calls=self._total_calls,
|
||||
)
|
||||
|
||||
def _new_call_id(self) -> str:
|
||||
"""Generate a readable call ID that is unique for this app."""
|
||||
call_id = generate_call_id()
|
||||
while call_id in self._issued_call_ids:
|
||||
call_id = generate_call_id()
|
||||
self._issued_call_ids.add(call_id)
|
||||
return call_id
|
||||
|
||||
async def run_code_stream(
|
||||
self,
|
||||
@@ -48,7 +77,8 @@ class JupyterApp:
|
||||
) -> AsyncIterator[KernelMessage | InputRequest]:
|
||||
"""Execute code and yield its subprocess messages."""
|
||||
_, shell = self._select_shell(shell_name)
|
||||
call_id = uuid4().hex
|
||||
call_id = self._new_call_id()
|
||||
self._total_calls += 1
|
||||
self._calls[call_id] = shell
|
||||
try:
|
||||
async for message in shell.run_cell_stream(code, call_id=call_id):
|
||||
|
||||
@@ -1,6 +1,9 @@
|
||||
from collections.abc import AsyncIterator
|
||||
from dataclasses import replace
|
||||
from datetime import UTC, datetime
|
||||
from uuid import uuid4
|
||||
|
||||
from ..models import ShellInfo, ShellStatus
|
||||
from .transport import InputRequest, JupyterTransport, KernelMessage
|
||||
|
||||
|
||||
@@ -9,6 +12,11 @@ class JupyterShell:
|
||||
|
||||
def __init__(self, name: str, transport: JupyterTransport | None = None) -> None:
|
||||
self.name = name
|
||||
self.shell_id = uuid4().hex
|
||||
self.created_at = datetime.now(UTC).isoformat()
|
||||
self.last_used_at: str | None = None
|
||||
self.status: ShellStatus = "ready"
|
||||
self._last_execution_count = 0
|
||||
self.transport = transport or JupyterTransport()
|
||||
self._active_call_id: str | None = None
|
||||
|
||||
@@ -26,6 +34,7 @@ class JupyterShell:
|
||||
"""Yield kernel messages while one cell executes."""
|
||||
await self.start()
|
||||
self._active_call_id = call_id
|
||||
self.status = "running"
|
||||
try:
|
||||
kernel_call_id = await self.transport.execute(code)
|
||||
async for message in self.transport.messages_for(kernel_call_id):
|
||||
@@ -33,10 +42,25 @@ class JupyterShell:
|
||||
yield replace(message, call_id=call_id)
|
||||
else:
|
||||
# Keep the kernel's execution ID private to the transport.
|
||||
yield replace(message, parent_id=call_id)
|
||||
public_message = replace(message, parent_id=call_id)
|
||||
if public_message.msg_type == "execute_reply":
|
||||
self._last_execution_count = int(
|
||||
public_message.content.get(
|
||||
"execution_count", self._last_execution_count
|
||||
)
|
||||
)
|
||||
self.status = (
|
||||
"error"
|
||||
if public_message.content.get("status") == "error"
|
||||
else "ready"
|
||||
)
|
||||
self.last_used_at = datetime.now(UTC).isoformat()
|
||||
yield public_message
|
||||
finally:
|
||||
if self._active_call_id == call_id:
|
||||
self._active_call_id = None
|
||||
if self.status == "running" and self._active_call_id is None:
|
||||
self.status = "ready"
|
||||
|
||||
async def reply_to_input(self, call_id: str, value: str) -> None:
|
||||
"""Send input to the cell currently waiting in this shell."""
|
||||
@@ -47,3 +71,23 @@ class JupyterShell:
|
||||
async def shutdown(self) -> None:
|
||||
"""Stop this shell's kernel process."""
|
||||
await self.transport.shutdown()
|
||||
|
||||
def describe(self) -> ShellInfo:
|
||||
"""Return metadata without starting an unused kernel."""
|
||||
return ShellInfo(
|
||||
shell_id=self.shell_id,
|
||||
name=self.name,
|
||||
status=self.status,
|
||||
execution_count=self._last_execution_count,
|
||||
created_at=self.created_at,
|
||||
last_used_at=self.last_used_at,
|
||||
capabilities=(
|
||||
"stdout",
|
||||
"stderr",
|
||||
"display_data",
|
||||
"execute_result",
|
||||
"error",
|
||||
"matplotlib-inline",
|
||||
"input",
|
||||
),
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user