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,456 @@
|
||||
"""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"
|
||||
)
|
||||
Reference in new issue
Block a user