464 lines
17 KiB
Python
464 lines
17 KiB
Python
"""Campaign runner — the side-effect shell around the pure scheduler.
|
|
|
|
Given a campaign and a clock position (real elapsed seconds), it asks
|
|
``campaign_scheduler.decide_schedule`` what is due, then durably claims each
|
|
plan occurrence before reusing the existing single-run execution path
|
|
(``EvalEngine.run`` with an ``existing_run``). Child-Run identities are the
|
|
primary idempotency authority; the legacy campaign summary remains a fallback
|
|
for pre-identity Runs created before this migration.
|
|
|
|
``CampaignRuntime`` owns the durable scheduler loop and exposes only lifecycle-
|
|
level operations. Clock mapping, reconciliation, ticking, claims and process-
|
|
local task handles remain implementation details. All authority lives in the
|
|
DB, so loops can be rebuilt on restart without losing or duplicating work.
|
|
"""
|
|
|
|
import asyncio
|
|
import logging
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime
|
|
from typing import Awaitable, Callable, Optional
|
|
|
|
from sqlmodel import Session
|
|
|
|
from agenteval.evaluation.analysis import resolve_analysis_model
|
|
from agenteval.evaluation.campaign_lifecycle import complete_campaign
|
|
from agenteval.evaluation.campaign_lifecycle import start_campaign as start_campaign_lifecycle
|
|
from agenteval.evaluation.campaign_scheduler import (
|
|
TickAction,
|
|
clock_offset,
|
|
decide_schedule,
|
|
decide_tick,
|
|
elapsed_seconds,
|
|
resolve_finalize,
|
|
)
|
|
from agenteval.evaluation.engine import EvalEngine
|
|
from agenteval.evaluation.intelligence_jobs import enqueue_campaign_analysis
|
|
from agenteval.models import (
|
|
Campaign,
|
|
CampaignStatus,
|
|
CampaignSummary,
|
|
EvalRun,
|
|
RunStatus,
|
|
RunTrigger,
|
|
Scenario,
|
|
SchedulerState,
|
|
)
|
|
from agenteval.storage.db import get_session, utc_now
|
|
from agenteval.storage.repository import (
|
|
CampaignRepository,
|
|
CampaignRunClaimStatus,
|
|
RunRepository,
|
|
ScenarioRepository,
|
|
TargetRepository,
|
|
)
|
|
from agenteval.task_registry import TaskRegistry
|
|
|
|
# Real wall-clock seconds between scheduler ticks. time_scale compresses the
|
|
# *window*, not the tick cadence — production 24h campaigns still tick slowly.
|
|
DEFAULT_TICK_SECONDS = 1.0
|
|
|
|
_logger = logging.getLogger("agenteval")
|
|
|
|
|
|
ChildRunExecutor = Callable[[Campaign, Scenario, EvalRun, Session], Awaitable[None]]
|
|
|
|
|
|
@dataclass
|
|
class AdvanceResult:
|
|
"""Outcome of advancing a campaign's clock once."""
|
|
|
|
spawned_run_ids: list[str] = field(default_factory=list)
|
|
finished: bool = False
|
|
|
|
|
|
@dataclass
|
|
class CampaignRecoveryResult:
|
|
"""Outcome of reconciling child Runs left by a previous process."""
|
|
|
|
resumed_run_ids: list[str] = field(default_factory=list)
|
|
failed_run_ids: list[str] = field(default_factory=list)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class CampaignRuntimeRecovery:
|
|
"""Durable Campaign work reconciled during process startup."""
|
|
|
|
interrupted_runs: int = 0
|
|
resumed_campaigns: int = 0
|
|
|
|
|
|
def _spawned_indices(campaign: Campaign, session: Optional[Session] = None) -> set[int]:
|
|
"""Return plan entries completed by durable identities or legacy progress.
|
|
|
|
Once any durable occurrence exists for an entry, identities replace the
|
|
legacy summary for that entry. This repairs summaries written by older
|
|
runners after only part of a multi-occurrence entry was claimed.
|
|
"""
|
|
legacy = set(campaign.summary.scheduler.spawned_indices) if campaign.summary else set()
|
|
if not campaign.id or session is None:
|
|
return legacy
|
|
|
|
occurrences: dict[int, set[int]] = {}
|
|
for run in RunRepository(session).list_by_campaign(campaign.id):
|
|
if run.campaign_plan_index is None or run.campaign_occurrence_index is None:
|
|
continue
|
|
occurrences.setdefault(run.campaign_plan_index, set()).add(run.campaign_occurrence_index)
|
|
|
|
return {
|
|
plan_index
|
|
for plan_index, entry in enumerate(campaign.plan)
|
|
if (plan_index in occurrences and set(range(entry.count)).issubset(occurrences[plan_index]))
|
|
or (plan_index not in occurrences and plan_index in legacy)
|
|
}
|
|
|
|
|
|
async def _execute_child_run(
|
|
campaign: Campaign,
|
|
scenario: Scenario,
|
|
run: EvalRun,
|
|
session: Session,
|
|
) -> None:
|
|
target = TargetRepository(session).get(campaign.target_id)
|
|
if not target:
|
|
raise ValueError(f"campaign target missing: {campaign.target_id}")
|
|
engine = EvalEngine(target=target, scenario=scenario, triggered_by=RunTrigger.CAMPAIGN)
|
|
await engine.run(existing_run=run)
|
|
|
|
|
|
async def _spawn_child_run(
|
|
campaign: Campaign,
|
|
scenario_id: str,
|
|
*,
|
|
plan_index: int,
|
|
occurrence_index: int,
|
|
session: Session,
|
|
execute_child_run: ChildRunExecutor = _execute_child_run,
|
|
) -> Optional[str]:
|
|
"""Claim one occurrence and drive its pending Run through the engine.
|
|
|
|
A rejected claim means the Campaign stopped being runnable between the
|
|
scheduler decision and this write; callers must stop spawning. An existing
|
|
pending Run is resumed, while an existing running/completed/failed Run is
|
|
returned without replaying external messages.
|
|
"""
|
|
scenario = ScenarioRepository(session).get(scenario_id)
|
|
if not scenario:
|
|
raise ValueError(f"campaign scenario missing: {scenario_id}")
|
|
|
|
claim = RunRepository(session).claim_campaign_run(
|
|
campaign_id=campaign.id or "",
|
|
scenario_id=scenario_id,
|
|
scenario_version=scenario.version or 1,
|
|
plan_index=plan_index,
|
|
occurrence_index=occurrence_index,
|
|
)
|
|
if claim.status in (CampaignRunClaimStatus.NOT_FOUND, CampaignRunClaimStatus.CONFLICT):
|
|
return None
|
|
run = claim.run
|
|
if run is None:
|
|
raise RuntimeError(f"Campaign child claim returned no Run: {claim.status.value}")
|
|
|
|
if run.status is RunStatus.PENDING:
|
|
await execute_child_run(campaign, scenario, run, session)
|
|
return run.id or ""
|
|
|
|
|
|
async def _advance_campaign(
|
|
*,
|
|
campaign_id: str,
|
|
elapsed_seconds: float,
|
|
session: Session,
|
|
cancel_event: Optional[asyncio.Event] = None,
|
|
execute_child_run: ChildRunExecutor = _execute_child_run,
|
|
) -> Optional[AdvanceResult]:
|
|
"""Advance the campaign clock to ``elapsed_seconds`` and spawn due Runs.
|
|
|
|
``elapsed_seconds`` is real wall-clock time since the window started; it is
|
|
mapped to a window offset via ``time_scale``. A due entry whose spawning
|
|
raises is recorded but remains due until every occurrence has a durable
|
|
child-Run identity. Returns ``None`` if the campaign does not exist.
|
|
"""
|
|
repo = CampaignRepository(session)
|
|
campaign = repo.get(campaign_id)
|
|
if not campaign:
|
|
return None
|
|
|
|
offset = clock_offset(elapsed_seconds=elapsed_seconds, time_scale=campaign.time_scale)
|
|
spawned = _spawned_indices(campaign, session)
|
|
decision = decide_schedule(
|
|
plan=campaign.plan,
|
|
window_seconds=campaign.window_seconds,
|
|
clock_offset_seconds=offset,
|
|
spawned_indices=spawned,
|
|
)
|
|
|
|
result = AdvanceResult(finished=decision.finished)
|
|
errors: list[dict] = list(campaign.summary.scheduler.errors) if campaign.summary else []
|
|
for due in decision.due:
|
|
# Cancellation stops further spawning; runs already in flight finish.
|
|
if cancel_event is not None and cancel_event.is_set():
|
|
break
|
|
claim_rejected = False
|
|
try:
|
|
for occurrence_index in range(due.entry.count):
|
|
run_id = await _spawn_child_run(
|
|
campaign,
|
|
due.entry.scenario_id,
|
|
plan_index=due.index,
|
|
occurrence_index=occurrence_index,
|
|
session=session,
|
|
execute_child_run=execute_child_run,
|
|
)
|
|
if run_id is None:
|
|
claim_rejected = True
|
|
break
|
|
result.spawned_run_ids.append(run_id)
|
|
except Exception as exc: # one entry failing must not wedge the campaign
|
|
errors.append({"entry_index": due.index, "error": str(exc)})
|
|
_logger.warning("活动 %s 计划条目 %d 派生失败（跳过）: %s", campaign_id, due.index, exc)
|
|
if claim_rejected:
|
|
break
|
|
spawned = _spawned_indices(campaign, session)
|
|
# Persist progress per entry: a failure partway through a multi-entry
|
|
# advance must never lose which entries already spawned, since restart
|
|
# recovery reads this back from the DB. Mutate the existing summary so
|
|
# any unknown top-level keys survive the read-modify-write.
|
|
summary = campaign.summary or CampaignSummary()
|
|
summary.scheduler = SchedulerState(spawned_indices=sorted(spawned), errors=errors)
|
|
campaign.summary = summary
|
|
repo.save_scheduler_state(campaign_id, summary)
|
|
|
|
return result
|
|
|
|
|
|
def _auto_start_analysis(campaign: Campaign, session: Session) -> None:
|
|
"""Enqueue 智能分析 when a 正式线 campaign completes.
|
|
|
|
Accelerated campaigns, an unresolvable analysis model, and any enqueue
|
|
failure all skip silently — the analysis is an enhancement and must never
|
|
block or break campaign completion.
|
|
"""
|
|
if campaign.time_scale != 1 or not campaign.id:
|
|
return
|
|
try:
|
|
if resolve_analysis_model(campaign, session) is None:
|
|
return
|
|
enqueue_campaign_analysis(campaign.id, triggered_by="auto")
|
|
except Exception as exc:
|
|
_logger.warning("活动 %s 自动分析触发失败（已跳过）: %s", campaign.id, exc)
|
|
|
|
|
|
async def _reconcile_campaign_child_runs(
|
|
campaign_id: str,
|
|
session: Session,
|
|
*,
|
|
cancel_event: Optional[asyncio.Event] = None,
|
|
execute_child_run: ChildRunExecutor = _execute_child_run,
|
|
) -> Optional[CampaignRecoveryResult]:
|
|
"""Resume safe pending claims and fail running Runs without replaying.
|
|
|
|
A pending claim has not crossed the Engine's send seam and can be resumed.
|
|
A running Run may already have reached the target, so recovery records an
|
|
explicit interruption instead of executing it again.
|
|
"""
|
|
campaign = CampaignRepository(session).get(campaign_id)
|
|
if campaign is None:
|
|
return None
|
|
if campaign.status is not CampaignStatus.RUNNING:
|
|
return CampaignRecoveryResult()
|
|
|
|
run_repo = RunRepository(session)
|
|
result = CampaignRecoveryResult(failed_run_ids=run_repo.mark_campaign_running_interrupted(campaign_id))
|
|
for run in run_repo.list_pending_campaign_children(campaign_id):
|
|
if cancel_event is not None and cancel_event.is_set():
|
|
break
|
|
plan_index = run.campaign_plan_index
|
|
occurrence_index = run.campaign_occurrence_index
|
|
valid_identity = (
|
|
plan_index is not None
|
|
and occurrence_index is not None
|
|
and 0 <= plan_index < len(campaign.plan)
|
|
and 0 <= occurrence_index < campaign.plan[plan_index].count
|
|
and run.scenario_id == campaign.plan[plan_index].scenario_id
|
|
)
|
|
if not valid_identity:
|
|
failed = run_repo.mark_pending_campaign_child_interrupted(run.id or "")
|
|
if failed and failed.id:
|
|
result.failed_run_ids.append(failed.id)
|
|
continue
|
|
|
|
try:
|
|
resumed_id = await _spawn_child_run(
|
|
campaign,
|
|
run.scenario_id,
|
|
plan_index=plan_index,
|
|
occurrence_index=occurrence_index,
|
|
session=session,
|
|
execute_child_run=execute_child_run,
|
|
)
|
|
if resumed_id is None:
|
|
break
|
|
result.resumed_run_ids.append(resumed_id)
|
|
except Exception as exc:
|
|
fresh = run_repo.get(run.id or "")
|
|
if fresh is not None and fresh.status is RunStatus.PENDING:
|
|
failed = run_repo.mark_pending_campaign_child_interrupted(run.id or "")
|
|
if failed and failed.id:
|
|
result.failed_run_ids.append(failed.id)
|
|
elif fresh is not None and fresh.status is RunStatus.FAILED and fresh.id:
|
|
result.failed_run_ids.append(fresh.id)
|
|
_logger.warning("活动 %s 子运行 %s 恢复失败: %s", campaign_id, run.id, exc)
|
|
|
|
return result
|
|
|
|
|
|
# ── durable scheduler loop ──────────────────────────────────────────────────
|
|
|
|
|
|
async def _run_campaign_loop(
|
|
campaign_id: str,
|
|
cancel: asyncio.Event,
|
|
*,
|
|
session_factory: Callable[[], Session],
|
|
now: Callable[[], datetime],
|
|
tick_seconds: float,
|
|
execute_child_run: ChildRunExecutor,
|
|
) -> None:
|
|
"""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
|
|
campaign, asks the pure ``decide_tick`` what to do, and only performs I/O.
|
|
It exits when the window finishes, the campaign leaves RUNNING (e.g. it was
|
|
cancelled), or the cooperative cancel event fires.
|
|
"""
|
|
session = session_factory()
|
|
try:
|
|
await _reconcile_campaign_child_runs(
|
|
campaign_id,
|
|
session,
|
|
cancel_event=cancel,
|
|
execute_child_run=execute_child_run,
|
|
)
|
|
while not cancel.is_set():
|
|
repo = CampaignRepository(session)
|
|
campaign = repo.get(campaign_id)
|
|
if not campaign:
|
|
return
|
|
|
|
current_time = now()
|
|
decision = decide_tick(
|
|
now=current_time,
|
|
started_at=campaign.started_at,
|
|
status=campaign.status,
|
|
window_seconds=campaign.window_seconds,
|
|
time_scale=campaign.time_scale,
|
|
plan=campaign.plan,
|
|
spawned_indices=_spawned_indices(campaign, session),
|
|
)
|
|
if decision.action is TickAction.STOP:
|
|
return
|
|
|
|
elapsed = elapsed_seconds(now=current_time, started_at=campaign.started_at)
|
|
await _advance_campaign(
|
|
campaign_id=campaign_id,
|
|
elapsed_seconds=elapsed,
|
|
session=session,
|
|
cancel_event=cancel,
|
|
execute_child_run=execute_child_run,
|
|
)
|
|
|
|
if decision.action is TickAction.COMPLETE:
|
|
current = repo.get(campaign_id)
|
|
# The lifecycle CAS is the authoritative cancel-race guard.
|
|
if current and resolve_finalize(current.status) is TickAction.COMPLETE:
|
|
try:
|
|
completed = complete_campaign(campaign_id=campaign_id, session=session)
|
|
except Exception as exc:
|
|
_logger.warning("活动 %s 完成迁移失败: %s", campaign_id, exc)
|
|
else:
|
|
_auto_start_analysis(completed, session)
|
|
return
|
|
|
|
try:
|
|
await asyncio.wait_for(cancel.wait(), timeout=tick_seconds)
|
|
except asyncio.TimeoutError:
|
|
pass
|
|
except Exception as exc: # a crashing loop must not take down the app
|
|
_logger.warning("活动调度循环 %s 异常退出: %s", campaign_id, exc)
|
|
finally:
|
|
session.close()
|
|
|
|
|
|
class CampaignRuntime:
|
|
"""Own durable Campaign execution behind one lifecycle-level interface."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
session_factory: Callable[[], Session] = get_session,
|
|
now: Callable[[], datetime] = utc_now,
|
|
tick_seconds: float = DEFAULT_TICK_SECONDS,
|
|
execute_child_run: ChildRunExecutor = _execute_child_run,
|
|
) -> None:
|
|
self._session_factory = session_factory
|
|
self._now = now
|
|
self._tick_seconds = tick_seconds
|
|
self._execute_child_run = execute_child_run
|
|
self._registry = TaskRegistry()
|
|
|
|
def start(self, campaign_id: str) -> bool:
|
|
"""Start or resume one durable Campaign, idempotently."""
|
|
session = self._session_factory()
|
|
try:
|
|
started = start_campaign_lifecycle(session, campaign_id)
|
|
finally:
|
|
session.close()
|
|
if started is None:
|
|
return False
|
|
self._launch(campaign_id)
|
|
return True
|
|
|
|
def cancel(self, campaign_id: str) -> None:
|
|
"""Signal the process-local loop after durable cancellation commits."""
|
|
self._registry.cancel(campaign_id)
|
|
|
|
def recover(self) -> CampaignRuntimeRecovery:
|
|
"""Repair interrupted Runs and relaunch every running Campaign."""
|
|
session = self._session_factory()
|
|
try:
|
|
interrupted_runs = RunRepository(session).mark_orphans_failed()
|
|
running_ids = [
|
|
campaign.id
|
|
for campaign in CampaignRepository(session).list_all()
|
|
if campaign.status is CampaignStatus.RUNNING and campaign.id
|
|
]
|
|
finally:
|
|
session.close()
|
|
for campaign_id in running_ids:
|
|
self._launch(campaign_id)
|
|
return CampaignRuntimeRecovery(
|
|
interrupted_runs=interrupted_runs,
|
|
resumed_campaigns=len(running_ids),
|
|
)
|
|
|
|
async def shutdown(self) -> None:
|
|
"""Gracefully stop all process-local Campaign loops."""
|
|
await self._registry.shutdown_all()
|
|
|
|
def _launch(self, campaign_id: str) -> None:
|
|
self._registry.launch(
|
|
campaign_id,
|
|
lambda cancel: _run_campaign_loop(
|
|
campaign_id,
|
|
cancel,
|
|
session_factory=self._session_factory,
|
|
now=self._now,
|
|
tick_seconds=self._tick_seconds,
|
|
execute_child_run=self._execute_child_run,
|
|
),
|
|
)
|
|
|
|
|
|
campaign_runtime = CampaignRuntime()
|