From c4cef70bef0c105d772fd255b22a9a38794a2017 Mon Sep 17 00:00:00 2001 From: David Brochart Date: Sun, 1 Nov 2026 14:46:09 +0100 Subject: [PATCH] Use AnyIO portal for asynchronous worker commands --- pyproject.toml | 2 +- src/akernel/kernel.py | 51 ++++++++++++++++++++--------------------- tests/test_interrupt.py | 11 +++++++++ 3 files changed, 37 insertions(+), 27 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 15f1643..efee562 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,7 @@ classifiers = [ ] keywords = [ "jupyter" ] dependencies = [ - "anyio >=4.14.2,<5.0.0", + "anyio >=4.16.0,<5.0.0", "python-dateutil", "rich >=14.0,<15", "comm >=0.1.3,<1", diff --git a/src/akernel/kernel.py b/src/akernel/kernel.py index d3627ac..28431a5 100644 --- a/src/akernel/kernel.py +++ b/src/akernel/kernel.py @@ -3,7 +3,6 @@ from collections.abc import Callable from contextvars import ContextVar from dataclasses import dataclass -from functools import partial from io import StringIO from types import TracebackType from typing import Any @@ -21,7 +20,8 @@ run, to_thread, ) -from anyio.lowlevel import EventLoopToken, checkpoint, current_token +from anyio.from_thread import BlockingPortal +from anyio.lowlevel import checkpoint import akernel.IPython from akernel.comm.manager import CommManager @@ -40,7 +40,7 @@ from .traceback import get_traceback ThreadResult = tuple[Any, BaseException | None, TracebackType | None] -WorkerCommand = tuple[EventLoopToken, Callable[..., Any], tuple[Any, ...]] +WorkerCommand = tuple[Callable[..., Any], tuple[Any, ...]] PARENT_VAR: ContextVar = ContextVar("parent") IDENTS_VAR: ContextVar = ContextVar("idents") @@ -137,7 +137,7 @@ def __init__( self.msg_cnt = 0 self.stop_event = Event() self._stopping = False - self._thread_token: EventLoopToken | None = None + self._thread_portal: BlockingPortal | None = None self._thread_jobs: dict[int, ThreadExecution] = {} self._cancelled_cells: set[int] = set() self._finishing_cells: dict[int, CancelScope] = {} @@ -198,25 +198,22 @@ def cancel_thread_execution(self, job: ThreadExecution) -> None: if job.cancelled(): return job.cancel() - if self._thread_token is not None and job.cancel_scope is not None: + if self._thread_portal is not None and job.cancel_scope is not None: self.send_worker_command(job.cancel_scope.cancel) def send_worker_command(self, func: Callable[..., Any], *args: Any) -> None: - token = self._thread_token - assert token is not None - self._worker_commands_send.send_nowait((token, func, args)) + self._worker_commands_send.send_nowait((func, args)) - async def forward_worker_commands(self) -> None: - # Cross-thread calls wait for the worker loop. Run them off the server - # thread and preserve their order, including the shutdown sentinel. + async def forward_worker_commands(self, portal: BlockingPortal) -> None: + # Preserve command order, including the shutdown sentinel, without + # blocking the server loop while the worker loop is busy. with CancelScope(shield=True): async with self._worker_commands_receive: - async for token, func, args in self._worker_commands_receive: + async for func, args in self._worker_commands_receive: try: - await to_thread.run_sync( - partial(from_thread.run_sync, func, *args, token=token) - ) - except RunFinishedError: + await portal.call_async(func, *args) + except (RunFinishedError, RuntimeError): + # Shutdown may close the portal before queued commands run. if not self._stopping: raise @@ -259,17 +256,18 @@ def _is_stopping(self) -> bool: return self._stopping async def thread_main(self) -> None: - self._thread_token = current_token() self._thread_send, self._thread_receive = create_memory_object_stream[ ThreadExecution | None ](float("inf")) - from_thread.run_sync(self._thread_ready.set) - try: - async with self._thread_send, self._thread_receive: - if not from_thread.run_sync(self._is_stopping): - await self.thread_execute() - finally: - self._thread_token = None + async with BlockingPortal() as portal: + self._thread_portal = portal + from_thread.run_sync(self._thread_ready.set) + try: + async with self._thread_send, self._thread_receive: + if not from_thread.run_sync(self._is_stopping): + await self.thread_execute() + finally: + self._thread_portal = None def run_thread(self) -> None: run(self.thread_main) @@ -282,10 +280,11 @@ async def start(self) -> None: self._worker_commands_send, self._worker_commands_receive = ( create_memory_object_stream[WorkerCommand](float("inf")) ) - self.task_group.start_soon(self.forward_worker_commands) self._thread_ready = Event() self.task_group.start_soon(to_thread.run_sync, self.run_thread) await self._thread_ready.wait() + assert self._thread_portal is not None + self.task_group.start_soon(self.forward_worker_commands, self._thread_portal) msg = self.create_message( "status", content={"execution_state": self.execution_state} ) @@ -303,7 +302,7 @@ async def start(self) -> None: self.interrupt() for scope in list(self._finishing_cells.values()): scope.cancel() - if self._thread_token is not None: + if self._thread_portal is not None: self.send_worker_command(self._thread_send.send_nowait, None) finally: if self.execute_in_thread: diff --git a/tests/test_interrupt.py b/tests/test_interrupt.py index ba59e62..4db38a9 100644 --- a/tests/test_interrupt.py +++ b/tests/test_interrupt.py @@ -247,6 +247,17 @@ async def test_shutdown_while_cell_is_running(self): class TestThreadInterrupt(TestInterrupt): execute_in_thread = True + async def test_worker_commands_do_not_need_another_thread(self): + limiter = anyio.to_thread.current_default_thread_limiter() + original_tokens = limiter.total_tokens + try: + # The user-code event loop already occupies the only worker slot. + limiter.total_tokens = 1 + msg_id = await self.execute("1 + 2") + await self.assert_finished(msg_id, "ok") + finally: + limiter.total_tokens = original_tokens + async def test_interrupt_before_worker_starts_cell(self): commands = [] original_send = self.kernel.send_worker_command