462 lines
17 KiB
Python
462 lines
17 KiB
Python
"""Revocable, bounded runtime evidence owned by an actual dequeue attempt."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from contextvars import ContextVar
|
|
from dataclasses import dataclass, replace
|
|
import threading
|
|
from typing import Callable, TYPE_CHECKING
|
|
from uuid import uuid4
|
|
import weakref
|
|
|
|
from .execution_context_budget import ContextBudget, PayloadLease, SealedContext
|
|
from .execution_context_graph import (
|
|
GraphProjection,
|
|
GraphRole,
|
|
SourceRef,
|
|
project_contexts,
|
|
)
|
|
from .execution_contexts import (
|
|
ContextBundle,
|
|
PromptSnapshot,
|
|
RecipeContext,
|
|
TextEvidence,
|
|
bundle_payload_size,
|
|
text_size,
|
|
MAX_ATTEMPT_BYTES,
|
|
)
|
|
from .execution_history_events import RunStart, encode_command
|
|
|
|
if TYPE_CHECKING:
|
|
from .execution_history_writer import HistoryWriterStatus
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AttemptCapability:
|
|
nonce: str
|
|
run_id: str
|
|
session_id: str
|
|
queue_item_id: int
|
|
job_id: str
|
|
storage_epoch: str
|
|
capture_generation: int
|
|
sources: tuple[SourceRef, ...]
|
|
|
|
|
|
@dataclass
|
|
class _Attempt:
|
|
capability: AttemptCapability
|
|
projection: GraphProjection
|
|
lease: PayloadLease | None
|
|
samples: dict
|
|
retained_bytes: int
|
|
reason: str = ""
|
|
executor_identity: int | None = None
|
|
executor_rejected: bool = False
|
|
|
|
|
|
_binding: ContextVar[tuple | None] = ContextVar("prompt_manager_attempt", default=None)
|
|
|
|
|
|
def active_capture() -> tuple[ContextCapture, AttemptCapability] | None:
|
|
binding = _binding.get()
|
|
if binding is not None:
|
|
manager = binding[0]()
|
|
if manager is not None and manager.valid(binding[1]):
|
|
return manager, binding[1]
|
|
return None
|
|
|
|
|
|
def unbind_current() -> None:
|
|
binding = _binding.get()
|
|
_binding.set(None)
|
|
if binding is not None:
|
|
manager = binding[0]()
|
|
if manager is not None:
|
|
manager.revoke(binding[1])
|
|
|
|
|
|
def _unknown(reason: str = "missing") -> PromptSnapshot:
|
|
field = TextEvidence(None, "unknown", "unknown", reason)
|
|
return PromptSnapshot(field, field, field, field)
|
|
|
|
|
|
def _snapshot_bytes(snapshot: PromptSnapshot) -> int:
|
|
# Reuse the scalar validator, including provenance and per-field bounds.
|
|
bundle = ContextBundle(
|
|
(RecipeContext("x", snapshot, None, "inferred"),), "partial", "partial"
|
|
)
|
|
return bundle_payload_size(bundle)
|
|
|
|
|
|
def _bounded_snapshot(snapshot: PromptSnapshot) -> tuple[PromptSnapshot, bool]:
|
|
if type(snapshot) is not PromptSnapshot:
|
|
raise ValueError("invalid_snapshot")
|
|
fields, limited = [], False
|
|
for field in (snapshot.raw, snapshot.effective, snapshot.prepend, snapshot.append):
|
|
if type(field) is not TextEvidence:
|
|
raise ValueError("invalid_field")
|
|
if field.value is not None:
|
|
try:
|
|
text_size(field.value)
|
|
except ValueError:
|
|
field = _unknown("limit_reached").raw
|
|
limited = True
|
|
fields.append(field)
|
|
if fields[1].value is not None and fields[1].provenance != "recorded":
|
|
fields[1] = _unknown("missing").effective
|
|
bounded = PromptSnapshot(*fields)
|
|
_snapshot_bytes(bounded)
|
|
return bounded, limited
|
|
|
|
|
|
class ContextCapture:
|
|
def __init__(
|
|
self,
|
|
queue_identity: int,
|
|
status_reader: Callable[[], HistoryWriterStatus],
|
|
budget: ContextBudget,
|
|
) -> None:
|
|
self.queue_identity = queue_identity
|
|
self._status_reader = status_reader
|
|
self._budget = budget
|
|
self._lock = threading.RLock()
|
|
self._attempts: dict[str, _Attempt] = {}
|
|
|
|
def _current(self, capability: AttemptCapability) -> _Attempt | None:
|
|
attempt = self._attempts.get(capability.nonce)
|
|
if attempt is None or attempt.capability is not capability:
|
|
return None
|
|
status = self._status_reader()
|
|
if (
|
|
not status.accepting
|
|
or status.epoch != capability.storage_epoch
|
|
or status.capture_generation != capability.capture_generation
|
|
):
|
|
self._drop(capability)
|
|
return None
|
|
return attempt
|
|
|
|
def valid(self, capability: AttemptCapability) -> bool:
|
|
with self._lock:
|
|
return self._current(capability) is not None
|
|
|
|
def begin(
|
|
self, start: RunStart, generation: int, graph: object, outputs: object
|
|
) -> AttemptCapability:
|
|
unbind_current()
|
|
encode_command(start)
|
|
with self._lock:
|
|
cap = AttemptCapability(
|
|
str(uuid4()),
|
|
start.run_id,
|
|
start.session_id,
|
|
start.queue_item_id,
|
|
start.job_id,
|
|
start.storage_epoch,
|
|
generation,
|
|
(),
|
|
)
|
|
if len(self._attempts) >= 128:
|
|
return cap
|
|
lease = self._budget.reserve(MAX_ATTEMPT_BYTES)
|
|
projection = GraphProjection((), (), "")
|
|
reason = "budget_exhausted" if lease is None else ""
|
|
retained = 0
|
|
if lease is not None:
|
|
try:
|
|
projection = project_contexts(graph, outputs)
|
|
retained = (
|
|
sum(
|
|
_snapshot_bytes(role.inferred) if role.inferred else 0
|
|
for recipe in projection.recipes
|
|
for role in (recipe.positive, recipe.negative)
|
|
)
|
|
+ sum(
|
|
_snapshot_bytes(snapshot)
|
|
for _, snapshot in projection.source_snapshots
|
|
)
|
|
+ sum(
|
|
text_size(s.node_id, 512) + text_size(s.node_type, 128) + 32
|
|
for s in projection.sources
|
|
)
|
|
)
|
|
if retained > MAX_ATTEMPT_BYTES:
|
|
raise ValueError("projection_limit")
|
|
except ValueError:
|
|
projection = GraphProjection((), (), "")
|
|
reason = "graph_limit"
|
|
retained = 0
|
|
except Exception:
|
|
projection = GraphProjection((), (), "")
|
|
reason = "invalid_evidence"
|
|
retained = 0
|
|
lease.release()
|
|
lease = None
|
|
cap = replace(cap, sources=projection.sources)
|
|
self._attempts[cap.nonce] = _Attempt(
|
|
cap, projection, lease, {}, retained, reason
|
|
)
|
|
_binding.set((weakref.ref(self), cap))
|
|
self._current(cap)
|
|
return cap
|
|
|
|
def bind_executor(self, capability: AttemptCapability, identity: int) -> None:
|
|
"""Bind one scalar host identity; recursive starts invalidate cache proof."""
|
|
with self._lock:
|
|
attempt = self._current(capability)
|
|
if attempt is not None:
|
|
if attempt.executor_identity is not None:
|
|
attempt.executor_rejected = True
|
|
else:
|
|
attempt.executor_identity = identity
|
|
|
|
def matches_executor(self, capability: AttemptCapability, identity: int) -> bool:
|
|
with self._lock:
|
|
attempt = self._current(capability)
|
|
return bool(
|
|
attempt is not None
|
|
and not attempt.executor_rejected
|
|
and attempt.executor_identity == identity
|
|
)
|
|
|
|
def source_snapshot(
|
|
self, capability: AttemptCapability, source: SourceRef
|
|
) -> PromptSnapshot | None:
|
|
with self._lock:
|
|
attempt = self._current(capability)
|
|
if attempt is not None:
|
|
for ref, snapshot in attempt.projection.source_snapshots:
|
|
if ref == source:
|
|
return snapshot
|
|
return None
|
|
|
|
def observe(
|
|
self,
|
|
capability: AttemptCapability,
|
|
source: SourceRef,
|
|
invocation_index: int,
|
|
snapshot: PromptSnapshot,
|
|
) -> None:
|
|
try:
|
|
with self._lock:
|
|
attempt = self._current(capability)
|
|
if (
|
|
attempt is None
|
|
or attempt.lease is None
|
|
or source not in capability.sources
|
|
or type(source) is not SourceRef
|
|
or type(invocation_index) is not int
|
|
or not 0 <= invocation_index <= 2**63 - 1
|
|
):
|
|
return
|
|
snapshot, limited = _bounded_snapshot(snapshot)
|
|
size = _snapshot_bytes(snapshot)
|
|
if (
|
|
limited
|
|
or any(
|
|
field.reason == "limit_reached"
|
|
for field in (
|
|
snapshot.raw,
|
|
snapshot.effective,
|
|
snapshot.prepend,
|
|
snapshot.append,
|
|
)
|
|
)
|
|
) and attempt.reason not in ("multiple_samples", "limit_reached"):
|
|
attempt.reason = "field_limit"
|
|
key = source, invocation_index
|
|
if key in attempt.samples:
|
|
if attempt.samples[key] != snapshot:
|
|
attempt.reason = "multiple_samples"
|
|
return
|
|
if len(attempt.samples) >= 64:
|
|
attempt.reason = "multiple_samples"
|
|
return
|
|
if attempt.retained_bytes + size > MAX_ATTEMPT_BYTES:
|
|
attempt.reason = "limit_reached"
|
|
return
|
|
attempt.samples[key] = snapshot
|
|
attempt.retained_bytes += size
|
|
except (ValueError, AttributeError, TypeError):
|
|
return
|
|
|
|
@staticmethod
|
|
def _discard_payload(attempt: _Attempt) -> None:
|
|
attempt.projection = GraphProjection((), (), "")
|
|
attempt.samples.clear()
|
|
attempt.retained_bytes = 0
|
|
|
|
def _drop(self, capability: AttemptCapability) -> _Attempt | None:
|
|
attempt = self._attempts.get(capability.nonce)
|
|
if attempt is not None and attempt.capability is capability:
|
|
self._attempts.pop(capability.nonce)
|
|
self._discard_payload(attempt)
|
|
if attempt.lease is not None:
|
|
attempt.lease.release()
|
|
return attempt
|
|
return None
|
|
|
|
def revoke(self, capability: AttemptCapability) -> None:
|
|
with self._lock:
|
|
self._drop(capability)
|
|
|
|
def revoke_all(self) -> None:
|
|
with self._lock:
|
|
for attempt in self._attempts.values():
|
|
self._discard_payload(attempt)
|
|
if attempt.lease is not None:
|
|
attempt.lease.release()
|
|
self._attempts.clear()
|
|
|
|
def seal(self, capability: AttemptCapability) -> SealedContext:
|
|
with self._lock:
|
|
attempt = self._current(capability)
|
|
if attempt is None:
|
|
return SealedContext(None, None, "unavailable", "unavailable")
|
|
self._attempts.pop(capability.nonce)
|
|
lease = attempt.lease
|
|
if lease is None:
|
|
return SealedContext(None, None, "unavailable", attempt.reason)
|
|
try:
|
|
ambiguous = attempt.reason == "multiple_samples"
|
|
|
|
def role_snapshot(role: GraphRole) -> PromptSnapshot:
|
|
nonlocal ambiguous
|
|
if role.reason == "source-transformed":
|
|
return _unknown("unsupported")
|
|
if not role.sources:
|
|
return role.inferred or _unknown("unsupported")
|
|
values = [
|
|
snapshot
|
|
for (source, _), snapshot in attempt.samples.items()
|
|
if source in role.sources
|
|
]
|
|
if len(values) > 1 or attempt.reason == "multiple_samples":
|
|
ambiguous = True
|
|
return _unknown("ambiguous")
|
|
if len(role.sources) == 1 and len(values) == 1:
|
|
return values[0]
|
|
inferred = role.inferred
|
|
if len(role.sources) == 1:
|
|
snapshots = attempt.projection.source_snapshots
|
|
inferred = next(
|
|
(
|
|
snapshot
|
|
for source, snapshot in snapshots
|
|
if source == role.sources[0]
|
|
),
|
|
inferred,
|
|
)
|
|
if inferred is not None:
|
|
s = inferred
|
|
return PromptSnapshot(
|
|
s.raw, _unknown().effective, s.prepend, s.append
|
|
)
|
|
return _unknown()
|
|
|
|
contexts = tuple(
|
|
RecipeContext(
|
|
r.anchor_id,
|
|
role_snapshot(r.positive),
|
|
role_snapshot(r.negative),
|
|
"inferred",
|
|
)
|
|
for r in attempt.projection.recipes
|
|
)
|
|
supported = any(
|
|
role.sources or role.inferred is not None
|
|
for r in attempt.projection.recipes
|
|
for role in (r.positive, r.negative)
|
|
)
|
|
complete = bool(contexts) and all(
|
|
field.value is not None
|
|
for r in contexts
|
|
for s in (r.positive, r.negative)
|
|
for field in (s.raw, s.effective, s.prepend, s.append)
|
|
)
|
|
state = (
|
|
"ambiguous" if ambiguous else "complete" if complete else "partial"
|
|
)
|
|
reason = (
|
|
"multiple_samples"
|
|
if ambiguous
|
|
else "captured" if complete else "partial"
|
|
)
|
|
if attempt.reason == "field_limit" and not ambiguous:
|
|
state, reason = "partial", "field_limit"
|
|
if not supported:
|
|
contexts, state, reason = (), "unavailable", "no_supported_context"
|
|
if attempt.reason in ("limit_reached", "graph_limit"):
|
|
contexts, state, reason = (), "unavailable", attempt.reason
|
|
bundle = ContextBundle(contexts, state, reason)
|
|
size = bundle_payload_size(bundle)
|
|
if size + attempt.retained_bytes > MAX_ATTEMPT_BYTES:
|
|
raise ValueError("assembly_limit")
|
|
self._discard_payload(attempt)
|
|
lease.shrink_to(size)
|
|
return SealedContext(
|
|
bundle, lease, state, reason, capability.capture_generation
|
|
)
|
|
except ValueError:
|
|
self._discard_payload(attempt)
|
|
lease.release()
|
|
return SealedContext(None, None, "unavailable", "limit_reached")
|
|
except Exception:
|
|
self._discard_payload(attempt)
|
|
lease.release()
|
|
return SealedContext(None, None, "unavailable", "invalid_evidence")
|
|
|
|
|
|
def record_runtime_context(
|
|
node_type: str,
|
|
raw: str,
|
|
prepend: str,
|
|
append: str,
|
|
effective: str,
|
|
*,
|
|
role: str | None = None,
|
|
) -> None:
|
|
"""Record only a successful node return in a trusted, bound host invocation."""
|
|
try:
|
|
binding = active_capture()
|
|
if binding is None:
|
|
return
|
|
# Optional host import is independent of our package namespace.
|
|
from comfy_execution.utils import get_executing_context
|
|
|
|
host = get_executing_context()
|
|
manager, cap = binding
|
|
if (
|
|
host is None
|
|
or type(host.prompt_id) is not str
|
|
or host.prompt_id != cap.job_id
|
|
):
|
|
return
|
|
if type(host.node_id) is not str or type(host.list_index) is not int:
|
|
return
|
|
if node_type == "PromptManagerPair":
|
|
slots = (
|
|
(0, 2) if role == "positive" else (1, 3) if role == "negative" else ()
|
|
)
|
|
else:
|
|
slots = (
|
|
(0,)
|
|
if node_type == "PromptManagerText"
|
|
else (0, 1) if role is None else ()
|
|
)
|
|
for slot in slots:
|
|
source = SourceRef(host.node_id, node_type, slot)
|
|
if source not in cap.sources:
|
|
continue
|
|
snapshot = PromptSnapshot(
|
|
*(
|
|
TextEvidence(value, "recorded", "runtime_return", None)
|
|
for value in (raw, effective, prepend, append)
|
|
)
|
|
)
|
|
manager.observe(cap, source, host.list_index, snapshot)
|
|
except Exception:
|
|
# Capture must never change the host result or expose input/exception text.
|
|
return
|