refactor(tasks): unify run/campaign task registries into TaskRegistry

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.
This commit is contained in:
sinohqb 2026-07-31 03:39:03 +08:00
parent e815298ce5
commit 0cca4963d1
6 changed files with 198 additions and 56 deletions

View File

@ -37,6 +37,7 @@ from agenteval.storage.repository import (
ScenarioRepository, ScenarioRepository,
TargetRepository, TargetRepository,
) )
from agenteval.task_registry import TaskRegistry
_SPAWNED_KEY = "spawned_indices" _SPAWNED_KEY = "spawned_indices"
_ERRORS_KEY = "errors" _ERRORS_KEY = "errors"
@ -48,9 +49,8 @@ DEFAULT_TICK_SECONDS = 1.0
_logger = logging.getLogger("agenteval") _logger = logging.getLogger("agenteval")
# Live loop tasks + cooperative cancel events, keyed by campaign id. Authority # Live loop tasks + cooperative cancel events, keyed by campaign id. Authority
# is the DB; these are only the in-process handles for the running loop. # is the DB; this registry only holds the in-process handles for the loop.
_tasks: dict[str, asyncio.Task] = {} campaign_registry = TaskRegistry()
_cancel_events: dict[str, asyncio.Event] = {}
@dataclass @dataclass
@ -182,7 +182,9 @@ async def advance_campaign(
# ── durable scheduler loop ────────────────────────────────────────────────── # ── durable scheduler loop ──────────────────────────────────────────────────
async def run_campaign_loop(campaign_id: str, *, tick_seconds: float = DEFAULT_TICK_SECONDS) -> None: async def run_campaign_loop(
campaign_id: str, cancel: asyncio.Event, *, tick_seconds: float = DEFAULT_TICK_SECONDS
) -> None:
"""Drive one campaign to completion, ticking on real wall-clock time. """Drive one campaign to completion, ticking on real wall-clock time.
The loop holds no authority and no decision logic: every tick it reloads the The loop holds no authority and no decision logic: every tick it reloads the
@ -190,7 +192,6 @@ async def run_campaign_loop(campaign_id: str, *, tick_seconds: float = DEFAULT_T
It exits when the window finishes, the campaign leaves RUNNING (e.g. it was It exits when the window finishes, the campaign leaves RUNNING (e.g. it was
cancelled), or the cooperative cancel event fires. cancelled), or the cooperative cancel event fires.
""" """
cancel = _cancel_events.setdefault(campaign_id, asyncio.Event())
session = get_session() session = get_session()
try: try:
while not cancel.is_set(): while not cancel.is_set():
@ -237,8 +238,6 @@ async def run_campaign_loop(campaign_id: str, *, tick_seconds: float = DEFAULT_T
_logger.warning("活动调度循环 %s 异常退出: %s", campaign_id, exc) _logger.warning("活动调度循环 %s 异常退出: %s", campaign_id, exc)
finally: finally:
session.close() session.close()
_cancel_events.pop(campaign_id, None)
_tasks.pop(campaign_id, None)
def start_campaign(campaign_id: str, session: Session, *, tick_seconds: float = DEFAULT_TICK_SECONDS) -> Optional[asyncio.Task]: def start_campaign(campaign_id: str, session: Session, *, tick_seconds: float = DEFAULT_TICK_SECONDS) -> Optional[asyncio.Task]:
@ -251,28 +250,21 @@ def start_campaign(campaign_id: str, session: Session, *, tick_seconds: float =
campaign = repo.get(campaign_id) campaign = repo.get(campaign_id)
if not campaign or campaign.status in (CampaignStatus.COMPLETED, CampaignStatus.CANCELLED, CampaignStatus.FAILED): if not campaign or campaign.status in (CampaignStatus.COMPLETED, CampaignStatus.CANCELLED, CampaignStatus.FAILED):
return None return None
if campaign_id in _tasks:
return _tasks[campaign_id]
if campaign.status == CampaignStatus.PLANNED: if campaign.status == CampaignStatus.PLANNED:
campaign.status = CampaignStatus.RUNNING campaign.status = CampaignStatus.RUNNING
campaign.started_at = utc_now() campaign.started_at = utc_now()
repo.update(campaign) repo.update(campaign)
_cancel_events.setdefault(campaign_id, asyncio.Event()) return campaign_registry.launch(
task = asyncio.create_task( campaign_id,
run_campaign_loop(campaign_id, tick_seconds=tick_seconds), lambda cancel: run_campaign_loop(campaign_id, cancel, tick_seconds=tick_seconds),
name=f"campaign-{campaign_id}",
) )
_tasks[campaign_id] = task
return task
def request_cancel(campaign_id: str) -> None: def request_cancel(campaign_id: str) -> None:
"""Signal the loop (if live) to stop spawning and exit promptly.""" """Signal the loop (if live) to stop spawning and exit promptly."""
event = _cancel_events.get(campaign_id) campaign_registry.cancel(campaign_id)
if event is not None:
event.set()
def resume_running_campaigns(session: Session, *, tick_seconds: float = DEFAULT_TICK_SECONDS) -> int: def resume_running_campaigns(session: Session, *, tick_seconds: float = DEFAULT_TICK_SECONDS) -> int:
@ -287,17 +279,4 @@ def resume_running_campaigns(session: Session, *, tick_seconds: float = DEFAULT_
async def shutdown_all() -> None: async def shutdown_all() -> None:
"""Gracefully stop all live loops (Web app shutdown).""" """Gracefully stop all live loops (Web app shutdown)."""
for event in list(_cancel_events.values()): await campaign_registry.shutdown_all()
event.set()
tasks = list(_tasks.values())
for task in tasks:
task.cancel()
for task in tasks:
# CancelledError is expected here (we just cancelled the task) and does
# not derive from Exception in 3.8+, so it must be listed explicitly.
try:
await task
except (asyncio.CancelledError, Exception):
pass
_tasks.clear()
_cancel_events.clear()

View File

@ -0,0 +1,87 @@
"""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)

View File

@ -38,11 +38,13 @@ async def lifespan(_: FastAPI):
except Exception as exc: except Exception as exc:
logging.getLogger("agenteval").warning("启动清理失败(忽略): %s", exc) logging.getLogger("agenteval").warning("启动清理失败(忽略): %s", exc)
yield yield
# 优雅停止所有活动调度循环 # 优雅停止所有进程内任务:先停活动调度循环,再停在跑的评测运行
try: try:
from agenteval.evaluation.campaign_runner import shutdown_all from agenteval.evaluation.campaign_runner import shutdown_all
from agenteval.web.routers.runs import run_registry
await shutdown_all() await shutdown_all()
await run_registry.shutdown_all()
except Exception as exc: except Exception as exc:
logging.getLogger("agenteval").warning("活动调度停止失败(忽略): %s", exc) logging.getLogger("agenteval").warning("活动调度停止失败(忽略): %s", exc)

View File

@ -1,7 +1,6 @@
"""API routes for evaluation runs.""" """API routes for evaluation runs."""
import asyncio import asyncio
from typing import Optional
from fastapi import APIRouter, Depends, HTTPException from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel from pydantic import BaseModel
@ -11,6 +10,7 @@ from agenteval.evaluation.engine import EvalEngine
from agenteval.models import EvalRun, RunStatus, RunTrigger from agenteval.models import EvalRun, RunStatus, RunTrigger
from agenteval.storage.db import get_session, iso_utc from agenteval.storage.db import get_session, iso_utc
from agenteval.storage.repository import RunRepository, ScenarioRepository, TargetRepository from agenteval.storage.repository import RunRepository, ScenarioRepository, TargetRepository
from agenteval.task_registry import TaskRegistry
from agenteval.utils.llm import extract_reply_text from agenteval.utils.llm import extract_reply_text
from agenteval.utils.webhook import send_run_webhook from agenteval.utils.webhook import send_run_webhook
from agenteval.web.deps import get_db from agenteval.web.deps import get_db
@ -28,15 +28,14 @@ class StartRunRequest(BaseModel):
# ── Task registry for live evaluation runs ───────────────────────────── # ── Task registry for live evaluation runs ─────────────────────────────
# Each running evaluation is an asyncio.Task keyed by run_id. The cancel # Each running evaluation is an asyncio.Task keyed by run_id. The cancel
# token is a cooperative ``asyncio.Event`` the engine checks between cases. # token is a cooperative ``asyncio.Event`` the engine checks between cases.
_tasks: dict[str, asyncio.Task] = {} run_registry = TaskRegistry()
_cancel_tokens: dict[str, asyncio.Event] = {}
async def _run_evaluation(run_id: str, target_id: str, scenario_id: str) -> None: async def _run_evaluation(
run_id: str, target_id: str, scenario_id: str, *, cancel_token: asyncio.Event
) -> None:
"""Background coroutine that drives one evaluation run to completion.""" """Background coroutine that drives one evaluation run to completion."""
session = get_session() session = get_session()
cancel_token = asyncio.Event()
_cancel_tokens[run_id] = cancel_token
try: try:
target = TargetRepository(session).get(target_id) target = TargetRepository(session).get(target_id)
scenario = ScenarioRepository(session).get(scenario_id) scenario = ScenarioRepository(session).get(scenario_id)
@ -64,8 +63,6 @@ async def _run_evaluation(run_id: str, target_id: str, scenario_id: str) -> None
) )
finally: finally:
session.close() session.close()
_cancel_tokens.pop(run_id, None)
_tasks.pop(run_id, None)
@router.get("") @router.get("")
@ -100,11 +97,12 @@ async def start_run(
) )
run = RunRepository(session).create(run) run = RunRepository(session).create(run)
task = asyncio.create_task( run_registry.launch(
_run_evaluation(run.id, request.target_id, request.scenario_id), run.id,
name=f"eval-run-{run.id}", lambda cancel_token: _run_evaluation(
run.id, request.target_id, request.scenario_id, cancel_token=cancel_token
),
) )
_tasks[run.id] = task
return run.model_dump() return run.model_dump()
@ -125,16 +123,8 @@ async def cancel_run(run_id: str, session: Session = Depends(get_db)) -> dict:
if run.status not in (RunStatus.PENDING, RunStatus.RUNNING): if run.status not in (RunStatus.PENDING, RunStatus.RUNNING):
raise HTTPException(status_code=400, detail="run is not in a cancellable state") raise HTTPException(status_code=400, detail="run is not in a cancellable state")
cancel_token = _cancel_tokens.get(run_id) signalled = run_registry.cancel(run_id)
task: Optional[asyncio.Task] = _tasks.get(run_id) if not signalled:
if cancel_token is not None:
# Cooperative cancel: the engine will catch CancelledError and mark
# the run as FAILED with code=cancelled_by_user.
cancel_token.set()
elif task is not None:
# Fallback: hard-cancel the task if no token exists (shouldn't happen).
task.cancel()
else:
# No live task (e.g. process restarted): mark the DB row directly. # No live task (e.g. process restarted): mark the DB row directly.
run.status = RunStatus.FAILED run.status = RunStatus.FAILED
run.summary = { run.summary = {

View File

@ -70,7 +70,7 @@ def _make_campaign(session, **overrides) -> Campaign:
async def _await_task(campaign_id, timeout=3.0): async def _await_task(campaign_id, timeout=3.0):
import asyncio import asyncio
task = campaign_runner._tasks.get(campaign_id) task = campaign_runner.campaign_registry.get(campaign_id)
if task is not None: if task is not None:
await asyncio.wait_for(task, timeout=timeout) await asyncio.wait_for(task, timeout=timeout)

View File

@ -0,0 +1,84 @@
"""Unit tests for TaskRegistry — the durable in-process task registry.
Covers the single seam both runs and campaigns share: launch (with an
injected cancel event), idempotency, cooperative cancel, done-callback
cleanup, and graceful shutdown_all.
"""
import asyncio
from agenteval.task_registry import TaskRegistry
async def test_launch_tracks_task_and_passes_cancel_event():
reg = TaskRegistry()
seen: dict = {}
async def coro(cancel: asyncio.Event):
seen["event"] = cancel
await cancel.wait()
task = reg.launch("a", lambda ev: coro(ev))
assert reg.is_running("a")
await asyncio.sleep(0) # let the coroutine start and capture the event
assert isinstance(seen["event"], asyncio.Event)
reg.cancel("a")
await task
assert not reg.is_running("a") # done-callback popped it
async def test_launch_is_idempotent_for_live_id():
reg = TaskRegistry()
async def coro(cancel: asyncio.Event):
await cancel.wait()
t1 = reg.launch("a", lambda ev: coro(ev))
t2 = reg.launch("a", lambda ev: coro(ev))
assert t1 is t2 # same live task, second coro never started
reg.cancel("a")
await t1
async def test_cancel_returns_true_when_live_false_otherwise():
reg = TaskRegistry()
async def coro(cancel: asyncio.Event):
await cancel.wait()
reg.launch("a", lambda ev: coro(ev))
assert reg.cancel("a") is True # live event signalled
assert reg.cancel("missing") is False # nothing to signal
await asyncio.sleep(0) # let the cancelled coro finish
async def test_done_callback_cleans_up_after_natural_completion():
reg = TaskRegistry()
async def quick(cancel: asyncio.Event):
return "done"
task = reg.launch("a", lambda ev: quick(ev))
await task
await asyncio.sleep(0) # allow done-callback to run
assert not reg.is_running("a")
async def test_shutdown_all_cancels_and_awaits_every_task():
reg = TaskRegistry()
started = []
async def coro(cancel: asyncio.Event):
started.append(1)
await cancel.wait()
reg.launch("a", lambda ev: coro(ev))
reg.launch("b", lambda ev: coro(ev))
await asyncio.sleep(0) # let both start
assert len(started) == 2
await reg.shutdown_all()
assert not reg.is_running("a")
assert not reg.is_running("b")