295 lines
11 KiB
Python
295 lines
11 KiB
Python
"""LLM 后台任务生命周期直测（架构保养第二轮候选 1）。
|
||
|
||
分析 / 对比 / judge 复核三条 LLM 任务链收进各自的 TaskRegistry：
|
||
强引用防 GC、按 id 幂等、shutdown 统一收敛；启动清扫把滞留的
|
||
generating 行标记为 failed。
|
||
"""
|
||
|
||
import asyncio
|
||
|
||
import pytest
|
||
from agenteval.evaluation import intelligence_jobs
|
||
from agenteval.exploration import judge
|
||
from agenteval.storage.repository import (
|
||
AsyncJobClaimStatus,
|
||
CampaignAnalysisRepository,
|
||
CampaignPeriodComparisonRepository,
|
||
)
|
||
from sqlmodel import Session, SQLModel, create_engine
|
||
|
||
|
||
@pytest.fixture()
|
||
def db_session(tmp_path):
|
||
from agenteval.storage.db import ( # noqa: F401
|
||
CampaignAnalysisDB,
|
||
CampaignDB,
|
||
CampaignPeriodComparisonDB,
|
||
)
|
||
|
||
engine = create_engine(
|
||
f"sqlite:///{tmp_path / 'llmtasks.db'}",
|
||
connect_args={"check_same_thread": False},
|
||
)
|
||
SQLModel.metadata.create_all(engine)
|
||
session = Session(engine)
|
||
try:
|
||
yield session
|
||
finally:
|
||
session.close()
|
||
engine.dispose()
|
||
|
||
|
||
async def test_start_analysis_registers_in_registry(db_session, monkeypatch):
|
||
monkeypatch.setattr(intelligence_jobs, "execute_campaign_analysis_job", lambda *a, **kw: asyncio.sleep(0))
|
||
intelligence_jobs.enqueue_campaign_analysis("c-1", triggered_by="manual", session=db_session)
|
||
assert intelligence_jobs.is_intelligence_job_running("analysis", "c-1")
|
||
await intelligence_jobs.shutdown_campaign_intelligence_jobs()
|
||
assert not intelligence_jobs.is_intelligence_job_running("analysis", "c-1")
|
||
|
||
|
||
async def test_start_comparison_registers_in_registry(db_session, monkeypatch):
|
||
monkeypatch.setattr(intelligence_jobs, "execute_campaign_comparison_job", lambda *a, **kw: asyncio.sleep(0))
|
||
intelligence_jobs.enqueue_campaign_comparison(
|
||
"c-2",
|
||
triggered_by="manual",
|
||
baseline_campaign_id="c-1",
|
||
session=db_session,
|
||
)
|
||
assert intelligence_jobs.is_intelligence_job_running("comparison", "c-2")
|
||
await intelligence_jobs.shutdown_campaign_intelligence_jobs()
|
||
|
||
|
||
async def test_start_judge_registers_in_registry(monkeypatch):
|
||
monkeypatch.setattr(judge, "execute_judge_review", lambda *a, **kw: asyncio.sleep(0))
|
||
judge.start_judge_review("s-1")
|
||
assert judge.judge_registry.is_running("s-1")
|
||
await judge.judge_registry.shutdown_all()
|
||
|
||
|
||
async def test_launch_is_idempotent_for_live_id(db_session, monkeypatch):
|
||
gate = asyncio.Event()
|
||
|
||
async def hang(*args, **kwargs):
|
||
await gate.wait()
|
||
|
||
monkeypatch.setattr(intelligence_jobs, "execute_campaign_analysis_job", hang)
|
||
intelligence_jobs.enqueue_campaign_analysis("c-dup", triggered_by="manual", session=db_session)
|
||
intelligence_jobs.enqueue_campaign_analysis("c-dup", triggered_by="manual", session=db_session)
|
||
assert intelligence_jobs.is_intelligence_job_running("analysis", "c-dup")
|
||
gate.set()
|
||
await intelligence_jobs.shutdown_campaign_intelligence_jobs()
|
||
|
||
|
||
async def test_shutdown_all_cancels_hanging_task(db_session, monkeypatch):
|
||
async def hang(*args, **kwargs):
|
||
await asyncio.Event().wait()
|
||
|
||
monkeypatch.setattr(intelligence_jobs, "execute_campaign_comparison_job", hang)
|
||
intelligence_jobs.enqueue_campaign_comparison(
|
||
"c-hang",
|
||
triggered_by="manual",
|
||
baseline_campaign_id="c-base",
|
||
session=db_session,
|
||
)
|
||
assert intelligence_jobs.is_intelligence_job_running("comparison", "c-hang")
|
||
await intelligence_jobs.shutdown_campaign_intelligence_jobs()
|
||
assert not intelligence_jobs.is_intelligence_job_running("comparison", "c-hang")
|
||
|
||
|
||
def test_mark_orphans_failed_flips_generating_analysis(db_session):
|
||
repo = CampaignAnalysisRepository(db_session)
|
||
repo.upsert("c-gen", status="generating", triggered_by="auto")
|
||
repo.upsert("c-done", status="completed", result={"ok": True}, triggered_by="auto")
|
||
|
||
count = repo.mark_orphans_failed()
|
||
|
||
assert count == 1
|
||
assert repo.get_by_campaign("c-gen").status == "failed"
|
||
assert repo.get_by_campaign("c-gen").error
|
||
assert repo.get_by_campaign("c-done").status == "completed"
|
||
|
||
|
||
def test_enqueue_analysis_persists_before_launch(db_session, monkeypatch):
|
||
launched = []
|
||
monkeypatch.setattr(
|
||
intelligence_jobs,
|
||
"_launch_analysis",
|
||
lambda cid, *, triggered_by: launched.append((cid, triggered_by)),
|
||
)
|
||
|
||
intelligence_jobs.enqueue_campaign_analysis("c-queued", triggered_by="auto", session=db_session)
|
||
|
||
assert launched == [("c-queued", "auto")]
|
||
row = CampaignAnalysisRepository(db_session).get_by_campaign("c-queued")
|
||
assert row.status == "queued"
|
||
assert row.triggered_by == "auto"
|
||
|
||
|
||
def test_only_one_session_can_claim_a_queued_analysis(db_session):
|
||
repo = CampaignAnalysisRepository(db_session)
|
||
repo.enqueue("c-claim", triggered_by="auto")
|
||
|
||
competing_session = Session(db_session.get_bind())
|
||
try:
|
||
assert repo.claim_queued("c-claim").status is AsyncJobClaimStatus.CLAIMED
|
||
assert (
|
||
CampaignAnalysisRepository(competing_session).claim_queued("c-claim").status
|
||
is AsyncJobClaimStatus.ALREADY_CLAIMED
|
||
)
|
||
finally:
|
||
competing_session.close()
|
||
|
||
row = repo.get_by_campaign("c-claim")
|
||
assert row.status == "generating"
|
||
assert row.triggered_by == "auto"
|
||
|
||
|
||
def test_recovery_relaunches_persisted_analysis_jobs(db_session, monkeypatch):
|
||
repo = CampaignAnalysisRepository(db_session)
|
||
repo.enqueue("c-queued", triggered_by="auto")
|
||
launched = []
|
||
monkeypatch.setattr(
|
||
intelligence_jobs,
|
||
"_launch_analysis",
|
||
lambda cid, *, triggered_by: launched.append((cid, triggered_by)),
|
||
)
|
||
|
||
assert intelligence_jobs.recover_campaign_intelligence_jobs(db_session) == (0, 1)
|
||
assert launched == [("c-queued", "auto")]
|
||
|
||
|
||
def test_mark_orphans_failed_flips_generating_comparison(db_session):
|
||
repo = CampaignPeriodComparisonRepository(db_session)
|
||
repo.upsert("c-gen", status="generating", baseline_campaign_id="b-1", triggered_by="auto")
|
||
repo.upsert("c-done", status="completed", baseline_campaign_id="b-2", result={"ok": True}, triggered_by="auto")
|
||
|
||
count = repo.mark_orphans_failed()
|
||
|
||
assert count == 1
|
||
assert repo.get_by_campaign("c-gen").status == "failed"
|
||
assert repo.get_by_campaign("c-gen").error
|
||
assert repo.get_by_campaign("c-done").status == "completed"
|
||
|
||
|
||
def test_enqueue_comparison_preserves_baseline_before_claim(db_session):
|
||
repo = CampaignPeriodComparisonRepository(db_session)
|
||
|
||
row = repo.enqueue(
|
||
"c-current",
|
||
baseline_campaign_id="c-baseline",
|
||
triggered_by="auto",
|
||
)
|
||
|
||
assert row.status == "queued"
|
||
assert row.baseline_campaign_id == "c-baseline"
|
||
assert row.triggered_by == "auto"
|
||
assert repo.claim_queued("c-current").status is AsyncJobClaimStatus.CLAIMED
|
||
assert repo.claim_queued("c-current").status is AsyncJobClaimStatus.ALREADY_CLAIMED
|
||
|
||
|
||
def test_enqueue_comparison_persists_before_launch(db_session, monkeypatch):
|
||
launched = []
|
||
monkeypatch.setattr(
|
||
intelligence_jobs,
|
||
"_launch_comparison",
|
||
lambda cid, *, triggered_by, baseline_campaign_id: launched.append(
|
||
(cid, triggered_by, baseline_campaign_id)
|
||
),
|
||
)
|
||
|
||
intelligence_jobs.enqueue_campaign_comparison(
|
||
"c-current",
|
||
baseline_campaign_id="c-baseline",
|
||
triggered_by="auto",
|
||
session=db_session,
|
||
)
|
||
|
||
assert launched == [("c-current", "auto", "c-baseline")]
|
||
row = CampaignPeriodComparisonRepository(db_session).get_by_campaign("c-current")
|
||
assert row.status == "queued"
|
||
assert row.baseline_campaign_id == "c-baseline"
|
||
|
||
|
||
def test_recovery_relaunches_persisted_comparison_jobs(db_session, monkeypatch):
|
||
CampaignPeriodComparisonRepository(db_session).enqueue(
|
||
"c-current",
|
||
baseline_campaign_id="c-baseline",
|
||
triggered_by="auto",
|
||
)
|
||
launched = []
|
||
monkeypatch.setattr(
|
||
intelligence_jobs,
|
||
"_launch_comparison",
|
||
lambda cid, *, triggered_by, baseline_campaign_id: launched.append(
|
||
(cid, triggered_by, baseline_campaign_id)
|
||
),
|
||
)
|
||
|
||
assert intelligence_jobs.recover_campaign_intelligence_jobs(db_session) == (0, 1)
|
||
assert launched == [("c-current", "auto", "c-baseline")]
|
||
|
||
|
||
def test_claim_reports_missing_and_terminal_jobs(db_session):
|
||
repo = CampaignAnalysisRepository(db_session)
|
||
|
||
assert repo.claim_queued("missing").status is AsyncJobClaimStatus.NOT_FOUND
|
||
repo.upsert("c-done", status="completed", result={"ok": True})
|
||
assert repo.claim_queued("c-done").status is AsyncJobClaimStatus.NOT_QUEUED
|
||
|
||
|
||
def test_reenqueue_does_not_move_generating_analysis_back_to_queued(db_session, monkeypatch):
|
||
repo = CampaignAnalysisRepository(db_session)
|
||
repo.upsert("c-live", status="generating", triggered_by="auto")
|
||
launched = []
|
||
monkeypatch.setattr(
|
||
intelligence_jobs,
|
||
"_launch_analysis",
|
||
lambda cid, *, triggered_by: launched.append((cid, triggered_by)),
|
||
)
|
||
|
||
intelligence_jobs.enqueue_campaign_analysis("c-live", triggered_by="manual", session=db_session)
|
||
|
||
row = repo.get_by_campaign("c-live")
|
||
assert row.status == "generating"
|
||
assert row.triggered_by == "auto"
|
||
assert launched == []
|
||
|
||
|
||
def test_reenqueue_preserves_queued_comparison_context(db_session):
|
||
repo = CampaignPeriodComparisonRepository(db_session)
|
||
repo.enqueue("c-current", baseline_campaign_id="c-original", triggered_by="auto")
|
||
|
||
row = repo.enqueue("c-current", baseline_campaign_id="c-replacement", triggered_by="manual")
|
||
|
||
assert row.status == "queued"
|
||
assert row.baseline_campaign_id == "c-original"
|
||
assert row.triggered_by == "auto"
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
("repository", "enqueue_kwargs"),
|
||
[
|
||
(CampaignAnalysisRepository, {}),
|
||
(CampaignPeriodComparisonRepository, {"baseline_campaign_id": "c-base"}),
|
||
],
|
||
)
|
||
def test_queued_recovery_is_bounded(db_session, repository, enqueue_kwargs):
|
||
repo = repository(db_session)
|
||
repo.enqueue("c-recover", triggered_by="auto", **enqueue_kwargs)
|
||
|
||
for expected_attempts in range(1, intelligence_jobs.MAX_QUEUED_RECOVERY_ATTEMPTS + 1):
|
||
recovered = repo.prepare_queued_recovery(
|
||
intelligence_jobs.MAX_QUEUED_RECOVERY_ATTEMPTS,
|
||
"recovery exhausted",
|
||
)
|
||
assert [row.campaign_id for row in recovered] == ["c-recover"]
|
||
assert repo.get_by_campaign("c-recover").recovery_attempts == expected_attempts
|
||
|
||
assert repo.prepare_queued_recovery(
|
||
intelligence_jobs.MAX_QUEUED_RECOVERY_ATTEMPTS,
|
||
"recovery exhausted",
|
||
) == []
|
||
exhausted = repo.get_by_campaign("c-recover")
|
||
assert exhausted.status == "failed"
|
||
assert exhausted.error == "recovery exhausted"
|