303 lines
12 KiB
Python
303 lines
12 KiB
Python
"""Real PostgreSQL admission serialization with deterministic offline capacity."""
|
|
|
|
import importlib
|
|
import threading
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from datetime import UTC, datetime
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock
|
|
from uuid import uuid4
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.orm import Session
|
|
|
|
from modelforge_api.domain.dispatch_context import DispatchContext
|
|
from modelforge_api.domain.external_creative import (
|
|
ExternalAllocationCreate,
|
|
ExternalResourceApproval,
|
|
)
|
|
from modelforge_api.persistence.models import (
|
|
Accelerator,
|
|
ExternalCreativeAllocation,
|
|
GatewayRequest,
|
|
ServiceClient,
|
|
ServingGpuLease,
|
|
)
|
|
from modelforge_api.services import external_creative as external_creative_module
|
|
from modelforge_api.services.external_creative import ExternalCreativeService
|
|
from modelforge_api.services.manifest_registry import ManifestRegistry
|
|
from modelforge_api.services.serving import ServingError, ServingService
|
|
from modelforge_api.settings import Settings
|
|
from tests.test_managed_identities_postgres import IdentityDatabase, identity_database
|
|
from tests.test_selected_comparison_postgres import _valid_selection
|
|
from tests.test_serving import MemoryPayloadStore
|
|
|
|
__all__ = ["identity_database"]
|
|
|
|
|
|
def test_competing_external_admissions_retain_one_capacity_reservation(
|
|
identity_database: IdentityDatabase,
|
|
) -> None:
|
|
database = identity_database
|
|
importlib.import_module("alembic.command").upgrade(database.migration(), "head")
|
|
with Session(database.runtime) as session:
|
|
_request, client_id = _valid_selection(session)
|
|
client = session.get(ServiceClient, client_id)
|
|
client.allowed_capabilities = ["creative.resource@1"]
|
|
session.commit()
|
|
accelerator_id = session.scalar(select(Accelerator.id))
|
|
service = ExternalCreativeService(session, Mock())
|
|
project_id = service._project(client)
|
|
service.upsert_resource(
|
|
ExternalResourceApproval(
|
|
project_id=project_id,
|
|
resource_ref="race-fixture",
|
|
workload="image",
|
|
provider_resource_id="offline-provider",
|
|
configured_vram_bytes=60,
|
|
),
|
|
"qualification",
|
|
)
|
|
|
|
barrier = threading.Barrier(2)
|
|
|
|
def contender(index: int) -> str:
|
|
with Session(database.runtime) as session:
|
|
|
|
def overview(*, commit):
|
|
assert commit is False
|
|
retained = sum(
|
|
session.scalars(
|
|
select(ExternalCreativeAllocation.required_vram_bytes).where(
|
|
ExternalCreativeAllocation.state.in_(["granted", "uncertain"])
|
|
)
|
|
)
|
|
)
|
|
return [
|
|
SimpleNamespace(
|
|
accelerator_id=accelerator_id,
|
|
attribution_confidence="HIGH",
|
|
pressure_state="LOW",
|
|
schedulable_free_vram_bytes=100 - retained,
|
|
policy_revision="offline-fixture",
|
|
observed_at=None,
|
|
)
|
|
]
|
|
|
|
serving = Mock()
|
|
serving.scheduler_overview.side_effect = overview
|
|
service = ExternalCreativeService(session, serving)
|
|
client = session.get(ServiceClient, client_id)
|
|
session.commit()
|
|
barrier.wait(timeout=30)
|
|
return service.acquire(
|
|
ExternalAllocationCreate(
|
|
request_id=f"race-{index}",
|
|
idempotency_key=f"race-{index}",
|
|
installation_id="fixture",
|
|
workload="image",
|
|
resource_ref="race-fixture",
|
|
provider_resource_id="offline-provider",
|
|
required_vram_bytes=60,
|
|
),
|
|
client,
|
|
).state
|
|
|
|
with ThreadPoolExecutor(max_workers=2) as pool:
|
|
assert sorted(pool.map(contender, [1, 2])) == ["granted", "queued"]
|
|
with Session(database.runtime) as session:
|
|
rows = list(session.scalars(select(ExternalCreativeAllocation)))
|
|
assert len(rows) == 2
|
|
assert sum(row.required_vram_bytes for row in rows if row.state == "granted") == 60
|
|
|
|
|
|
def test_native_and_external_admission_share_actual_postgres_capacity(
|
|
identity_database: IdentityDatabase,
|
|
) -> None:
|
|
database = identity_database
|
|
importlib.import_module("alembic.command").upgrade(database.migration(), "head")
|
|
settings = Settings(
|
|
_env_file=None,
|
|
env="test",
|
|
gateway_queue_timeout_seconds=1,
|
|
gateway_request_timeout_seconds=60,
|
|
)
|
|
|
|
def serving(session):
|
|
return ServingService(session, settings, ManifestRegistry(), MemoryPayloadStore())
|
|
|
|
with Session(database.runtime) as session:
|
|
selection, client_id = _valid_selection(session)
|
|
client = session.get(ServiceClient, client_id)
|
|
client.allowed_capabilities = ["rag.embedding@1", "creative.resource@1"]
|
|
session.commit()
|
|
service = serving(session)
|
|
deployment_id = selection.current_deployment_id
|
|
deployment = service.repo.deployment(deployment_id)
|
|
baseline = service._budget(deployment)
|
|
plan = service._placement_plan(
|
|
deployment, priority=client.workload_priority, dry_run=True, persist=False, commit=False
|
|
)
|
|
assert plan.verdict == "ADMIT"
|
|
available, native_required = baseline.schedulable_free_vram_bytes, plan.required_vram_bytes
|
|
assert 0 < native_required < available
|
|
external_required = available - native_required // 2
|
|
assert external_required <= available < external_required + native_required
|
|
external = ExternalCreativeService(session, service)
|
|
external.upsert_resource(
|
|
ExternalResourceApproval(
|
|
project_id=external._project(client),
|
|
resource_ref="mixed-race",
|
|
workload="image",
|
|
provider_resource_id="offline-provider",
|
|
configured_vram_bytes=external_required,
|
|
),
|
|
"qualification",
|
|
)
|
|
gateway = GatewayRequest(
|
|
request_id=uuid4(),
|
|
service_client_id=client.id,
|
|
capability_key="rag.embedding",
|
|
capability_version=1,
|
|
capability_deployment_id=deployment.id,
|
|
compute_node_id=deployment.compute_node_id,
|
|
status="queued",
|
|
priority=client.workload_priority,
|
|
input_sha256="a" * 64,
|
|
input_count=1,
|
|
)
|
|
session.add(gateway)
|
|
session.commit()
|
|
gateway_id = gateway.id
|
|
|
|
barrier = threading.Barrier(2)
|
|
|
|
def contender(native):
|
|
with Session(database.runtime) as session:
|
|
service = serving(session)
|
|
client = session.get(ServiceClient, client_id)
|
|
deployment = service.repo.deployment(deployment_id)
|
|
gateway = session.get(GatewayRequest, gateway_id)
|
|
context = DispatchContext.for_service_client(
|
|
client.id, client.name, client.workload_priority
|
|
)
|
|
session.commit()
|
|
barrier.wait(timeout=30)
|
|
if native:
|
|
try:
|
|
service._acquire_request_lease(deployment, gateway, context)
|
|
return "native_granted"
|
|
except ServingError as error:
|
|
assert error.code in {
|
|
"REQUEST_TIMEOUT",
|
|
"INSUFFICIENT_SCHEDULABLE_VRAM",
|
|
"INSUFFICIENT_CAPACITY",
|
|
}, error.code
|
|
return "native_blocked"
|
|
result = ExternalCreativeService(session, service).acquire(
|
|
ExternalAllocationCreate(
|
|
request_id="mixed-race",
|
|
idempotency_key="mixed-race",
|
|
installation_id="fixture",
|
|
resource_ref="mixed-race",
|
|
workload="image",
|
|
provider_resource_id="offline-provider",
|
|
required_vram_bytes=external_required,
|
|
ttl_seconds=60,
|
|
),
|
|
client,
|
|
)
|
|
return "external_" + result.state
|
|
|
|
with ThreadPoolExecutor(max_workers=2) as pool:
|
|
outcomes = list(pool.map(contender, [True, False]))
|
|
assert outcomes in (
|
|
["native_granted", "external_queued"],
|
|
["native_blocked", "external_granted"],
|
|
)
|
|
with Session(database.runtime) as session:
|
|
native = list(
|
|
session.scalars(
|
|
select(ServingGpuLease).where(ServingGpuLease.state.in_(["active", "granted"]))
|
|
)
|
|
)
|
|
external = list(
|
|
session.scalars(
|
|
select(ExternalCreativeAllocation).where(
|
|
ExternalCreativeAllocation.state.in_(["granted", "uncertain"])
|
|
)
|
|
)
|
|
)
|
|
assert len(native) + len(external) == 1
|
|
assert (
|
|
sum(row.reserved_vram_bytes for row in native)
|
|
+ sum(row.required_vram_bytes for row in external)
|
|
<= available
|
|
)
|
|
|
|
|
|
def test_expired_external_listing_keeps_postgres_count_and_rows_in_sync(
|
|
identity_database: IdentityDatabase, monkeypatch
|
|
) -> None:
|
|
database = identity_database
|
|
importlib.import_module("alembic.command").upgrade(database.migration(), "head")
|
|
with Session(database.runtime) as session:
|
|
_request, client_id = _valid_selection(session)
|
|
client = session.get(ServiceClient, client_id)
|
|
client.allowed_capabilities = ["creative.resource@1"]
|
|
session.commit()
|
|
|
|
accelerator_id = session.scalar(select(Accelerator.id))
|
|
serving = Mock()
|
|
serving.scheduler_overview.return_value = [
|
|
SimpleNamespace(
|
|
accelerator_id=accelerator_id,
|
|
attribution_confidence="HIGH",
|
|
pressure_state="LOW",
|
|
schedulable_free_vram_bytes=100,
|
|
policy_revision="postgres-boundary-fixture",
|
|
observed_at=None,
|
|
)
|
|
]
|
|
service = ExternalCreativeService(session, serving)
|
|
project_id = service._project(client)
|
|
service.upsert_resource(
|
|
ExternalResourceApproval(
|
|
project_id=project_id,
|
|
resource_ref="expiry-boundary-fixture",
|
|
workload="image",
|
|
provider_resource_id="offline-provider",
|
|
configured_vram_bytes=60,
|
|
),
|
|
"qualification",
|
|
)
|
|
acquired = service.acquire(
|
|
ExternalAllocationCreate(
|
|
request_id="expiry-boundary",
|
|
idempotency_key="expiry-boundary",
|
|
installation_id="fixture",
|
|
workload="image",
|
|
resource_ref="expiry-boundary-fixture",
|
|
provider_resource_id="offline-provider",
|
|
required_vram_bytes=60,
|
|
),
|
|
client,
|
|
)
|
|
row = session.get(ExternalCreativeAllocation, acquired.allocation_id)
|
|
boundary = datetime(2030, 1, 1, tzinfo=UTC)
|
|
row.expires_at = boundary
|
|
session.commit()
|
|
monkeypatch.setattr(external_creative_module, "_now", lambda: boundary)
|
|
|
|
uncertain = service.list_allocations("uncertain", 20, 0)
|
|
granted = service.list_allocations("granted", 20, 0)
|
|
|
|
assert uncertain.total == len(uncertain.items) == 1
|
|
assert uncertain.items[0].allocation_id == acquired.allocation_id
|
|
assert uncertain.items[0].state == "uncertain"
|
|
assert uncertain.items[0].decision_code == "ALLOCATION_EXPIRED_UNCERTAIN"
|
|
assert uncertain.items[0].uncertain_at is None
|
|
assert granted.total == len(granted.items) == 0
|
|
session.refresh(row)
|
|
assert row.state == "granted" and row.expires_at == boundary
|