nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'origin/main' into fix/460-case-history

Browse Source 
# Conflicts:
#	database/execution_contexts.py
#	database/history_reads.py
#	tests/comfyui/execution-history.spec.mjs
#	tests/test_execution_context_database.py
#	tests/test_execution_history_hooks.py
#	tests/test_hashing.py
#	tests/test_history_reads.py
This commit is contained in:
Vito Sansevero
2026-10-29 05:47:06 -07:00
parent 9d183a26f6 b5ca2eda94
commit fabd87e1bc
61 changed files with 4682 additions and 213 deletions
Download Patch File Download Diff File Expand all files Collapse all files
utils/foreign_runs.py
+142
View File
@@ -0,0 +1,142 @@
"""The monitor's path for outputs of workflows without a PromptManager node (#247).
A foreign output (no positive PromptManager node in its graph) is decided by
the ladder the output scan uses (``utils.file_prompt.decide_file_prompt``)
with the ``MONITOR`` policy: certain files are stored at once, files only
recorded history can decide wait for their run's outputs to be sealed, and
everything else stays unlinked. This module holds the parts that do not
need the monitor's waiting sets: the setting, the decision call and the
store. The waiting and the re-check stay on ``ImageGenerationHandler``.
"""
import os
from dataclasses import dataclass
from typing import Any, Callable, Mapping, Optional, Tuple
from .file_prompt import (
NOT_SAVED,
NOT_SAVED_WIRED,
PAIR,
POSITIVE,
Decision,
Policy,
decide_file_prompt,
)
from .hashing import generate_prompt_hash
from .usage_tracking import survive_switch
try:
from ..database.pair_ingestion import RUN_MARKS, ingest_resolved_pair
from ..database.run_attribution import recorded_output_lookup
from ..database.scanned_prompts import store_scanned_prompt_or_raise
except ImportError:
from database.pair_ingestion import RUN_MARKS, ingest_resolved_pair
from database.run_attribution import recorded_output_lookup
from database.scanned_prompts import store_scanned_prompt_or_raise
# Decisions that name a prompt to store
CERTAIN = (PAIR, POSITIVE)
# Decisions of a run whose prompt has Save to library off, or wired from a
# node the file cannot show (#90): final, nothing is stored, nothing waits
NOT_SAVED_KINDS = (NOT_SAVED, NOT_SAVED_WIRED)
# Sources certain without a recorded run: the graph alone, or its recorded saver
ANCHORED_SOURCES = ("inspector", "recorded-run")
@dataclass(frozen=True)
class ForeignRecheck:
"""``schedule`` snapshot: the one re-check of a foreign output after its
run's outputs were sealed. ``reference`` is the sealed output reference,
the spelling history recorded (``recorded_output_lookup(reference=)``)."""
reference: Mapping[str, Any]
def save_foreign_runs() -> bool:
"""Whether a foreign output may be stored now, read for every file.
``GalleryConfig.SAVE_FOREIGN_RUNS`` and ``MONITORING_ENABLED``: a
monitoring switch saved off stops foreign writes at once, although the
running observer watches until a restart. True when the configuration
cannot be imported (the defaults).
"""
try:
from ..py.config import GalleryConfig
except ImportError:
try:
from py.config import GalleryConfig
except ImportError:
return True
return (
GalleryConfig.SAVE_FOREIGN_RUNS is not False
and GalleryConfig.MONITORING_ENABLED is not False
)
def history_lookup(
db: Any, image_path: str, reference: Optional[Mapping[str, Any]] = None
) -> Callable[[], Any]:
"""The ladder's ``lookup`` for one file: its ``RecordedOutput`` or None.
With ``reference`` the file is spelled as history recorded it. Failures
are logged by ``recorded_output_lookup`` and name nothing.
"""
lookup = recorded_output_lookup(
db, reference=dict(reference) if reference is not None else None
)
return lambda: lookup(image_path)
def decide_foreign(
image_path: str,
metadata: Any,
*,
in_output: bool,
lookup: Optional[Callable[[], Any]],
) -> Decision:
"""The ``MONITOR`` ladder for one foreign output (the size is read only
when the ladder needs it)."""
def file_size() -> Optional[int]:
try:
return os.path.getsize(image_path)
except OSError:
return None
return decide_file_prompt(
metadata,
file_name=os.path.basename(image_path),
file_size=file_size,
in_output=in_output,
lookup=lookup,
policy=Policy.MONITOR,
)
def store_foreign(
db: Any, image_path: str, decision: Decision
) -> Tuple[str, Optional[int]]:
"""Store a certain decision: ``(status, prompt id)``.
One write through ``survive_switch``, so it runs again on the moved
library when a Move library switched while it waited (#249 S6b). Created
prompts carry ``RUN_MARKS``. A file linked to another prompt is a
``conflict`` and nothing is written. Database errors are raised.
"""
def store():
if decision.kind == PAIR:
result = ingest_resolved_pair(
db, image_path, decision.canonical, decision.resolution, marks=RUN_MARKS
)
return result["status"], result.get("image_id")
return store_scanned_prompt_or_raise(
db, image_path, decision.canonical, decision.text, RUN_MARKS
)
status, linked = survive_switch(db, store)
if decision.kind != PAIR or linked is None:
return status, linked
# The pair ingester reports the image row; its prompt is the positive's.
prompt = db.get_prompt_by_hash(generate_prompt_hash(decision.text))
return status, prompt["id"] if prompt else None
utils/image_monitor.py
+201 -41
View File
@@ -33,6 +33,17 @@ from typing import Callable, Optional, Dict, Any, List, Tuple
from watchdog.observers import Observer
from watchdog.events import FileSystemEventHandler
from .file_prompt import HISTORY, NOT_SAVED, NOT_SAVED_WIRED
from .foreign_runs import (
ANCHORED_SOURCES,
CERTAIN,
NOT_SAVED_KINDS,
ForeignRecheck,
decide_foreign,
history_lookup,
save_foreign_runs,
store_foreign,
)
from .hashing import generate_prompt_hash
from .metadata_extractor import ComfyUIMetadataExtractor
from .parallel import note_refusal, refused_at_shutdown
@@ -153,6 +164,11 @@ class ImageGenerationHandler(FileSystemEventHandler):
# Sealed re-checks ``schedule`` refused (the path still pending, or the
# pending set full): retried each time a pending image finishes.
self._history_due: "OrderedDict[str, None]" = OrderedDict()
# The same two sets for foreign outputs (#247), apart so a large
# foreign run never evicts a waiting keyless file; the due set keeps
# each file's sealed output reference.
self._awaiting_foreign: "OrderedDict[str, Tuple[str, float]]" = OrderedDict()
self._foreign_due: "OrderedDict[str, dict]" = OrderedDict()
self._awaiting_lock = threading.Lock()
self._sealed_subscribed = False
# Set by ``close``: no new work, no new worker, ever (#302).
@@ -253,6 +269,8 @@ class ImageGenerationHandler(FileSystemEventHandler):
with self._awaiting_lock:
self._awaiting.clear()
self._history_due.clear()
self._awaiting_foreign.clear()
self._foreign_due.clear()
with self._pending_lock:
self._pending.clear()
worker, self._worker = self._worker, None
@@ -438,6 +456,9 @@ class ImageGenerationHandler(FileSystemEventHandler):
if prompt_snapshot is HISTORY_ONLY:
self._recheck_history(image_path)
return
if isinstance(prompt_snapshot, ForeignRecheck):
self._recheck_foreign(image_path, prompt_snapshot.reference)
return
self.logger.info(f"Processing image: {image_path}")
if not os.path.exists(image_path):
@@ -457,17 +478,12 @@ class ImageGenerationHandler(FileSystemEventHandler):
# negative-prompt PromptManager nodes in the workflow.
state, current_prompt = self._find_prompt_from_metadata(metadata)
if state == "opted_out":
self.logger.info(
"Save to library is off for this run's prompt, skipping image: "
f"{image_path}"
)
self._log_not_saved(image_path)
return
if state == "foreign":
# Another workflow saved this image: it belongs to no prompt of
# ours, and must not consume a batch queue entry.
self.logger.info(
f"No PromptManager prompt in workflow, skipping image: {image_path}"
)
# Another workflow saved this image: never a batch queue entry
# of ours; its own graph may name its prompt (#247).
self._ingest_foreign(image_path, metadata)
return
if current_prompt:
self.logger.info(
@@ -589,36 +605,55 @@ class ImageGenerationHandler(FileSystemEventHandler):
resolved = os.path.abspath(image_path)
return os.path.normcase(resolved)
def _prune_awaiting(self) -> None:
"""Drop expired and surplus entries (``_awaiting_lock`` held)."""
def _prune_awaiting(self, store=None) -> List[Tuple[str, bool]]:
"""Drop expired and surplus entries of ``store`` (the keyless set by
default; ``_awaiting_lock`` held). Returns the dropped foreign files
as ``(path, evicted)`` for ``_log_dropped`` after the lock."""
store = self._awaiting if store is None else store
now = time.monotonic()
for key in [k for k, (_p, due) in self._awaiting.items() if due <= now]:
del self._awaiting[key]
while len(self._awaiting) > MAX_AWAITING_HISTORY:
self._awaiting.popitem(last=False)
expired = [k for k, (_p, due) in store.items() if due <= now]
dropped = [(store.pop(key)[0], False) for key in expired]
while len(store) > MAX_AWAITING_HISTORY:
dropped.append((store.popitem(last=False)[1][0], True))
return dropped if store is self._awaiting_foreign else []
def _await_history(self, image_path: str) -> None:
def _log_dropped(self, dropped: List[Tuple[str, bool]]) -> None:
"""Log foreign files that stopped waiting (outside ``_awaiting_lock``)."""
for image_path, evicted in dropped:
name = os.path.basename(image_path)
if evicted:
self.logger.info(
f"Too many files wait for their run; {name} left unlinked"
)
else:
self.logger.info(f"No recorded run for {name}; left unlinked")
def _await_history(self, image_path: str, store=None) -> None:
"""Re-check ``image_path`` once, history only, when its run is sealed.
A closed handler re-checks nothing.
``store`` is the waiting set (``_awaiting_foreign`` for a foreign
output). A closed handler re-checks nothing.
"""
if self._closed:
return
if not self._sealed_subscribed:
subscribe_outputs_sealed(self._on_outputs_sealed)
self._sealed_subscribed = True
store = self._awaiting if store is None else store
key = self._awaiting_key(image_path)
with self._awaiting_lock:
self._awaiting.pop(key, None)
self._awaiting[key] = (
image_path,
time.monotonic() + AWAIT_HISTORY_SECONDS,
)
self._prune_awaiting()
store.pop(key, None)
store[key] = (image_path, time.monotonic() + AWAIT_HISTORY_SECONDS)
dropped = self._prune_awaiting(store)
self._log_dropped(dropped)
def _stop_awaiting(self, image_path: str) -> None:
def _await_foreign(self, image_path: str) -> None:
self._await_history(image_path, self._awaiting_foreign)
def _stop_awaiting(self, image_path: str, store=None) -> None:
store = self._awaiting if store is None else store
with self._awaiting_lock:
self._awaiting.pop(self._awaiting_key(image_path), None)
store.pop(self._awaiting_key(image_path), None)
def awaiting_history(self) -> List[str]:
"""Normalized paths still waiting for their run's recorded outputs."""
@@ -626,6 +661,14 @@ class ImageGenerationHandler(FileSystemEventHandler):
self._prune_awaiting()
return list(self._awaiting)
def awaiting_foreign(self) -> List[str]:
"""Normalized paths of foreign outputs waiting for their run (#247)."""
with self._awaiting_lock:
dropped = self._prune_awaiting(self._awaiting_foreign)
waiting = list(self._awaiting_foreign)
self._log_dropped(dropped)
return waiting
def _on_outputs_sealed(self, run_id: str, references: List[dict]) -> None:
"""A run's outputs were recorded: re-check its waiting images once.
@@ -637,6 +680,7 @@ class ImageGenerationHandler(FileSystemEventHandler):
due = []
with self._awaiting_lock:
self._prune_awaiting()
dropped = self._prune_awaiting(self._awaiting_foreign)
for reference in references:
if reference.get("folder_type") != "output":
continue
@@ -647,39 +691,53 @@ class ImageGenerationHandler(FileSystemEventHandler):
)
waiting = self._awaiting.pop(key, None)
if waiting is not None:
due.append(waiting[0])
for image_path in due:
self._schedule_recheck(image_path)
due.append((waiting[0], None))
waiting = self._awaiting_foreign.pop(key, None)
if waiting is not None:
due.append((waiting[0], dict(reference)))
self._log_dropped(dropped)
for image_path, reference in due:
self._schedule_recheck(image_path, reference)
def _schedule_recheck(self, image_path: str) -> None:
"""Schedule the history-only re-check, or keep it due for later.
@staticmethod
def _recheck_snapshot(reference):
return HISTORY_ONLY if reference is None else ForeignRecheck(reference)
def _schedule_recheck(self, image_path: str, reference=None) -> None:
"""Schedule the history-only re-check (with the sealed ``reference``:
a foreign output's), or keep it due for later.
A closed handler drops it.
"""
if self._closed:
return
if self.schedule(image_path, HISTORY_ONLY):
if self.schedule(image_path, self._recheck_snapshot(reference)):
return
with self._awaiting_lock:
# ``close`` sets the flag before ``shutdown`` clears this set, so
# a close that raced the refusal above is seen here.
if self._closed:
return
self._history_due.pop(image_path, None)
self._history_due[image_path] = None
while len(self._history_due) > MAX_AWAITING_HISTORY:
self._history_due.popitem(last=False)
due = self._history_due if reference is None else self._foreign_due
due.pop(image_path, None)
due[image_path] = reference
evicted = []
while len(due) > MAX_AWAITING_HISTORY:
evicted.append(due.popitem(last=False)[0])
if reference is not None:
self._log_dropped([(path, True) for path in evicted])
def _drain_history_due(self) -> None:
"""Schedule due re-checks now that a pending slot or path is free."""
with self._awaiting_lock:
if not self._history_due:
return
due = list(self._history_due)
for image_path in due:
if self.schedule(image_path, HISTORY_ONLY):
due = [(self._history_due, p, None) for p in self._history_due]
due += [(self._foreign_due, *item) for item in self._foreign_due.items()]
for store, image_path, reference in due:
if self.schedule(image_path, self._recheck_snapshot(reference)):
with self._awaiting_lock:
self._history_due.pop(image_path, None)
# A newer seal may have replaced it meanwhile: keep that one
if image_path in store and store[image_path] is reference:
del store[image_path]
def _recheck_history(self, image_path: str) -> None:
"""The one history-only re-check of a keyless image (#342)."""
@@ -703,6 +761,108 @@ class ImageGenerationHandler(FileSystemEventHandler):
metadata or {"file_info": self.get_basic_file_info(image_path)},
)
def _ingest_foreign(self, image_path: str, metadata) -> None:
"""A foreign output at settle: store its certain prompt, or wait for
its run's recorded outputs when only history can decide (#247)."""
name = os.path.basename(image_path)
if not save_foreign_runs():
self.logger.info(
f"No PromptManager prompt in workflow, skipping image: {image_path}"
)
return
in_output = in_output_folder(image_path, host_output_directory())
decision = decide_foreign(
image_path, metadata, in_output=in_output, lookup=None
)
if decision.kind == HISTORY:
decision = self._settle_lookup(image_path, metadata)
if decision is None:
self.logger.debug(f"{name} waits for its run to end")
return
self._stop_awaiting(image_path, self._awaiting_foreign)
if decision.kind in NOT_SAVED_KINDS:
self._log_not_saved(image_path, decision.kind)
elif decision.kind in CERTAIN:
self._store_foreign(image_path, decision)
else:
self.logger.info(f"No certain prompt for {name}; left unlinked")
def _log_not_saved(self, image_path: str, kind: str = NOT_SAVED) -> None:
"""The log line of a file whose run's prompt does not save (#90)."""
if kind == NOT_SAVED_WIRED:
reason = "is wired from a node the image cannot show"
else:
reason = "is off"
self.logger.info(
f"Save to library {reason} for this run's prompt, skipping image: "
f"{image_path}"
)
def _settle_lookup(self, image_path: str, metadata):
"""The settle-time history lookup of a file only history can decide.
Registered first, so a run sealed before this lookup is not missed.
Returns None to keep waiting: no record yet, and the answer is not
the inspector's or the recorded saver's (the legacy floor waits for
the re-check). With a record, any decision ends the wait. A
``not_saved`` here is the undecided step's (an early one never asks
history, so it never reaches this lookup): a recorded saver may still
resolve another branch, so without a record it waits too (#90).
"""
self._await_foreign(image_path)
found = []
lookup = history_lookup(self.db_manager, image_path)
def recorded():
found.append(lookup())
return found[-1]
decision = decide_foreign(image_path, metadata, in_output=True, lookup=recorded)
if any(record is not None for record in found):
return decision
if decision.kind in CERTAIN and decision.source in ANCHORED_SOURCES:
return decision
return None
def _recheck_foreign(self, image_path: str, reference) -> None:
"""The one re-check of a foreign output after its run ended (#247)."""
name = os.path.basename(image_path)
if not os.path.exists(image_path):
return
if not save_foreign_runs():
self.logger.info(f"Saving prompts from other workflows is off: {name}")
return
try:
metadata = self.metadata_extractor.extract_metadata(image_path)
except Exception as error:
self.logger.warning(f"Metadata of {name} could not be read: {error}")
return
lookup = history_lookup(self.db_manager, image_path, reference)
decision = decide_foreign(image_path, metadata, in_output=True, lookup=lookup)
if decision.kind in NOT_SAVED_KINDS:
self._log_not_saved(image_path, decision.kind)
return
if decision.kind in CERTAIN:
self._store_foreign(image_path, decision)
return
self.logger.info(
f"No certain prompt for {name} after its run ended; left unlinked"
)
def _store_foreign(self, image_path: str, decision) -> None:
name = os.path.basename(image_path)
if not save_foreign_runs(): # switched off while this file was decided
self.logger.info(f"Saving prompts from other workflows is off: {name}")
return
try:
status, prompt_id = store_foreign(self.db_manager, image_path, decision)
except Exception as error:
self.logger.error(f"Failed to store the prompt of {name}: {error}")
return
self.logger.info(
f"Foreign run: {status} prompt {prompt_id} for {name} ({decision.source})"
)
def _history_prompt(self, image_path: str) -> Optional[int]:
"""The stored prompt recorded history attributes to this file, if one."""
try: