411 lines
16 KiB
Python
411 lines
16 KiB
Python
"""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()
|