Vito
d03d0ed1c5
* fix(history): count cached re-runs of runtime-linked text as tracked runs (#403) Since #347 a cached re-run of runtime-linked prompt text (PreviewAny or an LLM into a PromptManager text input) counts in the Library, but its history capture had no raw text, so the history page's "tracked runs" stayed at the number of jobs in which the node executed. The #347 receipt now also keeps the raw, prepend and append text the node gave the history capture. When a cached node's graph inference leaves a field unknown, the capture takes it from the receipt that matches the cache entry, labeled inferred/graph_semantics, so the job joins the prompt's history scope and gets its prompt hint. No schema change and no new evidence vocabulary; no backfill. Fails closed on no receipt, a disagreeing known field, equal outputs with different fields, oversized or invalid text and a full budget (least recently used receipts give up their fields first; counting stays). * docs(testing): link the whitespace history gap to #438
1056 lines
40 KiB
Python
1056 lines
40 KiB
Python
"""Cached re-runs of runtime-linked prompt text count once per job (#347).
|
|
|
|
Jobs are simulated end to end: the queue hook (submit + dequeue), the
|
|
``execution_cached`` replay, the real node executions and ComfyUI's output
|
|
cache (fakes of ``execution`` and ``comfy_execution.caching``, as in
|
|
tests/test_execution_context_host.py).
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import types
|
|
import unittest
|
|
from collections import namedtuple
|
|
from unittest import mock
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
try:
|
|
from tests.open_handles import assert_closed
|
|
except ImportError: # discovered with tests/ as the top-level directory
|
|
from open_handles import assert_closed
|
|
|
|
from database.operations import PromptDatabase
|
|
from utils import cached_runs as module
|
|
from utils.cached_runs import (
|
|
CachedRuns,
|
|
Facts,
|
|
Role,
|
|
Seen,
|
|
graph_role,
|
|
graph_seen,
|
|
own_counts,
|
|
register_cached_runs,
|
|
stored_facts,
|
|
)
|
|
from utils.execution_context_host import cached_string_outputs
|
|
from utils.hashing import generate_prompt_hash
|
|
from utils.run_keys import prompt_node_key, upstream_closure, upstream_nodes
|
|
from utils.usage_tracking import (
|
|
PendingFirstUse,
|
|
QueuedRuns,
|
|
find_pair_id,
|
|
find_pair_id_by_hashes,
|
|
handle_queued_prompt,
|
|
on_dequeue,
|
|
)
|
|
|
|
Entry = namedtuple("CacheEntry", "ui outputs")
|
|
LLM = "a lighthouse the LLM wrote"
|
|
|
|
|
|
class _Cache:
|
|
def __init__(self):
|
|
self.entries = {}
|
|
|
|
def get_local(self, node_id):
|
|
return self.entries.get(node_id)
|
|
|
|
|
|
class HierarchicalCache(_Cache):
|
|
pass
|
|
|
|
|
|
class LRUCache(_Cache):
|
|
pass
|
|
|
|
|
|
class RAMPressureCache(LRUCache):
|
|
pass
|
|
|
|
|
|
class NullCache(_Cache):
|
|
pass
|
|
|
|
|
|
def host_modules():
|
|
"""Fake ``execution`` and ``comfy_execution.caching`` for the cache reader."""
|
|
execution = types.ModuleType("execution")
|
|
execution.CacheEntry = Entry
|
|
|
|
class PromptExecutor:
|
|
def add_message(self, event, data, broadcast):
|
|
return None
|
|
|
|
execution.PromptExecutor = PromptExecutor
|
|
caching = types.ModuleType("comfy_execution.caching")
|
|
for cls in (HierarchicalCache, LRUCache, RAMPressureCache, NullCache):
|
|
setattr(caching, cls.__name__, cls)
|
|
package = types.ModuleType("comfy_execution")
|
|
package.caching = caching
|
|
return {
|
|
"execution": execution,
|
|
"comfy_execution": package,
|
|
"comfy_execution.caching": caching,
|
|
}
|
|
|
|
|
|
def graph(
|
|
kind="PromptManager",
|
|
text=None,
|
|
role="positive",
|
|
sampler=True,
|
|
negative_text="",
|
|
seed=1,
|
|
):
|
|
"""PrimitiveString -> PreviewAny -> ``kind`` -> KSampler (runtime text).
|
|
|
|
``text`` set: typed text instead of the runtime link. ``role`` "negative":
|
|
the node feeds the sampler's negative input. Pair: its negative editor holds
|
|
``negative_text`` and it feeds both sampler inputs.
|
|
"""
|
|
g = {
|
|
"str": {"class_type": "PrimitiveString", "inputs": {"value": "seed text"}},
|
|
"llm": {"class_type": "PreviewAny", "inputs": {"source": ["str", 0]}},
|
|
"ckpt": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt": "m"}},
|
|
"pm": {
|
|
"class_type": kind,
|
|
"inputs": {"text": ["llm", 0] if text is None else text},
|
|
},
|
|
"other": {
|
|
"class_type": "CLIPTextEncode",
|
|
"inputs": {"text": "", "clip": ["ckpt", 1]},
|
|
},
|
|
"save": {"class_type": "SaveImage", "inputs": {"images": ["ks", 0]}},
|
|
}
|
|
if kind != "PromptManagerText":
|
|
g["pm"]["inputs"]["clip"] = ["ckpt", 1]
|
|
if kind == "PromptManagerPair":
|
|
g["pm"]["inputs"]["negative_text"] = negative_text
|
|
source = ["pm", 0]
|
|
if kind == "PromptManagerText":
|
|
g["enc"] = {
|
|
"class_type": "CLIPTextEncode",
|
|
"inputs": {"text": ["pm", 0], "clip": ["ckpt", 1]},
|
|
}
|
|
source = ["enc", 0]
|
|
positive, negative = (source, ["other", 0])
|
|
if role == "negative":
|
|
positive, negative = negative, positive
|
|
if kind == "PromptManagerPair":
|
|
positive, negative = ["pm", 0], ["pm", 1]
|
|
if sampler:
|
|
g["ks"] = {
|
|
"class_type": "KSampler",
|
|
"inputs": {
|
|
"model": ["ckpt", 0],
|
|
"positive": positive,
|
|
"negative": negative,
|
|
"seed": seed,
|
|
},
|
|
}
|
|
else:
|
|
g["save"]["inputs"]["images"] = positive
|
|
return g
|
|
|
|
|
|
class JobCase(unittest.TestCase):
|
|
"""A database, the queue hook stashes and one ``CachedRuns`` per test."""
|
|
|
|
def setUp(self):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.tmp.cleanup)
|
|
self.addCleanup(assert_closed, self, self.tmp.name)
|
|
self.db = PromptDatabase(os.path.join(self.tmp.name, "runs.db"))
|
|
self.addCleanup(self.db.close_all)
|
|
patcher = mock.patch.dict(sys.modules, host_modules())
|
|
patcher.start()
|
|
self.addCleanup(patcher.stop)
|
|
self.queued, self.pair_queued = QueuedRuns(), QueuedRuns()
|
|
self.pending, self.pair_pending = PendingFirstUse(), PendingFirstUse()
|
|
self.context = (None, None)
|
|
self.runs = CachedRuns(context=lambda: self.context)
|
|
self.queue = object()
|
|
self.executor = types.SimpleNamespace(
|
|
server=types.SimpleNamespace(prompt_queue=self.queue),
|
|
caches=types.SimpleNamespace(outputs=HierarchicalCache()),
|
|
)
|
|
self.item_id = 0
|
|
self.trackers = []
|
|
|
|
# -- nodes ----------------------------------------------------------------
|
|
|
|
def node(self, kind):
|
|
if kind == "PromptManagerPair":
|
|
from prompt_manager_pair import PromptManagerPair as cls
|
|
elif kind == "PromptManagerText":
|
|
from prompt_manager_text import PromptManagerText as cls
|
|
else:
|
|
import prompt_manager
|
|
|
|
cls = getattr(prompt_manager, kind)
|
|
node = cls.__new__(cls)
|
|
node.db = self.db
|
|
node.logger = logging.getLogger("test.cached_runs")
|
|
node.prompt_tracker = mock.Mock()
|
|
self.trackers.append(node.prompt_tracker)
|
|
node.comfyui_integration = mock.Mock()
|
|
node._inject_lora_trigger_words = lambda text: text
|
|
node.pending_first_use = self.pending
|
|
node.pending_pair_first_use = self.pair_pending
|
|
node.cached_runs = self.runs
|
|
return node
|
|
|
|
def execute(self, g, node_id="pm", text=LLM, negative="", index=None):
|
|
"""Run the node as ComfyUI would; returns its STRING outputs."""
|
|
kind = g[node_id]["class_type"]
|
|
node = self.node(kind)
|
|
self.context = (self.context[0], index)
|
|
clip = mock.Mock()
|
|
# The node receives its switch as the queued graph holds it (#90)
|
|
save = g[node_id]["inputs"].get("save_to_library", True)
|
|
if kind == "PromptManagerText":
|
|
return node.process_text(
|
|
text, prompt=g, unique_id=node_id, save_to_library=save
|
|
)
|
|
if kind == "PromptManagerPair":
|
|
return node.encode_prompt(
|
|
clip,
|
|
text,
|
|
negative_text=negative,
|
|
prompt=g,
|
|
unique_id=node_id,
|
|
save_to_library=save,
|
|
)
|
|
return node.encode_prompt(
|
|
clip, text, prompt=g, unique_id=node_id, save_to_library=save
|
|
)
|
|
|
|
def cache(self, node_id, *results):
|
|
"""Store executed results (one per list index) as ComfyUI's cache entry."""
|
|
width = len(results[0])
|
|
outputs = [[result[slot] for result in results] for slot in range(width)]
|
|
self.executor.caches.outputs.entries[node_id] = Entry({}, outputs)
|
|
|
|
# -- jobs -----------------------------------------------------------------
|
|
|
|
def job(self, g, cached=(), run=None, outputs=("save",), prompt_id=None):
|
|
"""One queued job: hook, dequeue, ``execution_cached``, executions, done."""
|
|
self.item_id += 1
|
|
prompt_id = prompt_id or f"job-{self.item_id}"
|
|
request = handle_queued_prompt(
|
|
{"prompt": g, "prompt_id": prompt_id},
|
|
self.queued,
|
|
pair_queued=self.pair_queued,
|
|
)
|
|
on_dequeue(
|
|
request["prompt_id"],
|
|
self.queued,
|
|
lambda: self.db,
|
|
self.pending,
|
|
pair_queued=self.pair_queued,
|
|
pair_pending=self.pair_pending,
|
|
)
|
|
self.runs.begin_attempt(
|
|
((self.item_id, prompt_id, g, {}, list(outputs), {}), self.item_id)
|
|
)
|
|
self.context = (prompt_id, None)
|
|
self.runs.replay(self.executor, prompt_id, list(cached), lambda: self.db)
|
|
result = run() if run else None
|
|
self.runs.end_attempt(self.item_id)
|
|
self.context = (None, None)
|
|
return result
|
|
|
|
def first(self, g, text=LLM, negative="", node_id="pm"):
|
|
"""A job that executes the node and fills the cache."""
|
|
result = self.job(
|
|
g, run=lambda: self.execute(g, node_id, text=text, negative=negative)
|
|
)
|
|
self.cache(node_id, result)
|
|
return result
|
|
|
|
def count(self, text):
|
|
row = self.db.get_prompt_by_hash(generate_prompt_hash(text))
|
|
return row["run_count"] if row else None
|
|
|
|
def pair_count(self, positive, negative):
|
|
pair_id = find_pair_id(self.db, positive, negative)
|
|
return self.db.get_prompt_pair(pair_id)["run_count"] if pair_id else None
|
|
|
|
|
|
class TestCachedJobsCount(JobCase):
|
|
def test_runtime_text_counts_every_cached_job_for_each_node_kind(self):
|
|
for kind in (
|
|
"PromptManager",
|
|
"PromptManagerFocused",
|
|
"PromptManagerSplit",
|
|
"PromptManagerText",
|
|
):
|
|
with self.subTest(kind=kind):
|
|
text = f"{LLM} for {kind}"
|
|
g = graph(kind)
|
|
self.first(g, text=text)
|
|
self.job(graph(kind, seed=2), cached=["pm"])
|
|
self.job(graph(kind, seed=3), cached=["pm"])
|
|
self.assertEqual(self.count(text), 3)
|
|
|
|
def test_typed_text_still_counts_once_per_job(self):
|
|
g = graph(text="typed castle")
|
|
self.first(g, text="typed castle")
|
|
self.job(g, cached=["pm"])
|
|
self.job(g, cached=["pm"])
|
|
self.assertEqual(self.count("typed castle"), 3)
|
|
|
|
def test_rewired_to_negative_while_cached_counts_nothing(self):
|
|
self.first(graph())
|
|
self.job(graph(role="negative"), cached=["pm"])
|
|
self.assertEqual(self.count(LLM), 1)
|
|
# Rewired back to positive: the receipt still proves the text
|
|
self.job(graph(), cached=["pm"])
|
|
self.assertEqual(self.count(LLM), 2)
|
|
|
|
def test_no_sampler_counts_every_node_on_replay(self):
|
|
g = graph(sampler=False)
|
|
self.first(g)
|
|
self.job(g, cached=["pm"])
|
|
self.assertEqual(self.count(LLM), 2)
|
|
|
|
def test_empty_editor_replays_nothing(self):
|
|
g = graph()
|
|
self.first(g, text="")
|
|
self.job(g, cached=["pm"])
|
|
self.assertEqual(self.db.search_prompts(), [])
|
|
|
|
def test_disabled_node_then_cached_jobs_counts_0(self):
|
|
# #90 x #347: a node with Save to library off records no receipt and is
|
|
# no run prompt, so its cached jobs replay nothing. The prompt exists
|
|
# (an enabled run saved it), so any replayed count would show.
|
|
for kind in (
|
|
"PromptManager",
|
|
"PromptManagerFocused",
|
|
"PromptManagerSplit",
|
|
"PromptManagerText",
|
|
"PromptManagerPair",
|
|
):
|
|
with self.subTest(kind=kind):
|
|
text = f"{LLM} not saved by {kind}"
|
|
negative = "blurry" if kind == "PromptManagerPair" else ""
|
|
self.first(graph(kind, negative_text=negative), text, negative)
|
|
before = (self.count(text), self.pair_count(text, negative))
|
|
|
|
def off(seed):
|
|
g = graph(kind, negative_text=negative, seed=seed)
|
|
g["pm"]["inputs"]["save_to_library"] = False
|
|
return g
|
|
|
|
self.first(off(1), text, negative)
|
|
self.job(off(2), cached=["pm"])
|
|
self.job(off(3), cached=["pm"])
|
|
after = (self.count(text), self.pair_count(text, negative))
|
|
self.assertEqual(before[0], 1)
|
|
self.assertEqual(after, before)
|
|
|
|
def test_alternating_upstream_values_each_count_their_own_prompt(self):
|
|
g = graph()
|
|
self.first(g, text="value A")
|
|
self.first(g, text="value B")
|
|
self.cache("pm", self.execute_outputs("value A"))
|
|
self.job(g, cached=["pm"])
|
|
self.assertEqual((self.count("value A"), self.count("value B")), (2, 1))
|
|
|
|
def execute_outputs(self, text):
|
|
return (mock.ANY, text)
|
|
|
|
def test_list_items_replay_each_item_in_order(self):
|
|
g = graph()
|
|
results = []
|
|
|
|
def run():
|
|
for index, text in enumerate(("one", "two", "three")):
|
|
results.append(self.execute(g, text=text, index=index))
|
|
|
|
self.job(g, run=run)
|
|
self.cache("pm", *results)
|
|
self.job(g, cached=["pm"])
|
|
self.assertEqual([self.count(t) for t in ("one", "two", "three")], [2, 2, 2])
|
|
|
|
def test_a_repeated_id_or_message_never_counts_a_node_twice(self):
|
|
g = graph()
|
|
self.first(g)
|
|
self.job(g, cached=["pm", "pm"])
|
|
self.assertEqual(self.count(LLM), 2)
|
|
self.item_id += 1
|
|
self.runs.begin_attempt(((1, "again", g, {}, ["save"], {}), self.item_id))
|
|
for _ in range(2):
|
|
self.runs.replay(self.executor, "again", ["pm"], lambda: self.db)
|
|
self.assertEqual(self.count(LLM), 3)
|
|
|
|
def test_a_shorter_list_keeps_its_receipt_after_a_longer_one(self):
|
|
g = graph()
|
|
self.first(g, text="alpha")
|
|
results = []
|
|
|
|
def run():
|
|
for index, text in enumerate(("alpha", "beta")):
|
|
results.append(self.execute(g, text=text, index=index))
|
|
|
|
self.job(g, run=run)
|
|
self.cache("pm", (mock.ANY, "alpha"))
|
|
self.job(g, cached=["pm"])
|
|
self.assertEqual((self.count("alpha"), self.count("beta")), (3, 1))
|
|
|
|
def test_a_job_reusing_a_prompt_id_counts_and_executes_normally(self):
|
|
g = graph()
|
|
self.first(g)
|
|
self.job(g, cached=["pm"], prompt_id="same")
|
|
self.job(g, run=lambda: self.execute(g), prompt_id="same")
|
|
self.job(g, cached=["pm"], prompt_id="same")
|
|
self.assertEqual(self.count(LLM), 4)
|
|
|
|
|
|
class TestFailClosed(JobCase):
|
|
def setUp(self):
|
|
super().setUp()
|
|
self.g = graph()
|
|
self.first(self.g)
|
|
|
|
def assertNotCounted(self):
|
|
self.assertEqual(self.count(LLM), 1)
|
|
|
|
def test_changed_cached_text_is_not_counted(self):
|
|
self.cache("pm", (mock.ANY, LLM + " edited"))
|
|
self.job(self.g, cached=["pm"])
|
|
self.assertNotCounted()
|
|
|
|
def test_no_receipt_in_this_process_is_not_counted(self):
|
|
self.runs = CachedRuns(context=lambda: self.context)
|
|
self.job(self.g, cached=["pm"])
|
|
self.assertNotCounted()
|
|
|
|
def test_missing_cache_entry_is_not_counted(self):
|
|
self.executor.caches.outputs.entries.clear()
|
|
self.job(self.g, cached=["pm"])
|
|
self.assertNotCounted()
|
|
|
|
def test_unsupported_cache_is_not_counted(self):
|
|
entries = self.executor.caches.outputs.entries
|
|
self.executor.caches.outputs = NullCache()
|
|
self.executor.caches.outputs.entries = entries
|
|
self.job(self.g, cached=["pm"])
|
|
self.assertNotCounted()
|
|
|
|
def test_more_list_values_than_the_bound_are_not_counted(self):
|
|
self.cache("pm", *[(mock.ANY, LLM)] * 65)
|
|
self.job(self.g, cached=["pm"])
|
|
self.assertNotCounted()
|
|
|
|
def test_missing_list_index_is_not_counted(self):
|
|
self.cache("pm", (mock.ANY, LLM), (mock.ANY, LLM))
|
|
self.job(self.g, cached=["pm"])
|
|
self.assertNotCounted()
|
|
|
|
def test_node_the_job_does_not_run_is_not_counted(self):
|
|
self.job(self.g, cached=["pm"], outputs=("other",))
|
|
self.assertNotCounted()
|
|
|
|
def test_message_for_another_job_is_not_counted(self):
|
|
self.item_id += 1
|
|
self.runs.begin_attempt(((1, "this", self.g, {}, ["save"], {}), self.item_id))
|
|
self.runs.replay(self.executor, "another", ["pm"], lambda: self.db)
|
|
self.assertNotCounted()
|
|
|
|
def test_unkeyable_node_is_not_counted(self):
|
|
g = graph()
|
|
g["pm"]["inputs"]["extra"] = ["missing", 0]
|
|
self.job(g, cached=["pm"])
|
|
self.assertIsNone(prompt_node_key(g, "pm"))
|
|
self.assertNotCounted()
|
|
|
|
def test_deleted_prompt_is_skipped(self):
|
|
self.db.delete_prompt(
|
|
self.db.get_prompt_by_hash(generate_prompt_hash(LLM))["id"]
|
|
)
|
|
self.job(self.g, cached=["pm"])
|
|
self.assertIsNone(self.count(LLM))
|
|
|
|
def test_database_errors_never_raise(self):
|
|
def broken():
|
|
raise RuntimeError("database unavailable")
|
|
|
|
self.item_id += 1
|
|
self.runs.begin_attempt(((1, "j", self.g, {}, ["save"], {}), self.item_id))
|
|
self.runs.replay(self.executor, "j", ["pm"], broken)
|
|
self.assertNotCounted()
|
|
|
|
def test_malformed_messages_and_results_are_ignored(self):
|
|
self.runs.begin_attempt(None)
|
|
self.runs.begin_attempt((("x",), "not an int"))
|
|
self.runs.replay(self.executor, None, ["pm"], lambda: self.db)
|
|
self.item_id += 1
|
|
self.runs.begin_attempt(((1, "j", self.g, {}, ["save"], {}), self.item_id))
|
|
self.runs.replay(self.executor, "j", "pm", lambda: self.db)
|
|
self.runs.replay(self.executor, "j", ["pm"] * 501, lambda: self.db)
|
|
self.runs.replay(self.executor, "j", [7, None], lambda: self.db)
|
|
self.assertNotCounted()
|
|
|
|
|
|
class TestReplayThenExecution(JobCase):
|
|
def test_execution_after_replay_in_the_same_job_is_not_counted_again(self):
|
|
# RAM-pressure frees the entry after execution_cached: the node executes
|
|
g = graph()
|
|
self.first(g)
|
|
self.job(g, cached=["pm"], run=lambda: self.execute(g))
|
|
self.assertEqual(self.count(LLM), 2)
|
|
# The suppressed execution still leaves a receipt: the next job counts
|
|
self.job(g, cached=["pm"])
|
|
self.assertEqual(self.count(LLM), 3)
|
|
|
|
def test_suppression_keeps_the_pending_catch_up_and_the_tracker_queue(self):
|
|
g = graph()
|
|
self.first(g)
|
|
prompt_id = self.db.get_prompt_by_hash(generate_prompt_hash(LLM))["id"]
|
|
|
|
def run():
|
|
# The replay counted, then the prompt was deleted and re-saved
|
|
self.db.delete_prompt(prompt_id)
|
|
self.pending.add(generate_prompt_hash(LLM))
|
|
return self.execute(g)
|
|
|
|
self.job(g, cached=["pm"], run=run)
|
|
self.assertEqual(self.count(LLM), 1) # the catch-up, not the own count
|
|
push = self.trackers[-1].set_current_prompt.call_args.kwargs["push_to_queue"]
|
|
self.assertTrue(push)
|
|
|
|
def test_without_host_context_the_node_counts_as_before(self):
|
|
g = graph()
|
|
self.runs = CachedRuns(context=lambda: (None, None))
|
|
self.first(g)
|
|
self.job(g, cached=["pm"])
|
|
self.job(g, run=lambda: self.execute(g))
|
|
self.assertEqual(self.count(LLM), 2)
|
|
|
|
|
|
class TestPairReplay(JobCase):
|
|
def test_runtime_positive_counts_positive_and_pair_every_job(self):
|
|
g = graph("PromptManagerPair", negative_text="blurry")
|
|
self.first(g, negative="blurry")
|
|
self.job(graph("PromptManagerPair", negative_text="blurry", seed=2), ["pm"])
|
|
self.job(graph("PromptManagerPair", negative_text="blurry", seed=3), ["pm"])
|
|
self.assertEqual(self.count(LLM), 3)
|
|
self.assertEqual(self.pair_count(LLM, "blurry"), 3)
|
|
self.assertEqual(self.count("blurry"), 0)
|
|
|
|
def test_empty_negative_pair_replays_its_pair_count(self):
|
|
g = graph("PromptManagerPair")
|
|
self.first(g)
|
|
self.job(g, cached=["pm"])
|
|
self.assertEqual(self.pair_count(LLM, ""), 2)
|
|
self.assertEqual(self.count(LLM), 2)
|
|
|
|
def test_pair_execution_after_replay_is_not_counted_again(self):
|
|
g = graph("PromptManagerPair", negative_text="blurry")
|
|
self.first(g, negative="blurry")
|
|
self.job(g, ["pm"], run=lambda: self.execute(g, negative="blurry"))
|
|
self.assertEqual(self.count(LLM), 2)
|
|
self.assertEqual(self.pair_count(LLM, "blurry"), 2)
|
|
|
|
def test_pair_role_follows_the_current_graph(self):
|
|
g = graph("PromptManagerPair", negative_text="blurry")
|
|
self.first(g, negative="blurry")
|
|
rewired = graph("PromptManagerPair", negative_text="blurry")
|
|
rewired["ks"]["inputs"]["positive"] = ["other", 0]
|
|
self.job(rewired, cached=["pm"])
|
|
self.assertEqual(self.count(LLM), 1)
|
|
self.assertEqual(self.pair_count(LLM, "blurry"), 1)
|
|
|
|
|
|
class TestOwnCounts(unittest.TestCase):
|
|
"""The node path and replay decide through the same function."""
|
|
|
|
def test_single_node_table(self):
|
|
facts = Facts("p")
|
|
for used, seen, expected in (
|
|
({"text"}, False, [("prompt", "p")]),
|
|
({"text"}, True, []),
|
|
(set(), False, []),
|
|
):
|
|
with self.subTest(used=used, seen=seen):
|
|
counts = own_counts(
|
|
"PromptManager", facts, Role(frozenset(used)), Seen(text=seen)
|
|
)
|
|
self.assertEqual(counts, expected)
|
|
self.assertEqual(
|
|
own_counts("PromptManager", Facts(None), Role(frozenset({"text"})), Seen()),
|
|
[],
|
|
)
|
|
|
|
def test_pair_table(self):
|
|
both = Role(frozenset({"text", "negative_text"}), True)
|
|
self.assertEqual(
|
|
own_counts("PromptManagerPair", Facts("p", "n"), both, Seen()),
|
|
[("prompt", "p"), ("prompt", "n"), ("pair", "p", "n")],
|
|
)
|
|
self.assertEqual(
|
|
own_counts("PromptManagerPair", Facts("p", "p"), both, Seen()),
|
|
[("prompt", "p"), ("pair", "p", "p")],
|
|
)
|
|
self.assertEqual(
|
|
own_counts("PromptManagerPair", Facts("p"), both, Seen(text=True)),
|
|
[("pair", "p", None)],
|
|
)
|
|
self.assertEqual(
|
|
own_counts(
|
|
"PromptManagerPair", Facts("p", "n"), both, Seen(True, True, True)
|
|
),
|
|
[],
|
|
)
|
|
|
|
def test_node_and_replay_agree_on_role_and_seen(self):
|
|
for g, node_id in (
|
|
(graph(), "pm"),
|
|
(graph(text="typed"), "pm"),
|
|
(graph(role="negative"), "pm"),
|
|
(graph(sampler=False), "pm"),
|
|
(graph("PromptManagerPair", negative_text="n"), "pm"),
|
|
(graph("PromptManagerPair", text="typed", negative_text="n"), "pm"),
|
|
):
|
|
kind = g[node_id]["class_type"]
|
|
text = "typed" if isinstance(g[node_id]["inputs"]["text"], str) else LLM
|
|
facts = stored_facts(text, g[node_id]["inputs"].get("negative_text"))
|
|
with self.subTest(kind=kind, inputs=g[node_id]["inputs"]):
|
|
replay = own_counts(
|
|
kind,
|
|
facts,
|
|
graph_role(g, node_id, kind),
|
|
graph_seen(g, node_id, kind, facts),
|
|
)
|
|
self.assertEqual(replay, self.node_counts(g, node_id, kind, text))
|
|
|
|
@staticmethod
|
|
def node_counts(g, node_id, kind, text):
|
|
"""What the executing node counted itself, from a recording database."""
|
|
db = mock.Mock()
|
|
db.get_prompt_by_hash.side_effect = lambda h: {"id": h}
|
|
db.save_prompt_pair.return_value = 1
|
|
negative = g[node_id]["inputs"].get("negative_text", "")
|
|
db.get_prompt_pair.return_value = {
|
|
"positive_prompt_id": generate_prompt_hash(text),
|
|
"negative_prompt_id": (
|
|
generate_prompt_hash(negative) if negative.strip() else None
|
|
),
|
|
}
|
|
case = JobCase()
|
|
case.db = db
|
|
case.pending, case.pair_pending = PendingFirstUse(), PendingFirstUse()
|
|
case.runs = CachedRuns(context=lambda: (None, None))
|
|
case.trackers = []
|
|
case.context = (None, None)
|
|
if kind == "PromptManagerPair":
|
|
node = case.node(kind)
|
|
node.encode_prompt(
|
|
mock.Mock(), text, negative_text=negative, prompt=g, unique_id=node_id
|
|
)
|
|
else:
|
|
node = case.node(kind)
|
|
node._track_prompt_execution(text, text, None, None, {}, g, node_id)
|
|
counts = [("prompt", c.args[0]) for c in db.record_prompt_use.call_args_list]
|
|
if kind == "PromptManagerPair":
|
|
counts += [
|
|
("pair", generate_prompt_hash(text), db.get_prompt_pair()[k])
|
|
for k in ["negative_prompt_id"]
|
|
if db.record_prompt_pair_use.called
|
|
]
|
|
return counts
|
|
|
|
|
|
class ReceiptCase(unittest.TestCase):
|
|
def setUp(self):
|
|
self.context = ("j1", None)
|
|
self.runs = CachedRuns(
|
|
context=lambda: self.context, max_keys=2, max_receipts=2, max_items=3
|
|
)
|
|
self.g = {
|
|
name: {"class_type": "PromptManager", "inputs": {"text": name}}
|
|
for name in ("a", "b", "c")
|
|
}
|
|
|
|
def begin(self, item_id, prompt_id="j1"):
|
|
self.context = (prompt_id, None)
|
|
self.runs.begin_attempt(
|
|
((item_id, prompt_id, self.g, {}, ["a", "b", "c"], {}), item_id)
|
|
)
|
|
|
|
def keys(self):
|
|
return list(self.runs._receipts)
|
|
|
|
|
|
class TestReceipts(ReceiptCase):
|
|
def test_least_recently_used_run_keys_are_dropped(self):
|
|
self.begin(1)
|
|
for name in ("a", "b", "c"):
|
|
self.runs.record_item(name, (name,), Facts("h"))
|
|
self.assertEqual(self.keys(), [prompt_node_key(self.g, n) for n in ("b", "c")])
|
|
|
|
def test_receipts_per_key_are_bounded_and_equal_outputs_replace(self):
|
|
key = prompt_node_key(self.g, "a")
|
|
for item_id, text in ((1, "x"), (2, "y"), (3, "z"), (4, "y")):
|
|
self.begin(item_id)
|
|
self.runs.record_item("a", (text,), Facts(text))
|
|
self.runs.end_attempt(item_id)
|
|
receipts = self.runs._receipts[key]
|
|
self.assertEqual([r.item_id for r in receipts], [4, 3])
|
|
|
|
def test_list_indices_past_the_bound_are_not_recorded(self):
|
|
self.begin(1)
|
|
self.context = ("j1", 3)
|
|
self.runs.record_item("a", ("x",), Facts("x"))
|
|
self.assertEqual(self.keys(), [])
|
|
|
|
def test_non_string_outputs_and_unknown_nodes_are_not_recorded(self):
|
|
self.begin(1)
|
|
self.runs.record_item("a", (None,), Facts("x"))
|
|
self.runs.record_item("zzz", ("x",), Facts("x"))
|
|
self.assertEqual(self.keys(), [])
|
|
|
|
def test_task_done_clears_only_its_own_attempt(self):
|
|
self.begin(1)
|
|
self.runs.end_attempt(2)
|
|
self.runs.record_item("a", ("x",), Facts("x"))
|
|
self.runs.end_attempt(1)
|
|
self.runs.record_item("b", ("x",), Facts("x"))
|
|
self.assertEqual(self.keys(), [prompt_node_key(self.g, "a")])
|
|
|
|
def test_job_nodes_follow_the_outputs_to_execute(self):
|
|
self.runs.begin_attempt(((1, "j1", self.g, {}, ["b"], {}), 1))
|
|
self.assertEqual(list(self.runs._attempt.nodes), ["b"])
|
|
|
|
|
|
class TestRecordedFields(ReceiptCase):
|
|
"""The texts a receipt keeps for the history capture of a cached job (#403)."""
|
|
|
|
F = (("raw a", " pre ", ""),)
|
|
|
|
def record(self, item_id, name="a", output="x", fields=F, facts=None):
|
|
self.begin(item_id)
|
|
self.runs.record_item(name, (output,), facts or Facts("h"), fields)
|
|
self.runs.end_attempt(item_id)
|
|
|
|
def lookup(self, name="a", values=(("x",),), prompt_id="j9", item_id=9):
|
|
self.begin(item_id, prompt_id)
|
|
return self.runs.recorded_fields(prompt_id, name, values)
|
|
|
|
def test_a_cached_job_reads_the_fields_of_the_execution_it_reuses(self):
|
|
self.record(1)
|
|
self.assertEqual(self.lookup(), [self.F])
|
|
|
|
def test_list_items_return_their_own_fields_in_order(self):
|
|
self.begin(1)
|
|
for index, text in enumerate(("x", "y")):
|
|
self.context = ("j1", index)
|
|
self.runs.record_item("a", (text,), Facts("h"), ((text, "", ""),))
|
|
self.runs.end_attempt(1)
|
|
self.assertEqual(
|
|
self.lookup(values=(("x", "y"),)), [(("x", "", ""),), (("y", "", ""),)]
|
|
)
|
|
|
|
def test_fails_closed(self):
|
|
self.record(1)
|
|
for case, kwargs in (
|
|
("other output", dict(values=(("changed",),))),
|
|
("other node", dict(name="b")),
|
|
("unknown node", dict(name="zzz")),
|
|
("no values", dict(values=())),
|
|
("malformed values", dict(values="x")),
|
|
):
|
|
with self.subTest(case=case):
|
|
self.assertIsNone(self.lookup(**kwargs))
|
|
with self.subTest(case="another job"):
|
|
self.begin(9, "j9")
|
|
self.assertIsNone(self.runs.recorded_fields("j8", "a", (("x",),)))
|
|
with self.subTest(case="no job"):
|
|
self.runs.end_attempt(9)
|
|
self.assertIsNone(self.runs.recorded_fields("j9", "a", (("x",),)))
|
|
|
|
def test_an_execution_without_fields_gives_none(self):
|
|
for fields in (None, (("a", "b"),), ((1, "", ""),), (), "abc"):
|
|
with self.subTest(fields=fields):
|
|
self.runs = CachedRuns(context=lambda: self.context)
|
|
self.record(1, fields=fields)
|
|
self.assertIsNone(self.lookup())
|
|
|
|
def test_a_lookup_never_counts_or_marks_the_node_replayed(self):
|
|
self.record(1)
|
|
self.lookup()
|
|
self.assertFalse(self.runs.replayed("a"))
|
|
|
|
def test_equal_outputs_with_other_raw_text_conflict_and_still_count(self):
|
|
key = prompt_node_key(self.g, "a")
|
|
self.record(1, fields=(("x ", "", ""),))
|
|
self.record(2, fields=(("x", "", ""),))
|
|
# Counting dedupes by outputs as before; the fields cannot name either
|
|
self.assertEqual([r.item_id for r in self.runs._receipts[key]], [2])
|
|
self.assertIsNone(self.lookup())
|
|
self.assertIsNotNone(self.runs._match(key, (("x",),)))
|
|
# The conflict stays: a later execution cannot tell which text it was
|
|
self.record(3, fields=(("x", "", ""),))
|
|
self.assertIsNone(self.lookup())
|
|
|
|
def test_equal_outputs_with_other_facts_still_count_the_newest(self):
|
|
key = prompt_node_key(self.g, "a")
|
|
self.record(1, facts=Facts("old"), fields=(("a x", "", ""),))
|
|
self.record(2, facts=Facts("new"), fields=(("x", "a ", ""),))
|
|
items = self.runs._match(key, (("x",),))
|
|
self.assertEqual([item.facts for item in items], [Facts("new")])
|
|
|
|
def test_equal_outputs_and_fields_still_replace(self):
|
|
key = prompt_node_key(self.g, "a")
|
|
self.record(1)
|
|
self.record(2)
|
|
self.assertEqual([r.item_id for r in self.runs._receipts[key]], [2])
|
|
self.assertEqual(self.lookup(), [self.F])
|
|
|
|
def test_fields_dropped_by_the_budget_do_not_conflict(self):
|
|
self.runs = CachedRuns(context=lambda: self.context, max_field_bytes=0)
|
|
self.record(1)
|
|
self.runs._max_field_bytes = 100
|
|
self.record(2)
|
|
self.assertEqual(self.lookup(), [self.F])
|
|
|
|
def test_oversized_or_invalid_fields_are_not_kept(self):
|
|
for raw in ("x" * (module.MAX_FIELD_BYTES + 1), "\ud800"):
|
|
with self.subTest(size=len(raw)):
|
|
self.runs = CachedRuns(context=lambda: self.context)
|
|
self.record(1, fields=((raw, "", ""),))
|
|
self.assertIsNone(self.lookup())
|
|
self.assertEqual(self.runs._field_bytes, 0)
|
|
|
|
def test_old_receipts_give_up_their_fields_to_new_ones(self):
|
|
self.runs = CachedRuns(
|
|
context=lambda: self.context, max_keys=3, max_field_bytes=10
|
|
)
|
|
self.record(1, "a", fields=(("123456", "", ""),))
|
|
self.record(2, "b", fields=(("123456", "", ""),))
|
|
self.assertEqual(self.runs._field_bytes, 6)
|
|
self.assertIsNone(self.lookup("a"))
|
|
self.assertEqual(self.lookup("b"), [(("123456", "", ""),)])
|
|
# "a" still counts
|
|
self.assertIsNotNone(self.runs._match(prompt_node_key(self.g, "a"), (("x",),)))
|
|
# Larger than the whole budget: not kept, nothing else given up
|
|
self.record(3, "c", fields=(("12345678901", "", ""),))
|
|
self.assertEqual(self.runs._field_bytes, 6)
|
|
self.assertEqual(self.lookup("b"), [(("123456", "", ""),)])
|
|
|
|
def test_one_receipt_holds_at_most_one_runs_history(self):
|
|
with mock.patch.object(module, "MAX_RECEIPT_FIELD_BYTES", 8):
|
|
self.begin(1)
|
|
for index, raw in enumerate(("12345", "12345")):
|
|
self.context = ("j1", index)
|
|
self.runs.record_item("a", (raw,), Facts("h"), ((raw, "", ""),))
|
|
self.runs.end_attempt(1)
|
|
self.assertEqual(self.runs._field_bytes, 5)
|
|
self.assertIsNone(self.lookup(values=(("12345", "12345"),)))
|
|
|
|
def test_the_byte_count_follows_replacement_and_eviction(self):
|
|
self.begin(1)
|
|
self.runs.record_item("a", ("x",), Facts("h"), (("12345", "", ""),))
|
|
self.runs.record_item("a", ("x",), Facts("h"), (("123", "", ""),))
|
|
self.assertEqual(self.runs._field_bytes, 3)
|
|
self.runs.end_attempt(1)
|
|
for item_id, raw in ((2, "1"), (3, "22"), (4, "333")):
|
|
self.record(item_id, output=raw, fields=((raw, "", ""),))
|
|
# max_receipts=2: only the receipts of jobs 4 and 3 remain
|
|
self.assertEqual(self.runs._field_bytes, 5)
|
|
for item_id, name in ((5, "b"), (6, "c")): # max_keys=2: "a" is dropped
|
|
self.record(item_id, name, fields=(("1", "", ""),))
|
|
self.assertEqual(self.runs._field_bytes, 2)
|
|
self.assertEqual(
|
|
self.runs._field_bytes,
|
|
sum(module._bytes(r) for r in self.runs._receipts.values()),
|
|
)
|
|
|
|
|
|
class TestCachedJobsRecordFields(JobCase):
|
|
"""The real nodes keep the texts they gave the history capture (#403)."""
|
|
|
|
def lookup(self, g, node_id="pm", prompt_id="job-x"):
|
|
self.item_id += 1
|
|
self.runs.begin_attempt(
|
|
((self.item_id, prompt_id, g, {}, ["save"], {}), self.item_id)
|
|
)
|
|
values = cached_string_outputs(self.executor, node_id, g[node_id]["class_type"])
|
|
found = self.runs.recorded_fields(prompt_id, node_id, values)
|
|
self.runs.end_attempt(self.item_id)
|
|
return found
|
|
|
|
def test_each_node_kind_keeps_its_runtime_fields(self):
|
|
for kind in (
|
|
"PromptManager",
|
|
"PromptManagerFocused",
|
|
"PromptManagerSplit",
|
|
"PromptManagerText",
|
|
):
|
|
with self.subTest(kind=kind):
|
|
g = graph(kind)
|
|
self.first(g, text=f" {LLM} ")
|
|
self.assertEqual(
|
|
self.lookup(graph(kind, seed=2)), [((f" {LLM} ", "", ""),)]
|
|
)
|
|
|
|
def test_pair_keeps_positive_and_negative_fields(self):
|
|
g = graph("PromptManagerPair", negative_text="typed negative")
|
|
self.first(g, negative="typed negative")
|
|
self.assertEqual(
|
|
self.lookup(
|
|
graph("PromptManagerPair", negative_text="typed negative", seed=2)
|
|
),
|
|
[((LLM, "", ""), ("typed negative", "", ""))],
|
|
)
|
|
|
|
def test_a_disabled_node_keeps_no_fields(self):
|
|
g = graph()
|
|
g["pm"]["inputs"]["save_to_library"] = False
|
|
self.first(g)
|
|
g2 = graph(seed=2)
|
|
g2["pm"]["inputs"]["save_to_library"] = False
|
|
self.assertIsNone(self.lookup(g2))
|
|
|
|
|
|
class TestCacheReader(unittest.TestCase):
|
|
def setUp(self):
|
|
patcher = mock.patch.dict(sys.modules, host_modules())
|
|
patcher.start()
|
|
self.addCleanup(patcher.stop)
|
|
self.cache = LRUCache()
|
|
self.executor = types.SimpleNamespace(
|
|
caches=types.SimpleNamespace(outputs=self.cache)
|
|
)
|
|
|
|
def read(self, outputs, kind="PromptManager"):
|
|
self.cache.entries["n"] = Entry({}, outputs)
|
|
return cached_string_outputs(self.executor, "n", kind)
|
|
|
|
def test_reads_each_string_slot_in_order(self):
|
|
self.assertEqual(self.read([["c"], ["a", "b"]]), (("a", "b"),))
|
|
self.assertEqual(self.read([["t"]], "PromptManagerText"), (("t",),))
|
|
self.assertEqual(
|
|
self.read([[1], [2], ["p"], ["n"]], "PromptManagerPair"), (("p",), ("n",))
|
|
)
|
|
|
|
def test_fails_closed(self):
|
|
for outputs, kind in (
|
|
([["c"], ["a", 1]], "PromptManager"),
|
|
([["c"]], "PromptManager"),
|
|
([["c"], ["a"] * 65], "PromptManager"),
|
|
([[1], [2], ["p", "q"], ["n"]], "PromptManagerPair"),
|
|
([["c"], "text"], "PromptManager"),
|
|
([["c"], ["a"]], "CLIPTextEncode"),
|
|
([[]] * 5, "PromptManager"),
|
|
):
|
|
with self.subTest(outputs=outputs, kind=kind):
|
|
self.assertIsNone(self.read(outputs, kind))
|
|
self.cache.entries["m"] = ({}, [["c"], ["a"]]) # not a CacheEntry
|
|
self.assertIsNone(cached_string_outputs(self.executor, "m", "PromptManager"))
|
|
self.executor.caches.outputs = NullCache()
|
|
self.assertIsNone(self.read([["c"], ["a"]]))
|
|
|
|
|
|
class TestGraphHelpers(unittest.TestCase):
|
|
def test_upstream_nodes_and_closure(self):
|
|
g = graph()
|
|
self.assertEqual(
|
|
set(upstream_nodes(g, ["save"])),
|
|
{"save", "ks", "pm", "llm", "str", "ckpt", "other"},
|
|
)
|
|
self.assertEqual(upstream_nodes(g, ["missing"]), [])
|
|
self.assertEqual(upstream_nodes(None, ["save"]), [])
|
|
self.assertIsNone(upstream_closure(g, ["save"], limit=2))
|
|
self.assertIsNone(upstream_closure({"a": {"inputs": {"x": ["b", 0]}}}, ["a"]))
|
|
|
|
def test_find_pair_by_hashes_matches_the_text_lookup(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
db = PromptDatabase(os.path.join(tmp, "p.db"))
|
|
try:
|
|
pair_id = db.save_prompt_pair("cat", "")
|
|
self.assertEqual(
|
|
find_pair_id_by_hashes(db, generate_prompt_hash("cat"), None),
|
|
pair_id,
|
|
)
|
|
self.assertIsNone(
|
|
find_pair_id_by_hashes(
|
|
db, generate_prompt_hash("cat"), generate_prompt_hash("x")
|
|
)
|
|
)
|
|
finally:
|
|
db.close_all()
|
|
assert_closed(self, tmp)
|
|
|
|
|
|
class TestRegistration(unittest.TestCase):
|
|
def setUp(self):
|
|
self.modules = host_modules()
|
|
patcher = mock.patch.dict(sys.modules, self.modules)
|
|
patcher.start()
|
|
self.addCleanup(patcher.stop)
|
|
flag = mock.patch.object(module, "_registered", threading.Event())
|
|
flag.start()
|
|
self.addCleanup(flag.stop)
|
|
|
|
class Queue:
|
|
def __init__(self):
|
|
self.done = []
|
|
|
|
def get(self, timeout=None):
|
|
return self.result
|
|
|
|
def task_done(self, item_id, *args, **kwargs):
|
|
self.done.append(item_id)
|
|
|
|
def test_registers_once_and_follows_the_job(self):
|
|
queue = self.Queue()
|
|
server = types.SimpleNamespace(prompt_queue=queue)
|
|
runs = mock.Mock()
|
|
self.assertTrue(register_cached_runs(server, lambda: "db", runs))
|
|
self.assertFalse(register_cached_runs(server, lambda: "db", runs))
|
|
queue.result = (("item",), 5)
|
|
queue.get()
|
|
runs.begin_attempt.assert_called_once_with((("item",), 5))
|
|
executor = self.modules["execution"].PromptExecutor()
|
|
executor.server = server
|
|
executor.add_message("execution_start", {"prompt_id": "j"}, broadcast=False)
|
|
executor.add_message(
|
|
"execution_cached", {"prompt_id": "j", "nodes": ["n"]}, broadcast=False
|
|
)
|
|
runs.replay.assert_called_once()
|
|
self.assertEqual(runs.replay.call_args.args[1:3], ("j", ["n"]))
|
|
other = self.modules["execution"].PromptExecutor()
|
|
other.server = types.SimpleNamespace(prompt_queue=object())
|
|
other.add_message("execution_cached", {"prompt_id": "j", "nodes": []}, False)
|
|
self.assertEqual(runs.replay.call_count, 1)
|
|
queue.task_done(5, None, status=None)
|
|
runs.end_attempt.assert_called_once_with(5)
|
|
|
|
def test_hosts_without_the_interfaces_register_nothing(self):
|
|
class NoTaskDone:
|
|
def get(self):
|
|
return None
|
|
|
|
self.assertFalse(
|
|
register_cached_runs(
|
|
types.SimpleNamespace(prompt_queue=NoTaskDone()), lambda: None
|
|
)
|
|
)
|
|
self.assertFalse(register_cached_runs(types.SimpleNamespace(), lambda: None))
|
|
del self.modules["execution"].PromptExecutor.add_message
|
|
self.assertFalse(
|
|
register_cached_runs(
|
|
types.SimpleNamespace(prompt_queue=self.Queue()), lambda: None
|
|
)
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|