457 lines
18 KiB
Python
457 lines
18 KiB
Python
"""Defensive allocation authority and reservation regressions; no GPU invocation."""
|
|
|
|
from datetime import UTC, datetime, timedelta
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
from sqlalchemy import create_engine, event, select
|
|
from sqlalchemy.orm import Session
|
|
|
|
from modelforge_api.api.routes.external_creative import operator_allocations
|
|
from modelforge_api.domain.external_creative import (
|
|
ExternalActionAttestation,
|
|
ExternalAllocationCreate,
|
|
ExternalProfileObservation,
|
|
ExternalResourceApproval,
|
|
)
|
|
from modelforge_api.persistence.models import (
|
|
Accelerator,
|
|
Base,
|
|
ExternalCreativeAllocation,
|
|
Project,
|
|
ProjectBinding,
|
|
ServiceClient,
|
|
)
|
|
from modelforge_api.services import external_creative as external_creative_module
|
|
from modelforge_api.services.external_creative import (
|
|
ExternalCreativeService,
|
|
attestations,
|
|
observations,
|
|
resources,
|
|
)
|
|
from modelforge_api.services.serving import ServingError
|
|
from tests.test_serving import setup_m5
|
|
|
|
|
|
@pytest.fixture
|
|
def allocation_setup():
|
|
engine = create_engine("sqlite://")
|
|
|
|
@event.listens_for(engine, "connect")
|
|
def foreign_keys(connection, _record):
|
|
connection.execute("PRAGMA foreign_keys=ON")
|
|
|
|
Base.metadata.create_all(engine)
|
|
with Session(engine) as session:
|
|
serving, _candidate, node = setup_m5(session)
|
|
contract = serving.ensure_contract("rag.embedding", 1)
|
|
project = Project(key="creative-fixture", name="Creative", description="Offline fixture")
|
|
session.add(project)
|
|
session.flush()
|
|
binding = ProjectBinding(
|
|
project_id=project.id,
|
|
capability_contract_id=contract.id,
|
|
channel="stable",
|
|
priority="interactive",
|
|
)
|
|
session.add(binding)
|
|
session.flush()
|
|
client = ServiceClient(
|
|
name="Creative client",
|
|
project_binding_id=binding.id,
|
|
allowed_capabilities=["creative.resource@1"],
|
|
workload_priority="interactive",
|
|
)
|
|
session.add(client)
|
|
session.commit()
|
|
accelerator = session.scalar(
|
|
select(Accelerator).where(Accelerator.compute_node_id == node.id)
|
|
)
|
|
budget = SimpleNamespace(
|
|
accelerator_id=accelerator.id,
|
|
attribution_confidence="HIGH",
|
|
pressure_state="LOW",
|
|
schedulable_free_vram_bytes=100,
|
|
policy_revision="fixture",
|
|
observed_at=None,
|
|
)
|
|
|
|
# Contract matches root's shared accounting: subtract retained reservations once.
|
|
def overview(*, commit):
|
|
assert commit is False
|
|
retained = sum(
|
|
row.required_vram_bytes
|
|
for row in session.scalars(
|
|
select(ExternalCreativeAllocation).where(
|
|
ExternalCreativeAllocation.state.in_(["granted", "uncertain"])
|
|
)
|
|
)
|
|
)
|
|
return [
|
|
SimpleNamespace(
|
|
**{
|
|
**vars(budget),
|
|
"schedulable_free_vram_bytes": budget.schedulable_free_vram_bytes
|
|
- retained,
|
|
}
|
|
)
|
|
]
|
|
|
|
serving.scheduler_overview = Mock(side_effect=overview)
|
|
service = ExternalCreativeService(session, serving)
|
|
service.upsert_resource(
|
|
ExternalResourceApproval(
|
|
project_id=project.id,
|
|
resource_ref="chosen-model",
|
|
workload="image",
|
|
provider_resource_id="private-provider",
|
|
configured_vram_bytes=60,
|
|
),
|
|
"fixture",
|
|
)
|
|
request = ExternalAllocationCreate(
|
|
request_id="one",
|
|
idempotency_key="one",
|
|
installation_id="fixture",
|
|
workload="image",
|
|
resource_ref="chosen-model",
|
|
provider_resource_id="private-provider",
|
|
required_vram_bytes=60,
|
|
)
|
|
yield session, service, request, client, budget, binding, project
|
|
engine.dispose()
|
|
|
|
|
|
def test_attestation_requires_evidence_and_records_authorized_action(allocation_setup):
|
|
session, service, request, client, _, _, _ = allocation_setup
|
|
allocation = service.acquire(request, client)
|
|
assert allocation.state == "granted"
|
|
assert allocation.required_actions == []
|
|
# The normal free-capacity grant needs no eviction. Exercise the retained
|
|
# action attestation contract with an explicitly action-bearing fixture.
|
|
row = session.get(ExternalCreativeAllocation, allocation.allocation_id)
|
|
row.required_actions = ["release_idle_comfy"]
|
|
session.commit()
|
|
allocation = service.acquire(request, client)
|
|
|
|
with pytest.raises(ServingError, match="attestation is required"):
|
|
service.action(allocation.allocation_id, request.request_id, "attest", client)
|
|
assert session.execute(select(attestations)).all() == []
|
|
|
|
service.action(
|
|
allocation.allocation_id,
|
|
request.request_id,
|
|
"attest",
|
|
client,
|
|
attestation=ExternalActionAttestation(
|
|
action=allocation.required_actions[0],
|
|
observed_before=100,
|
|
observed_after=40,
|
|
completed_at=datetime.now(UTC),
|
|
),
|
|
)
|
|
recorded = session.execute(select(attestations)).one()
|
|
assert recorded.action == allocation.required_actions[0]
|
|
assert recorded.before_vram_bytes == 100
|
|
assert recorded.after_vram_bytes == 40
|
|
|
|
|
|
def test_priority_cannot_exceed_either_client_or_binding(allocation_setup):
|
|
session, service, request, client, _, binding, _ = allocation_setup
|
|
with pytest.raises(ServingError, match="priority"):
|
|
service.acquire(request.model_copy(update={"priority": "production"}), client)
|
|
client.workload_priority = "production"
|
|
session.commit()
|
|
with pytest.raises(ServingError, match="priority"):
|
|
service.acquire(request.model_copy(update={"priority": "production"}), client)
|
|
assert list(session.scalars(select(ExternalCreativeAllocation))) == []
|
|
binding.priority = "production"
|
|
session.commit()
|
|
assert (
|
|
service.acquire(request.model_copy(update={"priority": "production"}), client).priority
|
|
== "production"
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"revocation", ["client", "scope", "binding", "project", "resource", "priority"]
|
|
)
|
|
def test_queue_rechecks_current_authority(allocation_setup, revocation):
|
|
session, service, request, client, budget, binding, project = allocation_setup
|
|
budget.schedulable_free_vram_bytes = 0
|
|
created = service.acquire(request, client)
|
|
assert created.state == "queued"
|
|
if revocation == "client":
|
|
client.status = "disabled"
|
|
elif revocation == "scope":
|
|
client.allowed_capabilities = []
|
|
elif revocation == "binding":
|
|
binding.deprecated_at = datetime.now(UTC)
|
|
elif revocation == "project":
|
|
project.active = False
|
|
elif revocation == "priority":
|
|
client.workload_priority = "background"
|
|
else:
|
|
session.execute(resources.update().values(approved=False))
|
|
session.commit()
|
|
budget.schedulable_free_vram_bytes = 100
|
|
service._promote_queue()
|
|
session.commit()
|
|
row = session.get(ExternalCreativeAllocation, created.allocation_id)
|
|
assert row.state == "blocked" and row.acquired_at is None and row.accelerator_id is None
|
|
|
|
|
|
def test_queue_recalculates_resource_floor_and_accounts_new_grants_once(allocation_setup):
|
|
session, service, request, client, budget, _, _ = allocation_setup
|
|
budget.schedulable_free_vram_bytes = 0
|
|
first = service.acquire(request, client)
|
|
second = service.acquire(
|
|
request.model_copy(update={"idempotency_key": "two", "request_id": "two"}), client
|
|
)
|
|
session.execute(resources.update().values(configured_vram_bytes=80))
|
|
session.commit()
|
|
budget.schedulable_free_vram_bytes = 150
|
|
service._promote_queue()
|
|
session.commit()
|
|
pair = [session.get(ExternalCreativeAllocation, item.allocation_id) for item in (first, second)]
|
|
assert sorted(row.state for row in pair) == ["granted", "queued"]
|
|
assert all(row.required_vram_bytes == 80 for row in pair)
|
|
assert all(row.required_actions == [] for row in pair)
|
|
|
|
|
|
def test_permanently_oversized_queue_head_does_not_starve_fitting_request(allocation_setup):
|
|
_, service, request, client, _, _, _ = allocation_setup
|
|
oversized = service.acquire(
|
|
request.model_copy(
|
|
update={
|
|
"idempotency_key": "oversized",
|
|
"request_id": "oversized",
|
|
"required_vram_bytes": 120,
|
|
}
|
|
),
|
|
client,
|
|
)
|
|
fitting = service.acquire(
|
|
request.model_copy(update={"idempotency_key": "fitting", "request_id": "fitting"}),
|
|
client,
|
|
)
|
|
|
|
assert oversized.state == "queued"
|
|
assert fitting.state == "granted"
|
|
assert fitting.required_actions == []
|
|
assert fitting.decision_evidence["headroom_before_bytes"] == 100
|
|
|
|
|
|
def test_new_request_promotes_fitting_queue_head_before_competing_for_capacity(allocation_setup):
|
|
session, service, request, client, budget, _, _ = allocation_setup
|
|
budget.schedulable_free_vram_bytes = 0
|
|
waiting = service.acquire(request, client)
|
|
budget.schedulable_free_vram_bytes = 100
|
|
|
|
newcomer = service.acquire(
|
|
request.model_copy(update={"idempotency_key": "new", "request_id": "new"}), client
|
|
)
|
|
|
|
session.refresh(waiting_row := session.get(ExternalCreativeAllocation, waiting.allocation_id))
|
|
assert waiting_row.state == "granted"
|
|
assert newcomer.state == "queued"
|
|
|
|
|
|
def test_queued_decision_reports_observed_capacity_instead_of_zero_sentinel(allocation_setup):
|
|
_, service, request, client, budget, _, _ = allocation_setup
|
|
observed_at = datetime(2026, 9, 22, 8, 30, tzinfo=UTC)
|
|
budget.schedulable_free_vram_bytes = 50
|
|
budget.observed_at = observed_at
|
|
|
|
queued = service.acquire(request, client)
|
|
|
|
assert queued.state == "queued"
|
|
assert queued.decision_evidence["headroom_before_bytes"] == 50
|
|
assert queued.decision_evidence["observed_at"] == observed_at.isoformat()
|
|
|
|
|
|
def test_expired_heartbeat_cannot_resurrect_or_release_reservation(allocation_setup):
|
|
session, service, request, client, _, _, _ = allocation_setup
|
|
first = service.acquire(request, client)
|
|
row = session.get(ExternalCreativeAllocation, first.allocation_id)
|
|
row.expires_at = datetime.now(UTC) - timedelta(seconds=1)
|
|
session.commit()
|
|
with pytest.raises(ServingError, match="expired"):
|
|
service.action(row.id, request.request_id, "heartbeat", client)
|
|
session.expire_all()
|
|
assert row.state == "uncertain" and row.required_vram_bytes == 60 and row.accelerator_id
|
|
second = service.acquire(
|
|
request.model_copy(update={"idempotency_key": "two", "request_id": "two"}), client
|
|
)
|
|
assert second.state == "queued"
|
|
service.action(row.id, request.request_id, "release", client)
|
|
assert session.get(ExternalCreativeAllocation, second.allocation_id).state == "granted"
|
|
|
|
|
|
def test_confirming_expired_grant_released_promotes_waiting_allocation(allocation_setup):
|
|
session, service, request, client, _, _, _ = allocation_setup
|
|
granted = service.acquire(request, client)
|
|
expired = session.get(ExternalCreativeAllocation, granted.allocation_id)
|
|
expired.expires_at = datetime.now(UTC) - timedelta(seconds=1)
|
|
session.commit()
|
|
|
|
waiting = service.acquire(
|
|
request.model_copy(update={"idempotency_key": "waiting", "request_id": "waiting"}),
|
|
client,
|
|
)
|
|
session.refresh(expired)
|
|
assert expired.state == "uncertain"
|
|
assert expired.released_at is None
|
|
assert waiting.state == "queued"
|
|
|
|
service.reconcile(expired.id, "confirm_released", "provider verified idle", "operator")
|
|
|
|
waiting_row = session.get(ExternalCreativeAllocation, waiting.allocation_id)
|
|
session.refresh(waiting_row)
|
|
assert waiting_row.state == "granted"
|
|
|
|
|
|
@pytest.mark.parametrize("mutation", ["queued", "reserved", "expired"])
|
|
def test_profile_requires_actual_grant_and_authoritative_reservation(allocation_setup, mutation):
|
|
session, service, request, client, budget, _, _ = allocation_setup
|
|
if mutation == "queued":
|
|
budget.schedulable_free_vram_bytes = 0
|
|
first = service.acquire(request, client)
|
|
row = session.get(ExternalCreativeAllocation, first.allocation_id)
|
|
if mutation == "expired":
|
|
row.expires_at = datetime.now(UTC) - timedelta(seconds=1)
|
|
session.commit()
|
|
body = ExternalProfileObservation(
|
|
allocation_id=row.id,
|
|
resource_ref=request.resource_ref,
|
|
provider_resource_id=request.provider_resource_id,
|
|
workload="image",
|
|
observed_peak_vram_bytes=50,
|
|
reserved_vram_bytes=61 if mutation == "reserved" else 60,
|
|
duration_ms=1000,
|
|
outcome="succeeded",
|
|
)
|
|
with pytest.raises(ServingError, match="actual grant"):
|
|
service.observe(body, client)
|
|
assert list(session.execute(select(observations))) == []
|
|
|
|
|
|
def test_released_grant_can_report_exact_profile_without_reopening_capacity(allocation_setup):
|
|
session, service, request, client, _, _, _ = allocation_setup
|
|
first = service.acquire(request, client)
|
|
service.action(first.allocation_id, request.request_id, "release", client)
|
|
body = ExternalProfileObservation(
|
|
allocation_id=first.allocation_id,
|
|
resource_ref=request.resource_ref,
|
|
provider_resource_id=request.provider_resource_id,
|
|
workload="image",
|
|
observed_peak_vram_bytes=50,
|
|
reserved_vram_bytes=60,
|
|
duration_ms=1000,
|
|
outcome="succeeded",
|
|
)
|
|
assert service.observe(body, client).measurement_count == 1
|
|
assert service.observe(body, client).measurement_count == 1
|
|
assert session.get(ExternalCreativeAllocation, first.allocation_id).state == "released"
|
|
|
|
|
|
def test_history_read_projects_expiry_without_inventing_release(allocation_setup):
|
|
session, service, request, client, _, _, _ = allocation_setup
|
|
first = service.acquire(request, client)
|
|
row = session.get(ExternalCreativeAllocation, first.allocation_id)
|
|
row.expires_at = datetime.now(UTC) - timedelta(seconds=1)
|
|
session.commit()
|
|
page = service.list_allocations("uncertain", 20, 0)
|
|
assert page.total == 1 and page.items[0].allocation_id == row.id
|
|
assert page.items[0].state == "uncertain"
|
|
assert page.items[0].decision_code == "ALLOCATION_EXPIRED_UNCERTAIN"
|
|
assert page.items[0].expires_at is None
|
|
assert page.items[0].uncertain_at is None
|
|
session.expire_all()
|
|
stored = session.get(ExternalCreativeAllocation, first.allocation_id)
|
|
assert stored.state == "granted" and stored.released_at is None
|
|
assert stored.expires_at.replace(tzinfo=UTC) < datetime.now(UTC)
|
|
assert stored.uncertain_at is None
|
|
assert stored.required_vram_bytes == 60 and stored.accelerator_id is not None
|
|
|
|
|
|
def test_operator_allocation_get_performs_no_dml_or_commit(allocation_setup):
|
|
session, service, request, client, _, _, _ = allocation_setup
|
|
first = service.acquire(request, client)
|
|
row = session.get(ExternalCreativeAllocation, first.allocation_id)
|
|
row.expires_at = datetime.now(UTC) - timedelta(seconds=1)
|
|
session.commit()
|
|
client.name = "pending read-only probe"
|
|
|
|
statements = []
|
|
commits = []
|
|
|
|
def record_statement(_connection, _cursor, statement, _parameters, _context, _many):
|
|
operation = statement.lstrip().split(None, 1)[0].upper()
|
|
if operation in {"INSERT", "UPDATE", "DELETE", "REPLACE"}:
|
|
statements.append(operation)
|
|
|
|
def record_commit(_session):
|
|
commits.append(True)
|
|
|
|
bind = session.get_bind()
|
|
event.listen(bind, "before_cursor_execute", record_statement)
|
|
event.listen(session, "before_commit", record_commit)
|
|
try:
|
|
page = operator_allocations(service, object(), "uncertain", 20, 0)
|
|
finally:
|
|
event.remove(bind, "before_cursor_execute", record_statement)
|
|
event.remove(session, "before_commit", record_commit)
|
|
|
|
assert page.total == 1 and page.items[0].state == "uncertain"
|
|
assert statements == []
|
|
assert commits == []
|
|
session.rollback()
|
|
session.refresh(client)
|
|
assert client.name == "Creative client"
|
|
|
|
|
|
def test_history_read_expiry_boundary_is_inclusive(allocation_setup, monkeypatch):
|
|
session, service, request, client, _, _, _ = allocation_setup
|
|
first = service.acquire(request, client)
|
|
row = session.get(ExternalCreativeAllocation, first.allocation_id)
|
|
boundary = datetime(2030, 1, 1, tzinfo=UTC)
|
|
row.expires_at = boundary
|
|
session.commit()
|
|
monkeypatch.setattr(external_creative_module, "_now", lambda: boundary)
|
|
|
|
expired = service.list_allocations("uncertain", 20, 0)
|
|
assert expired.total == 1
|
|
assert expired.items[0].state == "uncertain"
|
|
assert expired.items[0].decision_code == "ALLOCATION_EXPIRED_UNCERTAIN"
|
|
assert expired.items[0].uncertain_at is None
|
|
|
|
row.expires_at = boundary + timedelta(microseconds=1)
|
|
session.commit()
|
|
before_expiry = service.list_allocations("uncertain", 20, 0)
|
|
active = service.list_allocations("granted", 20, 0)
|
|
|
|
assert before_expiry.total == 0
|
|
assert active.total == 1
|
|
assert active.items[0].state == "granted"
|
|
assert active.items[0].expires_at.replace(tzinfo=UTC) == boundary + timedelta(microseconds=1)
|
|
|
|
|
|
def test_restored_queued_uncertainty_requires_explicit_reconciliation(allocation_setup):
|
|
session, service, request, client, budget, _, _ = allocation_setup
|
|
budget.schedulable_free_vram_bytes = 0
|
|
first = service.acquire(request, client)
|
|
row = session.get(ExternalCreativeAllocation, first.allocation_id)
|
|
row.state, row.decision_code = "uncertain", "CREATIVE_RESTORE_INTERRUPTED"
|
|
session.commit()
|
|
budget.schedulable_free_vram_bytes = 100
|
|
assert service.get(row.id, client).state == "uncertain"
|
|
assert row.accelerator_id is None and row.acquired_at is None
|
|
assert (
|
|
service.reconcile(
|
|
row.id, "confirm_released", "External owner verified idle", "operator"
|
|
).state
|
|
== "released"
|
|
)
|