diff --git a/tests/test_run_attribution.py b/tests/test_run_attribution.py index 4ff9c718..bb4e69dd 100644 --- a/tests/test_run_attribution.py +++ b/tests/test_run_attribution.py @@ -982,7 +982,7 @@ class TestMonitorForeignRuns(RecordedRuns, unittest.TestCase): ) self.assertEqual(self.handler.awaiting_history(), [self.key(keyless)]) self.assertIn( - "No recorded run for IMG_GEN_10000_.png; left unlinked", + "Too many files wait for their run; IMG_GEN_10000_.png left unlinked", "\n".join(logs.output), ) @@ -1083,6 +1083,145 @@ class TestMonitorForeignRuns(RecordedRuns, unittest.TestCase): self.handler.process_new_image(path) size.assert_not_called() + def untraceable(self): + """Unresolved and only the legacy parser names a text (an unknown node).""" + graph = sampled_graph("unused", "blurry") + graph["1"]["inputs"]["text"] = ["30", 0] + graph["30"] = { + "class_type": "VendorPromptMaker", + "inputs": {"prompt": "a misty valley at dawn"}, + } + return {"prompt": graph} + + def test_no_record_at_settle_keeps_the_file_waiting_for_the_legacy_floor(self): + """Review: a lookup that found no record is not the run's answer yet.""" + self.metadata = self.untraceable() + path = self.image("vendor_00001_.png") + self.handler.process_new_image(path) + self.assertIsNone(self.linked_text(path)) + self.assertEqual(self.handler.awaiting_foreign(), [self.key(path)]) + self.record_run(self.db, self.output, [("dataset", "vendor_00001_.png")], []) + self.assertTrue(self.wait_for(lambda: self.linked_text(path))) + self.assertTrue(self.settled()) + self.assertEqual(self.linked_text(path), "a misty valley at dawn") + + def test_a_record_found_at_settle_without_a_certain_prompt_ends_the_wait(self): + """Review: sealed before settle, still ambiguous: no seal is coming.""" + path = self.image("IMG_GEN_00012_.png") + # One run reported it from two savers: a record, but no saver + reported = [("dataset", "IMG_GEN_00012_.png")] * 2 + self.record_run(self.db, self.output, reported, [], nodes=["13", "23"]) + with self.assertLogs("prompt_manager.image_monitor", "INFO") as logs: + self.handler.process_new_image(path) + self.assertIn( + "No certain prompt for IMG_GEN_00012_.png; left unlinked", + "\n".join(logs.output), + ) + self.assertEqual(self.handler.awaiting_foreign(), []) + self.assertIsNone(self.linked_text(path)) + + def test_monitoring_saved_off_stops_foreign_writes_at_settle_and_recheck(self): + """Review: the observer runs until a restart; nothing is written.""" + enabled = GalleryConfig.MONITORING_ENABLED + self.addCleanup(setattr, GalleryConfig, "MONITORING_ENABLED", enabled) + GalleryConfig.MONITORING_ENABLED = False + self.metadata = {"prompt": sampled_graph("a quiet harbor", "blurry")} + plain = self.image("plain_00005_.png") + self.handler.process_new_image(plain) + self.assertIsNone(self.linked_text(plain)) + GalleryConfig.MONITORING_ENABLED = True + self.metadata = {"prompt": shapes.shared_prefix()} + waiting = self.image("IMG_GEN_00013_.png") + self.handler.process_new_image(waiting) + GalleryConfig.MONITORING_ENABLED = False + reference = {"folder_type": "output", "subfolder": "dataset"} + reference["filename"] = "IMG_GEN_00013_.png" + self.record_run( + self.db, self.output, [("dataset", reference["filename"])], [], nodes=["13"] + ) + self.assertTrue(self.wait_for(lambda: not self.handler.awaiting_foreign())) + self.assertTrue(self.settled()) + self.handler._recheck_foreign(waiting, reference) + self.assertIsNone(self.linked_text(waiting)) + self.assertEqual(self.prompt_count(), 0) + + def test_the_setting_is_read_again_just_before_the_store(self): + self.metadata = {"prompt": sampled_graph("a quiet harbor", "blurry")} + path = self.image("plain_00006_.png") + real = __import__("utils.foreign_runs", fromlist=["x"]).decide_foreign + + def decide_then_switch_off(*args, **kwargs): + decision = real(*args, **kwargs) + GalleryConfig.SAVE_FOREIGN_RUNS = False + return decision + + with mock.patch( + "utils.image_monitor.decide_foreign", side_effect=decide_then_switch_off + ): + self.handler.process_new_image(path) + self.assertIsNone(self.linked_text(path)) + self.assertEqual(self.prompt_count(), 0) + + def test_a_seal_between_registering_and_the_settle_lookup_links_once(self): + """Regression guard (passes before the review fixes): one link only.""" + path = self.image("IMG_GEN_00014_.png") + real = __import__("utils.foreign_runs", fromlist=["x"]).history_lookup + + def sealed_first(db, image_path, reference=None): + lookup = real(db, image_path, reference) + + def run_then_look(): + if reference is None: # the settle lookup: the run seals now + self.record_run( + self.db, + self.output, + [("dataset", "IMG_GEN_00014_.png")], + [], + nodes=["13"], + ) + return lookup() + + return run_then_look + + with mock.patch("utils.image_monitor.history_lookup", side_effect=sealed_first): + self.assertTrue(self.handler.schedule(path)) + self.assertTrue(self.wait_for(lambda: self.linked_text(path))) + self.assertTrue(self.settled()) + self.assertEqual(self.linked_text(path), shapes.BRANCH_TEXTS[0]) + self.assertEqual(self.link_count(path), 1) + self.assertEqual(list(self.handler._foreign_due), []) + + def test_a_drained_recheck_never_drops_a_newer_reference(self): + path = self.image("IMG_GEN_00015_.png") + older, newer = {"filename": "old"}, {"filename": "new"} + self.handler._foreign_due[path] = older + real = self.handler.schedule + + def schedule(image_path, snapshot=None): + # The run is sealed again while the older one is scheduled + self.handler._foreign_due[image_path] = newer + return real(image_path, snapshot) + + self.handler.schedule = schedule + with mock.patch.object(self.handler, "process_new_image"): + self.handler._drain_history_due() + self.assertIs(self.handler._foreign_due[path], newer) + + def test_an_evicted_due_recheck_is_logged(self): + from utils import image_monitor + + self.handler.schedule = lambda *args, **kwargs: False + with ( + mock.patch.object(image_monitor, "MAX_AWAITING_HISTORY", 1), + self.assertLogs("prompt_manager.image_monitor", "INFO") as logs, + ): + for name in ("a_00001_.png", "b_00001_.png"): + self.handler._schedule_recheck(self.image(name), {"filename": name}) + self.assertIn( + "Too many files wait for their run; a_00001_.png left unlinked", + "\n".join(logs.output), + ) + 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") diff --git a/utils/foreign_runs.py b/utils/foreign_runs.py index b5f52b95..a48f079f 100644 --- a/utils/foreign_runs.py +++ b/utils/foreign_runs.py @@ -28,6 +28,8 @@ except ImportError: # Decisions that name a prompt to store CERTAIN = (PAIR, POSITIVE) +# Sources certain without a recorded run: the graph alone, or its recorded saver +ANCHORED_SOURCES = ("inspector", "recorded-run") @dataclass(frozen=True) @@ -40,8 +42,13 @@ class ForeignRecheck: def save_foreign_runs() -> bool: - """``GalleryConfig.SAVE_FOREIGN_RUNS``, read for every file; True when - the configuration cannot be imported (the default).""" + """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: @@ -49,7 +56,10 @@ def save_foreign_runs() -> bool: from py.config import GalleryConfig except ImportError: return True - return GalleryConfig.SAVE_FOREIGN_RUNS is not False + return ( + GalleryConfig.SAVE_FOREIGN_RUNS is not False + and GalleryConfig.MONITORING_ENABLED is not False + ) def history_lookup( diff --git a/utils/image_monitor.py b/utils/image_monitor.py index c7798988..7ca93223 100644 --- a/utils/image_monitor.py +++ b/utils/image_monitor.py @@ -35,6 +35,7 @@ from watchdog.events import FileSystemEventHandler from .file_prompt import HISTORY from .foreign_runs import ( + ANCHORED_SOURCES, CERTAIN, ForeignRecheck, decide_foreign, @@ -594,21 +595,28 @@ class ImageGenerationHandler(FileSystemEventHandler): resolved = os.path.abspath(image_path) return os.path.normcase(resolved) - def _prune_awaiting(self, store=None) -> None: + 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). A dropped foreign file is logged.""" + 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() - dropped = [k for k, (_p, due) in store.items() if due <= now] - dropped = [store.pop(key)[0] for key in dropped] + 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]) - if store is self._awaiting_foreign: - for image_path in dropped: + dropped.append((store.popitem(last=False)[1][0], True)) + return dropped if store is self._awaiting_foreign else [] + + 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"No recorded run for {os.path.basename(image_path)}; " - "left unlinked" + 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. @@ -626,7 +634,8 @@ class ImageGenerationHandler(FileSystemEventHandler): with self._awaiting_lock: store.pop(key, None) store[key] = (image_path, time.monotonic() + AWAIT_HISTORY_SECONDS) - self._prune_awaiting(store) + dropped = self._prune_awaiting(store) + self._log_dropped(dropped) def _await_foreign(self, image_path: str) -> None: self._await_history(image_path, self._awaiting_foreign) @@ -645,8 +654,10 @@ class ImageGenerationHandler(FileSystemEventHandler): 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) + 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. @@ -659,7 +670,7 @@ class ImageGenerationHandler(FileSystemEventHandler): due = [] with self._awaiting_lock: self._prune_awaiting() - self._prune_awaiting(self._awaiting_foreign) + dropped = self._prune_awaiting(self._awaiting_foreign) for reference in references: if reference.get("folder_type") != "output": continue @@ -674,6 +685,7 @@ class ImageGenerationHandler(FileSystemEventHandler): 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) @@ -699,8 +711,11 @@ class ImageGenerationHandler(FileSystemEventHandler): 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: - due.popitem(last=False) + 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.""" @@ -710,7 +725,9 @@ class ImageGenerationHandler(FileSystemEventHandler): for store, image_path, reference in due: if self.schedule(image_path, self._recheck_snapshot(reference)): with self._awaiting_lock: - store.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).""" @@ -748,15 +765,8 @@ class ImageGenerationHandler(FileSystemEventHandler): 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: + 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) @@ -765,6 +775,29 @@ class ImageGenerationHandler(FileSystemEventHandler): else: self.logger.info(f"No certain prompt for {name}; left unlinked") + 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. + """ + 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) @@ -789,6 +822,9 @@ class ImageGenerationHandler(FileSystemEventHandler): 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: