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
|
|
import unittest.mock as 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()
|