771 lines
30 KiB
Python
771 lines
30 KiB
Python
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)
|