Publish curated ITWorx-ModelForge source from 990401737f96c8d46a9f7b0d84994516cadcb6a0

backend/tests/test_serving_external_budget.py

@@ -0,0 +1,89 @@
+"""Native admission honors external reservations until their explicit release."""
+from datetime import UTC, datetime, timedelta
+from uuid import uuid4
+
+import pytest
+from sqlalchemy import create_engine, event
+from sqlalchemy.orm import Session
+
+from modelforge_api.persistence.models import (
+    Base,
+    ExternalCreativeAllocation,
+    ProjectBinding,
+    ServiceClient,
+)
+from modelforge_api.services.manifest_registry import ManifestRegistry
+from modelforge_api.services.serving import ServingService
+from modelforge_api.settings import Settings
+from tests.test_selected_comparison_postgres import _valid_selection
+from tests.test_serving import MemoryPayloadStore
+
+
+@pytest.fixture
+def budget_fixture():
+    engine = create_engine("sqlite://")
+    event.listen(engine, "connect", lambda connection, _: connection.execute("PRAGMA foreign_keys=ON"))
+    Base.metadata.create_all(engine)
+    with Session(engine) as session:
+        selection, client_id = _valid_selection(session)
+        service = ServingService(session, Settings(_env_file=None, env="test"),
+                                 ManifestRegistry(), MemoryPayloadStore())
+        deployment = service.repo.deployment(selection.current_deployment_id)
+        client = session.get(ServiceClient, client_id)
+        assert deployment is not None and client is not None
+        binding = session.get(ProjectBinding, client.project_binding_id)
+        assert binding is not None
+        yield session, service, deployment, (client.id, binding.project_id)
+    engine.dispose()
+
+
+def allocation(deployment, client, *, state, required):
+    return ExternalCreativeAllocation(
+        service_client_id=client[0], project_id=client[1],
+        accelerator_id=deployment.accelerator_id, request_id=str(uuid4()),
+        idempotency_key=str(uuid4()), request_sha256="a" * 64, workload="image",
+        resource_ref="fixture", provider_resource_id="fixture", required_vram_bytes=required,
+        priority="background", ttl_seconds=30, state=state, decision_code="FIXTURE",
+        decision_evidence={}, required_actions=[],
+        expires_at=datetime.now(UTC) - timedelta(seconds=1),
+    )
+
+
+@pytest.mark.parametrize(("state", "reserved"), [
+    ("granted", True), ("uncertain", True), ("queued", False), ("released", False),
+])
+def test_actual_native_budget_retains_expired_grants_and_uncertain_reservations(
+    budget_fixture, state, reserved,
+):
+    session, service, deployment, client = budget_fixture
+    before = service._budget(deployment)
+    required = 512 * 1024**2
+    session.add(allocation(deployment, client, state=state, required=required))
+    session.flush()
+    after = service._budget(deployment)
+    delta = required if reserved else 0
+    assert after.leased_vram_bytes == before.leased_vram_bytes + delta
+    assert after.schedulable_free_vram_bytes == before.schedulable_free_vram_bytes - delta
+    assert after.invariant_delta_bytes == before.invariant_delta_bytes == 0
+
+
+def test_native_placement_cannot_spend_external_reservation(budget_fixture):
+    session, service, deployment, client = budget_fixture
+    before = service._budget(deployment)
+    session.add(allocation(deployment, client, state="uncertain",
+                           required=before.schedulable_free_vram_bytes - 1))
+    session.flush()
+    plan = service._placement_plan(deployment, priority="production", dry_run=True, commit=False)
+    assert plan.headroom_before_bytes == 1
+    assert plan.verdict not in {"ADMIT", "ADMIT_AFTER_EVICTION"}
+
+
+def test_overview_does_not_commit_owning_admission_transaction(budget_fixture):
+    session, service, _deployment, _client = budget_fixture
+    commits = []
+    event.listen(session, "after_commit", lambda _: commits.append(True))
+    assert service.scheduler_overview(commit=False)
+    assert commits == []
+    assert session.in_transaction()
+    service.scheduler_overview()
+    assert commits == [True]