tzhouam
723c550209
Extract versioned stage contracts, model shared CPU/iGPU/NPU memory constraints, and add explicit graph routes, reservations, output acknowledgements, and worker retirement. Fix native Windows shared-buffer lifetime and host paths. Document qualified hardware paths and remaining rollout gates.
102 lines
4.1 KiB
Python
102 lines
4.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Atomic reservations shared by native and external stage startup.
|
|
|
|
Capacities are admission ceilings for this controller, captured before loads.
|
|
They are not fresh `available` samples: resident allocations remain charged,
|
|
so sampling available memory cannot subtract our own allocation a second time.
|
|
Refresh ceilings only with an explicit, reconciled controller budget.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
from collections.abc import Mapping
|
|
from dataclasses import dataclass
|
|
from types import MappingProxyType
|
|
|
|
|
|
class ResourceUnavailable(RuntimeError): # noqa: N818 - public admission result, like PlanNotAdmitted
|
|
pass
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Reservation:
|
|
owner: str
|
|
demands: Mapping[str, int]
|
|
|
|
|
|
class ResourceLedger:
|
|
def __init__(self, capacities: Mapping[str, int]) -> None:
|
|
if not capacities or any(type(n) is not int or n < 0 for n in capacities.values()):
|
|
raise ValueError("explicit nonnegative memory ceilings are required")
|
|
self.capacities = MappingProxyType(dict(capacities))
|
|
self._active: dict[str, Reservation] = {}
|
|
self._quarantined: set[str] = set()
|
|
self._lock = threading.RLock()
|
|
|
|
def reserve(self, owner: str, demands: Mapping[str, int]) -> Reservation:
|
|
with self._lock:
|
|
if not owner or owner in self._active:
|
|
raise ValueError(f"reservation owner missing or already active: {owner}")
|
|
if not demands or any(type(n) is not int or n < 0 for n in demands.values()):
|
|
raise ValueError("nonnegative memory demands are required")
|
|
for pool, count in demands.items():
|
|
if pool not in self.capacities:
|
|
raise ResourceUnavailable(f"unknown memory pool: {pool}")
|
|
used = sum(r.demands.get(pool, 0) for r in self._active.values())
|
|
if used + count > self.capacities[pool]:
|
|
raise ResourceUnavailable(f"{pool}: {used} reserved + {count} requested > {self.capacities[pool]}")
|
|
reservation = Reservation(owner, MappingProxyType(dict(demands)))
|
|
self._active[owner] = reservation
|
|
return reservation
|
|
|
|
def release(self, reservation: Reservation, *, drained: bool) -> bool:
|
|
with self._lock:
|
|
if self._active.get(reservation.owner) is not reservation:
|
|
return False # stale token must never release a newer allocation
|
|
if not drained:
|
|
self._quarantined.add(reservation.owner)
|
|
return False
|
|
del self._active[reservation.owner]
|
|
self._quarantined.discard(reservation.owner)
|
|
return True
|
|
|
|
def snapshot(self) -> dict:
|
|
with self._lock:
|
|
return {
|
|
"capacities": dict(self.capacities),
|
|
"reserved": {p: sum(r.demands.get(p, 0) for r in self._active.values()) for p in self.capacities},
|
|
"owners": list(self._active),
|
|
"quarantined": sorted(self._quarantined),
|
|
}
|
|
|
|
|
|
class GraphRequestGate:
|
|
"""One admitted pipeline request, held through the final consumer ACK.
|
|
|
|
Stage reservations already include its payload allowance. Rejecting excess
|
|
work before the ingress queue keeps slow consumers from accumulating inputs
|
|
or error outputs, and leaves the existing control queue available to abort.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._lock = threading.Lock()
|
|
self._active = None
|
|
|
|
def acquire(self, request_id):
|
|
with self._lock:
|
|
if self._active is not None:
|
|
raise ResourceUnavailable("graph pipeline capacity is one; consume or cancel the active request first")
|
|
ticket = (request_id, object())
|
|
self._active = ticket
|
|
return ticket
|
|
|
|
def current(self, request_id):
|
|
with self._lock:
|
|
return self._active if self._active is not None and self._active[0] == request_id else None
|
|
|
|
def release(self, ticket):
|
|
with self._lock:
|
|
if ticket is not None and self._active is ticket:
|
|
self._active = None
|