feat: expose Jupyter shell through web app
This commit is contained in:
@@ -0,0 +1,11 @@
|
||||
from .app import JupyterApp
|
||||
from .shell import JupyterShell
|
||||
from .transport import InputRequest, JupyterTransport, KernelMessage
|
||||
|
||||
__all__ = [
|
||||
"InputRequest",
|
||||
"JupyterApp",
|
||||
"JupyterShell",
|
||||
"JupyterTransport",
|
||||
"KernelMessage",
|
||||
]
|
||||
@@ -0,0 +1,71 @@
|
||||
from collections.abc import AsyncIterator
|
||||
from uuid import uuid4
|
||||
|
||||
from .shell import JupyterShell
|
||||
from .transport import InputRequest, KernelMessage
|
||||
from ..utils import generate_good_names
|
||||
|
||||
|
||||
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._calls: dict[str, JupyterShell] = {}
|
||||
|
||||
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
|
||||
|
||||
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]
|
||||
|
||||
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]
|
||||
|
||||
async def run_code_stream(
|
||||
self,
|
||||
code: str,
|
||||
shell_name: str = "last",
|
||||
) -> AsyncIterator[KernelMessage | InputRequest]:
|
||||
"""Execute code and yield its subprocess messages."""
|
||||
_, shell = self._select_shell(shell_name)
|
||||
call_id = uuid4().hex
|
||||
self._calls[call_id] = shell
|
||||
try:
|
||||
async for message in shell.run_cell_stream(code, call_id=call_id):
|
||||
yield message
|
||||
finally:
|
||||
self._calls.pop(call_id, None)
|
||||
|
||||
async def reply_to_input(self, call_id: str, value: str) -> None:
|
||||
"""Reply to an input request emitted by a running call."""
|
||||
try:
|
||||
shell = self._calls[call_id]
|
||||
except KeyError as error:
|
||||
raise KeyError(f"Call {call_id} does not exist") from error
|
||||
await shell.reply_to_input(call_id, value)
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
"""Stop all kernel processes owned by this app."""
|
||||
awaitables = [shell.shutdown() for shell in self.shells.values()]
|
||||
for awaitable in awaitables:
|
||||
await awaitable
|
||||
@@ -0,0 +1,49 @@
|
||||
from collections.abc import AsyncIterator
|
||||
from dataclasses import replace
|
||||
|
||||
from .transport import InputRequest, JupyterTransport, KernelMessage
|
||||
|
||||
|
||||
class JupyterShell:
|
||||
"""One persistent, subprocess-backed IPython shell."""
|
||||
|
||||
def __init__(self, name: str, transport: JupyterTransport | None = None) -> None:
|
||||
self.name = name
|
||||
self.transport = transport or JupyterTransport()
|
||||
self._active_call_id: str | None = None
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start this shell's kernel if it is not already running."""
|
||||
if self.transport.client is None:
|
||||
await self.transport.start()
|
||||
|
||||
async def run_cell_stream(
|
||||
self,
|
||||
code: str,
|
||||
*,
|
||||
call_id: str,
|
||||
) -> AsyncIterator[KernelMessage | InputRequest]:
|
||||
"""Yield kernel messages while one cell executes."""
|
||||
await self.start()
|
||||
self._active_call_id = call_id
|
||||
try:
|
||||
kernel_call_id = await self.transport.execute(code)
|
||||
async for message in self.transport.messages_for(kernel_call_id):
|
||||
if isinstance(message, InputRequest):
|
||||
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)
|
||||
finally:
|
||||
if self._active_call_id == call_id:
|
||||
self._active_call_id = None
|
||||
|
||||
async def reply_to_input(self, call_id: str, value: str) -> None:
|
||||
"""Send input to the cell currently waiting in this shell."""
|
||||
if self._active_call_id != call_id:
|
||||
raise RuntimeError(f"No active input request for call: {call_id}")
|
||||
await self.transport.reply_to_input(value)
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
"""Stop this shell's kernel process."""
|
||||
await self.transport.shutdown()
|
||||
@@ -0,0 +1,139 @@
|
||||
import asyncio
|
||||
from collections.abc import AsyncIterator
|
||||
from dataclasses import dataclass
|
||||
|
||||
from jupyter_client import AsyncKernelManager
|
||||
|
||||
|
||||
@dataclass
|
||||
class KernelMessage:
|
||||
"""A decoded Jupyter message kept at the transport boundary."""
|
||||
|
||||
msg_type: str
|
||||
parent_id: str | None
|
||||
content: dict[str, object]
|
||||
buffers: list[bytes]
|
||||
|
||||
|
||||
@dataclass
|
||||
class InputRequest:
|
||||
"""A request for user input emitted by a running kernel call."""
|
||||
|
||||
call_id: str
|
||||
prompt: str
|
||||
password: bool
|
||||
|
||||
|
||||
class JupyterTransport:
|
||||
"""Own one persistent ipykernel subprocess and its message channels."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.manager = AsyncKernelManager()
|
||||
self.client = None
|
||||
self._active_call: str | None = None
|
||||
self._call_lock = asyncio.Lock()
|
||||
self._waiting_for_input: str | None = None
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the kernel process and wait until it accepts requests."""
|
||||
await self.manager.start_kernel()
|
||||
self.client = self.manager.client()
|
||||
self.client.start_channels()
|
||||
await self.client.wait_for_ready()
|
||||
|
||||
async def execute(self, code: str) -> str:
|
||||
"""Submit one cell and return its Jupyter message ID."""
|
||||
if self.client is None:
|
||||
raise RuntimeError("JupyterTransport has not been started")
|
||||
|
||||
async with self._call_lock:
|
||||
if self._active_call is not None:
|
||||
raise RuntimeError("JupyterTransport already has an active call")
|
||||
self._active_call = self.client.execute(code, allow_stdin=True)
|
||||
return self._active_call
|
||||
|
||||
async def messages_for(
|
||||
self, call_id: str
|
||||
) -> AsyncIterator[KernelMessage | InputRequest]:
|
||||
"""Yield decoded output, input, and completion messages for one call."""
|
||||
if self.client is None:
|
||||
raise RuntimeError("JupyterTransport has not been started")
|
||||
if self._active_call != call_id:
|
||||
raise RuntimeError(f"Unknown or inactive call: {call_id}")
|
||||
|
||||
reply: KernelMessage | None = None
|
||||
idle = False
|
||||
try:
|
||||
while True:
|
||||
iopub_task = asyncio.create_task(self.client.get_iopub_msg())
|
||||
shell_task = asyncio.create_task(self.client.get_shell_msg())
|
||||
stdin_task = asyncio.create_task(self.client.get_stdin_msg())
|
||||
done, pending = await asyncio.wait(
|
||||
{iopub_task, shell_task, stdin_task},
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
)
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
await asyncio.gather(*pending, return_exceptions=True)
|
||||
|
||||
for task in done:
|
||||
message = task.result()
|
||||
parent_id = message.get("parent_header", {}).get("msg_id")
|
||||
if parent_id != call_id:
|
||||
continue
|
||||
|
||||
if message["msg_type"] == "input_request":
|
||||
self._waiting_for_input = call_id
|
||||
yield InputRequest(
|
||||
call_id=call_id,
|
||||
prompt=str(message.get("content", {}).get("prompt", "")),
|
||||
password=bool(
|
||||
message.get("content", {}).get("password", False)
|
||||
),
|
||||
)
|
||||
continue
|
||||
|
||||
decoded = KernelMessage(
|
||||
msg_type=message["msg_type"],
|
||||
parent_id=parent_id,
|
||||
content=dict(message.get("content", {})),
|
||||
buffers=[bytes(buffer) for buffer in message.get("buffers", [])],
|
||||
)
|
||||
if decoded.msg_type == "execute_reply":
|
||||
# The shell reply and IOPub messages use different
|
||||
# channels. Buffer the reply so it remains terminal
|
||||
# even when both channels become ready together.
|
||||
reply = decoded
|
||||
else:
|
||||
yield decoded
|
||||
if (
|
||||
decoded.msg_type == "status"
|
||||
and decoded.content.get("execution_state") == "idle"
|
||||
):
|
||||
idle = True
|
||||
|
||||
if reply is not None and idle:
|
||||
yield reply
|
||||
return
|
||||
finally:
|
||||
if self._waiting_for_input == call_id:
|
||||
self._waiting_for_input = None
|
||||
if self._active_call == call_id:
|
||||
self._active_call = None
|
||||
|
||||
async def reply_to_input(self, value: str) -> None:
|
||||
"""Reply to the active kernel input request without using process stdin."""
|
||||
if self.client is None:
|
||||
raise RuntimeError("JupyterTransport has not been started")
|
||||
if self._waiting_for_input is None:
|
||||
raise RuntimeError("JupyterTransport is not waiting for input")
|
||||
|
||||
self.client.input(value)
|
||||
self._waiting_for_input = None
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
"""Stop channels and terminate the kernel process."""
|
||||
if self.client is not None:
|
||||
self.client.stop_channels()
|
||||
self.client = None
|
||||
await self.manager.shutdown_kernel(now=True)
|
||||
Reference in New Issue
Block a user