Vito
4ef236aaf7
* feat(history): define bounded output reference contracts * feat(database): migrate execution output evidence to schema v3 * feat(history): persist immutable run output references * test(history): isolate mocked node gallery monitor * feat(history): capture bounded output descriptors with shared leases * fix(history): authenticate payload owners and release failed captures * feat(history): attach output evidence to exact finalized attempts * test(history): verify run output evidence on real ComfyUI * fix(history): release idle worker payload references * docs(history): clarify defensive lease cleanup * test(history): align output fixtures with native path styles * test(history): separate restore completion from writer drain * test(history): bound batch drains by progress and total time
130 lines
4.1 KiB
Python
130 lines
4.1 KiB
Python
"""Shared payload reservations with explicit, single-owner transport leases."""
|
|
|
|
from dataclasses import dataclass
|
|
import json
|
|
import threading
|
|
|
|
from .execution_contexts import BUNDLE_STATES, CAPTURE_REASONS, ContextBundle
|
|
|
|
|
|
def _byte_count(value: int) -> None:
|
|
if type(value) is not int or value < 0:
|
|
raise ValueError("Byte count must be a nonnegative integer")
|
|
|
|
|
|
class ContextBudget:
|
|
def __init__(self, limit_bytes: int = 8388608) -> None:
|
|
_byte_count(limit_bytes)
|
|
self._limit = limit_bytes
|
|
self._charged = 0
|
|
self._owners: set[PayloadLease] = set()
|
|
self._lock = threading.Lock()
|
|
|
|
@property
|
|
def charged_bytes(self) -> int:
|
|
with self._lock:
|
|
return self._charged
|
|
|
|
def reserve(self, byte_count: int) -> "PayloadLease | None":
|
|
_byte_count(byte_count)
|
|
with self._lock:
|
|
if byte_count > self._limit - self._charged:
|
|
return None
|
|
owner = PayloadLease._new(self, byte_count)
|
|
self._owners.add(owner)
|
|
self._charged += byte_count
|
|
return owner
|
|
|
|
def reserve_up_to(self, max_bytes: int) -> "PayloadLease | None":
|
|
"""Atomically own up to the remaining capacity in this same pool."""
|
|
_byte_count(max_bytes)
|
|
with self._lock:
|
|
byte_count = min(max_bytes, self._limit - self._charged)
|
|
if not byte_count:
|
|
return None
|
|
owner = PayloadLease._new(self, byte_count)
|
|
self._owners.add(owner)
|
|
self._charged += byte_count
|
|
return owner
|
|
|
|
|
|
class PayloadLease:
|
|
def __init__(self, pool: ContextBudget, byte_count: int) -> None:
|
|
raise ValueError("Reserve payload leases through ContextBudget")
|
|
|
|
@classmethod
|
|
def _new(cls, pool: ContextBudget, byte_count: int) -> "PayloadLease":
|
|
owner = object.__new__(cls)
|
|
owner._pool = pool
|
|
owner._bytes = byte_count
|
|
owner._active = True
|
|
return owner
|
|
|
|
@property
|
|
def byte_count(self) -> int:
|
|
with self._pool._lock:
|
|
self._require_owner()
|
|
return self._bytes
|
|
|
|
def _require_owner(self) -> None:
|
|
if self not in self._pool._owners or not self._active:
|
|
raise ValueError("Inactive or unissued payload lease")
|
|
|
|
def transfer(self) -> "PayloadLease":
|
|
with self._pool._lock:
|
|
self._require_owner()
|
|
owner = PayloadLease._new(self._pool, self._bytes)
|
|
self._pool._owners.add(owner)
|
|
self._pool._owners.remove(self)
|
|
self._active = False
|
|
return owner
|
|
|
|
def shrink_to(self, byte_count: int) -> None:
|
|
_byte_count(byte_count)
|
|
with self._pool._lock:
|
|
self._require_owner()
|
|
if byte_count > self._bytes:
|
|
raise ValueError("Cannot grow payload lease")
|
|
self._pool._charged -= self._bytes - byte_count
|
|
self._bytes = byte_count
|
|
|
|
def release(self) -> None:
|
|
with self._pool._lock:
|
|
if self in self._pool._owners:
|
|
self._pool._owners.remove(self)
|
|
self._pool._charged -= self._bytes
|
|
self._active = False
|
|
|
|
|
|
SHARED_CONTEXT_BUDGET = ContextBudget()
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SealedContext:
|
|
bundle: ContextBundle | None
|
|
lease: PayloadLease | None
|
|
state: str
|
|
reason: str
|
|
capture_generation: int | None = None
|
|
|
|
def validate_header(self) -> None:
|
|
if (
|
|
type(self.state) is not str
|
|
or self.state not in BUNDLE_STATES
|
|
or type(self.reason) is not str
|
|
or self.reason not in CAPTURE_REASONS
|
|
or (
|
|
self.capture_generation is not None
|
|
and (
|
|
type(self.capture_generation) is not int
|
|
or not 0 <= self.capture_generation <= 2**63 - 1
|
|
)
|
|
)
|
|
):
|
|
raise ValueError("Invalid capture status")
|
|
header = json.dumps(
|
|
[self.state, self.reason, self.capture_generation], separators=(",", ":")
|
|
).encode("utf-8")
|
|
if len(header) > 256:
|
|
raise ValueError("Oversized capture status")
|