feat(history): capture cached prompt strings without mutating host caches
This commit is contained in:
Vito Sansevero
@@ -244,3 +244,22 @@ class ProjectionTests(unittest.TestCase):
|
||||
g[str(i)] = node("KSampler", positive=["p", 0], negative=["p", 0])
|
||||
result = self.m.project_contexts(g, None)
|
||||
self.assertEqual(result.recipes, ())
|
||||
|
||||
def test_partial_source_snapshots_keep_literal_fields_without_stateful_substitution(
|
||||
self,
|
||||
):
|
||||
for upstream in ("RuntimeGenerator", "PromptManagerText"):
|
||||
g = graph()
|
||||
g["dynamic"] = node(upstream, text="looks literal")
|
||||
g["p"]["inputs"]["text"] = ["dynamic", 0]
|
||||
projection = self.m.project_contexts(g, None)
|
||||
positive = next(
|
||||
snapshot
|
||||
for ref, snapshot in projection.source_snapshots
|
||||
if ref.output_slot == 0
|
||||
)
|
||||
self.assertIsNone(positive.raw.value)
|
||||
self.assertIsNone(positive.effective.value)
|
||||
self.assertEqual(positive.prepend.value, "pre")
|
||||
self.assertEqual(positive.append.value, "post")
|
||||
self.assertEqual(positive.append.provenance, "inferred")
|
||||
|
||||
@@ -0,0 +1,397 @@
|
||||
"""Read-only cached STRING evidence and cooperative executor notifications."""
|
||||
|
||||
from collections import namedtuple
|
||||
from contextvars import Context
|
||||
from dataclasses import replace
|
||||
import gc
|
||||
import sys
|
||||
import weakref
|
||||
import threading
|
||||
import types
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
from utils.execution_context_budget import ContextBudget
|
||||
from utils.execution_context_capture import ContextCapture
|
||||
from utils.execution_history_events import RunStart
|
||||
from utils.execution_history_writer import HistoryWriterStatus
|
||||
|
||||
try:
|
||||
from utils import execution_context_host as host
|
||||
except ImportError:
|
||||
host = None
|
||||
|
||||
SESSION = "11111111-1111-4111-8111-111111111111"
|
||||
EPOCH = "22222222-2222-4222-8222-222222222222"
|
||||
JOB = "33333333-3333-4333-8333-333333333333"
|
||||
Entry = namedtuple("CacheEntry", "ui outputs")
|
||||
|
||||
|
||||
class HostTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.assertIsNotNone(host, "read-only host adapter is missing")
|
||||
self.queue = object()
|
||||
self.pool = ContextBudget()
|
||||
self.status = HistoryWriterStatus(EPOCH, True, 0, 0, True, 3)
|
||||
self.capture = ContextCapture(id(self.queue), lambda: self.status, self.pool)
|
||||
self.addCleanup(self.capture.revoke_all)
|
||||
self.calls = []
|
||||
calls = self.calls
|
||||
|
||||
class BasicCache:
|
||||
def __init__(self):
|
||||
self.entries = {}
|
||||
|
||||
def get_local(self, node_id):
|
||||
calls.append(("local", node_id))
|
||||
return self.entries.get(node_id)
|
||||
|
||||
async def get(self, node_id):
|
||||
raise AssertionError("async get must never be called")
|
||||
|
||||
class HierarchicalCache(BasicCache):
|
||||
pass
|
||||
|
||||
class LRUCache(BasicCache):
|
||||
pass
|
||||
|
||||
class RAMPressureCache(LRUCache):
|
||||
pass
|
||||
|
||||
class NullCache:
|
||||
def get_local(self, node_id):
|
||||
raise AssertionError("null cache is not usable")
|
||||
|
||||
class PromptExecutor:
|
||||
def add_message(self, event, data, broadcast):
|
||||
calls.append(("original", event, data, broadcast))
|
||||
if event == "raise":
|
||||
raise RuntimeError("original")
|
||||
return "original-return"
|
||||
|
||||
self.execution = types.ModuleType("execution")
|
||||
self.execution.PromptExecutor = PromptExecutor
|
||||
self.execution.CacheEntry = Entry
|
||||
self.caching = types.ModuleType("comfy_execution.caching")
|
||||
for cls in (
|
||||
BasicCache,
|
||||
HierarchicalCache,
|
||||
LRUCache,
|
||||
RAMPressureCache,
|
||||
NullCache,
|
||||
):
|
||||
setattr(self.caching, cls.__name__, cls)
|
||||
package = types.ModuleType("comfy_execution")
|
||||
package.caching = self.caching
|
||||
self.patch = patch.dict(
|
||||
sys.modules,
|
||||
{
|
||||
"comfy_execution": package,
|
||||
"execution": self.execution,
|
||||
"comfy_execution.caching": self.caching,
|
||||
},
|
||||
)
|
||||
self.patch.start()
|
||||
self.addCleanup(self.patch.stop)
|
||||
self.executor = PromptExecutor()
|
||||
self.executor.server = types.SimpleNamespace(prompt_queue=self.queue)
|
||||
self.executor.caches = types.SimpleNamespace(outputs=HierarchicalCache())
|
||||
|
||||
def begin(
|
||||
self, kind="PromptManager", raw=" raw ", prepend=" pre ", append=" post "
|
||||
):
|
||||
g = {
|
||||
"p": {
|
||||
"class_type": kind,
|
||||
"inputs": {
|
||||
"text": raw,
|
||||
"negative_text": "negative",
|
||||
"prepend_text": prepend,
|
||||
"append_text": append,
|
||||
},
|
||||
},
|
||||
"clip": {"class_type": "CLIPTextEncode", "inputs": {"text": ["p", 0]}},
|
||||
"s": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {
|
||||
"positive": (
|
||||
["clip", 0] if kind == "PromptManagerText" else ["p", 0]
|
||||
),
|
||||
"negative": (
|
||||
["p", 1] if kind == "PromptManagerPair" else ["clip", 0]
|
||||
),
|
||||
},
|
||||
},
|
||||
}
|
||||
if type(raw) is list or type(prepend) is list:
|
||||
g["dynamic"] = {"class_type": "RuntimeGenerator", "inputs": {}}
|
||||
g["stateful"] = {
|
||||
"class_type": "PromptManagerText",
|
||||
"inputs": {"text": "looks literal"},
|
||||
}
|
||||
self.cap = self.capture.begin(
|
||||
RunStart(SESSION, SESSION, 1, JOB, EPOCH, 0), 3, g, None
|
||||
)
|
||||
return self.cap
|
||||
|
||||
def observe(self, event, executor=None, **data):
|
||||
host.observe_cached_context(
|
||||
executor or self.executor, event, {"prompt_id": JOB, **data}
|
||||
)
|
||||
|
||||
def seal(self):
|
||||
sealed = self.capture.seal(self.cap)
|
||||
if sealed.lease:
|
||||
self.addCleanup(sealed.lease.release)
|
||||
return sealed.bundle.contexts[0]
|
||||
|
||||
def test_original_return_exception_and_wrapper_cleanup_preserved(self):
|
||||
cls = self.execution.PromptExecutor
|
||||
original = cls.add_message
|
||||
|
||||
def earlier(executor, *args, **kwargs):
|
||||
self.calls.append("earlier")
|
||||
return original(executor, *args, **kwargs)
|
||||
|
||||
cls.add_message = earlier
|
||||
seen = []
|
||||
stop_a = host.subscribe_executor_messages("a", lambda *args: seen.append(args))
|
||||
stop_b = host.subscribe_executor_messages("b", lambda *args: seen.append(args))
|
||||
wrapped = cls.add_message
|
||||
|
||||
def later(executor, *args, **kwargs):
|
||||
self.calls.append("later")
|
||||
return wrapped(executor, *args, **kwargs)
|
||||
|
||||
cls.add_message = later
|
||||
data = {"prompt_id": JOB}
|
||||
self.assertEqual(
|
||||
self.executor.add_message("ok", data, broadcast=False), "original-return"
|
||||
)
|
||||
self.assertEqual(len(seen), 2)
|
||||
self.assertIs(seen[0][0], self.executor)
|
||||
self.assertIs(seen[0][1][1], data)
|
||||
with self.assertRaisesRegex(RuntimeError, "original"):
|
||||
self.executor.add_message("raise", data, False)
|
||||
self.assertEqual(len(seen), 2)
|
||||
stop_a()
|
||||
self.executor.add_message("ok", data, False)
|
||||
self.assertEqual(len(seen), 3)
|
||||
stop_b()
|
||||
self.assertIs(cls.add_message, later)
|
||||
self.executor.add_message("ok", data, False)
|
||||
self.assertEqual(len(seen), 3)
|
||||
stop_c = host.subscribe_executor_messages("c", lambda *args: seen.append(args))
|
||||
self.executor.add_message("ok", data, False)
|
||||
self.assertEqual(len(seen), 4)
|
||||
stop_c()
|
||||
|
||||
def test_cache_reads_only_local_string_slots(self):
|
||||
for mode in ("HierarchicalCache", "LRUCache", "RAMPressureCache"):
|
||||
for kind, values in (
|
||||
("PromptManager", [object(), ["cached"]]),
|
||||
("PromptManagerFocused", [object(), ["cached"]]),
|
||||
("PromptManagerSplit", [object(), ["cached"]]),
|
||||
("PromptManagerText", [["cached"]]),
|
||||
("PromptManagerPair", [object(), object(), ["cached"], [""]]),
|
||||
):
|
||||
with self.subTest(mode=mode, kind=kind):
|
||||
self.begin(kind)
|
||||
cache = getattr(self.caching, mode)()
|
||||
entry = Entry({"secret": object()}, values)
|
||||
cache.entries["p"] = entry
|
||||
self.executor.caches.outputs = cache
|
||||
self.observe("execution_start")
|
||||
self.observe("execution_cached", nodes=["p"])
|
||||
result = self.seal()
|
||||
self.assertEqual(result.positive.effective.value, "cached")
|
||||
self.assertEqual(
|
||||
result.positive.effective.source, "host_cache_output"
|
||||
)
|
||||
if kind == "PromptManagerPair":
|
||||
self.assertEqual(result.negative.effective.value, "")
|
||||
self.assertIs(cache.entries["p"], entry)
|
||||
self.assertEqual(set(cache.__dict__), {"entries"})
|
||||
self.assertTrue(all(call[0] == "local" for call in self.calls))
|
||||
|
||||
def test_changed_trigger_configuration_keeps_cached_effective_text(self):
|
||||
self.begin()
|
||||
self.executor.caches.outputs.entries["p"] = Entry(
|
||||
None, [object(), ["old injected text"]]
|
||||
)
|
||||
self.observe("execution_start")
|
||||
self.observe("execution_cached", nodes=["p"])
|
||||
result = self.seal().positive
|
||||
self.assertEqual(result.effective.value, "old injected text")
|
||||
self.assertEqual(result.raw.value, " raw ")
|
||||
self.assertEqual(result.prepend.value, " pre ")
|
||||
self.assertEqual(result.append.value, " post ")
|
||||
self.assertEqual(result.raw.provenance, "inferred")
|
||||
|
||||
def test_dynamic_cached_raw_is_unknown(self):
|
||||
for dynamic in (["dynamic", 0], ["stateful", 0]):
|
||||
for field in ("raw", "prepend"):
|
||||
with self.subTest(dynamic=dynamic, field=field):
|
||||
self.begin(**{field: dynamic})
|
||||
self.executor.caches.outputs.entries["p"] = Entry(
|
||||
None, [object(), ["cached"]]
|
||||
)
|
||||
self.observe("execution_start")
|
||||
self.observe("execution_cached", nodes=["p"])
|
||||
result = self.seal().positive
|
||||
self.assertEqual(result.effective.value, "cached")
|
||||
self.assertIsNone(getattr(result, field).value)
|
||||
self.assertEqual(result.append.value, " post ")
|
||||
if field == "raw":
|
||||
self.assertEqual(result.prepend.value, " pre ")
|
||||
else:
|
||||
self.assertEqual(result.raw.value, " raw ")
|
||||
|
||||
def test_provider_local_reset_eviction_and_unknown_modes_fail_closed(self):
|
||||
for mode in (
|
||||
"NullCache",
|
||||
"unknown",
|
||||
"reset",
|
||||
"evicted",
|
||||
"malformed",
|
||||
"provider-local",
|
||||
):
|
||||
with self.subTest(mode=mode):
|
||||
self.begin()
|
||||
cache = self.caching.HierarchicalCache()
|
||||
if mode == "NullCache":
|
||||
cache = self.caching.NullCache()
|
||||
elif mode == "unknown":
|
||||
cache = types.SimpleNamespace(
|
||||
get_local=lambda _: self.fail("unknown cache read")
|
||||
)
|
||||
elif mode == "malformed":
|
||||
cache.entries["p"] = Entry(None, [object(), [object()]])
|
||||
elif mode in ("reset", "provider-local"):
|
||||
cache.entries["p"] = Entry(
|
||||
None, [object(), ["resident provider output"]]
|
||||
)
|
||||
self.executor.caches.outputs = cache
|
||||
self.observe("execution_start")
|
||||
if mode == "reset":
|
||||
self.executor.caches.outputs = self.caching.HierarchicalCache()
|
||||
self.observe("execution_cached", nodes=["p"])
|
||||
result = self.seal().positive
|
||||
self.assertEqual(
|
||||
result.effective.value,
|
||||
"resident provider output" if mode == "provider-local" else None,
|
||||
)
|
||||
|
||||
def test_worker_queue_job_executor_and_generation_fences(self):
|
||||
for failure in (
|
||||
"unbound",
|
||||
"thread",
|
||||
"queue",
|
||||
"job",
|
||||
"executor",
|
||||
"recursive",
|
||||
"generation",
|
||||
"no-start",
|
||||
):
|
||||
with self.subTest(failure=failure):
|
||||
self.begin()
|
||||
self.executor.server.prompt_queue = self.queue
|
||||
self.executor.caches.outputs.entries["p"] = Entry(
|
||||
None, [object(), ["cached"]]
|
||||
)
|
||||
if failure != "no-start":
|
||||
self.observe("execution_start")
|
||||
target = self.executor
|
||||
if failure == "queue":
|
||||
self.executor.server.prompt_queue = object()
|
||||
if failure in ("executor", "recursive"):
|
||||
target = self.execution.PromptExecutor()
|
||||
target.server = self.executor.server
|
||||
target.caches = self.executor.caches
|
||||
if failure == "recursive":
|
||||
self.observe("execution_start", target)
|
||||
target = self.executor
|
||||
data = {
|
||||
"prompt_id": "wrong" if failure == "job" else JOB,
|
||||
"nodes": ["p"],
|
||||
}
|
||||
if failure == "generation":
|
||||
self.status = replace(self.status, capture_generation=4)
|
||||
|
||||
def action():
|
||||
host.observe_cached_context(target, "execution_cached", data)
|
||||
|
||||
if failure == "unbound":
|
||||
Context().run(action)
|
||||
elif failure == "thread":
|
||||
thread = threading.Thread(target=action)
|
||||
thread.start()
|
||||
thread.join(2)
|
||||
else:
|
||||
action()
|
||||
if failure == "generation":
|
||||
self.assertIsNone(self.capture.seal(self.cap).bundle)
|
||||
self.status = replace(self.status, capture_generation=3)
|
||||
else:
|
||||
self.assertIsNone(self.seal().positive.effective.value)
|
||||
|
||||
def test_missing_cache_preserves_known_affixes_when_raw_is_dynamic(self):
|
||||
self.begin(raw=["dynamic", 0])
|
||||
self.observe("execution_start")
|
||||
self.observe("execution_cached", nodes=["p"])
|
||||
result = self.seal().positive
|
||||
self.assertIsNone(result.raw.value)
|
||||
self.assertIsNone(result.effective.value)
|
||||
self.assertEqual(result.prepend.value, " pre ")
|
||||
self.assertEqual(result.append.value, " post ")
|
||||
|
||||
def test_cache_entries_tensors_and_executor_are_not_retained(self):
|
||||
class Tensor:
|
||||
pass
|
||||
|
||||
tensor = Tensor()
|
||||
tensor_ref = weakref.ref(tensor)
|
||||
self.begin()
|
||||
self.executor.caches.outputs.entries["p"] = Entry(tensor, [tensor, ["cached"]])
|
||||
self.observe("execution_start")
|
||||
self.observe("execution_cached", nodes=["p"])
|
||||
self.executor.caches.outputs.entries.clear()
|
||||
executor_ref = weakref.ref(self.executor)
|
||||
del tensor
|
||||
del self.executor
|
||||
gc.collect()
|
||||
self.assertIsNone(tensor_ref())
|
||||
self.assertIsNone(executor_ref())
|
||||
self.assertEqual(self.seal().positive.effective.value, "cached")
|
||||
|
||||
def test_unknown_and_malformed_cache_payloads_are_bounded(self):
|
||||
for values in (["x"] * 65, ["x" * 65537], [object()], ["x", "x"]):
|
||||
with self.subTest(count=len(values)):
|
||||
self.begin()
|
||||
self.executor.caches.outputs.entries["p"] = Entry(
|
||||
None, [object(), values]
|
||||
)
|
||||
self.observe("execution_start")
|
||||
self.observe("execution_cached", nodes=["p"])
|
||||
result = self.seal().positive
|
||||
self.assertIsNone(result.effective.value)
|
||||
self.assertEqual(
|
||||
result.append.value, " post " if len(values) != 2 else None
|
||||
)
|
||||
|
||||
def test_disposed_subscriptions_do_not_retain_dynamic_executor_classes(self):
|
||||
def install_and_remove():
|
||||
class TemporaryExecutor:
|
||||
def add_message(self, *args, **kwargs):
|
||||
return None
|
||||
|
||||
self.execution.PromptExecutor = TemporaryExecutor
|
||||
ref = weakref.ref(TemporaryExecutor)
|
||||
stop = host.subscribe_executor_messages("temporary", lambda *args: None)
|
||||
stop()
|
||||
del self.execution.PromptExecutor
|
||||
return ref
|
||||
|
||||
ref = install_and_remove()
|
||||
gc.collect()
|
||||
self.assertIsNone(ref())
|
||||
@@ -625,3 +625,82 @@ class HistoryHooksTests(unittest.TestCase):
|
||||
self.queue.task_done(10, {}, status())
|
||||
self.drain(second)
|
||||
self.assertEqual(len(self.rows()), 1)
|
||||
|
||||
def test_cached_executor_subscription_persists_strings_and_disposes_cooperatively(
|
||||
self,
|
||||
):
|
||||
from utils.execution_context_host import subscribe_executor_messages
|
||||
|
||||
calls = []
|
||||
execution = types.ModuleType("execution")
|
||||
caching = types.ModuleType("comfy_execution.caching")
|
||||
package = types.ModuleType("comfy_execution")
|
||||
entry_type = namedtuple("CacheEntry", "ui outputs")
|
||||
|
||||
class PromptExecutor:
|
||||
def add_message(self, event, data, broadcast):
|
||||
calls.append(event)
|
||||
return "host-result"
|
||||
|
||||
class HierarchicalCache:
|
||||
def get_local(self, node_id):
|
||||
return entry_type(None, [object(), object(), ["old trigger"], [""]])
|
||||
|
||||
execution.PromptExecutor, execution.CacheEntry = PromptExecutor, entry_type
|
||||
caching.HierarchicalCache = HierarchicalCache
|
||||
package.caching = caching
|
||||
with patch.dict(
|
||||
sys.modules,
|
||||
{
|
||||
"execution": execution,
|
||||
"comfy_execution": package,
|
||||
"comfy_execution.caching": caching,
|
||||
},
|
||||
):
|
||||
observer_calls = []
|
||||
independent = subscribe_executor_messages(
|
||||
"independent", lambda *args: observer_calls.append(args)
|
||||
)
|
||||
try:
|
||||
registration = self.register()
|
||||
g = {
|
||||
"p": {
|
||||
"class_type": "PromptManagerPair",
|
||||
"inputs": {"text": "raw", "negative_text": ""},
|
||||
},
|
||||
"s": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {"positive": ["p", 0], "negative": ["p", 1]},
|
||||
},
|
||||
}
|
||||
self.queue.results.append(((0, JOB, g, {}, ["s"]), 1))
|
||||
self.queue.get()
|
||||
executor = PromptExecutor()
|
||||
executor.server = self.server
|
||||
executor.caches = SimpleNamespace(outputs=HierarchicalCache())
|
||||
self.assertEqual(
|
||||
executor.add_message("execution_start", {"prompt_id": JOB}, False),
|
||||
"host-result",
|
||||
)
|
||||
executor.add_message(
|
||||
event="execution_cached",
|
||||
data={"prompt_id": JOB, "nodes": ["p"]},
|
||||
broadcast=False,
|
||||
)
|
||||
self.queue.task_done(1, {}, status())
|
||||
self.drain(registration)
|
||||
result = self.db.get_execution_contexts(self.rows()[0]["run_id"])
|
||||
self.assertEqual(
|
||||
result["contexts"][0]["positive"]["effective"], "old trigger"
|
||||
)
|
||||
self.assertEqual(result["contexts"][0]["negative"]["effective"], "")
|
||||
self.assertTrue(registration.dispose())
|
||||
executor.add_message(
|
||||
"execution_cached", {"prompt_id": JOB, "nodes": ["p"]}, False
|
||||
)
|
||||
self.assertEqual(len(observer_calls), 3)
|
||||
finally:
|
||||
independent()
|
||||
self.assertEqual(
|
||||
calls, ["execution_start", "execution_cached", "execution_cached"]
|
||||
)
|
||||
|
||||
@@ -51,6 +51,8 @@ class _Attempt:
|
||||
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)
|
||||
@@ -164,13 +166,20 @@ class ContextCapture:
|
||||
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(
|
||||
text_size(s.node_id, 512) + text_size(s.node_type, 128) + 32
|
||||
for s in projection.sources
|
||||
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")
|
||||
@@ -186,6 +195,36 @@ class ContextCapture:
|
||||
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,
|
||||
@@ -292,8 +331,19 @@ class ContextCapture:
|
||||
return _unknown("ambiguous")
|
||||
if len(role.sources) == 1 and len(values) == 1:
|
||||
return values[0]
|
||||
if role.inferred is not None:
|
||||
s = role.inferred
|
||||
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
|
||||
)
|
||||
|
||||
@@ -32,6 +32,7 @@ class GraphProjection:
|
||||
recipes: tuple[DeclaredRecipe, ...]
|
||||
sources: tuple[SourceRef, ...]
|
||||
reason: str
|
||||
source_snapshots: tuple[tuple[SourceRef, PromptSnapshot], ...] = ()
|
||||
|
||||
|
||||
_SAMPLERS = {"KSampler", "KSamplerAdvanced", "SamplerCustom"}
|
||||
@@ -214,6 +215,26 @@ class _Projector:
|
||||
raise _Unknown("source_reference_limit_reached")
|
||||
return snapshot, sources
|
||||
|
||||
def source_snapshot(self, source: SourceRef) -> PromptSnapshot:
|
||||
name, inputs = self.node(source.node_id, 1)
|
||||
adapter = _DEFAULT_ADAPTERS[(name, source.output_slot)]
|
||||
unknown = TextEvidence(None, "unknown", "unknown", "missing")
|
||||
|
||||
def literal(value: object) -> TextEvidence:
|
||||
try:
|
||||
snapshot, sources = self.text(value, "string", 2)
|
||||
# Stateful PM outputs cannot prove their injected result from inputs.
|
||||
return unknown if sources else snapshot.effective
|
||||
except _Limit:
|
||||
raise
|
||||
except ValueError:
|
||||
return unknown
|
||||
|
||||
raw = literal(inputs.get(adapter.input_name))
|
||||
before = literal(inputs.get("prepend_text", "") if adapter.compose else "")
|
||||
after = literal(inputs.get("append_text", "") if adapter.compose else "")
|
||||
return self.retain(PromptSnapshot(raw, unknown, before, after))
|
||||
|
||||
def discover_sources(
|
||||
self, value: object, kind: str, depth: int = 1, active: frozenset = frozenset()
|
||||
) -> tuple[SourceRef, ...]:
|
||||
@@ -266,7 +287,7 @@ class _Projector:
|
||||
|
||||
def retain(self, snapshot: PromptSnapshot) -> PromptSnapshot:
|
||||
size = sum(
|
||||
text_size(field.value) + 64
|
||||
(text_size(field.value) if field.value is not None else 0) + 64
|
||||
for field in (
|
||||
snapshot.raw,
|
||||
snapshot.effective,
|
||||
@@ -388,7 +409,12 @@ def project_contexts(graph: object, outputs: object) -> GraphProjection:
|
||||
for role in (recipe.positive, recipe.negative)
|
||||
):
|
||||
reason = "some_roles_require_inferred_or_unresolved_text"
|
||||
return GraphProjection(tuple(recipes), sources, reason)
|
||||
return GraphProjection(
|
||||
tuple(recipes),
|
||||
sources,
|
||||
reason,
|
||||
tuple((source, resolver.source_snapshot(source)) for source in sources),
|
||||
)
|
||||
except ValueError as error:
|
||||
reason = str(error) if isinstance(error, _Unknown) else "invalid_text_value"
|
||||
return GraphProjection((), (), reason)
|
||||
|
||||
@@ -0,0 +1,155 @@
|
||||
"""Pinned-host v1 adapter for local cached STRING evidence only."""
|
||||
|
||||
import functools
|
||||
import threading
|
||||
from typing import Callable
|
||||
|
||||
from .execution_context_capture import active_capture
|
||||
from .execution_contexts import PromptSnapshot, TextEvidence
|
||||
|
||||
_STATE_ATTRIBUTE = "_prompt_manager_executor_messages_v1"
|
||||
_registration_lock = threading.Lock()
|
||||
|
||||
|
||||
class _Messages:
|
||||
def __init__(self, executor_class: type) -> None:
|
||||
self.executor_class = executor_class
|
||||
self.lock = threading.Lock()
|
||||
self.callbacks = {}
|
||||
self.wrapper = None
|
||||
self.original = None
|
||||
self.installed = False
|
||||
|
||||
def subscribe(self, key: str, callback: Callable) -> Callable[[], None]:
|
||||
with self.lock:
|
||||
if not self.installed:
|
||||
original = self.executor_class.add_message
|
||||
if not callable(original):
|
||||
raise ValueError("Executor message capability unavailable")
|
||||
|
||||
@functools.wraps(original)
|
||||
def observe(executor, *args, **kwargs):
|
||||
result = original(executor, *args, **kwargs)
|
||||
with self.lock:
|
||||
callbacks = tuple(entry[1] for entry in self.callbacks.values())
|
||||
for notify in callbacks:
|
||||
try:
|
||||
notify(executor, args, kwargs)
|
||||
except Exception:
|
||||
# Best-effort observers cannot change host results.
|
||||
pass
|
||||
return result
|
||||
|
||||
self.executor_class.add_message = observe
|
||||
self.original, self.wrapper, self.installed = original, observe, True
|
||||
existing = self.callbacks.get(key)
|
||||
token = existing[0] if existing is not None else object()
|
||||
if existing is None:
|
||||
self.callbacks[key] = (token, callback)
|
||||
|
||||
def unsubscribe() -> None:
|
||||
with self.lock:
|
||||
current = self.callbacks.get(key)
|
||||
if current is None or current[0] is not token:
|
||||
return
|
||||
del self.callbacks[key]
|
||||
if (
|
||||
not self.callbacks
|
||||
and self.executor_class.add_message is self.wrapper
|
||||
):
|
||||
self.executor_class.add_message = self.original
|
||||
self.installed = False
|
||||
|
||||
return unsubscribe
|
||||
|
||||
|
||||
def subscribe_executor_messages(key: str, callback: Callable) -> Callable[[], None]:
|
||||
"""Observe successful original calls; preserve earlier and later wrappers."""
|
||||
from execution import PromptExecutor
|
||||
|
||||
with _registration_lock:
|
||||
state = vars(PromptExecutor).get(_STATE_ATTRIBUTE)
|
||||
if state is None:
|
||||
state = _Messages(PromptExecutor)
|
||||
setattr(PromptExecutor, _STATE_ATTRIBUTE, state)
|
||||
return state.subscribe(key, callback)
|
||||
|
||||
|
||||
def observe_cached_context(
|
||||
executor: object, message_type: object, message_data: object
|
||||
) -> None:
|
||||
"""Copy bounded local strings without cache/provider execution or retention."""
|
||||
try:
|
||||
binding = active_capture()
|
||||
if binding is None or type(message_type) is not str:
|
||||
return
|
||||
manager, cap = binding
|
||||
if (
|
||||
type(message_data) is not dict
|
||||
or type(message_data.get("prompt_id")) is not str
|
||||
or message_data["prompt_id"] != cap.job_id
|
||||
or id(executor.server.prompt_queue) != manager.queue_identity
|
||||
):
|
||||
return
|
||||
if message_type == "execution_start":
|
||||
manager.bind_executor(cap, id(executor))
|
||||
return
|
||||
if message_type != "execution_cached" or not manager.matches_executor(
|
||||
cap, id(executor)
|
||||
):
|
||||
return
|
||||
nodes = message_data.get("nodes")
|
||||
if type(nodes) not in (list, tuple) or len(nodes) > 500:
|
||||
return
|
||||
from comfy_execution import caching
|
||||
from execution import CacheEntry
|
||||
|
||||
cache = executor.caches.outputs
|
||||
supported = tuple(
|
||||
cls
|
||||
for name in ("HierarchicalCache", "LRUCache", "RAMPressureCache")
|
||||
if isinstance(cls := getattr(caching, name, None), type)
|
||||
)
|
||||
if type(cache) not in supported:
|
||||
return
|
||||
local = getattr(cache, "get_local", None)
|
||||
if not callable(local):
|
||||
return
|
||||
# Only source IDs retained by the bounded submitted graph are queried.
|
||||
for source in cap.sources:
|
||||
if not any(type(node) is str and node == source.node_id for node in nodes):
|
||||
continue
|
||||
slot = (
|
||||
(3 if source.output_slot in (1, 3) else 2)
|
||||
if source.node_type == "PromptManagerPair"
|
||||
else 0 if source.node_type == "PromptManagerText" else 1
|
||||
)
|
||||
entry = local(source.node_id)
|
||||
if type(entry) is not CacheEntry:
|
||||
continue
|
||||
outputs = entry.outputs
|
||||
if type(outputs) not in (list, tuple) or not slot < len(outputs) <= 4:
|
||||
continue
|
||||
values = outputs[slot]
|
||||
if type(values) not in (list, tuple) or len(values) > 64:
|
||||
continue
|
||||
inferred = manager.source_snapshot(cap, source)
|
||||
if inferred is None:
|
||||
continue
|
||||
for index, value in enumerate(values):
|
||||
if type(value) is not str:
|
||||
continue
|
||||
manager.observe(
|
||||
cap,
|
||||
source,
|
||||
index,
|
||||
PromptSnapshot(
|
||||
inferred.raw,
|
||||
TextEvidence(value, "recorded", "host_cache_output", None),
|
||||
inferred.prepend,
|
||||
inferred.append,
|
||||
),
|
||||
)
|
||||
except Exception:
|
||||
# Unknown host versions and malformed entries leave evidence unavailable.
|
||||
return
|
||||
@@ -13,7 +13,8 @@ except ImportError:
|
||||
from database import PromptDatabase
|
||||
|
||||
from .execution_history_events import make_finish, make_start
|
||||
from .execution_context_capture import ContextCapture, unbind_current
|
||||
from .execution_context_capture import ContextCapture, active_capture, unbind_current
|
||||
from .execution_context_host import observe_cached_context, subscribe_executor_messages
|
||||
from .execution_history_writer import HistoryWriter
|
||||
from .logging_config import get_logger
|
||||
from .queue_observers import subscribe_dequeue, subscribe_task_done
|
||||
@@ -57,12 +58,28 @@ class HistoryRegistration:
|
||||
self._unsubscribers.append(
|
||||
subscribe_task_done(queue, "execution_history", self._task_done)
|
||||
)
|
||||
try:
|
||||
self._unsubscribers.append(
|
||||
subscribe_executor_messages(
|
||||
f"execution_history_{id(self)}", self._executor_message
|
||||
)
|
||||
)
|
||||
except (ImportError, AttributeError, TypeError, ValueError):
|
||||
pass
|
||||
self.writer.start()
|
||||
atexit.register(self.dispose)
|
||||
except Exception:
|
||||
self.dispose()
|
||||
raise
|
||||
|
||||
def _executor_message(self, executor: object, args: tuple, kwargs: dict) -> None:
|
||||
binding = active_capture()
|
||||
if binding is None or binding[0] is not self.capture:
|
||||
return
|
||||
event = args[0] if args else kwargs.get("event")
|
||||
data = args[1] if len(args) > 1 else kwargs.get("data")
|
||||
observe_cached_context(executor, event, data)
|
||||
|
||||
def _invalidate_capture(self) -> None:
|
||||
self.capture.revoke_all()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user