0cca4963d1
Both the single-run path and the campaign scheduler drove long-lived asyncio tasks through their own duplicated _tasks/_cancel_events dicts and shutdown loops. Collapse them into one deep TaskRegistry module, instantiated as run_registry and campaign_registry. launch() creates the cancel event before the task (so a cancel during startup is never lost), wires done-callback cleanup, and is idempotent per id; this makes runs.py's hard-cancel fallback provably dead, so it is removed. App shutdown now gracefully stops in-flight runs too, not just campaigns.
88 lines
3.4 KiB
Python
88 lines
3.4 KiB
Python
"""Durable in-process task registry.
|
|
|
|
Both the single-run path (``web.routers.runs``) and the campaign scheduler
|
|
(``evaluation.campaign_runner``) drive long-lived ``asyncio`` tasks whose
|
|
authority lives in the DB — the in-process handle is only there to cancel or
|
|
gracefully stop the task. This module is the one deep seam they share: it owns
|
|
the (task, cancel-event) pair per id, wires cleanup, and stops everything on
|
|
shutdown. Each caller instantiates its own registry, so id spaces never collide.
|
|
"""
|
|
|
|
import asyncio
|
|
import logging
|
|
from typing import Awaitable, Callable
|
|
|
|
_logger = logging.getLogger("agenteval")
|
|
|
|
# A coroutine factory: given the cancel event the registry owns, return the
|
|
# coroutine to run. The event exists before the task starts, so a cancel that
|
|
# arrives during startup is never lost.
|
|
CoroFactory = Callable[[asyncio.Event], Awaitable[None]]
|
|
|
|
|
|
class TaskRegistry:
|
|
"""Tracks live asyncio tasks and their cooperative cancel events by id."""
|
|
|
|
def __init__(self) -> None:
|
|
self._tasks: dict[str, asyncio.Task] = {}
|
|
self._cancel_events: dict[str, asyncio.Event] = {}
|
|
|
|
def launch(self, id: str, make_coro: CoroFactory) -> asyncio.Task:
|
|
"""Start a tracked task for ``id`` (idempotent for a live id).
|
|
|
|
Creates the cancel event first, hands it to ``make_coro`` to build the
|
|
coroutine, then schedules the task and wires a done-callback that pops
|
|
both handles when it ends. If ``id`` already has a live task, returns it
|
|
without starting a second coroutine.
|
|
"""
|
|
existing = self._tasks.get(id)
|
|
if existing is not None:
|
|
return existing
|
|
|
|
cancel = self._cancel_events.setdefault(id, asyncio.Event())
|
|
task = asyncio.create_task(make_coro(cancel), name=id)
|
|
self._tasks[id] = task
|
|
task.add_done_callback(lambda _t, _id=id: self._cleanup(_id))
|
|
return task
|
|
|
|
def cancel(self, id: str) -> bool:
|
|
"""Signal the cooperative cancel event for ``id``.
|
|
|
|
Returns True if a live cancel event was signalled, False if there was
|
|
nothing to signal (e.g. the task already ended or never existed) — the
|
|
caller uses this to decide any fallback such as marking the DB directly.
|
|
"""
|
|
event = self._cancel_events.get(id)
|
|
if event is None:
|
|
return False
|
|
event.set()
|
|
return True
|
|
|
|
def get(self, id: str) -> "asyncio.Task | None":
|
|
"""Return the live task for ``id``, or None if there is none."""
|
|
return self._tasks.get(id)
|
|
|
|
def is_running(self, id: str) -> bool:
|
|
return id in self._tasks
|
|
|
|
async def shutdown_all(self) -> None:
|
|
"""Cooperatively stop, then hard-cancel and await, every live task."""
|
|
for event in list(self._cancel_events.values()):
|
|
event.set()
|
|
tasks = list(self._tasks.values())
|
|
for task in tasks:
|
|
task.cancel()
|
|
for task in tasks:
|
|
# CancelledError is expected (we just cancelled the task) and does
|
|
# not derive from Exception in 3.8+, so list it explicitly.
|
|
try:
|
|
await task
|
|
except (asyncio.CancelledError, Exception) as exc: # noqa: BLE001
|
|
_logger.debug("task registry shutdown: task ended with %r", exc)
|
|
self._tasks.clear()
|
|
self._cancel_events.clear()
|
|
|
|
def _cleanup(self, id: str) -> None:
|
|
self._tasks.pop(id, None)
|
|
self._cancel_events.pop(id, None)
|