from __future__ import annotations import hashlib import json import uuid from datetime import UTC, datetime, timedelta from typing import cast import sqlalchemy as sa from sqlalchemy import func, select from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session from modelforge_api.domain.external_creative import ( CreativeAllocationState, CreativeRequiredAction, CreativeWorkload, ExternalActionAttestation, ExternalAllocationCreate, ExternalAllocationList, ExternalAllocationResponse, ExternalProfileObservation, ExternalResourceApproval, ExternalResourceList, ExternalResourceResponse, ) from modelforge_api.persistence.models import ( Accelerator, ExternalCreativeActionAttestation, ExternalCreativeAllocation, ExternalCreativeProfileObservation, ExternalCreativeResource, Project, ProjectBinding, ServiceClient, ) from modelforge_api.services.audit import AuditWriter from modelforge_api.services.serving import ServingError, ServingService ACTIVE = {"granted", "uncertain"} PRIORITY = {"production": 0, "interactive": 1, "background": 2, "benchmark": 3} resources = cast(sa.Table, ExternalCreativeResource.__table__) attestations = cast(sa.Table, ExternalCreativeActionAttestation.__table__) observations = cast(sa.Table, ExternalCreativeProfileObservation.__table__) def _now() -> datetime: return datetime.now(UTC) def _digest(value: ExternalAllocationCreate) -> str: return hashlib.sha256( json.dumps(value.model_dump(mode="json"), sort_keys=True, separators=(",", ":")).encode() ).hexdigest() class ExternalCreativeService: def __init__(self, session: Session, serving: ServingService): self.session, self.serving = session, serving def _lock_accelerators(self) -> None: self.session.scalars(select(Accelerator).order_by(Accelerator.id).with_for_update()).all() def _project(self, client: ServiceClient) -> uuid.UUID: fresh = self.session.get( ServiceClient, client.id, populate_existing=True, with_for_update=True ) binding = ( self.session.get( ProjectBinding, fresh.project_binding_id, populate_existing=True, with_for_update=True, ) if fresh and fresh.project_binding_id else None ) project = ( self.session.get( Project, binding.project_id, populate_existing=True, with_for_update=True ) if binding else None ) if ( fresh is None or fresh.status != "active" or fresh.disabled_at is not None or "creative.resource@1" not in fresh.allowed_capabilities or binding is None or binding.deprecated_at is not None or project is None or not project.active ): raise ServingError( 403, "PROJECT_BINDING_REQUIRED", "creative coordination requires an active project-bound client", ) return binding.project_id def _priority(self, client: ServiceClient, requested: str) -> str: binding = self.session.get(ProjectBinding, client.project_binding_id) allowed = max( PRIORITY.get(client.workload_priority, 99), PRIORITY.get(binding.priority, 99) if binding else 99, ) if PRIORITY.get(requested, -1) < allowed: raise ServingError( 403, "ALLOCATION_PRIORITY_FORBIDDEN", "requested priority exceeds client or project policy", ) return requested def _expire(self, row: ExternalCreativeAllocation) -> bool: expiry = row.expires_at if row.state == "granted" and ( expiry is None or (expiry.astimezone(UTC) if expiry.tzinfo else expiry.replace(tzinfo=UTC)) <= _now() ): row.state, row.uncertain_at, row.updated_at = "uncertain", _now(), _now() row.expires_at = None row.decision_code = "ALLOCATION_EXPIRED_UNCERTAIN" self._audit("system", "external-creative-expiry").write( "EXTERNAL_CREATIVE_ALLOCATION_EXPIRED", "external_creative_allocation", str(row.id), {"state": row.state, "reservation_retained_bytes": row.required_vram_bytes}, ) return True return False def _expire_grants(self) -> None: for row in self.session.scalars( select(ExternalCreativeAllocation) .where(ExternalCreativeAllocation.state == "granted") .order_by(ExternalCreativeAllocation.id) .with_for_update() .execution_options(populate_existing=True) ): self._expire(row) self.session.flush() def _queued_authority(self, row: ExternalCreativeAllocation) -> bool: client = self.session.get(ServiceClient, row.service_client_id, populate_existing=True) try: if client is None or self._project(client) != row.project_id: return False self._priority(client, row.priority) except ServingError: return False resource = ( self.session.execute( select(resources) .where( resources.c.project_id == row.project_id, resources.c.resource_ref == row.resource_ref, resources.c.workload == row.workload, resources.c.provider_resource_id == row.provider_resource_id, ) .with_for_update() ) .mappings() .one_or_none() ) if resource is None or not resource["approved"]: return False row.required_vram_bytes = max( row.required_vram_bytes, resource["configured_vram_bytes"], resource["measured_peak_vram_bytes"] or 0, ) return True def _audit( self, client_or_type: ServiceClient | str, actor_id: str | None = None ) -> AuditWriter: if isinstance(client_or_type, ServiceClient): return AuditWriter(self.session, "service_client", str(client_or_type.id)) return AuditWriter(self.session, client_or_type, actor_id or "unknown") def _response( self, row: ExternalCreativeAllocation, *, as_of: datetime | None = None ) -> ExternalAllocationResponse: expired = False if as_of is not None and row.state == "granted": expiry = row.expires_at expired = ( expiry is None or (expiry.astimezone(UTC) if expiry.tzinfo else expiry.replace(tzinfo=UTC)) <= as_of ) return ExternalAllocationResponse( allocation_id=row.id, project_id=row.project_id, request_id=row.request_id, installation_id=row.decision_evidence.get("installation_id", "unknown"), workload=cast(CreativeWorkload, row.workload), resource_ref=row.resource_ref, provider_resource_id=row.provider_resource_id, required_vram_bytes=row.required_vram_bytes, priority=row.priority, state=cast(CreativeAllocationState, "uncertain" if expired else row.state), decision_code=("ALLOCATION_EXPIRED_UNCERTAIN" if expired else row.decision_code), required_actions=cast(list[CreativeRequiredAction], row.required_actions or []), accelerator_id=row.accelerator_id, acquired_at=row.acquired_at, heartbeat_at=row.heartbeat_at, expires_at=None if expired else row.expires_at, uncertain_at=row.uncertain_at, released_at=row.released_at, decision_evidence=row.decision_evidence or {}, queue_position=(row.decision_evidence or {}).get("queue_position"), ) def _promote_queue(self) -> None: self._lock_accelerators() self._expire_grants() queued = list( self.session.scalars( select(ExternalCreativeAllocation) .where(ExternalCreativeAllocation.state == "queued") .order_by(ExternalCreativeAllocation.created_at, ExternalCreativeAllocation.id) .with_for_update() ) ) queued.sort( key=lambda item: (PRIORITY.get(item.priority, 9), item.created_at, str(item.id)) ) # Snapshot already includes all existing external reservations. reserved: dict[uuid.UUID, int] = {} budgets, now = self.serving.scheduler_overview(commit=False), _now() remaining_position = 0 for item in queued: if not self._queued_authority(item): item.state, item.decision_code = "blocked", "QUEUED_AUTHORITY_REVOKED" item.decision_evidence = {**item.decision_evidence, "queue_position": None} self._audit("system", "external-creative-admission").write( "EXTERNAL_CREATIVE_ALLOCATION_BLOCKED", "external_creative_allocation", str(item.id), {"decision_code": item.decision_code}, ) continue chosen = next( ( budget for budget in budgets if budget.attribution_confidence != "UNKNOWN" and budget.pressure_state not in {"HIGH", "CRITICAL"} and budget.schedulable_free_vram_bytes - reserved.get(budget.accelerator_id, 0) >= item.required_vram_bytes ), None, ) evidence = dict(item.decision_evidence or {}) if chosen is None: remaining_position += 1 evidence["queue_position"] = remaining_position item.decision_evidence = evidence continue item.state, item.decision_code, item.accelerator_id = ( "granted", "CAPACITY_GRANTED_FROM_QUEUE", chosen.accelerator_id, ) item.acquired_at = item.heartbeat_at = item.updated_at = now item.expires_at = now + timedelta(seconds=item.ttl_seconds) # This grant is based on observed schedulable free VRAM. A forced # unload would discard a reusable resident model despite admission # already succeeding. Preparation is only meaningful when a future # decision explicitly accounts for reclaimable provider memory. item.required_actions = [] evidence.update( { "queue_position": None, "promoted_at": now.isoformat(), "policy_revision": chosen.policy_revision, } ) item.decision_evidence = evidence self._audit("system", "external-creative-admission").write( "EXTERNAL_CREATIVE_ALLOCATION_GRANTED", "external_creative_allocation", str(item.id), { "accelerator_id": str(chosen.accelerator_id), "reserved_vram_bytes": item.required_vram_bytes, }, ) reserved[chosen.accelerator_id] = ( reserved.get(chosen.accelerator_id, 0) + item.required_vram_bytes ) def acquire( self, request: ExternalAllocationCreate, client: ServiceClient ) -> ExternalAllocationResponse: self._lock_accelerators() self._expire_grants() project_id = self._project(client) self._priority(client, request.priority) digest = _digest(request) previous = self.session.scalar( select(ExternalCreativeAllocation).where( ExternalCreativeAllocation.service_client_id == client.id, ExternalCreativeAllocation.idempotency_key == request.idempotency_key, ) ) if previous: if previous.request_sha256 != digest: raise ServingError( 409, "IDEMPOTENCY_CONFLICT", "idempotency key belongs to another allocation request", ) self.session.commit() return self._response(previous) binding = ( self.session.execute( select(resources) .where( resources.c.project_id == project_id, resources.c.resource_ref == request.resource_ref, resources.c.workload == request.workload, resources.c.provider_resource_id == request.provider_resource_id, ) .with_for_update() ) .mappings() .one_or_none() ) if binding is None or not binding["approved"]: raise ServingError( 403, "RESOURCE_NOT_APPROVED", "resource and provider binding is not approved for this project", ) profiled = binding["measured_peak_vram_bytes"] if request.use_measured_profile else None required_vram = max( request.required_vram_bytes, profiled or 0, binding["configured_vram_bytes"] ) # Give already-queued work first access to newly available capacity, while # allowing requests that fit to bypass queue entries that still cannot. self._promote_queue() budgets = self.serving.scheduler_overview(commit=False) chosen = next( ( b for b in budgets if b.attribution_confidence != "UNKNOWN" and b.pressure_state not in {"HIGH", "CRITICAL"} and b.schedulable_free_vram_bytes >= required_vram ), None, ) queued = list( self.session.scalars( select(ExternalCreativeAllocation).where( ExternalCreativeAllocation.state == "queued" ) ) ) queue_position = 1 + sum( PRIORITY.get(item.priority, 9) <= PRIORITY[request.priority] for item in queued ) state, code = ("granted", "CAPACITY_GRANTED") if chosen else ("queued", "CAPACITY_QUEUED") observed_budget = chosen or max( budgets, key=lambda budget: budget.schedulable_free_vram_bytes, default=None ) actions: list[str] = [] evidence = { "installation_id": request.installation_id, "policy_revision": observed_budget.policy_revision if observed_budget else "unavailable", "observed_at": observed_budget.observed_at.isoformat() if observed_budget and observed_budget.observed_at else None, "headroom_before_bytes": ( observed_budget.schedulable_free_vram_bytes if observed_budget else 0 ), "profile_source": "measured" if profiled else "configured", "requested_vram_bytes": request.required_vram_bytes, "queue_position": None if chosen else queue_position, } now = _now() row = ExternalCreativeAllocation( service_client_id=client.id, project_id=project_id, accelerator_id=chosen.accelerator_id if chosen else None, request_id=request.request_id, idempotency_key=request.idempotency_key, request_sha256=digest, workload=request.workload, resource_ref=request.resource_ref, provider_resource_id=request.provider_resource_id, required_vram_bytes=required_vram, priority=request.priority, ttl_seconds=request.ttl_seconds, state=state, decision_code=code, decision_evidence=evidence, required_actions=actions if chosen else [], acquired_at=now if chosen else None, heartbeat_at=now if chosen else None, expires_at=now + timedelta(seconds=request.ttl_seconds) if chosen else None, ) self.session.add(row) try: self.session.flush() self._audit(client).write( "EXTERNAL_CREATIVE_ALLOCATION_DECIDED", "external_creative_allocation", str(row.id), { "project_id": str(project_id), "workload": row.workload, "resource_ref": row.resource_ref, "state": row.state, "decision_code": row.decision_code, "accelerator_id": str(row.accelerator_id) if row.accelerator_id else None, }, ) self.session.commit() except IntegrityError as exc: self.session.rollback() raise ServingError( 409, "ALLOCATION_CONFLICT", "allocation was concurrently created; retry with the same idempotency key", ) from exc self.session.refresh(row) return self._response(row) def get(self, allocation_id: uuid.UUID, client: ServiceClient) -> ExternalAllocationResponse: self._lock_accelerators() self._project(client) self._promote_queue() row = self.session.get(ExternalCreativeAllocation, allocation_id) if row is None or row.service_client_id != client.id: raise ServingError(404, "ALLOCATION_NOT_FOUND", "allocation was not found") self.session.commit() return self._response(row) def action( self, allocation_id: uuid.UUID, request_id: str, action: str, client: ServiceClient, observed_peak_vram_bytes: int | None = None, attestation: ExternalActionAttestation | None = None, ) -> ExternalAllocationResponse: self._lock_accelerators() row = self.session.get( ExternalCreativeAllocation, allocation_id, with_for_update=True, populate_existing=True, ) if row is None or row.service_client_id != client.id or row.request_id != request_id: raise ServingError(404, "ALLOCATION_NOT_FOUND", "allocation was not found") if self._expire(row) and action == "heartbeat": self.session.commit() raise ServingError( 409, "ALLOCATION_EXPIRED", "expired allocation retains an uncertain reservation" ) if action == "heartbeat": self._project(client) if action == "attest" and attestation is None: raise ServingError(422, "ATTESTATION_REQUIRED", "attestation is required") now = _now() if attestation is not None: if attestation.action not in (row.required_actions or []): raise ServingError(409, "ACTION_NOT_REQUIRED", "attested action was not required") if attestation.observed_after > attestation.observed_before: raise ServingError( 422, "ATTESTATION_INVALID", "action attestation did not release VRAM" ) self.session.execute( attestations.insert().values( id=uuid.uuid4(), allocation_id=row.id, action=attestation.action, before_vram_bytes=attestation.observed_before, after_vram_bytes=attestation.observed_after, observed_at=attestation.completed_at, validated=True, created_at=now, ) ) if action == "release": if row.state != "released": row.state, row.released_at, row.updated_at = "released", now, now elif action == "uncertain": if row.state != "released": row.state, row.uncertain_at, row.expires_at, row.updated_at = ( "uncertain", now, None, now, ) elif action == "heartbeat": if row.state != "granted": raise ServingError( 409, "ALLOCATION_NOT_ACTIVE", "only granted allocations can be renewed" ) row.heartbeat_at, row.expires_at, row.updated_at = ( now, now + timedelta(seconds=row.ttl_seconds), now, ) self._audit(client).write( f"EXTERNAL_CREATIVE_ALLOCATION_{action.upper()}", "external_creative_allocation", str(row.id), {"state": row.state, "workload": row.workload}, ) if action == "release": self.session.flush() self._promote_queue() self.session.commit() self.session.refresh(row) return self._response(row) def observe( self, body: ExternalProfileObservation, client: ServiceClient ) -> ExternalResourceResponse: self._lock_accelerators() self._project(client) row = self.session.get( ExternalCreativeAllocation, body.allocation_id, with_for_update=True, populate_existing=True, ) if row is None or row.service_client_id != client.id: raise ServingError(404, "ALLOCATION_NOT_FOUND", "allocation was not found") if self._expire(row): self.session.commit() if ( row.acquired_at is None or row.accelerator_id is None or row.state not in {"granted", "released"} or body.reserved_vram_bytes != row.required_vram_bytes ): raise ServingError( 409, "PROFILE_GRANT_REQUIRED", "observation requires an actual grant and exact reserved bytes", ) if (row.resource_ref, row.provider_resource_id, row.workload) != ( body.resource_ref, body.provider_resource_id, body.workload, ): raise ServingError( 409, "PROFILE_BINDING_MISMATCH", "observation does not match its allocation" ) record = ( self.session.execute( select(resources) .where( resources.c.project_id == row.project_id, resources.c.resource_ref == row.resource_ref, resources.c.workload == row.workload, resources.c.provider_resource_id == row.provider_resource_id, ) .with_for_update() ) .mappings() .one() ) previous = self.session.execute( select(observations.c.id).where(observations.c.allocation_id == row.id) ).scalar_one_or_none() if previous is not None: return self._resource(record["id"]) try: self.session.execute( observations.insert().values( id=uuid.uuid4(), allocation_id=row.id, resource_ref=body.resource_ref, provider_resource_id=body.provider_resource_id, workload=body.workload, observed_peak_vram_bytes=body.observed_peak_vram_bytes, reserved_vram_bytes=row.required_vram_bytes, duration_ms=body.duration_ms, outcome=body.outcome, created_at=_now(), ) ) except IntegrityError as exc: self.session.rollback() raise ServingError( 409, "PROFILE_OBSERVATION_EXISTS", "allocation already has a profile observation" ) from exc self.session.execute( resources.update() .where(resources.c.id == record["id"]) .values( measured_peak_vram_bytes=max( record["measured_peak_vram_bytes"] or 0, body.observed_peak_vram_bytes ), measurement_count=record["measurement_count"] + 1, updated_at=_now(), ) ) self._audit(client).write( "EXTERNAL_CREATIVE_PROFILE_OBSERVED", "external_creative_resource", str(record["id"]), {"outcome": body.outcome, "observed_peak_vram_bytes": body.observed_peak_vram_bytes}, ) self.session.commit() return self._resource(record["id"]) def _resource(self, identifier: uuid.UUID) -> ExternalResourceResponse: record = ( self.session.execute(select(resources).where(resources.c.id == identifier)) .mappings() .one() ) return ExternalResourceResponse( resource_id=identifier, **{ key: record[key] for key in ExternalResourceResponse.model_fields if key != "resource_id" }, ) def upsert_resource( self, body: ExternalResourceApproval, actor_id: str ) -> ExternalResourceResponse: record = ( self.session.execute( select(resources) .where( resources.c.project_id == body.project_id, resources.c.resource_ref == body.resource_ref, resources.c.workload == body.workload, resources.c.provider_resource_id == body.provider_resource_id, ) .with_for_update() ) .mappings() .one_or_none() ) identifier, now = (record["id"] if record else uuid.uuid4()), _now() values = {**body.model_dump(), "updated_at": now} if record: self.session.execute( resources.update().where(resources.c.id == identifier).values(**values) ) else: self.session.execute( resources.insert().values( id=identifier, measured_peak_vram_bytes=None, measurement_count=0, **values ) ) self._audit("operator", actor_id).write( "EXTERNAL_CREATIVE_RESOURCE_CONFIGURED", "external_creative_resource", str(identifier), {"project_id": str(body.project_id), "approved": body.approved}, ) self.session.commit() return self._resource(identifier) def list_allocations( self, state: str | None, limit: int, offset: int ) -> ExternalAllocationList: observed_at = _now() expired_grant = sa.and_( ExternalCreativeAllocation.state == "granted", sa.or_( ExternalCreativeAllocation.expires_at.is_(None), ExternalCreativeAllocation.expires_at <= observed_at, ), ) effective_state = sa.case( (expired_grant, "uncertain"), else_=ExternalCreativeAllocation.state ) condition = effective_state == state if state else sa.true() with self.session.no_autoflush: total = ( self.session.scalar( select(func.count()).select_from(ExternalCreativeAllocation).where(condition) ) or 0 ) rows = self.session.scalars( select(ExternalCreativeAllocation) .where(condition) .order_by(ExternalCreativeAllocation.created_at.desc()) .offset(offset) .limit(limit) ).all() return ExternalAllocationList( items=[self._response(row, as_of=observed_at) for row in rows], total=total ) def list_resources( self, project_id: uuid.UUID | None, limit: int, offset: int ) -> ExternalResourceList: condition = resources.c.project_id == project_id if project_id else sa.true() total = ( self.session.scalar(select(func.count()).select_from(resources).where(condition)) or 0 ) records = ( self.session.execute( select(resources) .where(condition) .order_by(resources.c.updated_at.desc()) .offset(offset) .limit(limit) ) .mappings() .all() ) return ExternalResourceList( items=[ ExternalResourceResponse( resource_id=record["id"], **{ key: record[key] for key in ExternalResourceResponse.model_fields if key != "resource_id" }, ) for record in records ], total=total, ) def reconcile( self, allocation_id: uuid.UUID, action: str, reason: str, actor_id: str ) -> ExternalAllocationResponse: self._lock_accelerators() row = self.session.get( ExternalCreativeAllocation, allocation_id, with_for_update=True, populate_existing=True, ) if row is None: raise ServingError(404, "ALLOCATION_NOT_FOUND", "allocation was not found") self._expire(row) if row.state != "uncertain": raise ServingError( 409, "ALLOCATION_NOT_UNCERTAIN", "only uncertain allocations can be reconciled" ) if action == "confirm_released": row.state, row.released_at, row.updated_at = "released", _now(), _now() self.session.flush() self._promote_queue() self._audit("operator", actor_id).write( "EXTERNAL_CREATIVE_ALLOCATION_RECONCILED", "external_creative_allocation", str(row.id), {"action": action, "reason": reason}, ) self.session.commit() self.session.refresh(row) return self._response(row)