193 lines
6.6 KiB
Python
193 lines
6.6 KiB
Python
from datetime import UTC, datetime
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock
|
|
from uuid import uuid4
|
|
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
|
|
from modelforge_api.api.authorization import (
|
|
Principal,
|
|
PrincipalRole,
|
|
authenticate_operator_principal,
|
|
)
|
|
from modelforge_api.api.routes.lab import get_benchmark_run_service, router
|
|
from modelforge_api.domain.benchmarking import (
|
|
BenchmarkRunProgress,
|
|
BenchmarkRunResponse,
|
|
BenchmarkSuiteResponse,
|
|
)
|
|
from modelforge_api.domain.contracts import BenchmarkEnvironmentFingerprint, BenchmarkRunContract
|
|
|
|
|
|
def _responses() -> tuple[BenchmarkSuiteResponse, BenchmarkRunResponse]:
|
|
now = datetime.now(UTC)
|
|
suite = BenchmarkSuiteResponse(
|
|
id=uuid4(),
|
|
key="modelforge.rag-embedding.operator-performance",
|
|
revision="1",
|
|
capability="rag.embedding@1",
|
|
dataset_revision="builtin:rag-embedding-performance-v1",
|
|
definition={"schema_version": 1},
|
|
immutable_at=now,
|
|
)
|
|
environment = BenchmarkEnvironmentFingerprint(
|
|
runtime_type="onnx",
|
|
runtime_version="1.0",
|
|
runtime_image_digest=None,
|
|
launch_arguments={},
|
|
cuda_version=None,
|
|
driver_version=None,
|
|
accelerator_name=None,
|
|
accelerator_uuid=None,
|
|
context_length=None,
|
|
concurrency=1,
|
|
seed=None,
|
|
operating_system=None,
|
|
python_version=None,
|
|
)
|
|
run = BenchmarkRunResponse(
|
|
id=uuid4(),
|
|
capability="rag.embedding@1",
|
|
suite_id=suite.id,
|
|
capability_deployment_id=uuid4(),
|
|
status="accepted",
|
|
contract=BenchmarkRunContract(
|
|
deployment_id=str(uuid4()),
|
|
model_revision_sha="a" * 40,
|
|
artifact_sha256="b" * 64,
|
|
suite_key=suite.key,
|
|
suite_revision=suite.revision,
|
|
dataset_revision=suite.dataset_revision,
|
|
environment=environment,
|
|
started_at=now,
|
|
),
|
|
target_snapshot={},
|
|
progress=BenchmarkRunProgress(
|
|
completed_invocations=0, total_invocations=4, active_case=None
|
|
),
|
|
results={"metrics": {"gpu_vram": {"telemetry_status": "UNKNOWN"}}},
|
|
created_at=now,
|
|
)
|
|
return suite, run
|
|
|
|
|
|
def _harness(
|
|
role: PrincipalRole,
|
|
) -> tuple[TestClient, Mock, BenchmarkSuiteResponse, BenchmarkRunResponse]:
|
|
app = FastAPI()
|
|
app.include_router(router)
|
|
service = Mock()
|
|
suite, run = _responses()
|
|
service.suites.return_value = [suite]
|
|
service.list.return_value = [run]
|
|
service.list_page.return_value = SimpleNamespace(runs=[run], next_cursor=None)
|
|
service.start.return_value = run
|
|
service.get.return_value = run
|
|
service.cancel.return_value = run
|
|
app.dependency_overrides[authenticate_operator_principal] = lambda: Principal(
|
|
"operator-one", role, "test"
|
|
)
|
|
app.dependency_overrides[get_benchmark_run_service] = lambda: service
|
|
return TestClient(app), service, suite, run
|
|
|
|
|
|
def _request(suite: BenchmarkSuiteResponse) -> dict[str, object]:
|
|
return {
|
|
"capability": "rag.embedding@1",
|
|
"suite_id": str(suite.id),
|
|
"target": {
|
|
key: str(uuid4())
|
|
for key in (
|
|
"deployment_candidate_id",
|
|
"artifact_set_id",
|
|
"runtime_profile_id",
|
|
"compute_node_id",
|
|
"execution_approval_id",
|
|
)
|
|
},
|
|
"deployment_id": str(uuid4()),
|
|
"experiment_route_id": str(uuid4()),
|
|
}
|
|
|
|
|
|
def test_benchmark_endpoints_forward_operator_identity_and_disable_caching() -> None:
|
|
client, service, suite, run = _harness(PrincipalRole.OPERATOR)
|
|
|
|
suites = client.get("/api/v1/lab/benchmark-suites")
|
|
history = client.get("/api/v1/lab/benchmark-runs")
|
|
started = client.post(
|
|
"/api/v1/lab/benchmark-runs",
|
|
headers={"Idempotency-Key": "operator-benchmark"},
|
|
json=_request(suite),
|
|
)
|
|
status = client.get(f"/api/v1/lab/benchmark-runs/{run.id}")
|
|
cancelled = client.post(f"/api/v1/lab/benchmark-runs/{run.id}/cancel")
|
|
|
|
assert suites.status_code == history.status_code == 200
|
|
assert started.status_code == 202
|
|
assert status.status_code == cancelled.status_code == 200
|
|
assert all(
|
|
response.headers["cache-control"] == "no-store"
|
|
for response in (suites, history, started, status, cancelled)
|
|
)
|
|
assert service.start.call_args.kwargs == {
|
|
"actor_subject": "operator-one",
|
|
"idempotency_key": "operator-benchmark",
|
|
}
|
|
assert service.list_page.call_args.kwargs == {
|
|
"actor_subject": "operator-one",
|
|
"limit": 50,
|
|
"cursor": None,
|
|
}
|
|
assert service.get.call_args.kwargs == {"actor_subject": "operator-one"}
|
|
assert service.cancel.call_args.kwargs == {"actor_subject": "operator-one"}
|
|
|
|
|
|
def test_benchmark_history_cursor_preserves_the_existing_list_response_contract() -> None:
|
|
client, service, _suite, run = _harness(PrincipalRole.OPERATOR)
|
|
service.list_page.return_value = SimpleNamespace(runs=[run], next_cursor="opaque-next-page")
|
|
|
|
response = client.get(
|
|
"/api/v1/lab/benchmark-runs",
|
|
params={"limit": 1, "cursor": "opaque-current-page"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert isinstance(response.json(), list)
|
|
assert [item["id"] for item in response.json()] == [str(run.id)]
|
|
assert response.headers["x-next-cursor"] == "opaque-next-page"
|
|
assert response.headers["access-control-expose-headers"] == "X-Next-Cursor"
|
|
assert service.list_page.call_args.kwargs == {
|
|
"actor_subject": "operator-one",
|
|
"limit": 1,
|
|
"cursor": "opaque-current-page",
|
|
}
|
|
|
|
|
|
def test_benchmark_write_routes_require_operator_and_strict_declared_capability() -> None:
|
|
client, service, suite, _run = _harness(PrincipalRole.VIEWER)
|
|
|
|
denied = client.post(
|
|
"/api/v1/lab/benchmark-runs",
|
|
headers={"Idempotency-Key": "viewer-benchmark"},
|
|
json=_request(suite),
|
|
)
|
|
invalid = _request(suite)
|
|
invalid["capability"] = "rag.rerank@1"
|
|
rejected = _harness(PrincipalRole.OPERATOR)[0].post(
|
|
"/api/v1/lab/benchmark-runs",
|
|
headers={"Idempotency-Key": "bad-capability"},
|
|
json=invalid,
|
|
)
|
|
operator_client, operator_service, operator_suite, _operator_run = _harness(
|
|
PrincipalRole.OPERATOR
|
|
)
|
|
missing_key = operator_client.post("/api/v1/lab/benchmark-runs", json=_request(operator_suite))
|
|
|
|
assert denied.status_code == 403
|
|
assert rejected.status_code == 422
|
|
assert missing_key.status_code == 422
|
|
service.start.assert_not_called()
|
|
operator_service.start.assert_not_called()
|