Publish curated ITWorx-ModelForge source from 990401737f96c8d46a9f7b0d84994516cadcb6a0

backend/tests/test_serving.py

@@ -9,7 +9,7 @@ from unittest.mock import Mock
 
 import pytest
 from fastapi.testclient import TestClient
-from sqlalchemy import create_engine, select, update
+from sqlalchemy import create_engine, event, select, update
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import Session
 from sqlalchemy.pool import StaticPool
@@ -19,6 +19,7 @@ from modelforge_api.domain.runtime import AgentRuntimeProbeComplete
 from modelforge_api.domain.serving import (
     AgentResidentState,
     AgentServingJobComplete,
+    AgentServingJobFailure,
     AgentServingStateReport,
     CapabilityExperimentCreate,
     CapabilityPromotionCreate,
@@ -41,6 +42,7 @@ from modelforge_api.persistence.models import (
     ComputeNode,
     EmbeddingSpace,
     GatewayRequest,
+    SchedulerAcceleratorState,
     ServiceCredential,
     ServingGpuLease,
     ServingJob,
@@ -543,7 +545,111 @@ def test_scheduler_accounts_external_usage_reserve_and_workstation_rejection(
     assert "NODE_NOT_PRODUCTION_ELIGIBLE" in rejected["reasons"]
 
 
+@pytest.mark.parametrize(
+    "pressure_boundary",
+    [False, True],
+    ids=["first-read", "recovery-stable-window"],
+)
+def test_scheduler_get_projects_expiry_and_pressure_without_dml_or_commit(
+    session: Session,
+    pressure_boundary: bool,
+) -> None:
+    base, candidate, _node = setup_m5(session)
+    _approval, deployment = approve_and_promote(base, candidate.id)
+    baseline = base.scheduler_overview(commit=False, update_pressure=False)[0]
+
+    now = datetime.now(UTC)
+    since = now - timedelta(seconds=base.settings.scheduler_pressure_stable_seconds + 1)
+    pressure = (
+        SchedulerAcceleratorState(
+            accelerator_id=deployment.accelerator_id,
+            pressure_state="HIGH",
+            recovery_candidate="NORMAL",
+            recovery_candidate_since=since,
+            pressure_changed_at=since - timedelta(minutes=1),
+            last_observed_at=since,
+        )
+        if pressure_boundary
+        else None
+    )
+    expired_lease = ServingGpuLease(
+        accelerator_id=deployment.accelerator_id,
+        capability_deployment_id=deployment.id,
+        request_id=uuid.uuid4(),
+        reserved_vram_bytes=4 * 1024**3,
+        priority="production",
+        state="active",
+        owner="scheduler-read-only-test",
+        expires_at=now - timedelta(milliseconds=1),
+    )
+    if pressure is not None:
+        session.add(pressure)
+    session.add(expired_lease)
+    session.commit()
+
+    dml_statements: list[str] = []
+    commits: list[bool] = []
+
+    def record_statement(_conn, _cursor, statement, _parameters, _context, _many) -> None:
+        operation = statement.lstrip().split(maxsplit=1)[0].upper()
+        if operation in {"INSERT", "UPDATE", "DELETE", "REPLACE", "MERGE", "TRUNCATE"}:
+            dml_statements.append(statement)
+
+    def record_commit(_session: Session) -> None:
+        commits.append(True)
+
+    engine = session.get_bind()
+    event.listen(engine, "before_cursor_execute", record_statement)
+    event.listen(session, "after_commit", record_commit)
+    request_settings = Settings(
+        _env_file=None,
+        operator_api_key="scheduler-read-only-test-key",
+    )
+    app.dependency_overrides[get_settings] = lambda: request_settings
+    app.dependency_overrides[get_serving_service] = lambda: base
+    try:
+        response = TestClient(app).get(
+            "/api/v1/scheduler",
+            headers={"X-ModelForge-Admin-Token": "scheduler-read-only-test-key"},
+        )
+    finally:
+        app.dependency_overrides.clear()
+        event.remove(engine, "before_cursor_execute", record_statement)
+        event.remove(session, "after_commit", record_commit)
+
+    assert response.status_code == 200, response.text
+    budget = response.json()[0]
+    assert budget["pressure_state"] == "NORMAL"
+    assert budget["leased_vram_bytes"] == baseline.leased_vram_bytes
+    assert dml_statements == []
+    assert commits == []
+
+    persisted_lease = session.get(ServingGpuLease, expired_lease.id)
+    pressure_query = select(SchedulerAcceleratorState).where(
+        SchedulerAcceleratorState.accelerator_id == deployment.accelerator_id
+    )
+    persisted_pressure = session.scalar(pressure_query)
+    if pressure_boundary:
+        assert persisted_pressure is not None
+        assert persisted_pressure.pressure_state == "HIGH"
+        assert persisted_pressure.recovery_candidate == "NORMAL"
+        assert persisted_pressure.recovery_candidate_since is not None
+        assert persisted_pressure.recovery_candidate_since.replace(tzinfo=UTC) == since
+    else:
+        assert persisted_pressure is None
+    assert persisted_lease is not None and persisted_lease.state == "active"
+
+    # The default service path still stabilizes and commits pressure for write/admission callers.
+    base.scheduler_overview()
+    persisted_pressure = session.scalar(pressure_query)
+    assert persisted_pressure is not None
+    assert persisted_pressure.pressure_state == "NORMAL"
+    assert persisted_pressure.recovery_candidate is None
+
+
 class InlineWorkerServingService(ServingService):
+    synthetic_worker_metrics: dict[str, object]
+
     def _wait_job(self, job_id: uuid.UUID, deadline: float):
         del deadline
         job = self.repo.job(job_id)
@@ -577,7 +683,10 @@ class InlineWorkerServingService(ServingService):
                 "normalized": True,
                 "prompt_tokens": len(lease.input) * 4,
             }
-            metrics = {"inference_time_ms": 25.0}
+            metrics = {
+                "inference_time_ms": 25.0,
+                **getattr(self, "synthetic_worker_metrics", {}),
+            }
         else:
             assert lease.rerank_query is not None
             assert lease.rerank_documents is not None
@@ -667,6 +776,67 @@ def test_real_gateway_path_uses_scheduler_cold_load_then_warm_reuse(session: Ses
     assert history[0].latency_breakdown["worker_total"] >= 0
 
 
+def test_runtime_failure_fails_gateway_and_releases_request_resources(session: Session) -> None:
+    base, candidate, _node = setup_m5(session)
+    _approval, deployment = approve_and_promote(base, candidate.id)
+
+    class FailingRuntimeService(InlineWorkerServingService):
+        def _wait_job(self, job_id: uuid.UUID, deadline: float):
+            job = self.repo.job(job_id)
+            assert job is not None
+            if job.operation != "invoke":
+                return super()._wait_job(job_id, deadline)
+            node = self.repo.node(job.compute_node_id)
+            assert node is not None
+            lease = self.claim_next(node)
+            assert lease is not None and lease.job_id == job_id
+            self.fail_job(
+                job_id,
+                node,
+                AgentServingJobFailure(
+                    lease_token=lease.lease_token,
+                    worker_instance_id="failed-runtime-worker",
+                    failure_code="RUNTIME_UNAVAILABLE",
+                    failure_message="synthetic worker transport failure",
+                ),
+            )
+            return ServingService._wait_job(self, job_id, deadline)
+
+    service = FailingRuntimeService(session, base.settings, base.manifests, base.payload_store)
+    created = service.create_client(
+        ServiceClientCreate(name="runtime-failure-client", allowed_capabilities=["rag.embedding@1"])
+    )
+    app.dependency_overrides[get_settings] = lambda: service.settings
+    app.dependency_overrides[get_serving_service] = lambda: service
+    try:
+        response = TestClient(app).post(
+            "/api/v1/capabilities/rag.embedding@1/invoke",
+            headers={"Authorization": f"Bearer {created.credential}"},
+            json={"input": "must not be inferred elsewhere"},
+        )
+    finally:
+        app.dependency_overrides.clear()
+
+    assert response.status_code == 503
+    assert response.json()["error"]["code"] == "RUNTIME_UNAVAILABLE"
+    requests = list(session.scalars(select(GatewayRequest)))
+    assert len(requests) == 1
+    assert requests[0].status == "failed"
+    assert requests[0].failure_code == "RUNTIME_UNAVAILABLE"
+    assert requests[0].capability_deployment_id == deployment.id
+    jobs = list(session.scalars(select(ServingJob)))
+    assert [job.operation for job in jobs] == ["load", "invoke"]
+    assert jobs[-1].status == "failed" and jobs[-1].failure_code == "RUNTIME_UNAVAILABLE"
+    assert jobs[-1].payload_reference is None
+    assert service.payload_store.values == {}
+    leases = list(session.scalars(select(ServingGpuLease)))
+    assert len(leases) == 1
+    assert leases[0].state == "failed" and leases[0].released_at is not None
+    residency = service.repo.residency(deployment.id)
+    assert residency is not None and residency.active_requests == 0
+    assert residency.state == "cold"
+
+
 def test_lab_experiment_has_opaque_route_own_space_and_no_production_cutover(
     session: Session,
 ) -> None: