Publish curated ITWorx-ModelForge source from 990401737f96c8d46a9f7b0d84994516cadcb6a0
This commit is contained in:
Public Source Publisher committed
1 parent
7082ab955a
commit
848f0d8f8a
434 files changed
+67298
-4315
No files matched your search
@@ -0,0 +1,302 @@
|
||||
"""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
|
||||
Reference in new issue
Block a user