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.
This commit is contained in:
Vito Sansevero
2026-10-28 06:43:35 -07:00
parent 6a494c7c1f
commit b313224d46
5 changed files with 576 additions and 38 deletions
+24 -1
View File
@@ -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")
+5
View File
@@ -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):
+286
View File
@@ -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."""
+117
View File
@@ -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
+144 -37
View File
@@ -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: