586 lines
22 KiB
Python
586 lines
22 KiB
Python
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
|