feat: add persisted batch text testing
This commit is contained in:
Xin Wang
@@ -0,0 +1,410 @@
|
||||
"""In-process MVP scheduler for persisted batch test runs."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from loguru import logger
|
||||
from sqlalchemy import select
|
||||
|
||||
from db.models import ModelResource, TestCase, TestRun, TestRunCase
|
||||
from db.session import SessionLocal
|
||||
from models import RuntimeModelResource
|
||||
from services.config_resolver import resolve_runtime_config
|
||||
from services.test_runs.errors import TestExecutionError
|
||||
from services.test_runs.evaluator import EvaluationEngine
|
||||
from services.test_runs.text_runner import TextPipelineRunner
|
||||
from test_schemas import TestCaseDefinition
|
||||
|
||||
|
||||
@dataclass
|
||||
class _RunControl:
|
||||
stop_event: asyncio.Event = field(default_factory=asyncio.Event)
|
||||
reason: str | None = None
|
||||
case_tasks: set[asyncio.Task[None]] = field(default_factory=set)
|
||||
|
||||
|
||||
class TestRunOrchestrator:
|
||||
"""Run cases concurrently while keeping PostgreSQL as the source of truth."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._run_tasks: dict[str, asyncio.Task[None]] = {}
|
||||
self._controls: dict[str, _RunControl] = {}
|
||||
|
||||
def start(self, run_id: str) -> None:
|
||||
current = self._run_tasks.get(run_id)
|
||||
if current is not None and not current.done():
|
||||
return
|
||||
control = _RunControl()
|
||||
self._controls[run_id] = control
|
||||
task = asyncio.create_task(self._execute_run(run_id, control), name=f"test-run-{run_id}")
|
||||
self._run_tasks[run_id] = task
|
||||
task.add_done_callback(lambda _task: self._forget(run_id))
|
||||
|
||||
async def cancel(self, run_id: str) -> None:
|
||||
control = self._controls.get(run_id)
|
||||
if control is None:
|
||||
return
|
||||
control.reason = "manual"
|
||||
control.stop_event.set()
|
||||
for task in tuple(control.case_tasks):
|
||||
task.cancel()
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
tasks = [task for task in self._run_tasks.values() if not task.done()]
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
if tasks:
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
self._run_tasks.clear()
|
||||
self._controls.clear()
|
||||
|
||||
def _forget(self, run_id: str) -> None:
|
||||
self._run_tasks.pop(run_id, None)
|
||||
self._controls.pop(run_id, None)
|
||||
|
||||
async def _execute_run(self, run_id: str, control: _RunControl) -> None:
|
||||
try:
|
||||
async with SessionLocal() as session:
|
||||
run = await session.get(TestRun, run_id)
|
||||
if run is None:
|
||||
return
|
||||
run.status = "running"
|
||||
run.started_at = run.started_at or datetime.now(UTC)
|
||||
await session.commit()
|
||||
try:
|
||||
cfg = await resolve_runtime_config(session, str(run.assistant_id or ""))
|
||||
except Exception as exc:
|
||||
await self._fail_run_setup(run_id, str(exc))
|
||||
return
|
||||
config = dict(run.config or {})
|
||||
evaluator_resource_id = str(
|
||||
config.get("evaluatorModelResourceId") or ""
|
||||
)
|
||||
evaluator_row = await session.get(
|
||||
ModelResource,
|
||||
evaluator_resource_id,
|
||||
)
|
||||
if (
|
||||
evaluator_row is None
|
||||
or evaluator_row.capability != "LLM"
|
||||
or not evaluator_row.enabled
|
||||
):
|
||||
await self._fail_run_setup(
|
||||
run_id,
|
||||
"评估模型不存在、已停用或不是 LLM 资源",
|
||||
)
|
||||
return
|
||||
evaluator_resource = RuntimeModelResource(
|
||||
id=evaluator_row.id,
|
||||
name=evaluator_row.name,
|
||||
capability=evaluator_row.capability,
|
||||
interface_type=evaluator_row.interface_type,
|
||||
values=evaluator_row.values or {},
|
||||
secrets=evaluator_row.secrets or {},
|
||||
support_image_input=bool(evaluator_row.support_image_input),
|
||||
)
|
||||
rows = (
|
||||
await session.execute(
|
||||
select(TestRunCase)
|
||||
.where(TestRunCase.run_id == run_id)
|
||||
.order_by(TestRunCase.position)
|
||||
)
|
||||
).scalars().all()
|
||||
|
||||
semaphore = asyncio.Semaphore(int(config.get("concurrency") or 3))
|
||||
tasks = {
|
||||
asyncio.create_task(
|
||||
self._execute_case(
|
||||
run_id,
|
||||
row.id,
|
||||
cfg,
|
||||
evaluator_resource,
|
||||
config,
|
||||
semaphore,
|
||||
control,
|
||||
),
|
||||
name=f"test-run-case-{row.id}",
|
||||
)
|
||||
for row in rows
|
||||
}
|
||||
control.case_tasks = tasks
|
||||
if tasks:
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
await self._finish_run(run_id, control)
|
||||
except asyncio.CancelledError:
|
||||
await self._interrupt_run(run_id)
|
||||
raise
|
||||
except Exception as exc: # noqa: BLE001 - background task must persist failure
|
||||
logger.exception(f"批量测试运行失败: run_id={run_id}: {exc}")
|
||||
await self._fail_run_setup(run_id, str(exc))
|
||||
|
||||
async def _execute_case(
|
||||
self,
|
||||
run_id: str,
|
||||
run_case_id: str,
|
||||
cfg,
|
||||
evaluator_resource: RuntimeModelResource,
|
||||
config: dict,
|
||||
semaphore: asyncio.Semaphore,
|
||||
control: _RunControl,
|
||||
) -> None:
|
||||
async with semaphore:
|
||||
if control.stop_event.is_set():
|
||||
await self._mark_skipped(run_case_id)
|
||||
return
|
||||
|
||||
async with SessionLocal() as session:
|
||||
row = await session.get(TestRunCase, run_case_id)
|
||||
if row is None:
|
||||
return
|
||||
snapshot = dict(row.case_snapshot or {})
|
||||
definition = TestCaseDefinition.model_validate(snapshot)
|
||||
max_attempts = row.max_attempts
|
||||
|
||||
for attempt in range(1, max_attempts + 1):
|
||||
if control.stop_event.is_set():
|
||||
await self._mark_skipped(run_case_id)
|
||||
return
|
||||
await self._mark_running(run_case_id, attempt)
|
||||
try:
|
||||
timeout_seconds = int(
|
||||
config.get("timeoutSecs")
|
||||
or config.get("timeout_secs")
|
||||
or 60
|
||||
)
|
||||
async with asyncio.timeout(timeout_seconds):
|
||||
execution = await TextPipelineRunner().run(
|
||||
cfg.model_copy(deep=True),
|
||||
definition,
|
||||
timeout_seconds=timeout_seconds,
|
||||
)
|
||||
evaluated = await EvaluationEngine(evaluator_resource).evaluate(
|
||||
definition,
|
||||
execution,
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
await self._mark_skipped(run_case_id)
|
||||
raise
|
||||
except TimeoutError:
|
||||
error = TestExecutionError(
|
||||
code="CASE_TIMEOUT",
|
||||
message=f"测试用例执行与评估超过 {timeout_seconds} 秒",
|
||||
stage="pipeline",
|
||||
retryable=True,
|
||||
)
|
||||
if attempt < max_attempts and not control.stop_event.is_set():
|
||||
continue
|
||||
await self._mark_error(run_case_id, error)
|
||||
if str(
|
||||
config.get("errorStrategy")
|
||||
or config.get("error_strategy")
|
||||
) == "stop_on_error":
|
||||
await self._request_stop(run_id, control, "execution_error")
|
||||
return
|
||||
except TestExecutionError as exc:
|
||||
if exc.retryable and attempt < max_attempts and not control.stop_event.is_set():
|
||||
continue
|
||||
await self._mark_error(run_case_id, exc)
|
||||
if str(config.get("errorStrategy") or config.get("error_strategy")) == "stop_on_error":
|
||||
await self._request_stop(run_id, control, "execution_error")
|
||||
return
|
||||
except Exception as exc: # noqa: BLE001 - normalize unknown runtime errors
|
||||
error = TestExecutionError(
|
||||
code="UNEXPECTED_EXECUTION_ERROR",
|
||||
message=str(exc) or type(exc).__name__,
|
||||
stage="pipeline",
|
||||
retryable=False,
|
||||
)
|
||||
await self._mark_error(run_case_id, error)
|
||||
if str(config.get("errorStrategy") or config.get("error_strategy")) == "stop_on_error":
|
||||
await self._request_stop(run_id, control, "execution_error")
|
||||
return
|
||||
|
||||
status = "pass" if evaluated.passed else "fail"
|
||||
await self._mark_evaluated(run_case_id, status, evaluated)
|
||||
if status == "fail" and str(
|
||||
config.get("failureStrategy") or config.get("failure_strategy")
|
||||
) == "stop_on_fail":
|
||||
await self._request_stop(run_id, control, "assertion_failure")
|
||||
return
|
||||
|
||||
async def _request_stop(
|
||||
self,
|
||||
run_id: str,
|
||||
control: _RunControl,
|
||||
reason: str,
|
||||
) -> None:
|
||||
if control.reason is None:
|
||||
control.reason = reason
|
||||
control.stop_event.set()
|
||||
async with SessionLocal() as session:
|
||||
run = await session.get(TestRun, run_id)
|
||||
if run is not None and run.stop_reason is None:
|
||||
run.stop_reason = reason
|
||||
await session.commit()
|
||||
|
||||
async def _mark_running(self, run_case_id: str, attempt: int) -> None:
|
||||
async with SessionLocal() as session:
|
||||
row = await session.get(TestRunCase, run_case_id)
|
||||
if row is None:
|
||||
return
|
||||
row.status = "running"
|
||||
row.attempt_count = attempt
|
||||
row.started_at = row.started_at or datetime.now(UTC)
|
||||
row.finished_at = None
|
||||
row.execution_error = None
|
||||
await session.commit()
|
||||
|
||||
async def _mark_evaluated(self, run_case_id: str, status: str, evaluated) -> None:
|
||||
async with SessionLocal() as session:
|
||||
row = await session.get(TestRunCase, run_case_id)
|
||||
if row is None:
|
||||
return
|
||||
row.status = status
|
||||
row.result = {
|
||||
"turns": [
|
||||
item.model_dump(mode="json", by_alias=True)
|
||||
for item in evaluated.turns
|
||||
],
|
||||
"overallCriteria": [
|
||||
item.model_dump(mode="json", by_alias=True)
|
||||
for item in evaluated.overall_criteria
|
||||
],
|
||||
}
|
||||
row.execution_error = None
|
||||
row.finished_at = datetime.now(UTC)
|
||||
if row.test_case_id:
|
||||
test_case = await session.get(TestCase, row.test_case_id)
|
||||
if test_case is not None:
|
||||
test_case.last_result = status
|
||||
await session.commit()
|
||||
|
||||
async def _mark_error(
|
||||
self,
|
||||
run_case_id: str,
|
||||
error: TestExecutionError,
|
||||
) -> None:
|
||||
async with SessionLocal() as session:
|
||||
row = await session.get(TestRunCase, run_case_id)
|
||||
if row is None:
|
||||
return
|
||||
row.status = "error"
|
||||
row.execution_error = error.as_dict()
|
||||
row.finished_at = datetime.now(UTC)
|
||||
if row.test_case_id:
|
||||
test_case = await session.get(TestCase, row.test_case_id)
|
||||
if test_case is not None:
|
||||
test_case.last_result = "fail"
|
||||
await session.commit()
|
||||
|
||||
async def _mark_skipped(self, run_case_id: str) -> None:
|
||||
async with SessionLocal() as session:
|
||||
row = await session.get(TestRunCase, run_case_id)
|
||||
if row is None or row.status in {"pass", "fail", "error", "skipped"}:
|
||||
return
|
||||
row.status = "skipped"
|
||||
row.finished_at = datetime.now(UTC)
|
||||
await session.commit()
|
||||
|
||||
async def _finish_run(self, run_id: str, control: _RunControl) -> None:
|
||||
async with SessionLocal() as session:
|
||||
run = await session.get(TestRun, run_id)
|
||||
if run is None:
|
||||
return
|
||||
waiting = (
|
||||
await session.execute(
|
||||
select(TestRunCase).where(
|
||||
TestRunCase.run_id == run_id,
|
||||
TestRunCase.status.in_(["waiting", "running"]),
|
||||
)
|
||||
)
|
||||
).scalars().all()
|
||||
for row in waiting:
|
||||
row.status = "skipped"
|
||||
row.finished_at = datetime.now(UTC)
|
||||
manual = bool(run.cancel_requested or control.reason == "manual")
|
||||
run.status = "cancelled" if manual else "completed"
|
||||
run.stop_reason = "manual" if manual else control.reason or run.stop_reason
|
||||
run.finished_at = datetime.now(UTC)
|
||||
await session.commit()
|
||||
|
||||
async def _fail_run_setup(self, run_id: str, message: str) -> None:
|
||||
error = TestExecutionError(
|
||||
code="RUN_SETUP_FAILED",
|
||||
message=message or "批量测试初始化失败",
|
||||
stage="pipeline",
|
||||
retryable=False,
|
||||
)
|
||||
async with SessionLocal() as session:
|
||||
run = await session.get(TestRun, run_id)
|
||||
if run is None:
|
||||
return
|
||||
rows = (
|
||||
await session.execute(
|
||||
select(TestRunCase).where(
|
||||
TestRunCase.run_id == run_id,
|
||||
TestRunCase.status.in_(["waiting", "running"]),
|
||||
)
|
||||
)
|
||||
).scalars().all()
|
||||
now = datetime.now(UTC)
|
||||
for row in rows:
|
||||
row.status = "error"
|
||||
row.execution_error = error.as_dict()
|
||||
row.finished_at = now
|
||||
run.status = "completed"
|
||||
run.stop_reason = "execution_error"
|
||||
run.finished_at = now
|
||||
await session.commit()
|
||||
|
||||
async def _interrupt_run(self, run_id: str) -> None:
|
||||
await _mark_run_interrupted(run_id, "后端服务停止，运行已中断")
|
||||
|
||||
|
||||
async def _mark_run_interrupted(run_id: str, message: str) -> None:
|
||||
async with SessionLocal() as session:
|
||||
run = await session.get(TestRun, run_id)
|
||||
if run is None or run.status not in {"queued", "running"}:
|
||||
return
|
||||
rows = (
|
||||
await session.execute(
|
||||
select(TestRunCase).where(
|
||||
TestRunCase.run_id == run_id,
|
||||
TestRunCase.status.in_(["waiting", "running"]),
|
||||
)
|
||||
)
|
||||
).scalars().all()
|
||||
now = datetime.now(UTC)
|
||||
for row in rows:
|
||||
if row.status == "running":
|
||||
row.status = "error"
|
||||
row.execution_error = {
|
||||
"code": "SERVER_RESTARTED",
|
||||
"message": message,
|
||||
"stage": "pipeline",
|
||||
"retryable": True,
|
||||
}
|
||||
else:
|
||||
row.status = "skipped"
|
||||
row.finished_at = now
|
||||
run.status = "completed"
|
||||
run.stop_reason = "execution_error"
|
||||
run.finished_at = now
|
||||
await session.commit()
|
||||
|
||||
|
||||
async def recover_interrupted_test_runs() -> None:
|
||||
async with SessionLocal() as session:
|
||||
run_ids = (
|
||||
await session.execute(
|
||||
select(TestRun.id).where(TestRun.status.in_(["queued", "running"]))
|
||||
)
|
||||
).scalars().all()
|
||||
for run_id in run_ids:
|
||||
await _mark_run_interrupted(str(run_id), "后端服务重启，运行已中断")
|
||||
|
||||
|
||||
test_run_orchestrator = TestRunOrchestrator()
|
||||
Reference in New Issue
Block a user