diff --git a/TESTING.md b/TESTING.md index 1ce783d..4f2cfd1 100644 --- a/TESTING.md +++ b/TESTING.md @@ -5283,7 +5283,7 @@ which used to save `null` under a link. Primitive or through `PreviewAny`, for PromptManager, PromptManagerText, Focused, Split and the Pair (positive, negative, both, neither linked, and a Primitive negative). A wildcard Primitive on PromptManager (queued as written) and Focused (one choice) saves - and counts what was queued. The linked text is saved, counted once, recorded as + and counts what was queued. The linked text is saved, counted once per job, recorded as the run's raw positive/negative, and after a scan the run's image links to it. Text only known at run time (`PreviewAny`) cannot be read from the image; the case only requires that it is never linked to another prompt (its attribution from the recorded @@ -5297,6 +5297,39 @@ PM_COMFY_ROOT=... PM_TEST_PYTHON=... npm run --prefix tools/vendor test:comfyui PM_COMFY_ROOT=... PM_TEST_PYTHON=... npm run --prefix tools/vendor test:executions -- execution-text-link ``` +## Cached runtime-linked runs (#347) + +Text only known at run time (a `PreviewAny` or LLM link, a list node) is counted by the +node when it executes; the queue hook cannot read it. When ComfyUI serves such a node from +its output cache, `utils/cached_runs.py` counts the job instead: every execution records +a receipt (digests of its STRING outputs and hashes of the text it stores, per run key and +queue item), and on `execution_cached` the receipt whose digests equal the cache entry's +text is replayed once. The role (positive, negative, which Pair outputs are used) and what +the hook already saw are recomputed from the job's own graph, and the node and replay +decide through the same `own_counts`. + +- Fail closed: no receipt in this process (after a restart, an external cache provider), + changed cached text, more than 64 list values, a node the job's outputs do not need, or + a host without the queue/executor interfaces: nothing is counted, as before #347. +- A cached node counts at job start like the hook: a job failing later still counts it. + A node that executes after its replay in the same job (RAM-pressure freed the entry) + skips its own count but still consumes the pending first-use catch-up. +- No backfill: earlier cached jobs keep their counts. The history page's "tracked runs" + can stay below the Library count for these prompts (#403). +- Tests: `tests/test_cached_runs.py` (whole jobs with the real nodes, a real database, + fake host caches); `tests/comfyui/execution-text-link.spec.mjs` "#347" cases on the + managed host: PreviewAny into each node kind, 3 jobs changing only the sampler seed, + count 3 in both renderers; in one renderer the Classic, LRU, RAM-pressure and no-cache + modes, A, B, A from the owned `PMCycledText` under LRU, typed text, a downstream + failure, a "run selected output" job that does not need the node, and a rewire to the + negative input. The RAM-pressure eviction path cannot be forced on a real host and is + covered by the Python tests. + +```sh +python3 -m unittest discover -s tests -p "test_cached_runs.py" +PM_COMFY_ROOT=... PM_TEST_PYTHON=... npm run --prefix tools/vendor test:executions -- execution-text-link --grep "#347" +``` + ## Durable execution ledger `npm run --prefix tools/vendor test:executions` runs 52 real execution cases diff --git a/__init__.py b/__init__.py index 1365a2d..41c6fba 100644 --- a/__init__.py +++ b/__init__.py @@ -195,6 +195,26 @@ except Exception as e: f"Failed to register usage hook: {e}" ) +# A cached node counts what its last execution of the same text counted (#347). +try: + from .py.config import server_instance + from .utils.cached_runs import register_cached_runs + + register_cached_runs(server_instance, _db_factory) +except Exception as e: + try: + from .utils.logging_config import get_logger + + get_logger("prompt_manager.init").error( + f"Failed to register cached-run counting: {e}" + ) + except Exception: + import logging + + logging.getLogger("prompt_manager.init").error( + f"Failed to register cached-run counting: {e}" + ) + def _history_db_factory() -> "PromptDatabase": """Initialize history on its worker even if monitoring initialization failed.""" diff --git a/prompt_manager.py b/prompt_manager.py index b9b3f5e..2dcf003 100644 --- a/prompt_manager.py +++ b/prompt_manager.py @@ -268,6 +268,7 @@ class PromptManager(PromptManagerBase, ComfyNodeABC): record_runtime_context( type(self).__name__, text, prepend_text, append_text, encoding_text ) + self._record_cached_item(unique_id, (encoding_text,), storage_text) return result @classmethod diff --git a/prompt_manager_base.py b/prompt_manager_base.py index 3a65531..62977b1 100644 --- a/prompt_manager_base.py +++ b/prompt_manager_base.py @@ -26,6 +26,13 @@ try: from .utils.prompt_graph import resolve_text, run_prompt_nodes from .utils.run_keys import executing_run_id, prompt_node_key from .utils.usage_tracking import PENDING_FIRST_USE + from .utils.cached_runs import ( + CACHED_RUNS, + Role, + Seen, + own_counts, + stored_facts, + ) except ImportError: import sys @@ -38,6 +45,13 @@ except ImportError: from utils.prompt_graph import resolve_text, run_prompt_nodes from utils.run_keys import executing_run_id, prompt_node_key from utils.usage_tracking import PENDING_FIRST_USE + from utils.cached_runs import ( + CACHED_RUNS, + Role, + Seen, + own_counts, + stored_facts, + ) class PromptManagerBase: @@ -71,8 +85,10 @@ class PromptManagerBase: """Save the prompt to the SQLite database. Runs are normally counted by the queue hook (utils/usage_tracking.py), so a - save only counts one when asked to (batch items whose text the hook could not - see) or when the hook queued this prompt before it existed. + save only counts one when asked to (text the hook could not see, unless + this job already counted the node from ComfyUI's cache) or when the hook + queued this prompt before it existed. The two are independent: the + pending catch-up is written even when ``count_run`` is False. Args: text: The prompt text @@ -167,8 +183,19 @@ class PromptManagerBase: batch_item = ( role["known"] and role["is_run_prompt"] and not role["seen_at_queue"] ) + counts = own_counts( + type(self).__name__, + stored_facts(text), + Role( + frozenset({"text"} if role["known"] and role["is_run_prompt"] else ()) + ), + Seen(text=role["seen_at_queue"]), + ) prompt_id = self._save_prompt_to_database( - text=text, category=category, tags=tags, count_run=batch_item + text=text, + category=category, + tags=tags, + count_run=bool(counts) and not self._replayed(unique_id), ) if role["is_run_prompt"] and not prompt_id: # The save failed: an older entry of this key must not stand in (#342) @@ -182,6 +209,34 @@ class PromptManagerBase: ) return prompt_id + def _cached_runs(self) -> Any: + return getattr(self, "cached_runs", CACHED_RUNS) + + def _replayed(self, unique_id: Any) -> bool: + """This job already counted the node from ComfyUI's cache (#347). + + A RAM-pressure cache can free the entry after ``execution_cached``, so + the node executes in a job that replayed it: its own count is skipped. + """ + try: + return unique_id is not None and self._cached_runs().replayed(unique_id) + except Exception as error: + self.logger.debug(f"Could not check the cached-run replay: {error}") + return False + + def _record_cached_item( + self, unique_id: Any, outputs: tuple, text: Any, negative_text: Any = None + ) -> None: + """Remember this execution's outputs and stored text for cached jobs (#347).""" + if unique_id is None: + return + try: + self._cached_runs().record_item( + unique_id, outputs, stored_facts(text, negative_text) + ) + except Exception as error: + self.logger.debug(f"Could not record the cached-run receipt: {error}") + @staticmethod def _run_identity(prompt_graph: Any, unique_id: Any) -> Dict[str, Any]: """The tracker's proof that an image belongs to this execution (#342).""" diff --git a/prompt_manager_pair.py b/prompt_manager_pair.py index 5c90333..df62d63 100644 --- a/prompt_manager_pair.py +++ b/prompt_manager_pair.py @@ -12,6 +12,14 @@ else: try: from .prompt_manager import IO, InputTypeDict, PromptManagerFocused + from .utils.cached_runs import ( + PAIR, + Role, + Seen, + own_counts, + pair_components, + stored_facts, + ) from .utils.hashing import generate_prompt_hash from .utils.pair_execution import record_pair_execution from .utils.prompt_graph import intrinsic_pair_nodes, run_prompt_components @@ -23,6 +31,14 @@ try: ) except ImportError: from prompt_manager import IO, InputTypeDict, PromptManagerFocused + from utils.cached_runs import ( + PAIR, + Role, + Seen, + own_counts, + pair_components, + stored_facts, + ) from utils.hashing import generate_prompt_hash from utils.pair_execution import record_pair_execution from utils.prompt_graph import intrinsic_pair_nodes, run_prompt_components @@ -165,6 +181,9 @@ class PromptManagerPair(PromptManagerFocused): record_runtime_context( type(self).__name__, negative_text, "", "", negative_text, role="negative" ) + self._record_cached_item( + unique_id, (positive, negative_text), text, negative_text + ) return result def _save_and_track_pair( @@ -209,36 +228,35 @@ class PromptManagerPair(PromptManagerFocused): and generate_prompt_hash(resolved_negative) == generate_prompt_hash(negative_text) ) + negative_seen = resolved_negative is not None and generate_prompt_hash( + resolved_negative + ) == generate_prompt_hash(negative_text) positive_pending = getattr(self, "pending_first_use", PENDING_FIRST_USE) pair_pending = getattr(self, "pending_pair_first_use", PENDING_PAIR_FIRST_USE) - counted = set() - for field, component_text, resolved_text, component_id in ( - ("text", text, resolved_positive, prompt_id), - ( - "negative_text", - negative_text, - resolved_negative, - pair["negative_prompt_id"], - ), - ): - component_hash = generate_prompt_hash(component_text) - if ( - field not in used_fields - or not component_id - or component_hash in counted - ): + facts = stored_facts(text, negative_text) + role = Role(frozenset(used_fields), pair_used) + # Own counts are skipped when this job already counted the node from + # ComfyUI's cache (#347); the pending first-use catch-up never is. + counts = ( + [] + if self._replayed(unique_id) + else own_counts( + PAIR, facts, role, Seen(positive_seen, negative_seen, pair_seen) + ) + ) + component_ids = {"text": prompt_id, "negative_text": pair["negative_prompt_id"]} + for field, component_hash in pair_components(facts, role): + component_id = component_ids[field] + if not component_id: continue - counted.add(component_hash) - seen_at_queue = ( - resolved_text is not None - and generate_prompt_hash(resolved_text) == component_hash + runs = positive_pending.consume(component_hash) + int( + ("prompt", component_hash) in counts ) - runs = positive_pending.consume(component_hash) + int(not seen_at_queue) if runs: self.db.record_prompt_use(component_id, times=runs) if pair_used: runs = pair_pending.consume(pair_run_key(text, negative_text)) + int( - not pair_seen + ("pair", facts.positive_hash, facts.negative_hash) in counts ) if runs: self.db.record_prompt_pair_use(pair_id, times=runs) diff --git a/prompt_manager_text.py b/prompt_manager_text.py index b09945d..6e316f9 100644 --- a/prompt_manager_text.py +++ b/prompt_manager_text.py @@ -239,6 +239,7 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC): record_runtime_context( type(self).__name__, text, prepend_text, append_text, final_text ) + self._record_cached_item(unique_id, (final_text,), storage_text) return result @classmethod diff --git a/tests/comfyui/execution-text-link.spec.mjs b/tests/comfyui/execution-text-link.spec.mjs index 8eefb6d..20e60de 100644 --- a/tests/comfyui/execution-text-link.spec.mjs +++ b/tests/comfyui/execution-text-link.spec.mjs @@ -11,7 +11,7 @@ import { installHostHooks } from './host-fixture.mjs'; // The worker's owned host must be up before the shared host hooks open the page. test.beforeEach(async ({ execution }) => { expect(execution.databasePath).toBeTruthy(); }); -installHostHooks(test, ['PromptManager', 'PromptManagerText', 'PromptManagerFocused', 'PromptManagerPair', 'PMHistoryInputs']); +installHostHooks(test, ['PromptManager', 'PromptManagerText', 'PromptManagerFocused', 'PromptManagerPair', 'PMHistoryInputs', 'PMCycledText']); // `primitive`: a core string node the queue-time resolver follows. `legacy`: the frontend's virtual Primitive node, // whose value is queued as literal text. `preview`: PrimitiveString -> PreviewAny, the shape of the maintainer's @@ -37,6 +37,7 @@ function build(page, { kind, positive, negative, prefix }) { primitive.widgets[0].value = side.text; return; } + if (side.via === 'cycled') { connect(add('PMCycledText', 0, y), 0, prompt, name); return; } const text = add('PrimitiveString', 0, y); text.widgets.find(widget => widget.name === 'value').value = side.text; let source = text; @@ -75,9 +76,9 @@ function build(page, { kind, positive, negative, prefix }) { }, { kind, positive, negative, prefix }); } -async function run(execution, page, testCase) { +async function run(execution, page, testCase, built) { const prefix = `textlink/${randomUUID()}`; - const { output, workflow, promptNode } = await build(page, { ...testCase, prefix }); + const { output, workflow, promptNode } = built ?? await build(page, { ...testCase, prefix }); const id = randomUUID(); const response = await execution.request.post('/prompt', { data: { prompt: output, prompt_id: id, extra_data: { extra_pnginfo: { workflow } } } }); expect(response.ok(), await response.text()).toBe(true); @@ -193,3 +194,164 @@ for (const [label, linked, negativeVia] of [['positive', ['text'], 'primitive'], await expectImagesLinked(execution, evidence, sides.text.via === 'preview'); }); } + +// #347: a cached PromptManager node counts every job, as an executing one does. Jobs after the first change only +// the sampler seed, so ComfyUI serves the PromptManager node (and its PreviewAny source) from its output cache. +const COUNTS = `import json, sqlite3, sys, pathlib +db = sqlite3.connect(pathlib.Path(sys.argv[1]).resolve().as_uri() + "?mode=ro", uri=True) +one = lambda sql, key: (db.execute(sql, (int(key),)).fetchone() or [None])[0] if key != "-" else None +print(json.dumps({"positive": one("SELECT run_count FROM prompts WHERE id=?", sys.argv[2]), + "negative": one("SELECT run_count FROM prompts WHERE id=?", sys.argv[3]), + "pair": one("SELECT run_count FROM prompt_pairs WHERE id=?", sys.argv[4])}))`; +async function runCounts(execution, context) { + const key = value => value == null ? '-' : String(value); + const { stdout } = await promisify(execFile)(process.env.PM_TEST_PYTHON || 'python', + ['-c', COUNTS, execution.databasePath, key(context.positive_prompt_id), key(context.negative_prompt_id), key(context.pair_id)], { timeout: 5000 }); + return JSON.parse(stdout); +} + +async function setCache(execution, mode) { + expect((await execution.request.post('/pm_execution_fixture/cache', { data: { mode } })).ok()).toBe(true); +} + +// Queue an already built workflow again with another sampler seed; returns the job's history item and cached nodes. +async function queueAgain(execution, first, { seed, edit, targets, status = 'success' } = {}) { + const output = structuredClone(first.output); + const sampler = Object.values(output).find(node => node.class_type === 'KSampler'); + sampler.inputs.seed = seed; + edit?.(output); + const id = randomUUID(); + const data = { prompt: output, prompt_id: id, extra_data: { extra_pnginfo: { workflow: first.workflow } } }; + if (targets) data.partial_execution_targets = targets; + const response = await execution.request.post('/prompt', { data }); + expect(response.ok(), await response.text()).toBe(true); + let item; + await expect.poll(async () => { + item = (await (await execution.request.get(`/history/${id}`)).json())[id]; + return item?.status?.status_str ?? null; + }, { timeout: 20000 }).toBe(status); + const cached = item.status.messages.find(([event]) => event === 'execution_cached')?.[1]?.nodes ?? []; + return { id, item, cached }; +} + +// The recorded context of one job, once its capture reached SQLite. +async function capturedContext(execution, id) { + let result; + await expect.poll(async () => { + result = await readExecutionRuns(execution.databasePath, id); + const row = result.rows[0]; + return row ? result.captures[row.run_id].state : null; + }, { timeout: 10000 }).toMatch(/^(complete|partial)$/); + return result.captures[result.rows[0].run_id].contexts[0]; +} + +async function firstRun(execution, page, testCase) { + const prefix = `textlink/${randomUUID()}`; + const built = await build(page, { ...testCase, prefix }); + const evidence = await run(execution, page, testCase, built); + return { ...evidence, workflow: built.workflow }; +} + +for (const kind of ['PromptManager', 'PromptManagerFocused', 'PromptManagerSplit', 'PromptManagerText', 'PromptManagerPair']) { + test(`${kind} counts each cached re-run of text linked from PreviewAny (#347)`, async ({ execution, page }) => { + test.setTimeout(90000); + await setCache(execution, 'RAM_PRESSURE'); + const pair = kind === 'PromptManagerPair'; + const testCase = { kind, positive: { typed: 'Typed under the link', text: unique(`Cached ${kind}`), via: 'preview' }, + negative: pair ? { typed: unique('Typed negative') } : undefined }; + const first = await firstRun(execution, page, testCase); + for (const seed of [2, 3]) { + const again = await queueAgain(execution, first, { seed }); + expect(again.cached, 'ComfyUI served the PromptManager node from its cache').toContain(first.promptNode); + } + const counts = await runCounts(execution, first.context); + expect(counts.positive).toBe(3); + if (pair) expect(counts).toMatchObject({ negative: 0, pair: 3 }); + }); +} + +test.describe('cached runs in one renderer (#347)', () => { + // Cache behavior does not depend on the renderer: one project runs these. + test.beforeEach(({}, testInfo) => { test.skip(testInfo.project.name !== 'classic-start', 'covered in classic-start'); }); + const linked = () => ({ kind: 'PromptManager', positive: { typed: 'Typed under the link', text: unique('Cached'), via: 'preview' } }); + + for (const mode of ['CLASSIC', 'LRU', 'RAM_PRESSURE', 'NONE']) { + test(`counts 3 jobs under the ${mode} cache`, async ({ execution, page }) => { + test.setTimeout(90000); + await setCache(execution, mode); + try { + const first = await firstRun(execution, page, linked()); + for (const seed of [2, 3]) { + const again = await queueAgain(execution, first, { seed }); + if (mode === 'NONE') expect(again.cached).not.toContain(first.promptNode); + else expect(again.cached).toContain(first.promptNode); + } + expect((await runCounts(execution, first.context)).positive).toBe(3); + } finally { await setCache(execution, 'RAM_PRESSURE'); } + }); + } + + test('alternating upstream values A, B, A under LRU each count their own prompt', async ({ execution, page }) => { + test.setTimeout(90000); + await setCache(execution, 'LRU'); + try { + const [a, b] = [unique('Cycled A'), unique('Cycled B')]; + const cycle = async value => expect((await execution.request.post('/pm_execution_fixture/cycled', { data: { value } })).ok()).toBe(true); + await cycle(a); + const first = await firstRun(execution, page, { kind: 'PromptManager', positive: { typed: 'Typed under the link', text: a, via: 'cycled' } }); + expect(first.context.positive.raw).toBe(a); + await cycle(b); + const second = await queueAgain(execution, first, { seed: 2 }); + expect(second.cached).not.toContain(first.promptNode); + await cycle(a); + const third = await queueAgain(execution, first, { seed: 3 }); + expect(third.cached, 'LRU kept A').toContain(first.promptNode); + expect((await runCounts(execution, first.context)).positive).toBe(2); + const bContext = await capturedContext(execution, second.id); + expect(bContext.positive.raw).toBe(b); + expect((await runCounts(execution, bContext)).positive).toBe(1); + } finally { await setCache(execution, 'RAM_PRESSURE'); } + }); + + test('typed text served from the cache still counts once per job', async ({ execution, page }) => { + test.setTimeout(90000); + await setCache(execution, 'RAM_PRESSURE'); + const first = await firstRun(execution, page, { kind: 'PromptManager', positive: { typed: unique('Typed cached') } }); + for (const seed of [2, 3]) expect((await queueAgain(execution, first, { seed })).cached).toContain(first.promptNode); + expect((await runCounts(execution, first.context)).positive).toBe(3); + }); + + test('a job failing downstream of a cached node counts it once; a job that does not need it, never', async ({ execution, page }) => { + test.setTimeout(90000); + await setCache(execution, 'RAM_PRESSURE'); + const first = await firstRun(execution, page, linked()); + const failing = await queueAgain(execution, first, { seed: 2, status: 'error', edit: output => { + output.pm_failure = { class_type: 'PMHistoryOutput', inputs: { value: [Object.keys(output).find(id => output[id].class_type === 'KSampler'), 0], mode: 'failure' } }; + } }); + expect(failing.cached).toContain(first.promptNode); + expect((await runCounts(execution, first.context)).positive).toBe(2); + // "Run selected output": only an unrelated output node; ComfyUI still reports the cached PromptManager node + const unrelated = await queueAgain(execution, first, { seed: 2, targets: ['pm_unrelated'], edit: output => { + output.pm_unrelated = { class_type: 'PMHistoryCached', inputs: { value: unique('unrelated') } }; + } }); + expect(unrelated.cached).toContain(first.promptNode); + expect((await runCounts(execution, first.context)).positive).toBe(2); + }); + + test('a cached node rewired to the negative input does not count the positive prompt', async ({ execution, page }) => { + test.setTimeout(90000); + await setCache(execution, 'RAM_PRESSURE'); + const first = await firstRun(execution, page, linked()); + const rewired = await queueAgain(execution, first, { seed: 2, edit: output => { + const sampler = Object.values(output).find(node => node.class_type === 'KSampler'); + [sampler.inputs.positive, sampler.inputs.negative] = [sampler.inputs.negative, sampler.inputs.positive]; + } }); + expect(rewired.cached).toContain(first.promptNode); + expect((await runCounts(execution, first.context)).positive).toBe(1); + // The history capture of a cached job is unchanged: effective text from the cache, no stored-prompt hint + const cachedContext = await capturedContext(execution, (await queueAgain(execution, first, { seed: 3 })).id); + expect(cachedContext.positive.effective).toBe(first.context.positive.effective); + expect(cachedContext.positive_prompt_id).toBeNull(); + expect((await runCounts(execution, first.context)).positive).toBe(2); + }); +}); diff --git a/tests/comfyui/fixtures/ledger_nodes.py b/tests/comfyui/fixtures/ledger_nodes.py index 844249a..a0c5744 100644 --- a/tests/comfyui/fixtures/ledger_nodes.py +++ b/tests/comfyui/fixtures/ledger_nodes.py @@ -201,6 +201,39 @@ NODE_CLASS_MAPPINGS.update( {"PMHistoryImages": PMHistoryImages, "PMHistoryImageOutput": PMHistoryImageOutput} ) +# #347: text whose value (and cache signature) follows a value the test sets, so +# one queued graph can serve A, then B, then A again (each cached under LRU). +_CYCLED = {"value": "cycled"} + + +class PMCycledText: + @classmethod + def INPUT_TYPES(cls) -> dict: + return {"required": {}} + + RETURN_TYPES = ("STRING",) + FUNCTION = "execute" + CATEGORY = "testing/execution-history" + + @classmethod + def IS_CHANGED(cls, **kwargs: object) -> str: + return _CYCLED["value"] + + def execute(self) -> tuple: + return (_CYCLED["value"],) + + +NODE_CLASS_MAPPINGS["PMCycledText"] = PMCycledText + + +@PromptServer.instance.routes.post("/pm_execution_fixture/cycled") +async def cycled_text(request: web.Request) -> web.Response: + value = (await request.json())["value"] + if type(value) is not str: + raise web.HTTPBadRequest() + _CYCLED["value"] = value + return web.json_response({"value": value}) + @PromptServer.instance.routes.post("/pm_execution_fixture/metadata") async def metadata(request: web.Request) -> web.Response: diff --git a/tests/test_cached_runs.py b/tests/test_cached_runs.py new file mode 100644 index 0000000..d6ea495 --- /dev/null +++ b/tests/test_cached_runs.py @@ -0,0 +1,823 @@ +"""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() + if kind == "PromptManagerText": + return node.process_text(text, prompt=g, unique_id=node_id) + if kind == "PromptManagerPair": + return node.encode_prompt( + clip, text, negative_text=negative, prompt=g, unique_id=node_id + ) + return node.encode_prompt(clip, text, prompt=g, unique_id=node_id) + + 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_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 TestReceipts(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) + + 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 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() diff --git a/tests/test_pair_usage_tracking.py b/tests/test_pair_usage_tracking.py index 9039f29..f451cad 100644 --- a/tests/test_pair_usage_tracking.py +++ b/tests/test_pair_usage_tracking.py @@ -85,6 +85,23 @@ class TestPairUsageTracking(unittest.TestCase): self.db.get_prompt_by_id(pair["negative_prompt_id"])["run_count"], 0 ) + def test_execution_records_a_cached_run_receipt_for_both_editors(self): + from utils.cached_runs import Facts + + node = self.node() + node.cached_runs = Mock() + node.cached_runs.replayed.return_value = False + node.encode_prompt( + Mock(), + "cat", + negative_text="", + prompt=pair_graph(negative=""), + unique_id="1", + ) + node.cached_runs.record_item.assert_called_once_with( + "1", ("cat", ""), Facts(usage_tracking.generate_prompt_hash("cat"), None) + ) + def test_validation_rejection_never_counts_or_saves(self): self.submit(pair_graph()) self.assertEqual(self.db.list_prompt_pairs()["total"], 0) diff --git a/tests/test_prompt_usage.py b/tests/test_prompt_usage.py index c460664..c1dd27f 100644 --- a/tests/test_prompt_usage.py +++ b/tests/test_prompt_usage.py @@ -278,5 +278,49 @@ class TestNodeRecordsUse(UsageTestCase): self.assertEqual(self.db.get_prompt_by_id(pid)["run_count"], 2) +class TestNodeRecordsCachedRunReceipt(UsageTestCase): + """Each execution leaves a receipt for cached jobs (#347).""" + + GRAPH = { + "1": {"class_type": "PromptManagerText", "inputs": {"text": "typed text"}}, + } + + def _node(self): + from unittest import mock + + from prompt_manager_text import PromptManagerText + + node = PromptManagerText.__new__(PromptManagerText) + node.db = self.db + node.logger = __import__("logging").getLogger("test.prompt_usage") + node.prompt_tracker = mock.Mock() + node.comfyui_integration = mock.Mock() + node._inject_lora_trigger_words = lambda text: text + node.cached_runs = mock.Mock() + node.cached_runs.replayed.return_value = False + return node + + def test_receipt_holds_the_returned_text_and_the_stored_hash(self): + from utils.cached_runs import Facts + + node = self._node() + node.process_text( + "typed text", prepend_text="pre", prompt=self.GRAPH, unique_id="1" + ) + node.cached_runs.record_item.assert_called_once_with( + "1", ("pre typed text",), Facts(generate_prompt_hash("typed text")) + ) + # Hook-visible text: the node counts nothing itself + row = self.db.get_prompt_by_hash(generate_prompt_hash("typed text")) + self.assertEqual(row["run_count"], 0) + + def test_empty_editor_records_empty_facts(self): + from utils.cached_runs import Facts + + node = self._node() + node.process_text(" ", prompt=self.GRAPH, unique_id="1") + node.cached_runs.record_item.assert_called_once_with("1", (" ",), Facts(None)) + + if __name__ == "__main__": unittest.main() diff --git a/tests/vendor/comfyui-config.test.mjs b/tests/vendor/comfyui-config.test.mjs index 61b10ef..2cb9587 100644 --- a/tests/vendor/comfyui-config.test.mjs +++ b/tests/vendor/comfyui-config.test.mjs @@ -34,15 +34,16 @@ test('execution script discovers only the new matrix and preserves all 78 origin assert.match(executions, /\[classic-start\]/); assert.match(executions, /\[nodes2-start\]/); // + the fresh-host library location case, + #342 image attribution, + #343 linked text, - // + #332 image category routes on a generated output (image-attribution.spec.mjs, 2 projects) - assert.match(executions, /Total: 92 tests in 5 files/); + // + #332 image category routes on a generated output (image-attribution.spec.mjs, 2 projects), + // + #347 cached runtime-linked runs: five node kinds and eight one-renderer cases (execution-text-link.spec.mjs, 2 projects) + assert.match(executions, /Total: 118 tests in 5 files/); for (const project of ['classic-start', 'nodes2-start']) { const cases = executions.split('\n').filter(line => line.includes(`[${project}]`)); assert.equal(cases.filter(line => line.includes('execution-history.spec.mjs')).length, 15); assert.equal(cases.filter(line => line.includes('execution-outputs.spec.mjs')).length, 8); assert.equal(cases.filter(line => line.includes('library-location.spec.mjs')).length, 1); assert.equal(cases.filter(line => line.includes('image-attribution.spec.mjs')).length, 3); - assert.equal(cases.filter(line => line.includes('execution-text-link.spec.mjs')).length, 19); + assert.equal(cases.filter(line => line.includes('execution-text-link.spec.mjs')).length, 32); } assert.doesNotMatch(executions, /prompt-manager.spec/); }); diff --git a/utils/cached_runs.py b/utils/cached_runs.py new file mode 100644 index 0000000..fd73cd2 --- /dev/null +++ b/utils/cached_runs.py @@ -0,0 +1,532 @@ +"""Count a cached PromptManager node like the execution it reuses (#347). + +The queue hook (``usage_tracking``) counts text it can read from the queued +graph. Text only known at run time (a PreviewAny or LLM link, a list node) is +counted by the node itself when it executes. When ComfyUI later serves the node +from its output cache the node does not execute, so without this module that +job is not counted. + +When a PromptManager-family node executes it records a *receipt item*: the +role-independent facts its own count is made from (the stored text hashes) +and digests of the STRING outputs it returned. When ComfyUI reports the node +as cached for a later job (``execution_cached``), the receipt whose digests +equal the text in the cache entry is replayed once for that job. The node's +role (positive, negative, which Pair outputs are used) and what the hook +already saw are recomputed from that job's own graph, and both paths count +through the same ``own_counts`` so they cannot drift. + +Fail closed: when the cached text cannot be matched to a receipt recorded in +this process, nothing is counted (the behavior before #347). In memory only and +bounded; starts no threads. +""" + +import hashlib +import threading +from collections import OrderedDict +from typing import Any, Callable, Dict, FrozenSet, List, NamedTuple, Optional, Tuple + +try: + from .hashing import generate_prompt_hash + from .logging_config import get_logger + from .prompt_graph import ( + PROMPT_MANAGER_TYPES, + intrinsic_pair_nodes, + resolve_text, + run_prompt_components, + run_prompt_nodes, + ) + from .prompt_pair_resolver import resolve_text_input + from .queue_observers import subscribe_dequeue, subscribe_task_done + from .run_keys import prompt_node_key, upstream_nodes + from .usage_tracking import find_pair_id_by_hashes, survive_switch +except ImportError: + from utils.hashing import generate_prompt_hash + from utils.logging_config import get_logger + from utils.prompt_graph import ( + PROMPT_MANAGER_TYPES, + intrinsic_pair_nodes, + resolve_text, + run_prompt_components, + run_prompt_nodes, + ) + from utils.prompt_pair_resolver import resolve_text_input + from utils.queue_observers import subscribe_dequeue, subscribe_task_done + from utils.run_keys import prompt_node_key, upstream_nodes + from utils.usage_tracking import find_pair_id_by_hashes, survive_switch + +logger = get_logger("prompt_manager.cached_runs") + +PAIR = "PromptManagerPair" +MAX_KEYS = 256 # run keys remembered (the same bound as MAX_KEYED_RUNS) +MAX_RECEIPTS = 4 # receipts per run key: a recurring upstream value (A, B, A) +MAX_ITEMS = 64 # list items per receipt (the cache reader's bound) +MAX_CACHED_NODES = 500 # the execution_cached guard the history capture uses +_EMPTY_HASH = generate_prompt_hash("") + +Count = Tuple[Any, ...] # ("prompt", hash) or ("pair", positive, negative or None) + + +class Facts(NamedTuple): + """What one execution stored: hashes of the stored text, None when nothing. + + ``positive_hash`` is None for an empty editor (or saving disabled); + ``negative_hash`` (Pair only) is None for an empty negative, which the pair + stores as ``negative_prompt_id IS NULL``. + """ + + positive_hash: Optional[str] + negative_hash: Optional[str] = None + + +class Role(NamedTuple): + """How a job's graph uses the node: which editors reach a sampler's positive.""" + + used_fields: FrozenSet[str] + pair_used: bool = False + + +class Seen(NamedTuple): + """Which of the node's texts the queue hook resolved, and so already counted.""" + + text: bool = False + negative_text: bool = False + pair: bool = False + + +class _Item(NamedTuple): + outputs: Tuple[str, ...] + facts: Facts + + +class _Receipt: + __slots__ = ("item_id", "items") + + def __init__(self, item_id: int) -> None: + self.item_id = item_id + self.items: Dict[int, _Item] = {} + + def outputs(self) -> Dict[int, Tuple[str, ...]]: + return {index: item.outputs for index, item in self.items.items()} + + +class _Attempt(NamedTuple): + item_id: int + prompt_id: str + graph: Dict[str, Any] + nodes: Dict[str, Tuple[str, str]] # node id -> (class, run key) + replayed: set + + +def text_digest(value: str) -> str: + """Digest of one returned STRING value, exactly as returned.""" + return hashlib.sha256(value.encode("utf-8", "surrogatepass")).hexdigest() + + +def stored_facts(text: Any, negative_text: Any = None) -> Facts: + """The facts of an execution that stores ``text`` (and a Pair's negative).""" + positive = text if isinstance(text, str) and text.strip() else None + negative = ( + negative_text + if isinstance(negative_text, str) and negative_text.strip() + else None + ) + return Facts( + generate_prompt_hash(positive) if positive is not None else None, + generate_prompt_hash(negative) if negative is not None else None, + ) + + +def pair_components(facts: Facts, role: Role) -> List[Tuple[str, str]]: + """``(field, hash)`` of each used, stored Pair component; each hash once.""" + found, counted = [], set() + for field, prompt_hash in ( + ("text", facts.positive_hash), + ("negative_text", facts.negative_hash), + ): + if field not in role.used_fields or prompt_hash is None: + continue + if prompt_hash in counted: + continue + counted.add(prompt_hash) + found.append((field, prompt_hash)) + return found + + +def own_counts(node_type: str, facts: Facts, role: Role, seen: Seen) -> List[Count]: + """What one execution of a node counts itself: text the hook could not see. + + Shared by the executing node and by replay. Excludes the pending first-use + catch-up, which only an execution that saves the prompt can consume. + """ + if facts.positive_hash is None: + return [] + if node_type != PAIR: + if "text" in role.used_fields and not seen.text: + return [("prompt", facts.positive_hash)] + return [] + counts: List[Count] = [ + ("prompt", prompt_hash) + for field, prompt_hash in pair_components(facts, role) + if not getattr(seen, field) + ] + if role.pair_used and not seen.pair: + counts.append(("pair", facts.positive_hash, facts.negative_hash)) + return counts + + +def graph_role(graph: Any, node_id: str, node_type: str) -> Role: + """The node's role in a queued graph, by the functions the nodes use.""" + if node_type != PAIR: + used = {"text"} if node_id in run_prompt_nodes(graph) else set() + return Role(frozenset(used)) + used = { + field for source, field in run_prompt_components(graph) if source == node_id + } + return Role(frozenset(used), node_id in intrinsic_pair_nodes(graph)) + + +def _same(resolved: Optional[str], prompt_hash: Optional[str]) -> bool: + return resolved is not None and generate_prompt_hash(resolved) == prompt_hash + + +def graph_seen(graph: Any, node_id: str, node_type: str, facts: Facts) -> Seen: + """Which texts the queue hook resolved from ``graph`` (as the node decides).""" + if node_type != PAIR: + return Seen(text=_same(resolve_text(graph, node_id), facts.positive_hash)) + positive = _same(resolve_text_input(graph, node_id, "text"), facts.positive_hash) + negative = _same( + resolve_text_input(graph, node_id, "negative_text"), + facts.negative_hash or _EMPTY_HASH, + ) + return Seen(positive, negative, positive and negative) + + +def _host_context() -> Tuple[Optional[str], Optional[int]]: + """(prompt id, list index) of the node executing on this thread, if known.""" + try: + from comfy_execution.utils import get_executing_context + except Exception: + return None, None + try: + context = get_executing_context() + except Exception: + return None, None + prompt_id = getattr(context, "prompt_id", None) + index = getattr(context, "list_index", None) + return ( + prompt_id if isinstance(prompt_id, str) and prompt_id else None, + index if type(index) is int else None, + ) + + +def write_count(db: Any, count: Count) -> bool: + """Record one counted run; False when the prompt or pair no longer exists.""" + + def write() -> bool: + if count[0] == "prompt": + row = db.get_prompt_by_hash(count[1]) + if row: + db.record_prompt_use(row["id"]) + return bool(row) + pair_id = find_pair_id_by_hashes(db, count[1], count[2]) + if pair_id: + db.record_prompt_pair_use(pair_id) + return bool(pair_id) + + return bool(survive_switch(db, write)) + + +class CachedRuns: + """Receipts of executed nodes and the one job currently executing.""" + + def __init__( + self, + *, + max_keys: int = MAX_KEYS, + max_receipts: int = MAX_RECEIPTS, + max_items: int = MAX_ITEMS, + context: Callable[[], Tuple[Optional[str], Optional[int]]] = _host_context, + ) -> None: + self._max_keys = max_keys + self._max_receipts = max_receipts + self._max_items = max_items + self._context = context + self._receipts: "OrderedDict[str, List[_Receipt]]" = OrderedDict() + self._attempt: Optional[_Attempt] = None + self._lock = threading.Lock() + + # -- the job ------------------------------------------------------------ + + def begin_attempt(self, result: Any) -> None: + """A job was dequeued: ``result`` is ``PromptQueue.get``'s return value.""" + attempt = None + try: + item, item_id = result[0], result[1] + prompt_id, graph, outputs = item[1], item[2], item[4] + if type(item_id) is int and isinstance(prompt_id, str): + if isinstance(graph, dict): + attempt = _Attempt( + item_id, prompt_id, graph, _job_nodes(graph, outputs), set() + ) + except Exception as error: + logger.debug(f"Dequeued job not tracked for cached runs: {error}") + with self._lock: + self._attempt = attempt + + def end_attempt(self, item_id: Any) -> None: + """The job finished: its receipts are complete and replace equal older ones. + + Deduplicating only complete receipts keeps a list execution that is + still recording ([A] so far of [A, B]) from dropping an older [A]. + """ + with self._lock: + if self._attempt is not None and self._attempt.item_id == item_id: + self._attempt = None + for run_key, receipts in self._receipts.items(): + if receipts and receipts[0].item_id == item_id: + self._receipts[run_key] = _deduplicated(receipts) + + def _current(self) -> Tuple[Optional[_Attempt], Optional[int]]: + """The attempt the executing node belongs to, and its list index.""" + prompt_id, index = self._context() + with self._lock: + attempt = self._attempt + if attempt is None or prompt_id is None or prompt_id != attempt.prompt_id: + return None, None + return attempt, index + + # -- node execution ----------------------------------------------------- + + def replayed(self, node_id: Any) -> bool: + """Whether this job already counted the executing node by replay.""" + attempt, _ = self._current() + if attempt is None: + return False + with self._lock: + return str(node_id) in attempt.replayed + + def record_item(self, node_id: Any, outputs: Tuple[Any, ...], facts: Facts) -> None: + """Remember what this execution returned and was entitled to count.""" + try: + attempt, index = self._current() + node = attempt.nodes.get(str(node_id)) if attempt else None + if node is None: + return + index = 0 if index is None else index + if not 0 <= index < self._max_items: + return + if any(not isinstance(value, str) for value in outputs): + return + item = _Item(tuple(text_digest(value) for value in outputs), facts) + self._store(node[1], attempt.item_id, index, item) + except Exception as error: + logger.debug(f"Could not record a cached-run receipt: {error}") + + def _store(self, run_key: str, item_id: int, index: int, item: _Item) -> None: + with self._lock: + receipts = self._receipts.pop(run_key, []) + if receipts and receipts[0].item_id == item_id: + receipt = receipts[0] + else: + # The previous job's receipt is complete (its task_done may + # never have fired): it replaces older receipts it equals. + receipts = _deduplicated(receipts) + receipt = _Receipt(item_id) + receipts.insert(0, receipt) + receipt.items[index] = item + self._receipts[run_key] = receipts[: self._max_receipts] + while len(self._receipts) > self._max_keys: + self._receipts.popitem(last=False) + + # -- cached nodes ------------------------------------------------------- + + def replay( + self, + executor: Any, + prompt_id: Any, + nodes: Any, + db_factory: Callable[[], Any], + reader: Optional[Callable[[Any, str, str], Any]] = None, + ) -> None: + """Count the job's cached nodes whose cached text a receipt proves.""" + if reader is None: + reader = _host().cached_string_outputs + with self._lock: + attempt = self._attempt + if attempt is None or prompt_id != attempt.prompt_id: + return + if type(nodes) not in (list, tuple) or len(nodes) > MAX_CACHED_NODES: + return + db = None + for node_id in nodes: + node = attempt.nodes.get(node_id) if type(node_id) is str else None + if node is None: + continue + node_type, run_key = node + items = self._match(run_key, reader(executor, node_id, node_type)) + if items is None: + logger.debug( + f"Cached node {node_id} not counted: no matching execution" + ) + continue + with self._lock: + if node_id in attempt.replayed: + continue # a repeated message or id never counts twice + attempt.replayed.add(node_id) + counts: List[Count] = [] + role = graph_role(attempt.graph, node_id, node_type) + for item in items: + seen = graph_seen(attempt.graph, node_id, node_type, item.facts) + counts.extend(own_counts(node_type, item.facts, role, seen)) + if not counts: + continue + try: + db = db if db is not None else db_factory() + for count in counts: + if not write_count(db, count): + logger.info( + f"Cached node {node_id}: its {count[0]} no longer exists, " + "not counted" + ) + except Exception as error: + logger.warning(f"Could not count cached node {node_id}: {error}") + + def _match(self, run_key: str, values: Any) -> Optional[List[_Item]]: + """The items of the one receipt whose outputs equal ``values``, in order.""" + if not values: + return None + length = len(values[0]) + if length == 0 or length > self._max_items: + return None + if any(len(slot) != length for slot in values): + return None + wanted = { + index: tuple(text_digest(slot[index]) for slot in values) + for index in range(length) + } + with self._lock: + receipts = self._receipts.get(run_key) + if not receipts: + return None + matches = [receipt for receipt in receipts if receipt.outputs() == wanted] + if not matches: + return None + self._receipts.move_to_end(run_key) + facts = { + tuple(receipt.items[index].facts for index in range(length)) + for receipt in matches + } + if len(facts) != 1: + return None + return [matches[0].items[index] for index in range(length)] + + +def _deduplicated(receipts: List[_Receipt]) -> List[_Receipt]: + """The newest receipt replaces older ones with the same outputs.""" + if not receipts: + return receipts + outputs = receipts[0].outputs() + return [receipts[0]] + [r for r in receipts[1:] if r.outputs() != outputs] + + +def _job_nodes(graph: Dict[str, Any], outputs: Any) -> Dict[str, Tuple[str, str]]: + """PromptManager-family nodes the job's outputs need, with their run keys.""" + starts = ( + [str(node) for node in outputs] if isinstance(outputs, (list, tuple)) else [] + ) + needed = upstream_nodes(graph, starts) + nodes = {} + for node_id in needed: + node = graph.get(node_id) + node_type = node.get("class_type") if isinstance(node, dict) else None + if node_type not in PROMPT_MANAGER_TYPES: + continue + run_key = prompt_node_key(graph, node_id) + if run_key is not None: + nodes[node_id] = (node_type, run_key) + return nodes + + +def _host() -> Any: + """The host adapter module (imported late: it needs ComfyUI's modules).""" + try: + from . import execution_context_host + except ImportError: + from utils import execution_context_host + return execution_context_host + + +# Shared between the nodes and the host observers +CACHED_RUNS = CachedRuns() + +_registered = threading.Event() +_register_lock = threading.Lock() + + +def register_cached_runs( + server: Any, db_factory: Callable[[], Any], runs: Optional[CachedRuns] = None +) -> bool: + """Observe dequeue, ``execution_cached`` and ``task_done``; once per process. + + Registers nothing when the host lacks any of them: counting is then as + before #347. + """ + runs = CACHED_RUNS if runs is None else runs + with _register_lock: + if _registered.is_set(): + return False + queue = getattr(server, "prompt_queue", None) + if not callable(getattr(queue, "get", None)) or not callable( + getattr(queue, "task_done", None) + ): + logger.info("Cached runs not counted: the prompt queue is not observable") + return False + try: + from execution import PromptExecutor + except Exception: + logger.info("Cached runs not counted: no ComfyUI executor") + return False + if not callable(getattr(PromptExecutor, "add_message", None)): + logger.info("Cached runs not counted: executor messages unavailable") + return False + subscribe_executor_messages = _host().subscribe_executor_messages + + def on_message(executor: Any, args: tuple, kwargs: dict) -> None: + try: + event = args[0] if args else kwargs.get("event") + if event != "execution_cached": + return + server_of = getattr(executor, "server", None) + if getattr(server_of, "prompt_queue", None) is not queue: + return + data = args[1] if len(args) > 1 else kwargs.get("data") + if type(data) is not dict: + return + runs.replay( + executor, data.get("prompt_id"), data.get("nodes"), db_factory + ) + except Exception as error: + logger.warning(f"Could not count cached nodes: {error}") + + unsubscribers = [] + try: + unsubscribers.append( + subscribe_dequeue(queue, "cached_runs", runs.begin_attempt) + ) + unsubscribers.append( + subscribe_task_done( + queue, + "cached_runs", + lambda args, kwargs: runs.end_attempt( + args[0] if args else kwargs.get("item_id") + ), + ) + ) + unsubscribers.append(subscribe_executor_messages("cached_runs", on_message)) + except Exception: + for unsubscribe in unsubscribers: + unsubscribe() + raise + _registered.set() + logger.info("Counting cached runs of runtime-linked prompt text") + return True diff --git a/utils/execution_context_host.py b/utils/execution_context_host.py index c3d8ba3..fdf02af 100644 --- a/utils/execution_context_host.py +++ b/utils/execution_context_host.py @@ -75,6 +75,82 @@ def subscribe_executor_messages(key: str, callback: Callable) -> Callable[[], No return state.subscribe(key, callback) +# STRING output slots of each PromptManager-family node, in output order +STRING_SLOTS = { + "PromptManager": (1,), + "PromptManagerFocused": (1,), + "PromptManagerSplit": (1,), + "PromptManagerText": (0,), + "PromptManagerPair": (2, 3), +} +MAX_CACHED_VALUES = 64 + + +def _local_cache_reader(executor: object) -> Callable | None: + """``node_id -> outputs`` of the executor's local output cache, or None. + + Only the supported cache classes, only ``get_local`` (never a provider or + an async lookup), only a real ``CacheEntry`` with at most four outputs. + """ + 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 None + local = getattr(cache, "get_local", None) + if not callable(local): + return None + + def read(node_id: str) -> list | tuple | None: + entry = local(node_id) + if type(entry) is not CacheEntry: + return None + outputs = entry.outputs + if type(outputs) not in (list, tuple) or not 0 < len(outputs) <= 4: + return None + return outputs + + return read + + +def cached_string_outputs( + executor: object, node_id: str, node_type: str +) -> tuple[tuple[str, ...], ...] | None: + """One tuple per STRING slot of a cached node's local entry, or None (#347). + + None when the cache or entry is not supported, a slot is missing or holds + more than 64 values, any value is not a ``str``, or the slots differ in + length. Never runs a provider or an async lookup. + """ + try: + slots = STRING_SLOTS.get(node_type) + if slots is None or type(node_id) is not str: + return None + local = _local_cache_reader(executor) + outputs = local(node_id) if local is not None else None + if outputs is None or not max(slots) < len(outputs): + return None + found = [] + for slot in slots: + values = outputs[slot] + if type(values) not in (list, tuple) or len(values) > MAX_CACHED_VALUES: + return None + if any(type(value) is not str for value in values): + return None + found.append(tuple(values)) + if len({len(values) for values in found}) != 1: + return None + return tuple(found) + except Exception: + return None + + def observe_cached_context( executor: object, message_type: object, message_data: object ) -> None: @@ -101,19 +177,8 @@ def observe_cached_context( 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): + local = _local_cache_reader(executor) + if local is None: return # Only source IDs retained by the bounded submitted graph are queried. for source in cap.sources: @@ -124,11 +189,8 @@ def observe_cached_context( 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: + outputs = local(source.node_id) + if outputs is None or not slot < len(outputs): continue values = outputs[slot] if type(values) not in (list, tuple) or len(values) > 64: diff --git a/utils/run_keys.py b/utils/run_keys.py index 05d74a8..f53df44 100644 --- a/utils/run_keys.py +++ b/utils/run_keys.py @@ -51,6 +51,53 @@ def _links(inputs: Any) -> List[str]: return [value[0] for value in inputs.values() if is_link(value)] +def upstream_closure( + graph: Dict[str, Any], starts: List[str], limit: Optional[int] = None +) -> Optional[Dict[str, Any]]: + """``{node_id: [class_type, inputs]}`` of ``starts`` and everything upstream. + + None when a node or an upstream link's node is missing, or the closure + exceeds ``limit`` nodes. + """ + closure: Dict[str, Any] = {} + stack = list(starts) + while stack: + current = stack.pop() + if current in closure: + continue + node = graph.get(current) + if not isinstance(node, dict): + return None + if limit is not None and len(closure) >= limit: + return None + inputs = node.get("inputs") + closure[current] = [node.get("class_type"), inputs] + stack.extend(_links(inputs)) + return closure + + +def upstream_nodes(graph: Any, starts: List[str]) -> List[str]: + """Ids of ``starts`` and every node upstream of them that exists in ``graph``. + + Used to tell which nodes a job's outputs to execute need ("run selected + output" jobs leave the rest of the graph out). Missing nodes are skipped. + """ + if not isinstance(graph, dict): + return [] + found: Dict[str, None] = {} + stack = [str(node) for node in starts] + while stack: + current = stack.pop() + if current in found: + continue + node = graph.get(current) + if not isinstance(node, dict): + continue + found[current] = None + stack.extend(_links(node.get("inputs"))) + return list(found) + + def prompt_node_key(graph: Any, node_id: Any) -> Optional[str]: """Digest of ``node_id`` and its whole upstream closure, or None. @@ -62,20 +109,9 @@ def prompt_node_key(graph: Any, node_id: Any) -> Optional[str]: if not isinstance(graph, dict) or node_id is None: return None start = str(node_id) - closure: Dict[str, Any] = {} - stack = [start] - while stack: - current = stack.pop() - if current in closure: - continue - node = graph.get(current) - if not isinstance(node, dict): - return None - if len(closure) >= MAX_VISITS: - return None - inputs = node.get("inputs") - closure[current] = [node.get("class_type"), inputs] - stack.extend(_links(inputs)) + closure = upstream_closure(graph, [start], MAX_VISITS) + if closure is None: + return None try: encoded = json.dumps( [RUN_KEY_VERSION, start, closure], diff --git a/utils/usage_tracking.py b/utils/usage_tracking.py index 0365e25..6fb6d39 100644 --- a/utils/usage_tracking.py +++ b/utils/usage_tracking.py @@ -16,7 +16,9 @@ it never counts. Servers without a ``prompt_queue`` fall back to counting at submit time. Text from known pure string nodes is resolved at queue time too. Anything else -(PromptSearchList batches, unknown nodes) is counted by the node as it executes. +(PromptSearchList batches, unknown nodes) is counted by the node as it executes, +and when ComfyUI serves such a node from its cache, the job replays what that +execution counted (``cached_runs``, #347). """ import threading @@ -167,13 +169,22 @@ def queued_pair_runs(graph: Any) -> List[Tuple[str, str, str]]: def find_pair_id(db: Any, positive: str, negative: str) -> Optional[int]: """Find a saved pair without creating or overwriting its metadata.""" - positive_row = db.get_prompt_by_hash(generate_prompt_hash(positive)) - negative_row = ( - db.get_prompt_by_hash(generate_prompt_hash(negative)) - if negative.strip() - else None + return find_pair_id_by_hashes( + db, + generate_prompt_hash(positive), + generate_prompt_hash(negative) if negative.strip() else None, ) - if not positive_row or (negative.strip() and not negative_row): + + +def find_pair_id_by_hashes( + db: Any, positive_hash: str, negative_hash: Optional[str] +) -> Optional[int]: + """The saved pair of two prompt hashes; ``None`` negative is the empty one.""" + positive_row = db.get_prompt_by_hash(positive_hash) + negative_row = ( + db.get_prompt_by_hash(negative_hash) if negative_hash is not None else None + ) + if not positive_row or (negative_hash is not None and not negative_row): return None row = ( db.model.get_connection()