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