Compare commits
3
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0e41d21730 | ||
|
|
e752e9c218 | ||
|
|
a6505a07af |
+28
-5
@@ -5572,8 +5572,26 @@ decide through the same `own_counts`.
|
||||
- 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).
|
||||
- No backfill: earlier cached jobs keep their counts.
|
||||
- History (#403): the receipt also keeps the raw, prepend and append text the node gave
|
||||
the history capture. For a cached node whose graph inference leaves a field unknown,
|
||||
the capture takes it from the receipt matching the cache entry
|
||||
(`CachedRuns.recorded_fields`), as `inferred` / `graph_semantics`: the node did not run
|
||||
in this job, the value follows from the run key (the whole upstream closure) plus the
|
||||
cache proof. The job then joins the prompt's history scope, gets the
|
||||
`positive_prompt_id` (and Pair) hint, and an image of it can be attributed through its
|
||||
run, so the history page's "tracked runs" equals the Library count from the restart on.
|
||||
Fail closed (fields stay unknown, as before): no matching receipt, a receipt that
|
||||
disagrees with a field the graph knows, executions with equal outputs that recorded
|
||||
different fields (a conflict that stays for that run key), a field over 64 KiB or
|
||||
invalid Unicode, more than one run's history (1 MiB) in one receipt, or a full 8 MiB
|
||||
field budget (least recently used receipts give up their fields first; their counting
|
||||
stays). No new evidence vocabulary, so older builds still open the database.
|
||||
- Still different: jobs before the restart (no backfill); list batches (per item in the
|
||||
Library, once per job in history); text with leading or trailing whitespace (the
|
||||
Library stores it stripped, history matches the exact raw text, #438); a job in which the
|
||||
node executes after its replay (RAM-pressure freed the entry: two samples, an
|
||||
ambiguous capture); jobs while the history writer is fenced (restore, library move).
|
||||
- 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,
|
||||
@@ -5581,7 +5599,10 @@ decide through the same `own_counts`.
|
||||
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.
|
||||
covered by the Python tests. Since #403 these cases also check the history page's
|
||||
tracked runs equal the Library count, and `tests/test_execution_history_hooks.py`
|
||||
(`CachedRuntimeHistoryTests`) runs the real history hook and writer with a real
|
||||
`CachedRuns`.
|
||||
|
||||
```sh
|
||||
python3 -m unittest discover -s tests -p "test_cached_runs.py"
|
||||
@@ -5680,10 +5701,12 @@ Context cases verify A,A,B as three jobs/two combinations/one positive snapshot
|
||||
when every positive field matches; all five registered PM variants; exact raw
|
||||
whitespace and Unicode; known empty negatives; Text negative roles and crossed
|
||||
pairs; fully cached repeats after trigger configuration changes; successful source
|
||||
text before downstream failure/interruption; unknown cached dynamic raw; unsupported
|
||||
text before downstream failure/interruption; cached dynamic raw from the reused execution (#403); unsupported
|
||||
transforms and duplicate list invocation ambiguity. Fresh values use
|
||||
`runtime_return`, cached effective STRING outputs use `host_cache_output`, and
|
||||
submitted literals use graph inference. SQL NULL means unknown; empty text is
|
||||
submitted literals use graph inference. A cached dynamic raw is unknown unless the
|
||||
#347 receipt of the execution the cache entry comes from proves it; it is then
|
||||
`inferred` / `graph_semantics` (graph identity plus a same-process cache receipt, #403). SQL NULL means unknown; empty text is
|
||||
known. Text is never trimmed or normalized by capture. Cache availability and
|
||||
inferred sampler roles prove neither sampler execution nor consumption, and list
|
||||
indices do not identify generated images. External trigger configuration changes
|
||||
|
||||
+6
-1
@@ -275,7 +275,12 @@ 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)
|
||||
self._record_cached_item(
|
||||
unique_id,
|
||||
(encoding_text,),
|
||||
storage_text,
|
||||
fields=((text, prepend_text, append_text),),
|
||||
)
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
|
||||
+13
-3
@@ -243,14 +243,24 @@ class PromptManagerBase:
|
||||
return False
|
||||
|
||||
def _record_cached_item(
|
||||
self, unique_id: Any, outputs: tuple, text: Any, negative_text: Any = None
|
||||
self,
|
||||
unique_id: Any,
|
||||
outputs: tuple,
|
||||
text: Any,
|
||||
negative_text: Any = None,
|
||||
fields: Optional[tuple] = None,
|
||||
) -> None:
|
||||
"""Remember this execution's outputs and stored text for cached jobs (#347)."""
|
||||
"""Remember this execution's outputs and stored text for cached jobs (#347).
|
||||
|
||||
``fields`` are the ``(raw, prepend, append)`` texts this execution gave
|
||||
``record_runtime_context`` for each output, so a cached job's history
|
||||
can record them too (#403).
|
||||
"""
|
||||
if unique_id is None:
|
||||
return
|
||||
try:
|
||||
self._cached_runs().record_item(
|
||||
unique_id, outputs, stored_facts(text, negative_text)
|
||||
unique_id, outputs, stored_facts(text, negative_text), fields
|
||||
)
|
||||
except Exception as error:
|
||||
self.logger.debug(f"Could not record the cached-run receipt: {error}")
|
||||
|
||||
@@ -198,7 +198,11 @@ class PromptManagerPair(PromptManagerFocused):
|
||||
type(self).__name__, negative_text, "", "", negative_text, role="negative"
|
||||
)
|
||||
self._record_cached_item(
|
||||
unique_id, (positive, negative_text), text, negative_text
|
||||
unique_id,
|
||||
(positive, negative_text),
|
||||
text,
|
||||
negative_text,
|
||||
fields=((text, prepend_text, append_text), (negative_text, "", "")),
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
@@ -246,7 +246,12 @@ 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)
|
||||
self._record_cached_item(
|
||||
unique_id,
|
||||
(final_text,),
|
||||
storage_text,
|
||||
fields=((text, prepend_text, append_text),),
|
||||
)
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
|
||||
@@ -263,15 +263,18 @@ test.describe('public chronological history with the actual acknowledged host br
|
||||
});
|
||||
});
|
||||
|
||||
test('a cached dynamic unknown positive stays readable by run UUID and outside the known raw Library scope', async ({ execution }) => {
|
||||
const text = `unknown positive ${randomUUID()}`, graph = recipeGraph({ text, dynamic: true });
|
||||
// #403: the graph cannot show a dynamic positive's raw text; a cached job takes it from the execution it reuses.
|
||||
test('a cached dynamic positive takes its raw text from the execution it reuses and joins the Library scope', async ({ execution }) => {
|
||||
const text = `dynamic positive ${randomUUID()}`, graph = recipeGraph({ text, dynamic: true });
|
||||
const first = await submit(execution.request, graph);await history(execution.request, first);const known = (await capture(execution, first)).contexts[0];
|
||||
const second = await submit(execution.request, graph);await history(execution.request, second);const unknown = (await capture(execution, second)).contexts[0];
|
||||
expect(unknown.positive.raw).toBeNull();expect(unknown.positive.effective).toBe(known.positive.effective);
|
||||
const second = await submit(execution.request, graph);await history(execution.request, second);const reused = await capture(execution, second);const cached = reused.contexts[0];
|
||||
expect(cached.evidence.positive.effective.source).toBe('host_cache_output');expect(cached.positive.effective).toBe(known.positive.effective);
|
||||
expect(cached.positive.raw).toBe(known.positive.raw);expect(cached.evidence.positive.raw).toMatchObject({ provenance: 'inferred', source: 'graph_semantics' });
|
||||
expect(cached.positive_prompt_id).toBe(known.positive_prompt_id);
|
||||
const result = await readExecutionRuns(execution.databasePath, second), runId = result.rows[0].run_id;
|
||||
const detail = await execution.request.get(`/prompt_manager/history/runs/${runId}`);expect(detail.ok()).toBe(true);expect((await detail.json()).contexts[0].positive.raw).toBeNull();
|
||||
const detail = await execution.request.get(`/prompt_manager/history/runs/${runId}`);expect(detail.ok()).toBe(true);expect((await detail.json()).contexts[0].positive.raw).toBe(known.positive.raw);
|
||||
const response = await execution.request.get(`/prompt_manager/library/${known.positive_prompt_id}/history`);expect(response.ok()).toBe(true);const window = await response.json();
|
||||
expect(window.recorded_run_count).toBe(1);expect(window.items.map(item => item.run_id)).not.toContain(runId);
|
||||
expect(window.recorded_run_count).toBe(2);expect(window.items.map(item => item.run_id)).toContain(runId);
|
||||
});
|
||||
|
||||
test('actual PM A,A,B jobs persist two combinations and one identical positive snapshot', async ({ execution }) => {
|
||||
@@ -369,7 +372,7 @@ test('successful source capture survives downstream failure and real interruptio
|
||||
}
|
||||
});
|
||||
|
||||
test('cached dynamic raw stays unknown while literal affixes and effective cache text survive', async ({ execution }) => {
|
||||
test('cached dynamic raw comes from the reused execution while literal affixes and effective cache text survive (#403)', async ({ execution }) => {
|
||||
const graph = recipeGraph({ dynamic: true, text: 'runtime dynamic recipe', prepend: ' pre ', append: ' post ' });
|
||||
const first = await submit(execution.request, graph);
|
||||
await history(execution.request, first);
|
||||
@@ -377,8 +380,9 @@ test('cached dynamic raw stays unknown while literal affixes and effective cache
|
||||
const second = await submit(execution.request, graph);
|
||||
await history(execution.request, second);
|
||||
const value = (await capture(execution, second)).contexts[0];
|
||||
expect(value.positive).toEqual({ raw: null, prepend: ' pre ', append: ' post ', effective: 'pre runtime dynamic recipe post' });
|
||||
expect(value.evidence.positive.raw.provenance).toBe('unknown');
|
||||
expect(value.positive).toEqual({ raw: 'runtime dynamic recipe', prepend: ' pre ', append: ' post ', effective: 'pre runtime dynamic recipe post' });
|
||||
expect(value.evidence.positive.raw).toMatchObject({ provenance: 'inferred', source: 'graph_semantics' });
|
||||
expect(value.evidence.positive.prepend.source).toBe('graph_literal');
|
||||
expect(value.evidence.positive.effective.source).toBe('host_cache_output');
|
||||
});
|
||||
|
||||
|
||||
@@ -287,6 +287,14 @@ async function capturedContext(execution, id) {
|
||||
return result.captures[result.rows[0].run_id].contexts[0];
|
||||
}
|
||||
|
||||
// #403: the prompt's history "tracked runs", once every job's capture reached SQLite.
|
||||
async function expectTrackedRuns(execution, promptId, expected) {
|
||||
await expect.poll(async () => {
|
||||
const response = await execution.request.get(`/prompt_manager/library/${promptId}/history`);
|
||||
return response.ok() ? (await response.json()).recorded_run_count : null;
|
||||
}, { timeout: 10000, message: 'history tracked runs equal the Library run count' }).toBe(expected);
|
||||
}
|
||||
|
||||
async function firstRun(execution, page, testCase) {
|
||||
const prefix = `textlink/${randomUUID()}`;
|
||||
const built = await build(page, { ...testCase, prefix });
|
||||
@@ -309,6 +317,8 @@ for (const kind of ['PromptManager', 'PromptManagerFocused', 'PromptManagerSplit
|
||||
const counts = await runCounts(execution, first.context);
|
||||
expect(counts.positive).toBe(3);
|
||||
if (pair) expect(counts).toMatchObject({ negative: 0, pair: 3 });
|
||||
// #403: the cached jobs are tracked runs of the prompt too
|
||||
await expectTrackedRuns(execution, first.context.positive_prompt_id, 3);
|
||||
});
|
||||
}
|
||||
|
||||
@@ -329,6 +339,7 @@ test.describe('cached runs in one renderer (#347)', () => {
|
||||
else expect(again.cached).toContain(first.promptNode);
|
||||
}
|
||||
expect((await runCounts(execution, first.context)).positive).toBe(3);
|
||||
await expectTrackedRuns(execution, first.context.positive_prompt_id, 3);
|
||||
} finally { await setCache(execution, 'RAM_PRESSURE'); }
|
||||
});
|
||||
}
|
||||
@@ -352,6 +363,9 @@ test.describe('cached runs in one renderer (#347)', () => {
|
||||
const bContext = await capturedContext(execution, second.id);
|
||||
expect(bContext.positive.raw).toBe(b);
|
||||
expect((await runCounts(execution, bContext)).positive).toBe(1);
|
||||
// #403: the third job reuses A's execution, not the newest (B's)
|
||||
await expectTrackedRuns(execution, first.context.positive_prompt_id, 2);
|
||||
await expectTrackedRuns(execution, bContext.positive_prompt_id, 1);
|
||||
} finally { await setCache(execution, 'RAM_PRESSURE'); }
|
||||
});
|
||||
|
||||
@@ -390,10 +404,14 @@ test.describe('cached runs in one renderer (#347)', () => {
|
||||
} });
|
||||
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
|
||||
// #403: a cached job's history takes the raw text of the execution it reuses and names the prompt
|
||||
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(cachedContext.positive.raw).toBe(first.context.positive.raw);
|
||||
expect(cachedContext.evidence.positive.raw).toMatchObject({ provenance: 'inferred', source: 'graph_semantics' });
|
||||
expect(cachedContext.positive_prompt_id).toBe(first.context.positive_prompt_id);
|
||||
expect((await runCounts(execution, first.context)).positive).toBe(2);
|
||||
// The rewired job used the prompt as its negative: neither number counts it
|
||||
await expectTrackedRuns(execution, first.context.positive_prompt_id, 2);
|
||||
});
|
||||
});
|
||||
|
||||
+193
-1
@@ -666,7 +666,7 @@ class TestOwnCounts(unittest.TestCase):
|
||||
return counts
|
||||
|
||||
|
||||
class TestReceipts(unittest.TestCase):
|
||||
class ReceiptCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.context = ("j1", None)
|
||||
self.runs = CachedRuns(
|
||||
@@ -686,6 +686,8 @@ class TestReceipts(unittest.TestCase):
|
||||
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"):
|
||||
@@ -726,6 +728,196 @@ class TestReceipts(unittest.TestCase):
|
||||
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())
|
||||
|
||||
@@ -100,14 +100,19 @@ class HostTests(unittest.TestCase):
|
||||
self.executor.caches = types.SimpleNamespace(outputs=HierarchicalCache())
|
||||
|
||||
def begin(
|
||||
self, kind="PromptManager", raw=" raw ", prepend=" pre ", append=" post "
|
||||
self,
|
||||
kind="PromptManager",
|
||||
raw=" raw ",
|
||||
prepend=" pre ",
|
||||
append=" post ",
|
||||
negative="negative",
|
||||
):
|
||||
g = {
|
||||
"p": {
|
||||
"class_type": kind,
|
||||
"inputs": {
|
||||
"text": raw,
|
||||
"negative_text": "negative",
|
||||
"negative_text": negative,
|
||||
"prepend_text": prepend,
|
||||
"append_text": append,
|
||||
},
|
||||
@@ -125,7 +130,7 @@ class HostTests(unittest.TestCase):
|
||||
},
|
||||
},
|
||||
}
|
||||
if type(raw) is list or type(prepend) is list:
|
||||
if any(type(value) is list for value in (raw, prepend, negative)):
|
||||
g["dynamic"] = {"class_type": "RuntimeGenerator", "inputs": {}}
|
||||
g["stateful"] = {
|
||||
"class_type": "PromptManagerText",
|
||||
@@ -445,3 +450,138 @@ class HostTests(unittest.TestCase):
|
||||
ref = install_and_remove()
|
||||
gc.collect()
|
||||
self.assertIsNone(ref())
|
||||
|
||||
# -- #403: a cached runtime text takes its fields from the reused execution
|
||||
|
||||
class Receipts:
|
||||
"""A fake ``CachedRuns.recorded_fields``: fixed fields, calls recorded."""
|
||||
|
||||
def __init__(self, fields):
|
||||
self.fields, self.calls = fields, []
|
||||
|
||||
def recorded_fields(self, prompt_id, node_id, values):
|
||||
self.calls.append((prompt_id, node_id, values))
|
||||
return self.fields
|
||||
|
||||
def observe_with(self, receipts, outputs, **begin):
|
||||
self.begin(**begin)
|
||||
self.executor.caches.outputs.entries["p"] = Entry(None, outputs)
|
||||
self.observe("execution_start")
|
||||
host.observe_cached_context(
|
||||
self.executor,
|
||||
"execution_cached",
|
||||
{"prompt_id": JOB, "nodes": ["p"]},
|
||||
receipts=receipts,
|
||||
)
|
||||
return self.seal()
|
||||
|
||||
def test_cached_runtime_raw_comes_from_the_matching_receipt(self):
|
||||
receipts = self.Receipts([(("LLM raw", " pre ", " post "),)])
|
||||
result = self.observe_with(
|
||||
receipts, [object(), ["cached"]], raw=["dynamic", 0]
|
||||
).positive
|
||||
self.assertEqual(receipts.calls, [(JOB, "p", (("cached",),))])
|
||||
self.assertEqual(result.raw.value, "LLM raw")
|
||||
self.assertEqual(
|
||||
(result.raw.provenance, result.raw.source), ("inferred", "graph_semantics")
|
||||
)
|
||||
self.assertEqual(result.effective.value, "cached")
|
||||
self.assertEqual(result.effective.source, "host_cache_output")
|
||||
# Fields the graph already knows are never replaced
|
||||
self.assertEqual(result.prepend.value, " pre ")
|
||||
self.assertEqual(result.prepend.source, "graph_literal")
|
||||
self.assertEqual(result.append.value, " post ")
|
||||
|
||||
def test_cached_runtime_prepend_comes_from_the_matching_receipt(self):
|
||||
receipts = self.Receipts([((" raw ", "LLM pre", " post "),)])
|
||||
result = self.observe_with(
|
||||
receipts, [object(), ["cached"]], prepend=["dynamic", 0]
|
||||
).positive
|
||||
self.assertEqual(result.prepend.value, "LLM pre")
|
||||
self.assertEqual(result.prepend.source, "graph_semantics")
|
||||
self.assertEqual(result.raw.value, " raw ")
|
||||
self.assertEqual(result.raw.source, "graph_literal")
|
||||
|
||||
def test_a_receipt_disagreeing_with_a_known_field_fills_nothing(self):
|
||||
receipts = self.Receipts([(("LLM raw", "other pre", " post "),)])
|
||||
result = self.observe_with(
|
||||
receipts, [object(), ["cached"]], raw=["dynamic", 0]
|
||||
).positive
|
||||
self.assertIsNone(result.raw.value)
|
||||
self.assertEqual(result.prepend.value, " pre ")
|
||||
|
||||
def test_pair_runtime_negative_takes_the_negative_fields(self):
|
||||
receipts = self.Receipts(
|
||||
[((" raw ", " pre ", " post "), ("LLM negative", "", ""))]
|
||||
)
|
||||
result = self.observe_with(
|
||||
receipts,
|
||||
[object(), object(), ["cached positive"], ["LLM negative"]],
|
||||
kind="PromptManagerPair",
|
||||
negative=["dynamic", 0],
|
||||
)
|
||||
self.assertEqual(result.negative.raw.value, "LLM negative")
|
||||
self.assertEqual(result.negative.raw.source, "graph_semantics")
|
||||
self.assertEqual(result.positive.raw.value, " raw ")
|
||||
self.assertEqual(result.positive.raw.source, "graph_literal")
|
||||
|
||||
def test_known_typed_fields_are_never_replaced_by_a_receipt(self):
|
||||
receipts = self.Receipts([(("other", "other", "other"),)])
|
||||
result = self.observe_with(receipts, [object(), ["cached"]]).positive
|
||||
self.assertEqual(
|
||||
(result.raw.value, result.prepend.value, result.append.value),
|
||||
(" raw ", " pre ", " post "),
|
||||
)
|
||||
self.assertEqual(result.raw.source, "graph_literal")
|
||||
|
||||
def test_no_matching_receipt_keeps_the_cached_raw_unknown(self):
|
||||
for fields in (None, [None], [], "bad", [(("a", "b"),)], [((1, "", ""),)]):
|
||||
with self.subTest(fields=fields):
|
||||
result = self.observe_with(
|
||||
self.Receipts(fields), [object(), ["cached"]], raw=["dynamic", 0]
|
||||
).positive
|
||||
self.assertIsNone(result.raw.value)
|
||||
self.assertEqual(result.effective.value, "cached")
|
||||
|
||||
def test_a_failing_receipt_lookup_keeps_the_capture(self):
|
||||
class Broken:
|
||||
def recorded_fields(self, *args):
|
||||
raise RuntimeError("boom")
|
||||
|
||||
result = self.observe_with(
|
||||
Broken(), [object(), ["cached"]], raw=["dynamic", 0]
|
||||
).positive
|
||||
self.assertIsNone(result.raw.value)
|
||||
self.assertEqual(result.effective.value, "cached")
|
||||
|
||||
def test_pair_slots_take_their_own_receipt_fields(self):
|
||||
receipts = self.Receipts(
|
||||
[(("LLM positive", " pre ", " post "), ("LLM negative", "", ""))]
|
||||
)
|
||||
result = self.observe_with(
|
||||
receipts,
|
||||
[object(), object(), ["cached positive"], ["cached negative"]],
|
||||
kind="PromptManagerPair",
|
||||
raw=["dynamic", 0],
|
||||
)
|
||||
self.assertEqual(
|
||||
receipts.calls[0], (JOB, "p", (("cached positive",), ("cached negative",)))
|
||||
)
|
||||
self.assertEqual(result.positive.raw.value, "LLM positive")
|
||||
self.assertEqual(result.positive.effective.value, "cached positive")
|
||||
# The negative editor is typed: the graph's value stays
|
||||
self.assertEqual(result.negative.raw.value, "negative")
|
||||
self.assertEqual(result.negative.raw.source, "graph_literal")
|
||||
|
||||
def test_default_receipts_are_the_process_cached_runs(self):
|
||||
from utils import cached_runs
|
||||
|
||||
receipts = self.Receipts([(("LLM raw", " pre ", " post "),)])
|
||||
with patch.object(cached_runs, "CACHED_RUNS", receipts):
|
||||
self.begin(raw=["dynamic", 0])
|
||||
self.executor.caches.outputs.entries["p"] = Entry(
|
||||
None, [object(), ["cached"]]
|
||||
)
|
||||
self.observe("execution_start")
|
||||
self.observe("execution_cached", nodes=["p"])
|
||||
self.assertEqual(self.seal().positive.raw.value, "LLM raw")
|
||||
|
||||
@@ -54,7 +54,7 @@ def status(kind="execution_success", job=JOB):
|
||||
return Status("success", True, [(kind, {"prompt_id": job, "timestamp": 20})])
|
||||
|
||||
|
||||
class HistoryHooksTests(unittest.TestCase):
|
||||
class HooksCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.directory = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self.directory.cleanup)
|
||||
@@ -122,6 +122,8 @@ class HistoryHooksTests(unittest.TestCase):
|
||||
def drain(self, registration):
|
||||
drain_writer(registration.writer, clock=time.monotonic)
|
||||
|
||||
|
||||
class HistoryHooksTests(HooksCase):
|
||||
def test_package_history_initializes_when_monitor_and_usage_setup_fail(self):
|
||||
try:
|
||||
from tests.test_init_package import (
|
||||
@@ -767,3 +769,150 @@ class HistoryHooksTests(unittest.TestCase):
|
||||
self.assertEqual(
|
||||
calls, ["execution_start", "execution_cached", "execution_cached"]
|
||||
)
|
||||
|
||||
|
||||
class CachedRuntimeHistoryTests(HooksCase):
|
||||
"""A cached runtime-linked prompt is a tracked run, like an executed one (#403).
|
||||
|
||||
The real history hook and writer, a real ``CachedRuns`` and the real host
|
||||
adapter over fakes of ComfyUI's executor, cache and executing context.
|
||||
"""
|
||||
|
||||
TEXT = "a lighthouse the LLM wrote"
|
||||
GRAPH = {
|
||||
"str": {"class_type": "PrimitiveString", "inputs": {"value": "seed"}},
|
||||
"llm": {"class_type": "PreviewAny", "inputs": {"source": ["str", 0]}},
|
||||
"p": {"class_type": "PromptManager", "inputs": {"text": ["llm", 0]}},
|
||||
"s": {"class_type": "KSampler", "inputs": {"positive": ["p", 0]}},
|
||||
}
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
from utils import cached_runs
|
||||
from utils.cached_runs import CachedRuns
|
||||
|
||||
self.context = None
|
||||
self.runs = CachedRuns(context=lambda: self.context or (None, None))
|
||||
self.cache_entries = {}
|
||||
entries = self.cache_entries
|
||||
execution = types.ModuleType("execution")
|
||||
caching = types.ModuleType("comfy_execution.caching")
|
||||
package = types.ModuleType("comfy_execution")
|
||||
utils_module = types.ModuleType("comfy_execution.utils")
|
||||
entry_type = namedtuple("CacheEntry", "ui outputs")
|
||||
|
||||
class PromptExecutor:
|
||||
def add_message(self, event, data, broadcast):
|
||||
return None
|
||||
|
||||
class HierarchicalCache:
|
||||
def get_local(self, node_id):
|
||||
return entries.get(node_id)
|
||||
|
||||
execution.PromptExecutor, execution.CacheEntry = PromptExecutor, entry_type
|
||||
caching.HierarchicalCache = HierarchicalCache
|
||||
package.caching, package.utils = caching, utils_module
|
||||
utils_module.get_executing_context = lambda: (
|
||||
SimpleNamespace(prompt_id=self.context[0], node_id="p", list_index=0)
|
||||
if self.context
|
||||
else None
|
||||
)
|
||||
self.entry_type = entry_type
|
||||
modules = patch.dict(
|
||||
sys.modules,
|
||||
{
|
||||
"execution": execution,
|
||||
"comfy_execution": package,
|
||||
"comfy_execution.caching": caching,
|
||||
"comfy_execution.utils": utils_module,
|
||||
},
|
||||
)
|
||||
modules.start()
|
||||
self.addCleanup(modules.stop)
|
||||
swap = patch.object(cached_runs, "CACHED_RUNS", self.runs)
|
||||
swap.start()
|
||||
self.addCleanup(swap.stop)
|
||||
self.executor = PromptExecutor()
|
||||
self.executor.server = self.server
|
||||
self.executor.caches = SimpleNamespace(outputs=HierarchicalCache())
|
||||
self.prompt_id = self.db.save_prompt(text=self.TEXT)
|
||||
self.item_id = 0
|
||||
|
||||
def job(self, cached):
|
||||
"""One job; the node executes unless ``cached`` (ComfyUI's cache serves it)."""
|
||||
from uuid import uuid4
|
||||
|
||||
from utils.cached_runs import stored_facts
|
||||
from utils.execution_context_capture import record_runtime_context
|
||||
|
||||
self.item_id += 1
|
||||
job = str(uuid4())
|
||||
result = ((0, job, self.GRAPH, {}, ["s"]), self.item_id)
|
||||
self.queue.results.append(result)
|
||||
self.queue.get()
|
||||
self.runs.begin_attempt(result)
|
||||
self.context = (job, 0)
|
||||
self.executor.add_message("execution_start", {"prompt_id": job}, False)
|
||||
if cached:
|
||||
self.executor.add_message(
|
||||
"execution_cached", {"prompt_id": job, "nodes": ["p"]}, False
|
||||
)
|
||||
else:
|
||||
# What PromptManager.encode_prompt records for this text
|
||||
record_runtime_context("PromptManager", self.TEXT, "", "", self.TEXT)
|
||||
self.runs.record_item(
|
||||
"p", (self.TEXT,), stored_facts(self.TEXT), ((self.TEXT, "", ""),)
|
||||
)
|
||||
self.cache_entries["p"] = self.entry_type(None, [object(), [self.TEXT]])
|
||||
self.context = None
|
||||
self.queue.task_done(self.item_id, {}, status(job=job))
|
||||
self.runs.end_attempt(self.item_id)
|
||||
|
||||
def tracked_runs(self):
|
||||
with self.db.recorded_history_snapshot():
|
||||
return self.db.list_history_window(self.prompt_id)["recorded_run_count"]
|
||||
|
||||
def test_three_jobs_of_runtime_text_are_three_tracked_runs(self):
|
||||
registration = self.register()
|
||||
for cached in (False, True, True):
|
||||
self.job(cached)
|
||||
self.drain(registration)
|
||||
self.assertEqual(self.tracked_runs(), 3)
|
||||
rows = self.rows()
|
||||
for row in rows:
|
||||
contexts = self.db.get_execution_contexts(row["run_id"])["contexts"]
|
||||
self.assertEqual(contexts[0]["positive"]["raw"], self.TEXT)
|
||||
self.assertEqual(contexts[0]["positive_prompt_id"], self.prompt_id)
|
||||
cached = self.db.get_execution_contexts(rows[1]["run_id"])["contexts"][0]
|
||||
self.assertEqual(
|
||||
cached["evidence"]["positive"]["raw"],
|
||||
{"provenance": "inferred", "source": "graph_semantics", "reason": None},
|
||||
)
|
||||
self.assertEqual(
|
||||
cached["evidence"]["positive"]["effective"]["source"], "host_cache_output"
|
||||
)
|
||||
# The stored history passes the whole-database check older builds run too
|
||||
from database.execution_history_schema import history_schema_error
|
||||
|
||||
self.assertIsNone(history_schema_error(self.db.model.get_connection()))
|
||||
# An image of a cached job can now be attributed through its run (#342)
|
||||
from database.run_attribution import named_prompt
|
||||
|
||||
self.assertEqual(
|
||||
[named_prompt(self.db, row["run_id"]) for row in rows],
|
||||
[self.prompt_id] * 3,
|
||||
)
|
||||
|
||||
def test_without_a_receipt_a_cached_job_stays_untracked(self):
|
||||
registration = self.register()
|
||||
self.job(False)
|
||||
self.runs._receipts.clear() # as after a restart with a provider entry
|
||||
self.job(True)
|
||||
self.drain(registration)
|
||||
self.assertEqual(self.tracked_runs(), 1)
|
||||
cached = self.db.get_execution_contexts(self.rows()[1]["run_id"])["contexts"]
|
||||
self.assertIsNone(cached[0]["positive"]["raw"])
|
||||
self.assertIsNone(cached[0]["positive_prompt_id"])
|
||||
from database.run_attribution import named_prompt
|
||||
|
||||
self.assertIsNone(named_prompt(self.db, self.rows()[1]["run_id"]))
|
||||
|
||||
@@ -99,7 +99,10 @@ class TestPairUsageTracking(unittest.TestCase):
|
||||
unique_id="1",
|
||||
)
|
||||
node.cached_runs.record_item.assert_called_once_with(
|
||||
"1", ("cat", ""), Facts(usage_tracking.generate_prompt_hash("cat"), None)
|
||||
"1",
|
||||
("cat", ""),
|
||||
Facts(usage_tracking.generate_prompt_hash("cat"), None),
|
||||
(("cat", "", ""), ("", "", "")),
|
||||
)
|
||||
|
||||
def test_validation_rejection_never_counts_or_saves(self):
|
||||
|
||||
@@ -308,7 +308,11 @@ class TestNodeRecordsCachedRunReceipt(UsageTestCase):
|
||||
"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"))
|
||||
"1",
|
||||
("pre typed text",),
|
||||
Facts(generate_prompt_hash("typed text")),
|
||||
# #403: the texts the history capture got, for a cached job's history
|
||||
(("typed text", "pre", ""),),
|
||||
)
|
||||
# Hook-visible text: the node counts nothing itself
|
||||
row = self.db.get_prompt_by_hash(generate_prompt_hash("typed text"))
|
||||
@@ -319,7 +323,9 @@ class TestNodeRecordsCachedRunReceipt(UsageTestCase):
|
||||
|
||||
node = self._node()
|
||||
node.process_text(" ", prompt=self.GRAPH, unique_id="1")
|
||||
node.cached_runs.record_item.assert_called_once_with("1", (" ",), Facts(None))
|
||||
node.cached_runs.record_item.assert_called_once_with(
|
||||
"1", (" ",), Facts(None), ((" ", "", ""),)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
+201
-28
@@ -15,9 +15,14 @@ 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.
|
||||
|
||||
The receipt also keeps the raw, prepend and append text the node gave the
|
||||
execution history (#403), so the history capture of a cached job can record
|
||||
the raw text the graph cannot show (``recorded_fields``) and the job becomes a
|
||||
tracked run of its prompt, as it is a counted one.
|
||||
|
||||
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.
|
||||
this process, nothing is counted (the behavior before #347) and the history
|
||||
fields stay unknown. In memory only and bounded; starts no threads.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
@@ -39,6 +44,7 @@ try:
|
||||
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
|
||||
from .execution_contexts import MAX_ATTEMPT_BYTES, MAX_FIELD_BYTES, text_size
|
||||
except ImportError:
|
||||
from utils.hashing import generate_prompt_hash
|
||||
from utils.logging_config import get_logger
|
||||
@@ -53,6 +59,7 @@ except ImportError:
|
||||
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
|
||||
from utils.execution_contexts import MAX_ATTEMPT_BYTES, MAX_FIELD_BYTES, text_size
|
||||
|
||||
logger = get_logger("prompt_manager.cached_runs")
|
||||
|
||||
@@ -61,9 +68,15 @@ 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
|
||||
MAX_FIELD_TEXT_BYTES = 8 * 1024 * 1024 # history fields all receipts keep
|
||||
MAX_RECEIPT_FIELD_BYTES = MAX_ATTEMPT_BYTES # one receipt: what one run's history holds
|
||||
_EMPTY_HASH = generate_prompt_hash("")
|
||||
|
||||
Count = Tuple[Any, ...] # ("prompt", hash) or ("pair", positive, negative or None)
|
||||
# (raw, prepend, append) per STRING slot, as given to record_runtime_context
|
||||
Fields = Tuple[Tuple[str, str, str], ...]
|
||||
# Executions with equal outputs recorded different fields: none can be named
|
||||
_CONFLICT: Fields = ()
|
||||
|
||||
|
||||
class Facts(NamedTuple):
|
||||
@@ -96,6 +109,8 @@ class Seen(NamedTuple):
|
||||
class _Item(NamedTuple):
|
||||
outputs: Tuple[str, ...]
|
||||
facts: Facts
|
||||
fields: Optional[Fields] = None
|
||||
field_bytes: int = 0
|
||||
|
||||
|
||||
class _Receipt:
|
||||
@@ -108,6 +123,18 @@ class _Receipt:
|
||||
def outputs(self) -> Dict[int, Tuple[str, ...]]:
|
||||
return {index: item.outputs for index, item in self.items.items()}
|
||||
|
||||
def field_bytes(self) -> int:
|
||||
return sum(item.field_bytes for item in self.items.values())
|
||||
|
||||
def strip_fields(self) -> int:
|
||||
"""Drop the history fields (counting is unchanged); the bytes freed."""
|
||||
freed = 0
|
||||
for index, item in list(self.items.items()):
|
||||
if item.field_bytes:
|
||||
freed += item.field_bytes
|
||||
self.items[index] = item._replace(fields=None, field_bytes=0)
|
||||
return freed
|
||||
|
||||
|
||||
class _Attempt(NamedTuple):
|
||||
item_id: int
|
||||
@@ -117,6 +144,22 @@ class _Attempt(NamedTuple):
|
||||
replayed: set
|
||||
|
||||
|
||||
def _field_size(fields: Any, slots: int) -> Optional[int]:
|
||||
"""UTF-8 size of well-formed fields for ``slots`` outputs, else None."""
|
||||
if type(fields) is not tuple or len(fields) != slots:
|
||||
return None
|
||||
size = 0
|
||||
for slot in fields:
|
||||
if type(slot) is not tuple or len(slot) != 3:
|
||||
return None
|
||||
for value in slot:
|
||||
try:
|
||||
size += text_size(value, MAX_FIELD_BYTES) # the capture's own check
|
||||
except ValueError:
|
||||
return None
|
||||
return size
|
||||
|
||||
|
||||
def text_digest(value: str) -> str:
|
||||
"""Digest of one returned STRING value, exactly as returned."""
|
||||
return hashlib.sha256(value.encode("utf-8", "surrogatepass")).hexdigest()
|
||||
@@ -245,11 +288,14 @@ class CachedRuns:
|
||||
max_keys: int = MAX_KEYS,
|
||||
max_receipts: int = MAX_RECEIPTS,
|
||||
max_items: int = MAX_ITEMS,
|
||||
max_field_bytes: int = MAX_FIELD_TEXT_BYTES,
|
||||
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._max_field_bytes = max_field_bytes
|
||||
self._field_bytes = 0 # UTF-8 bytes of the fields all receipts keep
|
||||
self._context = context
|
||||
self._receipts: "OrderedDict[str, List[_Receipt]]" = OrderedDict()
|
||||
self._attempt: Optional[_Attempt] = None
|
||||
@@ -284,7 +330,10 @@ class CachedRuns:
|
||||
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)
|
||||
before = _bytes(receipts)
|
||||
deduplicated = _deduplicated(receipts)
|
||||
self._field_bytes += _bytes(deduplicated) - before
|
||||
self._receipts[run_key] = deduplicated
|
||||
|
||||
def _current(self) -> Tuple[Optional[_Attempt], Optional[int]]:
|
||||
"""The attempt the executing node belongs to, and its list index."""
|
||||
@@ -305,8 +354,19 @@ class CachedRuns:
|
||||
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."""
|
||||
def record_item(
|
||||
self,
|
||||
node_id: Any,
|
||||
outputs: Tuple[Any, ...],
|
||||
facts: Facts,
|
||||
fields: Optional[Fields] = None,
|
||||
) -> None:
|
||||
"""Remember what this execution returned and was entitled to count.
|
||||
|
||||
``fields`` are the texts the node gave the execution history, one
|
||||
``(raw, prepend, append)`` per output; kept only when well formed,
|
||||
each within the history's field limit and the store within its budget.
|
||||
"""
|
||||
try:
|
||||
attempt, index = self._current()
|
||||
node = attempt.nodes.get(str(node_id)) if attempt else None
|
||||
@@ -317,14 +377,20 @@ class CachedRuns:
|
||||
return
|
||||
if any(not isinstance(value, str) for value in outputs):
|
||||
return
|
||||
item = _Item(tuple(text_digest(value) for value in outputs), facts)
|
||||
size = _field_size(fields, len(outputs)) if fields is not None else None
|
||||
item = _Item(
|
||||
tuple(text_digest(value) for value in outputs),
|
||||
facts,
|
||||
fields if size is not None else None,
|
||||
size or 0,
|
||||
)
|
||||
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, [])
|
||||
receipts = self._take(run_key)
|
||||
if receipts and receipts[0].item_id == item_id:
|
||||
receipt = receipts[0]
|
||||
else:
|
||||
@@ -334,9 +400,44 @@ class CachedRuns:
|
||||
receipt = _Receipt(item_id)
|
||||
receipts.insert(0, receipt)
|
||||
receipt.items[index] = item
|
||||
self._receipts[run_key] = receipts[: self._max_receipts]
|
||||
receipts = receipts[: self._max_receipts]
|
||||
if item.field_bytes and (
|
||||
receipt.field_bytes() > MAX_RECEIPT_FIELD_BYTES
|
||||
or not self._make_room(_bytes(receipts))
|
||||
):
|
||||
# Counting is as before; this job's history fields stay unknown
|
||||
receipt.items[index] = item._replace(fields=None, field_bytes=0)
|
||||
self._put(run_key, receipts)
|
||||
while len(self._receipts) > self._max_keys:
|
||||
self._receipts.popitem(last=False)
|
||||
_, dropped = self._receipts.popitem(last=False)
|
||||
self._field_bytes -= _bytes(dropped)
|
||||
|
||||
def _make_room(self, needed: int) -> bool:
|
||||
"""Free the fields of the least recently used receipts (under the lock).
|
||||
|
||||
Old receipts of workflows never run again must not hold the budget:
|
||||
their counting stays, only their history fields go. False when
|
||||
``needed`` does not fit even then.
|
||||
"""
|
||||
if needed > self._max_field_bytes:
|
||||
return False
|
||||
for receipts in self._receipts.values(): # least recently used first
|
||||
if self._field_bytes + needed <= self._max_field_bytes:
|
||||
break
|
||||
for receipt in receipts:
|
||||
self._field_bytes -= receipt.strip_fields()
|
||||
return self._field_bytes + needed <= self._max_field_bytes
|
||||
|
||||
def _take(self, run_key: str) -> List[_Receipt]:
|
||||
"""Remove a key's receipts and their bytes from the store (under the lock)."""
|
||||
receipts = self._receipts.pop(run_key, [])
|
||||
self._field_bytes -= _bytes(receipts)
|
||||
return list(receipts)
|
||||
|
||||
def _put(self, run_key: str, receipts: List[_Receipt]) -> None:
|
||||
"""Store a key's receipts as the most recently used (under the lock)."""
|
||||
self._receipts[run_key] = receipts
|
||||
self._field_bytes += _bytes(receipts)
|
||||
|
||||
# -- cached nodes -------------------------------------------------------
|
||||
|
||||
@@ -393,25 +494,11 @@ class CachedRuns:
|
||||
|
||||
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]
|
||||
matches = self._matching(run_key, values)
|
||||
if not matches:
|
||||
return None
|
||||
self._receipts.move_to_end(run_key)
|
||||
length = len(values[0])
|
||||
facts = {
|
||||
tuple(receipt.items[index].facts for index in range(length))
|
||||
for receipt in matches
|
||||
@@ -420,13 +507,99 @@ class CachedRuns:
|
||||
return None
|
||||
return [matches[0].items[index] for index in range(length)]
|
||||
|
||||
def _matching(self, run_key: str, values: Any) -> List[_Receipt]:
|
||||
"""Receipts of ``run_key`` whose outputs equal ``values`` (under the lock)."""
|
||||
if type(values) not in (list, tuple) or not values:
|
||||
return []
|
||||
if any(type(slot) not in (list, tuple) for slot in values):
|
||||
return []
|
||||
length = len(values[0])
|
||||
if length == 0 or length > self._max_items:
|
||||
return []
|
||||
if any(len(slot) != length for slot in values):
|
||||
return []
|
||||
if any(type(value) is not str for slot in values for value in slot):
|
||||
return []
|
||||
wanted = {
|
||||
index: tuple(text_digest(slot[index]) for slot in values)
|
||||
for index in range(length)
|
||||
}
|
||||
receipts = self._receipts.get(run_key)
|
||||
if not receipts:
|
||||
return []
|
||||
matches = [receipt for receipt in receipts if receipt.outputs() == wanted]
|
||||
if matches:
|
||||
self._receipts.move_to_end(run_key)
|
||||
return matches
|
||||
|
||||
# -- the history capture (#403) -------------------------------------------
|
||||
|
||||
def recorded_fields(
|
||||
self, prompt_id: Any, node_id: Any, values: Any
|
||||
) -> Optional[List[Fields]]:
|
||||
"""The history fields of the execution a cached node reuses, per list item.
|
||||
|
||||
``values`` are the node's cached STRING outputs (``cached_string_outputs``).
|
||||
Only for a node of the job currently executing, and only when every
|
||||
receipt matching ``values`` kept the same fields; None otherwise. Never
|
||||
counts and never marks the node replayed.
|
||||
"""
|
||||
try:
|
||||
with self._lock:
|
||||
attempt = self._attempt
|
||||
if (
|
||||
attempt is None
|
||||
or type(prompt_id) is not str
|
||||
or prompt_id != attempt.prompt_id
|
||||
or type(node_id) is not str
|
||||
or node_id not in attempt.nodes
|
||||
):
|
||||
return None
|
||||
matches = self._matching(attempt.nodes[node_id][1], values)
|
||||
if not matches:
|
||||
return None
|
||||
found = []
|
||||
for index in range(len(values[0])):
|
||||
fields = {receipt.items[index].fields for receipt in matches}
|
||||
if len(fields) != 1 or None in fields or _CONFLICT in fields:
|
||||
return None
|
||||
found.append(fields.pop())
|
||||
return found
|
||||
except Exception as error:
|
||||
logger.debug(f"Could not read a cached run's history fields: {error}")
|
||||
return None
|
||||
|
||||
|
||||
def _deduplicated(receipts: List[_Receipt]) -> List[_Receipt]:
|
||||
"""The newest receipt replaces older ones with the same outputs."""
|
||||
"""The newest receipt replaces older ones with the same outputs.
|
||||
|
||||
Counting needs only the outputs. When a replaced receipt recorded other
|
||||
history fields for equal outputs (a raw text the outputs do not show),
|
||||
the newest item's fields become a conflict, so no cached job of it is
|
||||
given either text (#403). A conflict stays until the run key is dropped.
|
||||
"""
|
||||
if not receipts:
|
||||
return receipts
|
||||
outputs = receipts[0].outputs()
|
||||
return [receipts[0]] + [r for r in receipts[1:] if r.outputs() != outputs]
|
||||
newest, outputs = receipts[0], receipts[0].outputs()
|
||||
kept = [newest]
|
||||
for receipt in receipts[1:]:
|
||||
if receipt.outputs() != outputs:
|
||||
kept.append(receipt)
|
||||
continue
|
||||
for index, item in list(newest.items.items()):
|
||||
older = receipt.items[index].fields
|
||||
if older == _CONFLICT or (
|
||||
older is not None and item.fields is not None and older != item.fields
|
||||
):
|
||||
newest.items[index] = item._replace(fields=_CONFLICT, field_bytes=0)
|
||||
return kept
|
||||
|
||||
|
||||
def _bytes(receipts: List[_Receipt]) -> int:
|
||||
"""UTF-8 bytes of the history fields ``receipts`` keep."""
|
||||
return sum(
|
||||
item.field_bytes for receipt in receipts for item in receipt.items.values()
|
||||
)
|
||||
|
||||
|
||||
def _job_nodes(graph: Dict[str, Any], outputs: Any) -> Dict[str, Tuple[str, str]]:
|
||||
|
||||
@@ -151,10 +151,77 @@ def cached_string_outputs(
|
||||
return None
|
||||
|
||||
|
||||
def _receipt_fields(
|
||||
receipts: object, executor: object, job_id: str, source: object
|
||||
) -> list | None:
|
||||
"""Per list index, the ``(raw, prepend, append)`` of the execution reused (#403).
|
||||
|
||||
From the #347 receipt whose digests equal every STRING slot of the cache
|
||||
entry; None when there is none or the answer is malformed.
|
||||
"""
|
||||
try:
|
||||
if receipts is None:
|
||||
try:
|
||||
from .cached_runs import CACHED_RUNS as receipts
|
||||
except ImportError:
|
||||
from utils.cached_runs import CACHED_RUNS as receipts
|
||||
values = cached_string_outputs(executor, source.node_id, source.node_type)
|
||||
if values is None:
|
||||
return None
|
||||
found = receipts.recorded_fields(job_id, source.node_id, values)
|
||||
if type(found) is not list or len(found) != len(values[0]):
|
||||
return None
|
||||
position = (
|
||||
(1 if source.output_slot in (1, 3) else 0)
|
||||
if source.node_type == "PromptManagerPair"
|
||||
else 0
|
||||
)
|
||||
fields = []
|
||||
for slots in found:
|
||||
slot = slots[position] if type(slots) is tuple else None
|
||||
if not (
|
||||
type(slot) is tuple
|
||||
and len(slot) == 3
|
||||
and all(type(value) is str for value in slot)
|
||||
):
|
||||
return None
|
||||
fields.append(slot)
|
||||
return fields
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _agrees(inferred: PromptSnapshot, fields: tuple) -> bool:
|
||||
"""Every field the graph knows equals the receipt's; else fill nothing."""
|
||||
known = (inferred.raw, inferred.prepend, inferred.append)
|
||||
return all(
|
||||
field.value is None or field.value == value
|
||||
for field, value in zip(known, fields)
|
||||
)
|
||||
|
||||
|
||||
def _filled(inferred: TextEvidence, value: str) -> TextEvidence:
|
||||
"""A field the graph cannot show, taken from the reused execution (#403)."""
|
||||
if inferred.value is not None:
|
||||
return inferred
|
||||
return TextEvidence(value, "inferred", "graph_semantics", None)
|
||||
|
||||
|
||||
def observe_cached_context(
|
||||
executor: object, message_type: object, message_data: object
|
||||
executor: object,
|
||||
message_type: object,
|
||||
message_data: object,
|
||||
*,
|
||||
receipts: object = None,
|
||||
) -> None:
|
||||
"""Copy bounded local strings without cache/provider execution or retention."""
|
||||
"""Copy bounded local strings without cache/provider execution or retention.
|
||||
|
||||
Raw, prepend and append come from graph inference; a field it leaves
|
||||
unknown (runtime-only text) is taken from the receipt of the execution the
|
||||
cache entry is proven to come from (``CachedRuns.recorded_fields``), so the
|
||||
cached job is a tracked run of its prompt (#403). ``receipts`` defaults to
|
||||
the process ``CACHED_RUNS``.
|
||||
"""
|
||||
try:
|
||||
binding = active_capture()
|
||||
if binding is None or type(message_type) is not str:
|
||||
@@ -198,18 +265,29 @@ def observe_cached_context(
|
||||
inferred = manager.source_snapshot(cap, source)
|
||||
if inferred is None:
|
||||
continue
|
||||
fields = None
|
||||
if any(
|
||||
field.value is None
|
||||
for field in (inferred.raw, inferred.prepend, inferred.append)
|
||||
):
|
||||
fields = _receipt_fields(receipts, executor, cap.job_id, source)
|
||||
for index, value in enumerate(values):
|
||||
if type(value) is not str:
|
||||
continue
|
||||
raw, prepend, append = inferred.raw, inferred.prepend, inferred.append
|
||||
if fields is not None and _agrees(inferred, fields[index]):
|
||||
raw = _filled(raw, fields[index][0])
|
||||
prepend = _filled(prepend, fields[index][1])
|
||||
append = _filled(append, fields[index][2])
|
||||
manager.observe(
|
||||
cap,
|
||||
source,
|
||||
index,
|
||||
PromptSnapshot(
|
||||
inferred.raw,
|
||||
raw,
|
||||
TextEvidence(value, "recorded", "host_cache_output", None),
|
||||
inferred.prepend,
|
||||
inferred.append,
|
||||
prepend,
|
||||
append,
|
||||
),
|
||||
)
|
||||
except Exception:
|
||||
|
||||
Reference in New Issue
Block a user