Publish curated ITWorx-ModelForge source from 990401737f96c8d46a9f7b0d84994516cadcb6a0

backend/tests/test_lab_execution_runs.py

@@ -0,0 +1,585 @@
+import hashlib
+import json
+from collections.abc import Iterator
+from datetime import UTC, datetime, timedelta
+
+import pytest
+from sqlalchemy import create_engine, func, select
+from sqlalchemy.orm import Session
+
+from modelforge_api.domain.lab import LabExecutionStartRequest
+from modelforge_api.domain.serving import (
+    AgentServingJobCancellationAck,
+    AgentServingJobCancellationCheck,
+    AgentServingJobComplete,
+)
+from modelforge_api.persistence.models import (
+    Base,
+    GatewayRequest,
+    LabExecutionRun,
+    ServingJob,
+)
+from modelforge_api.services.lab_runs import LabRunService
+from modelforge_api.services.serving import ServingError
+from tests.test_dispatch_context import operator_fixture
+
+
+@pytest.fixture
+def session() -> Iterator[Session]:
+    engine = create_engine("sqlite+pysqlite:///:memory:")
+    Base.metadata.create_all(engine)
+    with Session(engine) as value:
+        yield value
+    engine.dispose()
+
+
+def make_start_request(request, selected) -> LabExecutionStartRequest:
+    return LabExecutionStartRequest(
+        target=request.target,
+        input=request.input,
+        limits=request.limits,
+        deployment_id=selected.capability_deployment_id,
+        experiment_route_id=selected.id,
+    )
+
+
+def test_same_actor_key_and_fingerprint_returns_one_durable_run(session: Session) -> None:
+    serving, preflight, selected = operator_fixture(session)
+    service = LabRunService(serving)
+    request = make_start_request(preflight, selected)
+
+    first = service.start(request, actor_subject="operator-test", idempotency_key="run-001")
+    retry = service.start(request, actor_subject="operator-test", idempotency_key="run-001")
+
+    assert retry.id == first.id
+    assert retry.status == "accepted"
+    assert session.scalar(select(func.count()).select_from(LabExecutionRun)) == 1
+    assert "private input content" not in first.model_dump_json()
+    assert "payload_reference" not in first.model_dump()
+
+
+def test_reused_actor_key_with_different_input_is_a_conflict(session: Session) -> None:
+    serving, preflight, selected = operator_fixture(session)
+    service = LabRunService(serving)
+    request = make_start_request(preflight, selected)
+    changed = LabExecutionStartRequest.model_validate(
+        {
+            **request.model_dump(mode="json"),
+            "input": {"capability": "rag.embedding", "payload": {"input": "different"}},
+        }
+    )
+    service.start(request, actor_subject="operator-test", idempotency_key="run-002")
+
+    with pytest.raises(ServingError) as raised:
+        service.start(changed, actor_subject="operator-test", idempotency_key="run-002")
+
+    assert raised.value.status_code == 409
+    assert raised.value.code == "LAB_IDEMPOTENCY_CONFLICT"
+    assert session.scalar(select(func.count()).select_from(LabExecutionRun)) == 1
+
+
+def test_run_status_and_cancel_are_scoped_to_authenticated_actor(session: Session) -> None:
+    serving, preflight, selected = operator_fixture(session)
+    service = LabRunService(serving)
+    run = service.start(
+        make_start_request(preflight, selected),
+        actor_subject="operator-owner",
+        idempotency_key="run-003",
+    )
+
+    for action in (
+        lambda: service.get(run.id, actor_subject="operator-other"),
+        lambda: service.cancel(run.id, actor_subject="operator-other"),
+    ):
+        with pytest.raises(ServingError) as raised:
+            action()
+        assert raised.value.status_code == 404
+        assert raised.value.code == "LAB_RUN_NOT_FOUND"
+
+    assert service.get(run.id, actor_subject="operator-owner").id == run.id
+    cancelled = service.cancel(run.id, actor_subject="operator-owner")
+    assert cancelled.status == "cancelled"
+    assert service.cancel(run.id, actor_subject="operator-owner").status == "cancelled"
+
+
+def _claimed_run(session: Session, actor: str, key: str):
+    serving, preflight, selected = operator_fixture(session)
+    run_service = LabRunService(serving)
+    run_response = run_service.start(
+        make_start_request(preflight, selected), actor_subject=actor, idempotency_key=key
+    )
+    run = session.get(LabExecutionRun, run_response.id)
+    assert run is not None
+    gateway = session.get(GatewayRequest, run.gateway_request_id)
+    assert gateway is not None
+    token = "t" * 32
+    job = ServingJob(
+        capability_deployment_id=selected.capability_deployment_id,
+        compute_node_id=selected.deployment.compute_node_id,
+        gateway_request_id=gateway.id,
+        lab_execution_run_id=run.id,
+        operation="invoke",
+        status="running",
+        priority="lab",
+        idempotency_key=f"test-{key}",
+        payload_reference=run.payload_reference,
+        attempt_count=1,
+        lease_token_hash=hashlib.sha256(token.encode()).hexdigest(),
+        lease_expires_at=datetime.now(UTC) + timedelta(minutes=2),
+        result_summary={},
+        started_at=datetime.now(UTC),
+    )
+    session.add(job)
+    run.status = "running"
+    run.started_at = datetime.now(UTC)
+    run.heartbeat_at = datetime.now(UTC)
+    gateway.status = "running"
+    session.commit()
+    return serving, run_service, run, job, selected, token
+
+
+def _queued_run_with_job(session: Session, actor: str, key: str):
+    serving, preflight, selected = operator_fixture(session)
+    run_service = LabRunService(serving)
+    response = run_service.start(
+        make_start_request(preflight, selected), actor_subject=actor, idempotency_key=key
+    )
+    run = session.get(LabExecutionRun, response.id)
+    assert run is not None
+    gateway = session.get(GatewayRequest, run.gateway_request_id)
+    assert gateway is not None
+    job = ServingJob(
+        capability_deployment_id=selected.capability_deployment_id,
+        compute_node_id=selected.deployment.compute_node_id,
+        gateway_request_id=gateway.id,
+        lab_execution_run_id=run.id,
+        operation="invoke",
+        status="queued",
+        priority="lab",
+        idempotency_key=f"queued-{key}",
+        payload_reference=run.payload_reference,
+        attempt_count=0,
+        result_summary={},
+    )
+    run.status = "queued"
+    gateway.status = "queued"
+    session.add(job)
+    session.commit()
+    return serving, run_service, run, job, selected, preflight
+
+
+def test_cancel_racing_with_worker_completion_discards_vectors_and_acknowledges(
+    session: Session,
+) -> None:
+    serving, run_service, run, job, selected, token = _claimed_run(
+        session, "operator-race", "run-race"
+    )
+    requested = run_service.cancel(run.id, actor_subject="operator-race")
+    assert requested.status == "cancel_requested"
+
+    deployment = serving.repo.deployment(selected.capability_deployment_id)
+    assert deployment is not None
+    node = serving.repo.node(deployment.compute_node_id)
+    space = serving.repo.embedding_space(deployment.embedding_space_id)
+    assert node is not None and space is not None
+    vector = [1.0] + [0.0] * (space.dimension - 1)
+    serving.complete_job(
+        job.id,
+        node,
+        AgentServingJobComplete(
+            lease_token=token,
+            worker_instance_id="worker-cancel-race",
+            state="warm",
+            result={"vectors": [vector], "count": 1, "normalized": True},
+            metrics={"inference_time_ms": 12.0, "private_metric": "must-not-persist"},
+            health={},
+            runtime_facts={"private_runtime_fact": "must-not-persist"},
+        ),
+    )
+
+    session.refresh(run)
+    session.refresh(job)
+    assert run.status == "cancelled"
+    assert run.cancellation_acknowledged_at is not None
+    assert run.output_metadata is None
+    assert job.status == "cancelled"
+    assert job.result_reference is None
+    persisted = f"{run.worker_observations!r} {job.result_summary!r}"
+    assert "private_metric" not in persisted
+    assert "private_runtime_fact" not in persisted
+    assert "must-not-persist" not in persisted
+    assert serving.payload_store.get(run.payload_reference) is None
+
+
+def test_successful_lab_completion_persists_hash_and_allowlisted_metadata_only(
+    session: Session,
+) -> None:
+    serving, _run_service, run, job, selected, token = _claimed_run(
+        session, "operator-success", "run-success"
+    )
+    deployment = serving.repo.deployment(selected.capability_deployment_id)
+    assert deployment is not None
+    node = serving.repo.node(deployment.compute_node_id)
+    space = serving.repo.embedding_space(deployment.embedding_space_id)
+    assert node is not None and space is not None
+    vector = [1.0] + [0.0] * (space.dimension - 1)
+    serving.complete_job(
+        job.id,
+        node,
+        AgentServingJobComplete(
+            lease_token=token,
+            worker_instance_id="worker-success",
+            state="warm",
+            result={"vectors": [vector], "count": 1, "normalized": True},
+            metrics={
+                "inference_time_ms": 12.0,
+                "gpu_uuid": "GPU-test",
+                "private_metric": "must-not-persist",
+            },
+            health={},
+            runtime_facts={"private_runtime_fact": "must-not-persist"},
+        ),
+    )
+
+    session.refresh(run)
+    session.refresh(job)
+    assert run.status == "completed"
+    assert run.output_metadata is not None
+    assert run.output_metadata["dimension"] == space.dimension
+    assert run.output_metadata["item_count"] == 1
+    assert len(run.output_metadata["sha256"]) == 64
+    assert (
+        run.output_metadata["sha256"]
+        == hashlib.sha256(
+            json.dumps([vector], sort_keys=True, separators=(",", ":")).encode()
+        ).hexdigest()
+    )
+    assert run.output_metadata["content_logged"] is False
+    assert run.worker_observations is not None
+    assert run.worker_observations["source"] == "runtime-worker"
+    assert "private_metric" not in str(run.worker_observations)
+    assert "private_runtime_fact" not in str(run.worker_observations)
+    assert job.status == "completed"
+    assert job.result_reference is None
+    assert "vectors" not in str(job.result_summary)
+    gateway = session.get(GatewayRequest, run.gateway_request_id)
+    assert gateway is not None
+    persisted_evidence = gateway.decision_evidence["lab_run_evidence"]
+    assert persisted_evidence["status"] == "completed"
+    assert persisted_evidence["output"] == run.output_metadata
+    assert "private_metric" not in str(gateway.decision_evidence)
+    assert "private_runtime_fact" not in str(gateway.decision_evidence)
+    assert "vectors" not in str(gateway.decision_evidence)
+    assert "private input content" not in str(gateway.decision_evidence)
+    assert serving.payload_store.get(run.payload_reference) is None
+
+
+def test_reconciler_dispatches_exact_cold_target_through_load_and_invoke_jobs(
+    session: Session,
+) -> None:
+    serving, preflight, selected = operator_fixture(session)
+    run_service = LabRunService(serving)
+    created = run_service.start(
+        make_start_request(preflight, selected),
+        actor_subject="operator-reconcile",
+        idempotency_key="run-reconcile",
+    )
+
+    serving.reconcile()
+    run = session.get(LabExecutionRun, created.id)
+    assert run is not None
+    assert run.status == "preparing"
+    load_job = session.scalar(
+        select(ServingJob).where(
+            ServingJob.lab_execution_run_id == run.id,
+            ServingJob.operation == "load",
+        )
+    )
+    assert load_job is not None and load_job.status == "queued"
+    node = serving.repo.node(selected.deployment.compute_node_id)
+    assert node is not None
+    load_lease = serving.claim_next(node)
+    assert load_lease is not None and load_lease.job_id == load_job.id
+    assert load_lease.lab_run_id == run.id
+    serving.complete_job(
+        load_job.id,
+        node,
+        AgentServingJobComplete(
+            lease_token=load_lease.lease_token,
+            worker_instance_id="synthetic-worker",
+            state="warm",
+            result={},
+            metrics={
+                "baseline_vram_bytes": 1_000,
+                "loaded_vram_bytes": 3_000,
+            },
+            health={"model": "healthy"},
+            runtime_facts={"offline_local_only": True},
+        ),
+    )
+
+    serving.reconcile()
+    session.refresh(run)
+    assert run.status == "queued"
+    invoke_job = session.scalar(
+        select(ServingJob).where(
+            ServingJob.lab_execution_run_id == run.id,
+            ServingJob.operation == "invoke",
+        )
+    )
+    assert invoke_job is not None and invoke_job.status == "queued"
+    invoke_lease = serving.claim_next(node)
+    assert invoke_lease is not None and invoke_lease.job_id == invoke_job.id
+    assert invoke_lease.deployment_id == selected.capability_deployment_id
+    assert invoke_lease.lab_run_id == run.id
+    assert invoke_lease.input == ["private input content"]
+    assert invoke_lease.expected_dimension is not None
+    deployment = serving.repo.deployment(selected.capability_deployment_id)
+    assert deployment is not None
+    accelerator = serving.repo.accelerator(deployment.accelerator_id)
+    assert accelerator is not None
+    vector = [1.0] + [0.0] * (invoke_lease.expected_dimension - 1)
+    serving.complete_job(
+        invoke_job.id,
+        node,
+        AgentServingJobComplete(
+            lease_token=invoke_lease.lease_token,
+            worker_instance_id="synthetic-worker",
+            state="warm",
+            result={"vectors": [vector], "count": 1, "normalized": True},
+            metrics={
+                "inference_time_ms": 12.0,
+                "gpu_uuid": accelerator.device_uuid,
+                "used_vram_bytes": 3_000,
+                "process_used_vram_bytes": 2_000,
+                "free_vram_bytes": 7_000,
+                "gpu_utilization_percent": 35,
+                "temperature_c": 60,
+                "power_draw_w": 90,
+                "private_metric": "must-not-persist",
+            },
+            health={"model": "healthy"},
+            runtime_facts={"private_runtime_fact": "must-not-persist"},
+        ),
+    )
+
+    session.refresh(run)
+    session.refresh(invoke_job)
+    assert run.status == "completed"
+    assert run.output_metadata is not None
+    assert run.output_metadata["item_count"] == 1
+    assert len(run.output_metadata["sha256"]) == 64
+    assert run.worker_observations is not None
+    assert run.worker_observations["device_matches_target"] is True
+    assert "private_metric" not in str(run.worker_observations)
+    assert "private_runtime_fact" not in str(run.worker_observations)
+    assert invoke_job.result_reference is None
+    assert "vectors" not in str(invoke_job.result_summary)
+    assert (
+        "private input content"
+        not in run_service.get(run.id, actor_subject="operator-reconcile").model_dump_json()
+    )
+
+
+def test_expired_claim_is_interrupted_and_never_reclaimed(session: Session) -> None:
+    serving, _run_service, run, job, selected, _token = _claimed_run(
+        session, "operator-stale", "run-stale"
+    )
+    job.lease_expires_at = datetime.now(UTC) - timedelta(seconds=1)
+    session.commit()
+
+    serving.reconcile()
+
+    session.refresh(run)
+    session.refresh(job)
+    assert run.status == "interrupted"
+    assert run.failure_code == "LAB_WORKER_LEASE_STALE"
+    assert job.status == "failed"
+    assert job.failure_code == "LAB_WORKER_LEASE_STALE"
+    node = serving.repo.node(selected.deployment.compute_node_id)
+    assert node is not None
+    assert serving.claim_next(node) is None
+
+
+@pytest.mark.parametrize("expiry_mode", ["redis_key_missing", "durable_expiry_elapsed"])
+def test_expired_or_missing_redis_input_fails_closed_without_dispatch(
+    session: Session, expiry_mode: str
+) -> None:
+    serving, preflight, selected = operator_fixture(session)
+    run_service = LabRunService(serving)
+    created = run_service.start(
+        make_start_request(preflight, selected),
+        actor_subject=f"operator-expired-{expiry_mode}",
+        idempotency_key=f"expired-{expiry_mode}",
+    )
+    run = session.get(LabExecutionRun, created.id)
+    assert run is not None and run.payload_reference is not None
+    if expiry_mode == "redis_key_missing":
+        # Redis TTL expiry is represented by the payload store returning no key.
+        serving.payload_store.delete(run.payload_reference)
+    else:
+        run.input_payload_expires_at = datetime.now(UTC) - timedelta(seconds=1)
+        session.commit()
+
+    serving.reconcile()
+
+    session.refresh(run)
+    gateway = session.get(GatewayRequest, run.gateway_request_id)
+    assert gateway is not None
+    assert run.status == "failed"
+    assert run.failure_code == "REQUEST_PAYLOAD_EXPIRED"
+    assert run.payload_reference is None
+    assert gateway.status == "failed"
+    assert session.scalar(
+        select(func.count()).select_from(ServingJob).where(
+            ServingJob.lab_execution_run_id == run.id
+        )
+    ) == 0
+
+
+def test_claim_with_missing_redis_key_fails_run_instead_of_leaving_it_queued(
+    session: Session,
+) -> None:
+    serving, _run_service, run, job, selected, _preflight = _queued_run_with_job(
+        session, "operator-claim-expired", "claim-expired"
+    )
+    assert run.payload_reference is not None
+    serving.payload_store.delete(run.payload_reference)
+    node = serving.repo.node(selected.deployment.compute_node_id)
+    assert node is not None
+
+    lease = serving.claim_next(node)
+
+    session.refresh(run)
+    session.refresh(job)
+    assert lease is None
+    assert run.status == "failed"
+    assert run.failure_code == "REQUEST_PAYLOAD_EXPIRED"
+    assert job.status == "failed"
+    assert job.failure_code == "REQUEST_PAYLOAD_EXPIRED"
+
+
+def test_lab_claim_rejects_another_runs_live_redis_payload(session: Session) -> None:
+    serving, run_service, run, job, selected, preflight = _queued_run_with_job(
+        session, "operator-payload-isolation", "payload-isolation"
+    )
+    changed_request = LabExecutionStartRequest.model_validate(
+        {
+            **make_start_request(preflight, selected).model_dump(mode="json"),
+            "input": {
+                "capability": "rag.embedding",
+                "payload": {"input": "another run's private input"},
+            },
+        }
+    )
+    other = run_service.start(
+        changed_request,
+        actor_subject="operator-payload-isolation",
+        idempotency_key="payload-isolation-other",
+    )
+    assert other.id != run.id
+    assert other.input_summary.sha256 != run.input_sha256
+    other_run_payload = session.get(LabExecutionRun, other.id)
+    assert other_run_payload is not None and other_run_payload.payload_reference is not None
+    job.payload_reference = other_run_payload.payload_reference
+    session.commit()
+    node = serving.repo.node(selected.deployment.compute_node_id)
+    assert node is not None
+
+    lease = serving.claim_next(node)
+
+    session.refresh(run)
+    session.refresh(job)
+    session.refresh(other_run_payload)
+    assert lease is None
+    assert run.status == "failed"
+    assert run.failure_code == "LAB_PAYLOAD_REFERENCE_MISMATCH"
+    assert job.failure_code == "LAB_PAYLOAD_REFERENCE_MISMATCH"
+    assert other_run_payload.status == "accepted"
+
+
+def test_worker_poll_renews_lease_and_ack_completes_cancel_transition(session: Session) -> None:
+    serving, run_service, run, job, selected, token = _claimed_run(
+        session, "operator-ack", "run-ack"
+    )
+    run.timeout_seconds = 300
+    job.lease_expires_at = datetime.now(UTC) + timedelta(seconds=30)
+    session.commit()
+    node = serving.repo.node(selected.deployment.compute_node_id)
+    assert node is not None
+    before = job.lease_expires_at
+    requested = run_service.cancel(run.id, actor_subject="operator-ack")
+    assert requested.status == "cancel_requested"
+
+    state = serving.check_lab_cancellation(
+        job.id, node, AgentServingJobCancellationCheck(lease_token=token)
+    )
+    assert state.cancel_requested is True
+    session.refresh(job)
+    assert job.lease_expires_at is not None and job.lease_expires_at >= before
+    serving.acknowledge_lab_cancellation(
+        job.id,
+        node,
+        AgentServingJobCancellationAck(
+            lease_token=token,
+            worker_instance_id="worker-ack",
+            state="warm",
+        ),
+    )
+
+    session.refresh(run)
+    session.refresh(job)
+    assert run.status == "cancelled"
+    assert run.cancellation_acknowledged_at is not None
+    assert run.worker_observations == {
+        "source": "runtime-worker",
+        "phase": "after_inference",
+        "availability": "cancelled",
+        "safe_boundary_reached": True,
+        "worker_instance_id": "worker-ack",
+        "worker_state": "warm",
+    }
+    assert job.status == "cancelled"
+    assert job.result_reference is None
+
+
+def test_queued_run_cancel_prevents_worker_claim_and_deletes_ttl_input(session: Session) -> None:
+    serving, preflight, selected = operator_fixture(session)
+    run_service = LabRunService(serving)
+    response = run_service.start(
+        make_start_request(preflight, selected),
+        actor_subject="operator-queued-cancel",
+        idempotency_key="queued-cancel",
+    )
+    run = session.get(LabExecutionRun, response.id)
+    assert run is not None
+    payload_reference = run.payload_reference
+    gateway = session.get(GatewayRequest, run.gateway_request_id)
+    assert gateway is not None
+    job = ServingJob(
+        capability_deployment_id=selected.capability_deployment_id,
+        compute_node_id=selected.deployment.compute_node_id,
+        gateway_request_id=gateway.id,
+        lab_execution_run_id=run.id,
+        operation="invoke",
+        status="queued",
+        priority="lab",
+        idempotency_key="queued-cancel-job",
+        payload_reference=run.payload_reference,
+        attempt_count=0,
+        result_summary={},
+    )
+    session.add(job)
+    run.status = "queued"
+    gateway.status = "queued"
+    session.commit()
+
+    cancelled = run_service.cancel(run.id, actor_subject="operator-queued-cancel")
+
+    session.refresh(job)
+    assert cancelled.status == "cancelled"
+    assert job.status == "cancelled"
+    assert job.lease_token_hash is None
+    assert "input_payload_expires_at" in cancelled.model_dump()
+    assert payload_reference is not None
+    assert serving.payload_store.get(payload_reference) is None