From b313224d464ea48cb5007f2b069b95f5af13cc33 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Wed, 28 Oct 2026 06:43:35 -0700 Subject: [PATCH] feat(monitor): store the certain prompt of a foreign workflow's output (#247) A foreign output (no positive PromptManager node) is decided by the shared ladder with the MONITOR policy. Certain files are stored at settle with RUN_MARKS through survive_switch; files only recorded history can decide wait in their own bounded set and are re-checked once after the run's outputs are sealed, with the sealed reference. The setting GalleryConfig.SAVE_FOREIGN_RUNS is read for every file. --- tests/test_image_linking.py | 25 ++- tests/test_image_monitor_singleton.py | 5 + tests/test_run_attribution.py | 286 ++++++++++++++++++++++++++ utils/foreign_runs.py | 117 +++++++++++ utils/image_monitor.py | 181 ++++++++++++---- 5 files changed, 576 insertions(+), 38 deletions(-) create mode 100644 utils/foreign_runs.py diff --git a/tests/test_image_linking.py b/tests/test_image_linking.py index 3cc8586a..7d7d23be 100644 --- a/tests/test_image_linking.py +++ b/tests/test_image_linking.py @@ -172,17 +172,40 @@ class TestForeignWorkflows(LinkingTestCase): self.tracker.set_current_prompt(text, {"prompt_id": prompt_id}) return prompt_id + def _linked_text(self): + row = ( + self.db.model.get_connection() + .execute( + "SELECT p.text FROM generated_images g " + "JOIN prompts p ON p.id = g.prompt_id WHERE g.image_path = ?", + (self.image,), + ) + .fetchone() + ) + return row[0] if row else None + def test_foreign_image_leaves_the_batch_queue_aligned(self): + """#247: it links to its own prompt and never takes a queue entry.""" first = self._queue_batch_item("item one") self.assertIsNone(self._process({"prompt": foreign_graph()})) self.assertEqual(self.linked, []) self.assertEqual(self.tracker.pop_next_prompt()["id"], first) self.assertIsNone(self.tracker.pop_next_prompt()) + self.assertEqual(self._linked_text(), "a cat") + prompt = self.db.get_prompt_by_hash(generate_prompt_hash("a cat")) + self.assertEqual(prompt["run_count"] or 0, 0) - def test_foreign_image_is_not_linked_by_the_later_fallbacks_either(self): + def test_a_foreign_image_is_linked_only_by_the_ladder(self): self._queue_batch_item("item one") # live tracker and recent prompt exist + untraceable = foreign_graph() + untraceable["9"]["inputs"]["text"] = ["30", 0] + untraceable["30"] = {"class_type": "VendorPromptMaker", "inputs": {}} + self.assertIsNone(self._process({"prompt": untraceable})) + self.assertEqual(self.linked, []) + self.assertIsNone(self._linked_text()) self.assertIsNone(self._process({"prompt": foreign_graph()})) self.assertEqual(self.linked, []) + self.assertEqual(self._linked_text(), "a cat") def test_negative_only_workflow_is_foreign(self): self._save("blurry") diff --git a/tests/test_image_monitor_singleton.py b/tests/test_image_monitor_singleton.py index 320c70be..658404c7 100644 --- a/tests/test_image_monitor_singleton.py +++ b/tests/test_image_monitor_singleton.py @@ -597,6 +597,11 @@ class TestHandlerScheduling(unittest.TestCase): self.handler._schedule_recheck(path) self.assertIsNone(self.handler._worker) self.assertEqual(list(self.handler._history_due), []) + # #247: a foreign file's sealed re-check is dropped the same way + self.handler._await_foreign(path) + self.assertEqual(self.handler.awaiting_foreign(), []) + self.handler._schedule_recheck(path, {"filename": "late.png"}) + self.assertEqual(list(self.handler._foreign_due), []) self.assertEqual(self._live_workers() - before, set()) def test_a_sealed_run_starts_a_worker_on_an_open_handler(self): diff --git a/tests/test_run_attribution.py b/tests/test_run_attribution.py index b041bd14..b2f29da1 100644 --- a/tests/test_run_attribution.py +++ b/tests/test_run_attribution.py @@ -33,6 +33,7 @@ from utils.execution_outputs import ( snapshot_output_roots, ) from utils.hashing import generate_prompt_hash +from py.config import GalleryConfig from utils.image_monitor import ImageGenerationHandler STYLE = "windows" if os.name == "nt" else "posix" @@ -730,6 +731,291 @@ class TestMonitorHistoryRecheck(RecordedRuns, unittest.TestCase): run_attribution.notify_outputs_sealed("run", []) # no dead callback raises +class TestMonitorForeignRuns(RecordedRuns, unittest.TestCase): + """#247: outputs of workflows without a PromptManager node, in the monitor.""" + + def setUp(self): + history_tests.HistoryDatabaseTests.setUp(self) + self.db = history_tests.HistoryDatabaseTests.database(self) + self.output = str(Path(self.temp.name) / "output") + os.makedirs(os.path.join(self.output, "dataset")) + self.tracker = make_tracker(self.db) + self.handler = ImageGenerationHandler(self.db, self.tracker) + self.handler.processing_delay = 0 + self.handler.settle_interval = 0.01 + self.addCleanup(close_handler, self, self.handler) + self.metadata = {"prompt": shapes.shared_prefix()} + saved = GalleryConfig.SAVE_FOREIGN_RUNS + self.addCleanup(setattr, GalleryConfig, "SAVE_FOREIGN_RUNS", saved) + GalleryConfig.SAVE_FOREIGN_RUNS = True + patches = ( + mock.patch( + "utils.image_monitor.host_output_directory", return_value=self.output + ), + mock.patch( + "database.run_attribution.host_output_directory", + return_value=self.output, + ), + mock.patch.object( + self.handler.metadata_extractor, + "extract_metadata", + side_effect=lambda path: self.metadata, + ), + ) + for patch in patches: + patch.start() + self.addCleanup(patch.stop) + + close_databases = history_tests.HistoryDatabaseTests.close_databases + wait_for = TestMonitorHistoryRecheck.wait_for + + def image(self, name, subfolder="dataset"): + path = os.path.join(self.output, subfolder, name) + Path(path).write_bytes(b"png") + return path + + def linked_text(self, path): + row = ( + self.db.model.get_connection() + .execute( + "SELECT p.text FROM generated_images g " + "JOIN prompts p ON p.id = g.prompt_id WHERE g.image_path = ?", + (str(path),), + ) + .fetchone() + ) + return row[0] if row else None + + def key(self, path): + return self.handler._awaiting_key(path) + + def prompt_count(self): + connection = self.db.model.get_connection() + return connection.execute("SELECT COUNT(*) FROM prompts").fetchone()[0] + + def settled(self): + return self.wait_for(lambda: self.handler.pending_count() == 0) + + def test_each_dataset_file_links_to_its_branch_once_the_run_ends(self): + """CUJ 1: unlinked while the run is open, its own branch after the seal.""" + paths = [self.image(f"IMG_GEN_0000{n}_.png") for n in (1, 2)] + for path in paths: + self.handler.process_new_image(path) + self.assertEqual([self.linked_text(path) for path in paths], [None, None]) + self.assertEqual( + self.handler.awaiting_foreign(), [self.key(path) for path in paths] + ) + self.assertEqual(self.handler.awaiting_history(), []) + self.record_run( + self.db, + self.output, + [("dataset", os.path.basename(path)) for path in paths], + [], + nodes=["13", "23"], + ) + self.assertTrue(self.wait_for(lambda: all(map(self.linked_text, paths)))) + self.assertTrue(self.settled()) + self.assertEqual( + [self.linked_text(path) for path in paths], list(shapes.BRANCH_TEXTS[:2]) + ) + self.assertEqual(self.handler.awaiting_foreign(), []) + prompt = self.db.get_prompt_by_hash( + generate_prompt_hash(shapes.BRANCH_TEXTS[0]) + ) + self.assertEqual(prompt["category"], "scanned") + self.assertIn("Saved from a workflow run: IMG_GEN_00001_.png", prompt["notes"]) + + def test_a_run_sealed_before_the_file_settles_links_it_at_once(self): + path = self.image("IMG_GEN_00003_.png") + self.record_run( + self.db, self.output, [("dataset", "IMG_GEN_00003_.png")], [], nodes=["33"] + ) + self.handler.process_new_image(path) + self.assertEqual(self.linked_text(path), shapes.BRANCH_TEXTS[2]) + self.assertEqual(self.handler.awaiting_foreign(), []) + + def test_a_name_two_runs_reported_stays_unlinked(self): + """CUJ 3: the file was overwritten; no run proves its saver.""" + path = self.image("IMG_GEN_00004_.png") + reported = [("dataset", "IMG_GEN_00004_.png")] + self.record_run(self.db, self.output, reported, [], nodes=["13"]) + self.record_run(self.db, self.output, reported, [], item=1, nodes=["23"]) + self.handler.process_new_image(path) + self.assertIsNone(self.linked_text(path)) + reference = {"folder_type": "output", "subfolder": "dataset"} + reference["filename"] = "IMG_GEN_00004_.png" + with self.assertLogs("prompt_manager.image_monitor", "INFO") as logs: + self.handler._recheck_foreign(path, reference) + self.assertIn( + "No certain prompt for IMG_GEN_00004_.png after its run ended; " + "left unlinked", + "\n".join(logs.output), + ) + self.assertIsNone(self.linked_text(path)) + self.assertEqual(self.prompt_count(), 0) + + def test_the_setting_off_stores_nothing_at_settle_or_after_the_run(self): + """CUJ 4: read for every file, so it applies without a restart.""" + GalleryConfig.SAVE_FOREIGN_RUNS = False + skipped = self.image("IMG_GEN_00005_.png") + with self.assertLogs("prompt_manager.image_monitor", "INFO") as logs: + self.handler.process_new_image(skipped) + self.assertIn( + "No PromptManager prompt in workflow, skipping image", + "\n".join(logs.output), + ) + self.assertEqual(self.handler.awaiting_foreign(), []) + GalleryConfig.SAVE_FOREIGN_RUNS = True + waiting = self.image("IMG_GEN_00006_.png") + self.handler.process_new_image(waiting) + self.assertEqual(self.handler.awaiting_foreign(), [self.key(waiting)]) + GalleryConfig.SAVE_FOREIGN_RUNS = False + self.record_run( + self.db, self.output, [("dataset", "IMG_GEN_00006_.png")], [], nodes=["13"] + ) + self.assertTrue(self.wait_for(lambda: not self.handler.awaiting_foreign())) + self.assertTrue(self.settled()) + self.assertEqual( + [self.linked_text(p) for p in (skipped, waiting)], [None, None] + ) + self.assertEqual(self.prompt_count(), 0) + + def test_a_plain_foreign_workflow_links_at_settle_and_counts_no_run(self): + """CUJ 2: the default profile decides; the tracker queue is untouched.""" + text = "a lighthouse on a cliff at dusk" + self.metadata = {"prompt": sampled_graph(text, "blurry")} + queued = self.db.save_prompt( + text="queued item", prompt_hash=generate_prompt_hash("queued item") + ) + self.tracker.set_current_prompt("queued item", {"prompt_id": queued}) + first, second = self.image("plain_00001_.png"), self.image("plain_00002_.png") + self.handler.process_new_image(first) + self.handler.process_new_image(second) + self.assertEqual( + [self.linked_text(first), self.linked_text(second)], [text, text] + ) + prompt = self.db.get_prompt_by_hash(generate_prompt_hash(text)) + self.assertEqual(prompt["run_count"] or 0, 0) + self.assertIsNone(prompt.get("last_used_at")) + self.assertEqual(self.tracker.pop_next_prompt()["id"], queued) + self.assertEqual(self.handler.awaiting_foreign(), []) + + def test_the_sealed_reference_spells_the_file_as_history_recorded_it(self): + """Windows: the watcher's spelling differs in case from the run's report.""" + Path(self.output, "Dataset").mkdir() + recorded = self.image("IMG_GEN_00007_.png", subfolder="Dataset") + seen = os.path.join(self.output, "dataset", "img_gen_00007_.png") + Path(seen).write_bytes(b"png") + folded = staticmethod(lambda path: os.path.realpath(path).lower()) + with mock.patch.object(ImageGenerationHandler, "_awaiting_key", folded): + self.handler.process_new_image(seen) + self.assertIsNone(self.linked_text(seen)) + self.record_run( + self.db, + self.output, + [("Dataset", os.path.basename(recorded))], + [], + nodes=["43"], + ) + self.assertTrue(self.wait_for(lambda: self.linked_text(seen))) + self.assertTrue(self.settled()) + self.assertEqual(self.linked_text(seen), shapes.BRANCH_TEXTS[3]) + + def test_a_full_foreign_set_never_evicts_a_waiting_keyless_file(self): + from utils import image_monitor + + keyless = self.image("bare_00001_.png", subfolder="") + with mock.patch.object( + self.handler.metadata_extractor, "extract_metadata", return_value=None + ): + self.handler.process_new_image(keyless) + paths = [self.image(f"IMG_GEN_1000{n}_.png") for n in range(3)] + with ( + mock.patch.object(image_monitor, "MAX_AWAITING_HISTORY", 2), + self.assertLogs("prompt_manager.image_monitor", "INFO") as logs, + ): + for path in paths: + self.handler.process_new_image(path) + self.assertEqual( + self.handler.awaiting_foreign(), [self.key(p) for p in paths[1:]] + ) + self.assertEqual(self.handler.awaiting_history(), [self.key(keyless)]) + self.assertIn( + "No recorded run for IMG_GEN_10000_.png; left unlinked", + "\n".join(logs.output), + ) + + def test_an_expired_foreign_entry_is_logged_and_left_unlinked(self): + from utils import image_monitor + + path = self.image("IMG_GEN_00008_.png") + self.handler.process_new_image(path) + later = time.monotonic() + image_monitor.AWAIT_HISTORY_SECONDS + 1 + with ( + mock.patch.object(image_monitor.time, "monotonic", return_value=later), + self.assertLogs("prompt_manager.image_monitor", "INFO") as logs, + ): + self.assertEqual(self.handler.awaiting_foreign(), []) + self.assertIn( + "No recorded run for IMG_GEN_00008_.png; left unlinked", + "\n".join(logs.output), + ) + self.assertIsNone(self.linked_text(path)) + + def test_a_seal_while_the_pending_set_is_full_is_rechecked_later(self): + from utils import image_monitor + + path = self.image("IMG_GEN_00009_.png") + self.handler.process_new_image(path) + blocker = self.image("blocker.png", subfolder="") + release = threading.Event() + original = self.handler.process_new_image + + def slow(image_path, prompt_snapshot=None): + if image_path == blocker: + release.wait(5) + return None + return original(image_path, prompt_snapshot=prompt_snapshot) + + self.handler.process_new_image = slow + with mock.patch.object(image_monitor, "MAX_PENDING_IMAGES", 1): + self.assertTrue(self.handler.schedule(blocker)) + self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 1)) + self.record_run( + self.db, + self.output, + [("dataset", "IMG_GEN_00009_.png")], + [], + nodes=["13"], + ) + self.assertEqual(list(self.handler._foreign_due), [path]) + release.set() + self.assertTrue(self.wait_for(lambda: self.linked_text(path))) + self.assertTrue(self.settled()) + self.assertEqual(self.linked_text(path), shapes.BRANCH_TEXTS[0]) + + def test_shutdown_forgets_waiting_and_due_foreign_files(self): + path = self.image("IMG_GEN_00010_.png") + self.handler.process_new_image(path) + self.handler._foreign_due[path] = {"filename": "IMG_GEN_00010_.png"} + self.handler.shutdown() + self.assertEqual(self.handler.awaiting_foreign(), []) + self.assertEqual(list(self.handler._foreign_due), []) + + def test_a_failed_store_is_logged_and_leaves_the_file_unlinked(self): + self.metadata = {"prompt": sampled_graph("a quiet harbor", "blurry")} + path = self.image("plain_00003_.png") + with ( + mock.patch( + "utils.foreign_runs.ingest_resolved_pair", + side_effect=sqlite3.OperationalError("locked"), + ), + self.assertLogs("prompt_manager.image_monitor", "ERROR"), + ): + self.handler.process_new_image(path) + self.assertIsNone(self.linked_text(path)) + + class TestBoundedLookup(RecordedRuns, unittest.TestCase): """M5: history lookups are batched, and the history schema is unchanged.""" diff --git a/utils/foreign_runs.py b/utils/foreign_runs.py new file mode 100644 index 00000000..83518762 --- /dev/null +++ b/utils/foreign_runs.py @@ -0,0 +1,117 @@ +"""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 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) + + +@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: + """``GalleryConfig.SAVE_FOREIGN_RUNS``, read for every file; True when + the configuration cannot be imported (the default).""" + 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 + + +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.""" + try: + file_size = os.path.getsize(image_path) + except OSError: + file_size = 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 diff --git a/utils/image_monitor.py b/utils/image_monitor.py index 7c70d876..c7798988 100644 --- a/utils/image_monitor.py +++ b/utils/image_monitor.py @@ -33,6 +33,15 @@ from typing import Callable, Optional, Dict, Any, List, Tuple from watchdog.observers import Observer from watchdog.events import FileSystemEventHandler +from .file_prompt import HISTORY +from .foreign_runs import ( + CERTAIN, + 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 @@ -147,6 +156,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). @@ -247,6 +261,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 @@ -432,6 +448,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): @@ -451,11 +470,9 @@ class ImageGenerationHandler(FileSystemEventHandler): # negative-prompt PromptManager nodes in the workflow. state, current_prompt = self._find_prompt_from_metadata(metadata) 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( @@ -577,36 +594,47 @@ 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) -> None: + """Drop expired and surplus entries of ``store`` (the keyless set by + default; ``_awaiting_lock`` held). A dropped foreign file is logged.""" + 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) + dropped = [k for k, (_p, due) in store.items() if due <= now] + dropped = [store.pop(key)[0] for key in dropped] + while len(store) > MAX_AWAITING_HISTORY: + dropped.append(store.popitem(last=False)[1][0]) + if store is self._awaiting_foreign: + for image_path in dropped: + self.logger.info( + f"No recorded run for {os.path.basename(image_path)}; " + "left unlinked" + ) - def _await_history(self, image_path: str) -> None: + 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) + self._prune_awaiting(store) - 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.""" @@ -614,6 +642,12 @@ 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: + self._prune_awaiting(self._awaiting_foreign) + return list(self._awaiting_foreign) + def _on_outputs_sealed(self, run_id: str, references: List[dict]) -> None: """A run's outputs were recorded: re-check its waiting images once. @@ -625,6 +659,7 @@ class ImageGenerationHandler(FileSystemEventHandler): due = [] with self._awaiting_lock: self._prune_awaiting() + self._prune_awaiting(self._awaiting_foreign) for reference in references: if reference.get("folder_type") != "output": continue @@ -635,39 +670,47 @@ 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))) + 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 + while len(due) > MAX_AWAITING_HISTORY: + due.popitem(last=False) 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) + store.pop(image_path, None) def _recheck_history(self, image_path: str) -> None: """The one history-only re-check of a keyless image (#342).""" @@ -691,6 +734,70 @@ 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: + # Registered first, so a run sealed before this lookup is not missed + self._await_foreign(image_path) + decision = decide_foreign( + image_path, + metadata, + in_output=True, + lookup=history_lookup(self.db_manager, image_path), + ) + if decision.kind not in CERTAIN: + self.logger.debug(f"{name} waits for its run to end") + return + self._stop_awaiting(image_path, self._awaiting_foreign) + if decision.kind in CERTAIN: + self._store_foreign(image_path, decision) + else: + self.logger.info(f"No certain prompt for {name}; left unlinked") + + 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 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) + 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: