90 lines
3.8 KiB
Python
90 lines
3.8 KiB
Python
"""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]
|