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:
@@ -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")
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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."""
|
||||
|
||||
|
||||
@@ -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
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user