from __future__ import annotations from datetime import UTC, datetime, timedelta from uuid import UUID, uuid4 import pytest from sqlalchemy import create_engine, event, func, select, update from sqlalchemy.orm import Session from modelforge_api.domain.benchmarking import BenchmarkRunStartRequest from modelforge_api.persistence.models import ( Base, BenchmarkRun, BenchmarkSuite, GatewayRequest, LabExecutionRun, RuntimeProfile, ) from modelforge_api.services.benchmark_runs import BenchmarkRunService from modelforge_api.services.lab_runs import LabRunService from modelforge_api.services.serving import ServingError, ServingService from tests.test_dispatch_context import operator_fixture @pytest.fixture def session(): engine = create_engine("sqlite://") Base.metadata.create_all(engine) with Session(engine) as value: yield value engine.dispose() def _request(serving, target_request, selected) -> BenchmarkRunStartRequest: suite = BenchmarkRunService(serving).suites()[0] return BenchmarkRunStartRequest( capability="rag.embedding@1", suite_id=suite.id, target=target_request.target, deployment_id=selected.capability_deployment_id, experiment_route_id=selected.id, ) def _run(serving, target_request, selected, *, actor="operator-one", key="benchmark-once"): request = _request(serving, target_request, selected) return BenchmarkRunService(serving).start(request, actor_subject=actor, idempotency_key=key) def test_suite_listing_is_read_only_and_advertises_declared_contract(session: Session) -> None: serving, _target, _selected = operator_fixture(session) writes: list[str] = [] def capture(_connection, _cursor, statement, _parameters, _context, _executemany): if statement.lstrip().split(maxsplit=1)[0].upper() in {"INSERT", "UPDATE", "DELETE"}: writes.append(statement) event.listen(session.get_bind(), "before_cursor_execute", capture) suites = BenchmarkRunService(serving).suites() assert len(suites) == 1 assert suites[0].capability == "rag.embedding@1" assert suites[0].immutable_at.year == 2026 assert not writes assert session.scalar(select(func.count()).select_from(BenchmarkSuite)) == 0 def test_operator_run_is_idempotent_actor_scoped_and_links_real_lab_gateway( session: Session, ) -> None: serving, target_request, selected = operator_fixture(session) request = _request(serving, target_request, selected) service = BenchmarkRunService(serving) first = service.start(request, actor_subject="operator-one", idempotency_key="same-run") replay = service.start(request, actor_subject="operator-one", idempotency_key="same-run") assert first.id == replay.id assert first.status == "running" assert first.progress.active_case == "request-1" assert first.contract.deployment_id == str(selected.capability_deployment_id) assert first.contract.model_revision_sha assert first.contract.artifact_sha256 assert first.target_snapshot["deployment_id"] == str(selected.capability_deployment_id) assert session.scalar(select(func.count()).select_from(BenchmarkRun)) == 1 child = session.scalar( select(LabExecutionRun).where(LabExecutionRun.actor_subject == "operator-one") ) assert child is not None gateway = session.get(GatewayRequest, child.gateway_request_id) assert gateway is not None assert gateway.capability_deployment_id == selected.capability_deployment_id assert gateway.status == "queued" assert gateway.service_client_id is None assert service.list(actor_subject="someone-else") == [] with pytest.raises(ServingError) as raised: service.get(first.id, actor_subject="someone-else") assert raised.value.status_code == 404 changed = request.model_copy(update={"experiment_route_id": uuid4()}) with pytest.raises(ServingError) as raised: service.start(changed, actor_subject="operator-one", idempotency_key="same-run") assert raised.value.code == "BENCHMARK_IDEMPOTENCY_CONFLICT" def test_history_keyset_pages_tied_timestamps_without_skipping_when_new_run_arrives( session: Session, ) -> None: serving, target_request, selected = operator_fixture(session) service = BenchmarkRunService(serving) original = [ _run( serving, target_request, selected, key=f"history-page-{index}", ) for index in range(5) ] _run(serving, target_request, selected, actor="operator-two", key="other-actor-run") tied_created_at = datetime(2026, 9, 24, 12, tzinfo=UTC) session.execute( update(BenchmarkRun) .where(BenchmarkRun.actor_subject == "operator-one") .values(created_at=tied_created_at) ) session.commit() expected = sorted((run.id for run in original), key=lambda run_id: run_id.int, reverse=True) first_page = service.list_page(actor_subject="operator-one", limit=2) assert [run.id for run in first_page.runs] == expected[:2] assert first_page.next_cursor is not None inserted = _run( serving, target_request, selected, key="history-page-arrived-after-first-page", ) session.execute( update(BenchmarkRun) .where(BenchmarkRun.id == inserted.id) .values(created_at=tied_created_at + timedelta(days=1)) ) session.commit() second_page = service.list_page( actor_subject="operator-one", limit=2, cursor=first_page.next_cursor ) third_page = service.list_page( actor_subject="operator-one", limit=2, cursor=second_page.next_cursor ) assert [run.id for run in second_page.runs] == expected[2:4] assert [run.id for run in third_page.runs] == expected[4:] assert third_page.next_cursor is None assert service.list_page(actor_subject="operator-two", limit=2).runs[0].id != expected[0] with pytest.raises(ServingError) as raised: service.list_page( actor_subject="operator-two", limit=2, cursor=first_page.next_cursor ) assert raised.value.code == "BENCHMARK_CURSOR_INVALID" with pytest.raises(ServingError) as raised: service.list_page(actor_subject="operator-one", limit=2, cursor="") assert raised.value.code == "BENCHMARK_CURSOR_INVALID" def test_target_mismatch_fails_before_persisting_parent(session: Session) -> None: serving, target_request, selected = operator_fixture(session) request = _request(serving, target_request, selected) request = request.model_copy( update={"target": request.target.model_copy(update={"compute_node_id": uuid4()})} ) with pytest.raises(ServingError) as raised: BenchmarkRunService(serving).start( request, actor_subject="operator-one", idempotency_key="bad-target" ) assert raised.value.code == "BENCHMARK_TARGET_MISMATCH" assert session.scalar(select(func.count()).select_from(BenchmarkRun)) == 0 assert session.scalar(select(func.count()).select_from(LabExecutionRun)) == 0 def test_secret_like_launch_arguments_are_not_persisted_in_provenance( session: Session, monkeypatch: pytest.MonkeyPatch ) -> None: serving, target_request, selected = operator_fixture(session) request = _request(serving, target_request, selected) service = BenchmarkRunService(serving) service._suite(request.suite_id) original_snapshot = serving._operator_lab_target_snapshot def snapshot_with_secret(*args, **kwargs): snapshot = original_snapshot(*args, **kwargs) serving.session.execute( update(RuntimeProfile) .where(RuntimeProfile.id == target_request.target.runtime_profile_id) .values(launch_arguments={"environment": ["HF_TOKEN=private-canary"]}) ) serving.session.commit() return snapshot monkeypatch.setattr(serving, "_operator_lab_target_snapshot", snapshot_with_secret) with pytest.raises(ServingError) as raised: service.start(request, actor_subject="operator-one", idempotency_key="secret-arguments") assert raised.value.code == "BENCHMARK_RUNTIME_ARGUMENTS_UNSAFE" assert "private-canary" not in str(raised.value) assert session.scalar(select(func.count()).select_from(BenchmarkRun)) == 0 def test_lost_parent_link_reuses_same_child_idempotency_key(session: Session) -> None: serving, target_request, selected = operator_fixture(session) response = _run(serving, target_request, selected, key="lost-link") parent = session.get(BenchmarkRun, response.id) assert parent is not None execution = {**parent.results["execution"], "active_lab_run_id": None, "active_case": None} parent.results = {**parent.results, "execution": execution} parent.status = "accepted" session.commit() # A process crash after child commit but before parent relink replays the deterministic # actor-scoped LAB key. The child uniqueness constraint returns the same durable run. resumed = BenchmarkRunService(serving) assert resumed.reconcile_one(parent.id) parent = session.get(BenchmarkRun, response.id) assert parent is not None assert UUID(parent.results["execution"]["active_lab_run_id"]) == UUID( response.results["execution"]["active_lab_run_id"] ) assert ( session.scalar( select(func.count()) .select_from(LabExecutionRun) .where( LabExecutionRun.actor_subject == "operator-one", LabExecutionRun.idempotency_key == f"benchmark-{response.id.hex}-00", ) ) == 1 ) def test_vector_throughput_requires_valid_count_and_timing_from_same_attempt() -> None: results = BenchmarkRunService._initial_results(comparable=True) results["execution"]["attempts"] = [ {"status": "completed", "gateway": {"input_count": 4, "inference_time_ms": None}}, {"status": "completed", "gateway": {"input_count": True, "inference_time_ms": 20.0}}, {"status": "completed", "gateway": {"input_count": -2, "inference_time_ms": 10.0}}, {"status": "completed", "gateway": {"input_count": 1, "inference_time_ms": 10.0}}, ] refreshed = BenchmarkRunService._refresh_metrics(results, {}) throughput = refreshed["metrics"]["throughput"] assert throughput["successful_vectors_per_inference_second"] == 100.0 assert throughput["inference_samples"] == 1 assert throughput["measured_vectors"] == 1 assert throughput["successful_vectors_coverage"] == pytest.approx(1 / 5) def test_unmeasured_or_mismatched_worker_resources_remain_unknown() -> None: results = BenchmarkRunService._initial_results(comparable=True) results["execution"]["attempts"] = [ { "worker_observations": { "device_matches_target": True, "metrics": {"temperature_c": 61}, } }, { "worker_observations": { "device_matches_target": False, "metrics": {"used_vram_bytes": 40_000}, } }, ] metrics = BenchmarkRunService._refresh_metrics(results, {"accelerator_uuid": "GPU-0"})[ "metrics" ]["gpu_vram"] assert metrics["telemetry_status"] == "UNKNOWN" assert metrics["vram_status"] == "UNKNOWN" assert metrics["gpu_utilization_status"] == "UNKNOWN" assert metrics["peak_used_vram_bytes"] is None assert metrics["target_accelerator_uuid"] == "GPU-0" def test_cancellation_racing_lab_admission_preserves_terminal_parent_and_stops_child( session: Session, monkeypatch: pytest.MonkeyPatch ) -> None: serving, target_request, selected = operator_fixture(session) request = _request(serving, target_request, selected) original_start = LabRunService.start def start_then_cancel(lab_service, lab_request, *, actor_subject, idempotency_key): child_response = original_start( lab_service, lab_request, actor_subject=actor_subject, idempotency_key=idempotency_key, ) parent = session.scalar( select(BenchmarkRun).where( BenchmarkRun.actor_subject == "operator-one", BenchmarkRun.idempotency_key == "racing-cancel", ) ) assert parent is not None BenchmarkRunService(serving).cancel(parent.id, actor_subject="operator-one") return child_response monkeypatch.setattr(LabRunService, "start", start_then_cancel) # Start commits the parent before admission; the injected cancellation lands after the # child is durable but before the benchmark reconciler can link it. response = BenchmarkRunService(serving).start( request, actor_subject="operator-one", idempotency_key="racing-cancel" ) parent = session.get(BenchmarkRun, response.id) assert parent is not None and parent.status == "cancelled" execution = parent.results["execution"] child = session.get(LabExecutionRun, UUID(execution["active_lab_run_id"])) assert child is not None and child.status == "cancelled" assert response.status == "cancelled" assert response.progress.active_case is None def test_normal_reconciliation_resumes_children_and_aggregates_actual_gateway_fields( session: Session, ) -> None: serving, target_request, selected = operator_fixture(session) response = _run(serving, target_request, selected, key="resume-run") parent_id = response.id baseline = datetime.now(UTC).replace(microsecond=0) for index in range(4): parent = session.get(BenchmarkRun, parent_id) assert parent is not None execution = parent.results["execution"] child_id = UUID(execution["active_lab_run_id"]) child = session.get(LabExecutionRun, child_id) assert child is not None gateway = session.get(GatewayRequest, child.gateway_request_id) assert gateway is not None gateway.created_at = baseline + timedelta(seconds=index * 2) gateway.finished_at = gateway.created_at + timedelta(milliseconds=250) gateway.cold = index == 0 gateway.load_time_ms = 75.0 if index == 0 else None gateway.inference_time_ms = 100.0 if index == 0 else (None if index == 1 else 200.0) gateway.total_latency_ms = 125.0 + index * 10 gateway.input_count = 1 if index != 1 else 4 child.status = "completed" child.failure_code = None child.finished_at = gateway.finished_at child.worker_observations = ( { "device_matches_target": True, "metrics": { "used_vram_bytes": 3_000, "process_used_vram_bytes": 2_500, "free_vram_bytes": 9_000, "gpu_utilization_percent": 30, }, } if index == 0 else {"device_matches_target": False, "metrics": {"used_vram_bytes": 99}} ) session.commit() # Recreate the service object to prove accepted progress is recoverable from rows, # then use the same normal reconciliation entry point as the background tick. resumed = ServingService( session, serving._deployment_settings, serving.manifests, serving.payload_store ) summary = resumed.reconcile() assert "benchmark_runs_driven" in summary parent = session.get(BenchmarkRun, parent_id) assert parent is not None and parent.status == "completed" metrics = parent.results["metrics"] assert metrics["cold_load_latency_ms"] == {"p50": 75.0, "p95": 75.0, "samples": 1} assert metrics["warm_latency_ms"]["samples"] == 3 assert metrics["failures"] == {"attempted": 4, "failed": 0, "rate": 0.0} assert metrics["throughput"]["inference_samples"] == 3 assert metrics["throughput"]["measured_vectors"] == 3 assert metrics["throughput"]["successful_vectors_coverage"] == pytest.approx(3 / 7) assert metrics["gpu_vram"]["telemetry_status"] == "MEASURED" assert metrics["gpu_vram"]["peak_used_vram_bytes"] == 3_000 assert metrics["quality"]["status"] == "UNKNOWN"