perf(images): refresh the listing when the monitor sees a new output (#348, PR 2)

The image monitor now reports every settled output (linked or not, also
when linking raised) with its resolved path through set_output_listener.
The API's notify_output_settled hops from the monitor's worker thread onto
the server's loop (call_soon_threadsafe) and does nothing before the loop is
known or after cleanup. mark_root_dirty matches the path against each
walk's resolved_root (normcase(realpath), computed in the executor by
walk_root, now its fourth value) and marks the root dirty.

A dirty root gets one background walk, no sooner than
rewalk_interval = min(15, max(5, 5 x last walk's duration)) after the
previous walk started (last_started_at / last_duration are recorded on
every walk). _schedule_walk is shared with PR 1's delayed stale refresh:
one deferred walk per root, never armed while a walk runs (the walk's end
re-checks dirty: the trailing walk), at least 0.5 s away, the earlier
deadline wins, and every walk start cancels it. A failed walk clears dirty.
Rescan semantics are unchanged.
This commit is contained in:
Vito Sansevero
2026-10-26 12:28:56 -07:00
parent 9ac471eb72
commit e84e56aa3b
10 changed files with 857 additions and 26 deletions
+18
View File
@@ -333,6 +333,24 @@ hold a walk open with a gated scan instead of sleeping.
listing gets a 500. Rescan waits for a walk that started after the click
(by walk number, not by clock), Rescans clicked during one walk share the
next, and a failed walk is a 500 even when a previous listing exists.
- **New outputs (monitor-triggered walks, PR 2).** The image monitor reports
every settled output it accepts (by default `.png .jpg .jpeg .webp .gif`;
videos only when their extension is in `monitoring.extensions`), linked or
not, with its `realpath`, through `set_output_listener`. On its worker
thread the API's `notify_output_settled` only does
`loop.call_soon_threadsafe(mark_root_dirty, path)`, and does nothing before
the loop is known or after `cleanup_gallery_listings`. `mark_root_dirty`
matches the path against each walk's `resolved_root` (`normcase(realpath)`,
computed in the executor; `root_contains` ends the match at a separator)
and marks the root `dirty`. A dirty root gets one background walk, no
sooner than `rewalk_interval` = `min(15, max(5, 5 × last walk's duration))`
seconds after the previous walk started. `_schedule_walk` never arms a timer
while a walk runs (the walk's end re-checks `dirty`: the trailing walk),
arms at most one per root and never sooner than 0.5 s, and keeps the
earlier deadline when the stale refresh and a dirty root both want one;
every walk start (read, Rescan, warm-up or timer) cancels it. A failed walk
clears `dirty`. Tests patch `REWALK_FLOOR`/`REWALK_CAP` and run the pending
timer's callback by hand (`_fire`) instead of waiting for it.
- **Measuring it.** `tools/vendor/perf/images_load_probe.mjs` is a maintainer
tool, never run in CI. After `npm ci --prefix tools/vendor` and
`npx --prefix tools/vendor playwright install chromium`, run
+4
View File
@@ -149,6 +149,9 @@ try:
# Every settled output's preview is queued for the API's background
# worker (#249): the workspace never asks to generate thumbnails.
_global_image_monitor.set_thumbnail_queue(_review_api.queue_capture_thumbnail)
# Every settled output walks its Images listing again (#348): the
# API hops from the monitor's thread onto the server's loop.
_global_image_monitor.set_output_listener(_review_api.notify_output_settled)
_review_host = getattr(config.server_instance, "app", None)
if _review_host is not None and hasattr(_review_host, "on_cleanup"):
@@ -157,6 +160,7 @@ try:
):
monitor.set_review_recorder(None)
monitor.set_thumbnail_queue(None)
monitor.set_output_listener(None)
api.close_auto_thumbnails(final=True)
_review_host.on_cleanup.insert(0, _stop_review_recording)
+24 -1
View File
@@ -362,8 +362,11 @@ class PromptManagerAPI(
# background walks refreshing them (py/api/gallery_listing.py, #348).
self._root_listings: dict[str, RootListing] = {}
self._listing_tasks: set = set()
# The loop the walks run on, set by the first walk.
# The loop the walks run on, set by the startup warm-up or the first
# walk; the monitor's thread hops onto it (notify_output_settled).
self._listing_loop = None
# Set by cleanup_gallery_listings: no background walk starts after it.
self._listings_closed = False
# The workspace caches rebuilt behind a stored walk that changed the
# listing, once the host's startup warm-up has run (#348).
self._workspace_warm_enabled = False
@@ -539,15 +542,35 @@ class PromptManagerAPI(
# Each stored walk that changes a listing then rebuilds the workspace's
# per-file values, orders and favorites behind it (_warm_workspace).
self._workspace_warm_enabled = True
self._listing_loop = asyncio.get_running_loop()
self._listings_closed = False
try:
for root in self._get_all_output_dirs():
self._schedule_walk(root)
except Exception as e:
self.logger.warning(f"Could not warm the Images listing: {e}")
def notify_output_settled(self, path):
"""The image monitor settled an output at *path* (any thread).
Called on the monitor's worker thread with the resolved path. It only
hands the path to the event loop (``mark_root_dirty``) and never
touches a listing from this thread; before the loop is known, or once
the listings are cleaned up, it does nothing.
"""
loop = self._listing_loop
if loop is None or self._listings_closed or not loop.is_running():
return
try:
loop.call_soon_threadsafe(self.mark_root_dirty, path)
except RuntimeError:
pass # the loop closed meanwhile (shutdown)
async def cleanup_gallery_listings(self, app=None):
"""Stop the background listing walks (aiohttp on_cleanup)."""
self._workspace_warm_enabled = False
# No mark, deferred walk or trailing walk starts a walk from now on.
self._listings_closed = True
# Before any await, so no delayed refresh can start a walk meanwhile.
for record in self._root_listings.values():
handle, record.rewalk_handle = record.rewalk_handle, None
+59 -4
View File
@@ -6,6 +6,10 @@ loop; the walk itself, the comparison with the previous listing and both
folder sets run in the executor through ``walk_root``. Routes live in
``py/api/images.py`` (``_listing``); this module has no aiohttp import so
tests can exercise it without a server.
A settled output the image monitor reports marks its root dirty
(``mark_root_dirty``); a dirty root gets one background walk, no sooner than
``rewalk_interval`` after its previous walk started (PR 2).
"""
from __future__ import annotations
@@ -26,6 +30,18 @@ REFRESH_AFTER = 30
# (#348 live test: reads during the walk took 2.2 s instead of 0.07 s).
REFRESH_DELAY = 1.0
# No background walk of a root starts sooner than rewalk_interval(last walk's
# duration) after its previous walk started: FACTOR times that duration, at
# least FLOOR and at most CAP seconds. Background walks then hold the GIL at
# most ~20% of the time; the cap keeps a new output listed within ~20 s on a
# slow machine. Read through this module so tests can patch them.
REWALK_FLOOR = 5
REWALK_FACTOR = 5
REWALK_CAP = 15
# A deferred walk is never armed sooner than this: no zero-delay loop.
MIN_REWALK_DELAY = 0.5
# A reader with no listing gives up (500) after this many walks in a row were
# discarded by invalidations made while they ran.
MAX_DISCARDED_WALKS = 3
@@ -80,8 +96,17 @@ class RootListing:
flight: asyncio.Future | None = None
# True while walks keep failing: one warning per streak.
failing: bool = False
# The delayed refresh not started yet (loop.call_later), or None.
# The one deferred background walk not started yet (loop.call_later), or
# None: a stale refresh or a dirty root's rewalk, whichever is due first.
rewalk_handle: asyncio.TimerHandle | None = None
# normcase(realpath(root)), computed by the walk in the executor: monitor
# paths are matched against it without file I/O on the loop.
resolved_root: str | None = None
# Monotonic start and duration of the last walk (REWALK_MIN_INTERVAL).
last_started_at: float | None = None
last_duration: float | None = None
# A settled output was reported under this root; cleared when a walk starts.
dirty: bool = False
def drop(self) -> None:
self.files = None
@@ -121,6 +146,34 @@ def folder_sets(root, files) -> tuple[frozenset, frozenset]:
return frozenset(folders), frozenset(with_ancestors)
def rewalk_interval(last_duration) -> float:
"""REWALK_MIN_INTERVAL after a walk that took *last_duration* seconds."""
scaled = REWALK_FACTOR * (last_duration or 0.0)
return min(REWALK_CAP, max(REWALK_FLOOR, scaled))
def root_contains(resolved_root: str, path, pathmod=os.path) -> bool:
"""True when the resolved *path* lies under *resolved_root* (no file I/O).
*resolved_root* is already ``normcase``d; *path* is folded the same way,
so Windows matches case-insensitively and with either separator. The
match ends at a separator: ``/out2/x.png`` is not under ``/out``.
"""
target = pathmod.normcase(str(path))
prefix = resolved_root
if not prefix.endswith(pathmod.sep):
prefix += pathmod.sep
return target.startswith(prefix)
def resolve_root(root) -> str:
"""``normcase(realpath(root))``, the spelling monitor paths are matched in.
File I/O (realpath follows links): call it in the executor.
"""
return os.path.normcase(os.path.realpath(str(root)))
def same_listing(old, new) -> bool:
"""True when two listings hold the same ``(path, mtime)`` pairs in order."""
if old is None or new is None:
@@ -130,8 +183,10 @@ def same_listing(old, new) -> bool:
def walk_root(
scan_fn: Callable[[PurePath], list], root, previous_files
) -> tuple[list, frozenset, frozenset]:
"""The single executor job of a walk: ``(files, folders, with_ancestors)``.
) -> tuple[list, frozenset, frozenset, str]:
"""The single executor job of a walk.
Returns ``(files, folders, with_ancestors, resolved_root)``.
Opens the root first, so a missing or unreadable root raises ``OSError``
instead of listing nothing. Returns *previous_files* itself when nothing
@@ -143,4 +198,4 @@ def walk_root(
if same_listing(previous_files, files):
files = previous_files
folders, with_ancestors = folder_sets(root, files)
return files, folders, with_ancestors
return files, folders, with_ancestors, resolve_root(root)
+79 -15
View File
@@ -1433,28 +1433,81 @@ class ImageRoutesMixin:
)
def _refresh_if_stale(self, output_path, record):
"""Walk a stale listing again, REFRESH_DELAY after this read (on the loop).
"""Ask for a background walk when the listing is stale or dirty.
At most one pending refresh per root, none while a walk runs.
A stale listing (older than REFRESH_AFTER) is walked REFRESH_DELAY
after this read, so the walk does not share the GIL with the page's
own requests; a dirty one as soon as its interval allows.
"""
age = _time.monotonic() - (record.walked_at or 0.0)
if age < gallery_listing.REFRESH_AFTER:
if age >= gallery_listing.REFRESH_AFTER:
self._schedule_walk(output_path, gallery_listing.REFRESH_DELAY)
elif record.dirty:
self._schedule_walk(output_path)
def mark_root_dirty(self, path):
"""A settled output at *path*: walk every root holding it again (on the loop).
*path* is resolved by the monitor's thread (``notify_output_settled``),
so this does no file I/O: it compares ``normcase`` strings with each
walk's ``resolved_root``. Only roots with a listing, or a walk running,
are marked; a root without either is walked by its next read anyway.
"""
if self._listings_closed:
return
if record.flight is not None or record.rewalk_handle is not None:
for key, record in list(self._root_listings.items()):
if record.files is None and record.flight is None:
continue
resolved = record.resolved_root or os.path.normcase(key)
if gallery_listing.root_contains(resolved, path):
record.dirty = True
self._schedule_walk(Path(key))
def _schedule_walk(self, output_path, delay=0.0):
"""Start or defer a background walk of *output_path* (on the loop).
- A walk is running: nothing; its end re-checks ``dirty``.
- Otherwise the walk is due *delay* from now, and never sooner than
``rewalk_interval`` after the previous walk started. Due now: start
it. Later: one deferred walk per root, at least MIN_REWALK_DELAY
away; when one is already pending, the earlier deadline wins.
Every walk start cancels the pending one (``_start_walk``).
"""
if self._listings_closed:
return
record.rewalk_handle = asyncio.get_running_loop().call_later(
gallery_listing.REFRESH_DELAY, self._delayed_walk, output_path, record
record = self._root_listing(output_path)
if record.flight is not None:
return
now = _time.monotonic()
due = now + delay
if record.last_started_at is not None:
interval = gallery_listing.rewalk_interval(record.last_duration)
due = max(due, record.last_started_at + interval)
if due <= now:
self._start_walk(output_path, record)
return
wait = max(gallery_listing.MIN_REWALK_DELAY, due - now)
loop = asyncio.get_running_loop()
handle = record.rewalk_handle
if handle is not None:
if handle.when() - loop.time() <= wait:
return # the pending walk comes first
handle.cancel()
record.rewalk_handle = loop.call_later(
wait, self._delayed_walk, output_path, record
)
def _delayed_walk(self, output_path, record):
record.rewalk_handle = None
self._schedule_walk(output_path)
"""The deferred walk's timer: clear the handle first, then walk.
def _schedule_walk(self, output_path):
"""Start a background walk of *output_path* unless one is running."""
record = self._root_listing(output_path)
if record.flight is None:
self._start_walk(output_path, record)
It was armed for a stale or dirty listing at the earliest moment the
interval allowed. Only a walk start clears either reason, and every
start cancels this timer, so when it fires the walk is still due.
"""
record.rewalk_handle = None
if self._listings_closed or record.flight is not None:
return
self._start_walk(output_path, record)
def _start_walk(self, output_path, record):
"""Start the root's one walk; returns its flight future (on the loop)."""
@@ -1462,7 +1515,10 @@ class ImageRoutesMixin:
self._listing_loop = loop
handle, record.rewalk_handle = record.rewalk_handle, None
if handle is not None:
handle.cancel() # this walk is the refresh it was waiting to start
handle.cancel() # this walk is the one it was waiting to start
# This walk lists every output reported so far.
record.dirty = False
record.last_started_at = _time.monotonic()
record.walk_seq += 1
flight = loop.create_future()
record.flight = flight
@@ -1482,8 +1538,9 @@ class ImageRoutesMixin:
error never reaches a future nobody awaits.
"""
outcome = None
started = record.last_started_at
try:
files, folders, with_ancestors = await self._run_in_executor(
files, folders, with_ancestors, resolved = await self._run_in_executor(
walk_root, self._scan_gallery_files_sync, output_path, record.files
)
except asyncio.CancelledError:
@@ -1496,9 +1553,12 @@ class ImageRoutesMixin:
f"Could not list output folder {output_path}: {exc}"
)
record.failing = True
# Retried by the 30 s rule, not by a trailing walk.
record.dirty = False
outcome = WalkOutcome.failed(exc)
else:
record.failing = False
record.resolved_root = resolved
if record.generation == generation:
changed = files is not record.files
record.files = files
@@ -1512,10 +1572,14 @@ class ImageRoutesMixin:
else:
outcome = WalkOutcome.discarded()
finally:
record.last_duration = _time.monotonic() - started
if record.flight is flight:
record.flight = None
if not flight.done():
flight.set_result(outcome)
if record.dirty:
# An output settled during this walk: the trailing walk.
self._schedule_walk(output_path)
async def get_output_images(self, request):
"""Get all images from ComfyUI output folder(s)."""
+422 -1
View File
@@ -2540,7 +2540,10 @@ class TestSharedListing(ImageAPITestCase):
from py.api import gallery_listing
record = self._record()
record.walked_at -= (seconds or gallery_listing.REFRESH_AFTER) + 1
shift = (seconds or gallery_listing.REFRESH_AFTER) + 1
record.walked_at -= shift
# An old listing's walk also started long ago (REWALK_MIN_INTERVAL).
record.last_started_at -= shift
async def _settle(self):
"""Wait for every background walk this API started."""
@@ -3073,6 +3076,401 @@ class TestSharedListing(ImageAPITestCase):
self.assertTrue(task.done())
self.assertIsNone(self.api._workspace_warm_task)
# -- monitor-triggered walks (#348, PR 2) ---------------------------
def _interval(self, seconds):
"""Patch REWALK_MIN_INTERVAL to *seconds*, whatever the walk took."""
for name in ("REWALK_FLOOR", "REWALK_CAP"):
patcher = patch.object(gallery_listing, name, seconds)
patcher.start()
self.addCleanup(patcher.stop)
def _output(self, name):
"""A new output on disk, spelled as the monitor reports it (resolved)."""
return os.path.realpath(make_png(self.output_dir / name))
async def _listed(self):
"""A stored listing of the root, built by a read."""
await self._json("GET", self.URL)
await self._settle()
return self._record()
def _fire(self, record):
"""Run the pending deferred walk now, as its timer would."""
handle = record.rewalk_handle
self.assertIsNotNone(handle, "no deferred walk is pending")
handle.cancel()
self.api._delayed_walk(self.output_dir, record)
def _remaining(self, handle):
return handle.when() - asyncio.get_running_loop().time()
async def test_a_settled_output_is_listed_without_a_reload(self):
make_png(self.output_dir / "a.png")
record = await self._listed()
self._interval(0)
self.api.mark_root_dirty(self._output("b.png"))
self.assertIsNotNone(record.flight) # the interval allows: walks now
self.assertFalse(record.dirty)
await self._settle()
self.assertEqual(self.scan.calls, 2)
self.assertEqual(len(record.files), 2)
_, data = await self._json("GET", self.URL)
self.assertEqual(data["total"], 2)
self.assertEqual(self.scan.calls, 2)
async def test_the_monitor_thread_hops_onto_the_loop(self):
make_png(self.output_dir / "a.png")
record = await self._listed()
self._interval(0)
path = self._output("b.png")
calls = []
real = self.api.mark_root_dirty
def mark(seen):
calls.append(threading.get_ident())
real(seen)
self.api.mark_root_dirty = mark
loop = asyncio.get_running_loop()
await loop.run_in_executor(None, self.api.notify_output_settled, path)
# The hop was queued before the executor's result: one loop turn runs it.
await asyncio.sleep(0)
self.assertEqual(calls, [threading.get_ident()])
await self._settle()
self.assertEqual(len(record.files), 2)
async def test_notify_before_the_loop_is_known_does_nothing(self):
self.assertIsNone(self.api._listing_loop)
with patch.object(self.api, "mark_root_dirty") as mark:
self.api.notify_output_settled(self._output("b.png"))
await asyncio.sleep(0)
mark.assert_not_called()
async def test_notify_after_cleanup_does_nothing(self):
make_png(self.output_dir / "a.png")
record = await self._listed()
await self.api.cleanup_gallery_listings(web.Application())
self._interval(0)
with patch.object(self.api, "mark_root_dirty") as mark:
self.api.notify_output_settled(self._output("b.png"))
await asyncio.sleep(0)
mark.assert_not_called()
# A hop queued before the cleanup is dropped when it runs.
self.api.mark_root_dirty(self._output("c.png"))
self.assertFalse(record.dirty)
self.assertIsNone(record.flight)
self.assertIsNone(record.rewalk_handle)
async def test_notify_on_a_closed_loop_does_nothing(self):
loop = asyncio.new_event_loop()
loop.close()
self.api._listing_loop = loop
self.api.notify_output_settled(self._output("b.png")) # no exception
async def test_notify_racing_the_loop_closing_does_nothing(self):
class ClosingLoop:
def is_running(self):
return True
def call_soon_threadsafe(self, *_args):
raise RuntimeError("Event loop is closed")
self.api._listing_loop = ClosingLoop()
self.api.notify_output_settled(self._output("b.png")) # no exception
async def test_an_output_outside_every_root_is_ignored(self):
make_png(self.output_dir / "a.png")
record = await self._listed()
self._interval(0)
elsewhere = make_png(self.comfy_dir / "output2" / "x.png")
self.api.mark_root_dirty(os.path.realpath(elsewhere))
self.assertFalse(record.dirty)
self.assertIsNone(record.flight)
self.assertIsNone(record.rewalk_handle)
async def test_a_root_never_listed_is_not_walked_by_a_mark(self):
# No listing, no walk running: the first read walks anyway.
self.api.mark_root_dirty(self._output("b.png"))
self.assertEqual(self.api._root_listings, {})
self.assertEqual(self.scan.calls, 0)
async def test_a_dropped_listing_is_left_to_the_next_read(self):
make_png(self.output_dir / "a.png")
record = await self._listed()
self.api.invalidate_gallery_cache()
self._interval(0)
self.api.mark_root_dirty(self._output("b.png"))
self.assertIsNone(record.flight)
self.assertEqual(self.scan.calls, 1)
async def test_the_walk_records_the_resolved_root_and_its_timing(self):
make_png(self.output_dir / "a.png")
before = images_module._time.monotonic()
record = await self._listed()
self.assertEqual(
record.resolved_root, os.path.normcase(os.path.realpath(self.output_dir))
)
self.assertGreaterEqual(record.last_started_at, before)
self.assertLessEqual(record.last_started_at, record.walked_at)
# Measured when the walk ended, after its listing was stored.
self.assertGreaterEqual(
record.last_duration, record.walked_at - record.last_started_at
)
async def test_a_failed_walk_records_its_timing_and_clears_dirty(self):
make_png(self.output_dir / "a.png")
record = await self._listed()
self._interval(30)
self.api.mark_root_dirty(self._output("b.png"))
self.assertTrue(record.dirty)
self.api._scan_gallery_files_sync = _raise
started = record.last_started_at
with self.assertLogs(self.api.logger, "WARNING"):
self._fire(record)
await self._settle()
self.assertGreater(record.last_started_at, started)
self.assertIsNotNone(record.last_duration)
self.assertFalse(record.dirty)
self.assertIsNone(record.rewalk_handle) # no trailing walk after a failure
self.assertEqual(len(record.files), 1) # the previous listing is kept
async def test_a_burst_of_marks_coalesces_into_one_walk(self):
make_png(self.output_dir / "a.png")
record = await self._listed()
self._interval(30)
for i in range(5):
self.api.mark_root_dirty(self._output(f"b{i}.png"))
handle = record.rewalk_handle
self.assertIsNotNone(handle)
self.assertIsNone(record.flight)
self.assertEqual(self.scan.calls, 1)
# Due one interval after the last walk started, not 30 s from now.
due = record.last_started_at + 30 - images_module._time.monotonic()
self.assertAlmostEqual(self._remaining(handle), due, delta=1)
self._fire(record)
await self._settle()
self.assertEqual(self.scan.calls, 2)
self.assertEqual(len(record.files), 6)
self.assertFalse(record.dirty)
self.assertIsNone(record.rewalk_handle)
async def test_a_mark_during_a_walk_gives_exactly_one_trailing_walk(self):
make_png(self.output_dir / "a.png")
record = await self._listed()
self._interval(30)
gate = self.scan.hold_next()
rescan = asyncio.ensure_future(
self._json("POST", "/prompt_manager/gallery/rescan")
)
await self.scan.wait_entered() # listed before b.png exists
for name in ("b.png", "c.png"):
self.api.mark_root_dirty(self._output(name))
# Never armed while a walk runs: the walk's end re-checks dirty.
self.assertIsNone(record.rewalk_handle)
self.assertTrue(record.dirty)
gate.set()
status, _ = await rescan
self.assertEqual(status, 200)
self.assertEqual(len(record.files), 1)
handle = record.rewalk_handle
self.assertIsNotNone(handle) # the trailing walk, deferred by the interval
self.assertGreaterEqual(self._remaining(handle), 0.5)
self._fire(record)
await self._settle()
self.assertEqual(self.scan.calls, 3)
self.assertEqual(len(record.files), 3)
self.assertIsNone(record.rewalk_handle)
self.assertFalse(record.dirty)
async def test_a_walk_slower_than_the_interval_trails_at_once_and_stops(self):
make_png(self.output_dir / "a.png")
record = await self._listed()
self._interval(0) # the walk outlasts its interval
gate = self.scan.hold_next()
self.api.mark_root_dirty(self._output("b.png"))
await self.scan.wait_entered()
self.api.mark_root_dirty(self._output("c.png"))
self.assertIsNone(record.rewalk_handle)
gate.set()
await self._settle()
# One walk for the first mark, one trailing walk for the second, then
# nothing: no mark arrived during the trailing walk.
self.assertEqual(self.scan.calls, 3)
self.assertEqual(len(record.files), 3)
self.assertIsNone(record.rewalk_handle)
self.assertIsNone(record.flight)
async def test_a_read_triggered_walk_cancels_a_pending_rewalk(self):
make_png(self.output_dir / "a.png")
record = await self._listed()
self._interval(30)
self.api.mark_root_dirty(self._output("b.png"))
handle = record.rewalk_handle
status, data = await self._json("POST", "/prompt_manager/gallery/rescan")
self.assertEqual((status, data["total"]), (200, 2))
self.assertTrue(handle.cancelled())
self.assertIsNone(record.rewalk_handle)
self.assertFalse(record.dirty)
await self._settle()
self.assertEqual(self.scan.calls, 2)
async def test_after_the_trailing_walk_a_new_mark_arms_again(self):
make_png(self.output_dir / "a.png")
record = await self._listed()
self._interval(30)
self.api.mark_root_dirty(self._output("b.png"))
self._fire(record)
await self._settle()
self.assertIsNone(record.rewalk_handle)
self.api.mark_root_dirty(self._output("c.png"))
self.assertIsNotNone(record.rewalk_handle)
self._fire(record)
await self._settle()
self.assertEqual(self.scan.calls, 3)
self.assertEqual(len(record.files), 3)
async def test_a_deferred_walk_waits_at_least_half_a_second(self):
make_png(self.output_dir / "a.png")
record = await self._listed()
self._interval(30)
# The interval ends in a few milliseconds: still never a zero delay.
record.last_started_at = images_module._time.monotonic() - 29.99
self.api.mark_root_dirty(self._output("b.png"))
handle = record.rewalk_handle
self.assertIsNotNone(handle)
self.assertGreaterEqual(self._remaining(handle), 0.45)
handle.cancel()
async def test_a_mark_moves_a_later_stale_refresh_earlier(self):
make_png(self.output_dir / "a.png")
record = await self._listed()
self._age()
self._interval(5)
with patch.object(gallery_listing, "REFRESH_DELAY", 30):
await self._json("GET", self.URL)
stale = record.rewalk_handle
self.assertGreater(self._remaining(stale), 25)
# The interval has long passed: the mark's walk starts now, and that
# walk is the refresh the stale read was waiting for.
self.api.mark_root_dirty(self._output("b.png"))
self.assertTrue(stale.cancelled())
self.assertIsNotNone(record.flight)
self.assertIsNone(record.rewalk_handle)
await self._settle()
self.assertEqual(self.scan.calls, 2)
self.assertEqual(len(record.files), 2)
async def test_a_stale_read_keeps_an_earlier_pending_rewalk(self):
make_png(self.output_dir / "a.png")
record = await self._listed()
self._interval(20)
self.api.mark_root_dirty(self._output("b.png"))
pending = record.rewalk_handle
when = pending.when()
# Stale by its age, while the interval since the last walk still runs.
record.walked_at -= gallery_listing.REFRESH_AFTER + 1
with patch.object(gallery_listing, "REFRESH_DELAY", 30):
await self._json("GET", self.URL)
self.assertIs(record.rewalk_handle, pending)
self.assertEqual(pending.when(), when)
self.assertFalse(pending.cancelled())
async def test_a_stale_read_moves_a_later_pending_rewalk_earlier(self):
make_png(self.output_dir / "a.png")
record = await self._listed()
self._interval(15)
self.api.mark_root_dirty(self._output("b.png"))
pending = record.rewalk_handle
self.assertGreater(self._remaining(pending), 10)
# A listing older than REFRESH_AFTER whose interval is over.
self._age()
with patch.object(gallery_listing, "REFRESH_DELAY", 2):
await self._json("GET", self.URL)
self.assertTrue(pending.cancelled())
replaced = record.rewalk_handle
self.assertIsNotNone(replaced)
self.assertLess(self._remaining(replaced), 3)
self.assertEqual(self.scan.calls, 1)
replaced.cancel()
async def test_a_dirty_read_rearms_a_lost_rewalk(self):
make_png(self.output_dir / "a.png")
record = await self._listed()
self._interval(30)
record.dirty = True # e.g. its timer was cancelled by hand
await self._json("GET", self.URL)
self.assertIsNotNone(record.rewalk_handle)
self.assertEqual(self.scan.calls, 1)
async def test_cleanup_cancels_a_pending_rewalk(self):
make_png(self.output_dir / "a.png")
record = await self._listed()
self._interval(30)
self.api.mark_root_dirty(self._output("b.png"))
handle = record.rewalk_handle
await self.api.cleanup_gallery_listings(web.Application())
self.assertTrue(handle.cancelled())
self.assertIsNone(record.rewalk_handle)
self.assertEqual(self.scan.calls, 1)
async def test_a_walk_cancelled_by_cleanup_starts_no_trailing_walk(self):
make_png(self.output_dir / "a.png")
record = await self._listed()
self._interval(0)
gate = self.scan.hold_next()
self.api.mark_root_dirty(self._output("b.png"))
await self.scan.wait_entered()
self.api.mark_root_dirty(self._output("c.png"))
self.assertTrue(record.dirty)
await self.api.cleanup_gallery_listings(web.Application())
gate.set()
self.assertIsNone(record.flight)
self.assertIsNone(record.rewalk_handle)
self.assertEqual(self.api._listing_tasks, set())
self.assertEqual(self.scan.calls, 2)
async def test_a_warm_up_after_cleanup_listens_again(self):
make_png(self.output_dir / "a.png")
self._interval(0)
await self._warmed()
await self.api.cleanup_gallery_listings(web.Application())
await self._warmed()
record = self._record()
self.api.mark_root_dirty(self._output("b.png"))
await self._settle()
self.assertEqual(len(record.files), 2)
async def test_monitor_walks_during_a_workspace_build_add_one_more_build(self):
make_png(self.output_dir / "a.png")
release, entered = threading.Event(), threading.Event()
built = []
real = self.api._build_workspace_caches
def build(listings):
built.append(len(listings[0][2]))
if len(built) == 1:
entered.set()
release.wait(_HANG_LIMIT)
real(listings)
self.addCleanup(release.set)
self._interval(0)
with patch.object(self.api, "_build_workspace_caches", side_effect=build):
await self.api.warm_gallery_listings(web.Application())
loop = asyncio.get_running_loop()
await loop.run_in_executor(None, entered.wait, _HANG_LIMIT)
for name in ("b.png", "c.png", "d.png"): # three monitor walks
self.api.mark_root_dirty(self._output(name))
while self.api._listing_tasks:
await asyncio.gather(*list(self.api._listing_tasks))
await asyncio.sleep(0)
release.set()
await self._settle()
self.assertEqual(built, [1, 4])
self.assertEqual(self.scan.calls, 4)
self.assertTrue(self._workspace_ready())
class TestGalleryListingHostHooks(unittest.IsolatedAsyncioTestCase):
async def test_root_registers_the_warm_up_and_its_cleanup(self):
@@ -3092,3 +3490,26 @@ class TestGalleryListingHostHooks(unittest.IsolatedAsyncioTestCase):
) as (package, host):
self.assertIn(package.api.warm_gallery_listings, host.app.on_startup)
self.assertIn(package.api.cleanup_gallery_listings, host.app.on_cleanup)
async def test_root_connects_the_monitor_to_the_listing(self):
try:
from test_init_package import fake_server_module, load_package
except ImportError:
from tests.test_init_package import fake_server_module, load_package
with tempfile.TemporaryDirectory() as directory:
server = fake_server_module()
server.PromptServer.instance.app = web.Application()
with load_package(
"gallery_listing_monitor_pkg",
str(Path(directory) / "db.sqlite"),
directory,
server,
) as (package, host):
monitor = package._global_image_monitor
self.assertEqual(
monitor.output_listener, package.api.notify_output_settled
)
# The first cleanup hook stops the monitor's hand-offs.
await host.app.on_cleanup[0](host.app)
self.assertIsNone(monitor.output_listener)
+1
View File
@@ -188,6 +188,7 @@ class TestWorkspaceImages(ImageAPITestCase):
newest = self.api._workspace_order_cache[key][1]["newest"]
favorites = self.api._favorite_listing_cache[str(self.output_dir)]
record.walked_at -= gallery_listing.REFRESH_AFTER + 1
record.last_started_at -= gallery_listing.REFRESH_AFTER + 1
with patch.object(
self.api,
"_scan_gallery_files_sync",
+109 -3
View File
@@ -1,6 +1,8 @@
"""Pure helpers of the shared Images listing (py/api/gallery_listing.py, #348)."""
import ntpath
import os
import posixpath
import sys
import tempfile
import unittest
@@ -17,6 +19,8 @@ from py.api.gallery_listing import ( # noqa: E402
RootListing,
WalkOutcome,
folder_sets,
rewalk_interval,
root_contains,
same_listing,
walk_root,
)
@@ -123,24 +127,36 @@ class WalkRootTests(unittest.TestCase):
def test_returns_the_listing_and_both_folder_sets(self):
files = [(self.root / "sub" / "a.png", 2.0), (self.root / "b.png", 1.0)]
scan, calls = self._scan(files)
listed, folders, with_ancestors = walk_root(scan, self.root, None)
listed, folders, with_ancestors, resolved = walk_root(scan, self.root, None)
self.assertEqual(calls, [self.root])
self.assertEqual(listed, files)
self.assertEqual(folders, frozenset({"sub"}))
self.assertEqual(with_ancestors, frozenset({"sub"}))
self.assertEqual(resolved, os.path.normcase(os.path.realpath(self.root)))
def test_the_resolved_root_follows_a_symlinked_root(self):
link = Path(self._tmp.name) / "linked"
try:
os.symlink(self.root, link, target_is_directory=True)
except (OSError, NotImplementedError) as error:
self.skipTest(f"symlinks unavailable: {error}")
scan, _calls = self._scan([])
resolved = walk_root(scan, link, None)[3]
self.assertEqual(resolved, os.path.normcase(os.path.realpath(self.root)))
self.assertNotEqual(resolved, os.path.normcase(str(link)))
def test_an_unchanged_walk_returns_the_previous_list_object(self):
files = [(self.root / "sub" / "a.png", 2.0)]
previous = list(files)
scan, _calls = self._scan(files)
listed, _folders, _ancestors = walk_root(scan, self.root, previous)
listed, _folders, _ancestors, _resolved = walk_root(scan, self.root, previous)
self.assertIs(listed, previous)
def test_a_changed_walk_returns_the_new_list(self):
previous = [(self.root / "sub" / "a.png", 2.0)]
files = [(self.root / "sub" / "a.png", 3.0)]
scan, _calls = self._scan(files)
listed, _folders, _ancestors = walk_root(scan, self.root, previous)
listed, _folders, _ancestors, _resolved = walk_root(scan, self.root, previous)
self.assertIsNot(listed, previous)
self.assertEqual(listed, files)
@@ -180,6 +196,96 @@ class RecordTests(unittest.TestCase):
self.assertEqual(gallery_listing.REFRESH_DELAY, 1.0)
self.assertIsNone(RootListing().rewalk_handle)
def test_a_new_record_is_clean_and_never_walked(self):
record = RootListing()
self.assertFalse(record.dirty)
self.assertIsNone(record.resolved_root)
self.assertIsNone(record.last_started_at)
self.assertIsNone(record.last_duration)
def test_drop_keeps_the_resolved_root_and_the_walk_timing(self):
# A drop clears the listing, not what is known about the root itself.
record = RootListing(resolved_root="/o", last_started_at=1.0)
record.last_duration = 0.5
record.dirty = True
record.drop()
self.assertEqual(record.resolved_root, "/o")
self.assertEqual((record.last_started_at, record.last_duration), (1.0, 0.5))
self.assertTrue(record.dirty)
class RewalkIntervalTests(unittest.TestCase):
def test_constants(self):
self.assertEqual(
(
gallery_listing.REWALK_FLOOR,
gallery_listing.REWALK_FACTOR,
gallery_listing.REWALK_CAP,
gallery_listing.MIN_REWALK_DELAY,
),
(5, 5, 15, 0.5),
)
def test_short_walks_wait_the_floor(self):
self.assertEqual(rewalk_interval(0.2), 5)
self.assertEqual(rewalk_interval(0), 5)
self.assertEqual(rewalk_interval(None), 5) # never walked
def test_the_interval_scales_with_the_last_walk(self):
self.assertEqual(rewalk_interval(1.0), 5)
self.assertEqual(rewalk_interval(2.0), 10)
self.assertAlmostEqual(rewalk_interval(2.5), 12.5)
def test_slow_walks_are_capped(self):
self.assertEqual(rewalk_interval(3.0), 15)
self.assertEqual(rewalk_interval(60.0), 15)
def test_the_constants_are_read_when_called(self):
with (
patch.object(gallery_listing, "REWALK_FLOOR", 0),
patch.object(gallery_listing, "REWALK_CAP", 0),
):
self.assertEqual(rewalk_interval(2.0), 0)
class RootContainsTests(unittest.TestCase):
def test_posix_paths_under_the_root_match(self):
root = posixpath.normcase("/srv/out")
self.assertTrue(root_contains(root, "/srv/out/a.png", posixpath))
self.assertTrue(root_contains(root, "/srv/out/sub/deep/a.png", posixpath))
def test_a_sibling_sharing_the_prefix_does_not_match(self):
self.assertFalse(root_contains("/srv/out", "/srv/out2/a.png", posixpath))
self.assertFalse(root_contains("/srv/out", "/srv/outside.png", posixpath))
self.assertFalse(root_contains("/srv/out", "/srv/other/a.png", posixpath))
def test_posix_keeps_case(self):
self.assertFalse(root_contains("/srv/out", "/srv/OUT/a.png", posixpath))
def test_the_root_itself_is_not_an_output(self):
self.assertFalse(root_contains("/srv/out", "/srv/out", posixpath))
def test_windows_folds_case_and_separators(self):
root = ntpath.normcase("C:\\ComfyUI\\Output")
self.assertTrue(root_contains(root, "c:\\comfyui\\output\\a.png", ntpath))
self.assertTrue(root_contains(root, "C:/ComfyUI/OUTPUT/sub/a.PNG", ntpath))
self.assertFalse(root_contains(root, "C:\\ComfyUI\\Output2\\a.png", ntpath))
self.assertFalse(root_contains(root, "D:\\ComfyUI\\Output\\a.png", ntpath))
def test_a_windows_drive_root_and_a_unc_share(self):
drive = ntpath.normcase("E:\\")
self.assertTrue(root_contains(drive, "e:\\x.png", ntpath))
self.assertFalse(root_contains(drive, "F:\\x.png", ntpath))
share = ntpath.normcase("\\\\Server\\Share\\Out")
self.assertTrue(root_contains(share, "\\\\server\\share\\out\\a.png", ntpath))
self.assertFalse(root_contains(share, "\\\\server\\share\\a.png", ntpath))
def test_the_default_is_this_platforms_rule(self):
root = os.path.normcase(os.path.realpath(tempfile.gettempdir()))
inside = os.path.join(root, "x.png")
self.assertTrue(root_contains(root, inside))
self.assertTrue(root_contains(root, inside.upper()) == (os.name == "nt"))
if __name__ == "__main__":
unittest.main()
+103
View File
@@ -449,6 +449,109 @@ class TestHandlerScheduling(unittest.TestCase):
self.assertEqual(self.processed, [first, second])
class TestOutputListener(unittest.TestCase):
"""Every settled output is reported to the Images listing (#348, PR 2)."""
def setUp(self):
self.tmp = tempfile.TemporaryDirectory()
self.addCleanup(self.tmp.cleanup)
self.handler = im_mod.ImageGenerationHandler(MagicMock(), MagicMock())
self.addCleanup(self.handler.shutdown)
self.handler.processing_delay = 0
self.handler.settle_interval = 0.01
self.handler.process_new_image = lambda path, prompt_snapshot=None: None
self.reported = []
self.handler.output_listener = self.reported.append
def _image(self, name, content=b"x" * 16):
path = os.path.join(self.tmp.name, name)
with open(path, "wb") as f:
f.write(content)
return path
def test_a_settled_output_is_reported_once_with_its_resolved_path(self):
path = self._image("a.png")
self.handler._process_pending(path)
self.assertEqual(self.reported, [os.path.realpath(path)])
def test_a_symlinked_spelling_is_reported_resolved(self):
real_dir = os.path.join(self.tmp.name, "real")
os.makedirs(real_dir)
link_dir = os.path.join(self.tmp.name, "link")
try:
os.symlink(real_dir, link_dir, target_is_directory=True)
except (OSError, NotImplementedError) as error:
self.skipTest(f"symlinks unavailable: {error}")
with open(os.path.join(real_dir, "b.png"), "wb") as f:
f.write(b"x" * 16)
self.handler._process_pending(os.path.join(link_dir, "b.png"))
self.assertEqual(
self.reported, [os.path.realpath(os.path.join(real_dir, "b.png"))]
)
def test_it_is_reported_even_when_linking_raises(self):
def explode(path, prompt_snapshot=None):
raise RuntimeError("boom")
self.handler.process_new_image = explode
path = self._image("c.png")
with self.assertRaises(RuntimeError):
self.handler._process_pending(path)
self.assertEqual(self.reported, [os.path.realpath(path)])
self.assertEqual(self.handler.pending_count(), 0)
def test_an_output_that_never_settles_is_not_reported(self):
self.handler.settle_attempts = 3
self.handler._process_pending(os.path.join(self.tmp.name, "gone.png"))
empty = self._image("empty.png", content=b"")
self.handler._process_pending(empty)
self.assertEqual(self.reported, [])
def test_a_video_is_reported_when_its_extension_is_monitored(self):
self.handler.supported_extensions = (".png", ".mp4")
path = self._image("clip.mp4")
self.handler.on_created(FakeEvent(path))
self.assertTrue(wait_for(lambda: self.reported != []))
self.assertEqual(self.reported, [os.path.realpath(path)])
def test_a_failing_listener_is_logged_and_the_worker_carries_on(self):
def broken(_path):
raise RuntimeError("listing gone")
self.handler.output_listener = broken
first = self._image("d.png")
with self.assertLogs("prompt_manager.image_monitor", "WARNING") as logs:
self.handler._process_pending(first) # no exception
self.assertIn("Images listing", "\n".join(logs.output))
self.handler.output_listener = self.reported.append
second = self._image("e.png")
self.handler.schedule(second)
self.assertTrue(wait_for(lambda: self.reported != []))
self.assertEqual(self.reported, [os.path.realpath(second)])
def test_no_listener_reports_nothing(self):
self.handler.output_listener = None
self.handler._process_pending(self._image("f.png")) # no exception
self.assertEqual(self.reported, [])
def test_the_monitor_wires_running_and_later_handlers(self):
monitor = im_mod.ImageMonitor(MagicMock(), MagicMock())
listener = self.reported.append
monitor.set_output_listener(listener)
with (
patch.object(im_mod, "Observer", FakeObserver),
patch.object(im_mod.ImageMonitor, "_gallery_config", return_value=None),
):
monitor.start_monitoring([self.tmp.name])
try:
self.assertIs(monitor.handler.output_listener, listener)
monitor.set_output_listener(None)
self.assertIsNone(monitor.output_listener)
self.assertIsNone(monitor.handler.output_listener)
finally:
monitor.stop_monitoring()
class TestModuleIsolation(unittest.TestCase):
"""This module must not leave mocks behind for later test modules."""
+38 -2
View File
@@ -106,6 +106,10 @@ class ImageGenerationHandler(FileSystemEventHandler):
# ``queue_capture_thumbnail``, wired at start-up; #249). It only
# queues: the preview is written on its own worker. None: no queue.
self.thumbnail_queue: Optional[Callable[[str], str]] = None
# Told about every settled output, with its resolved path (the API's
# ``notify_output_settled``, wired at start-up; #348): the Images
# listing walks that root again. None: nobody is told.
self.output_listener: Optional[Callable[[str], None]] = None
# Read from GalleryConfig if available, otherwise use defaults
try:
@@ -259,8 +263,12 @@ class ImageGenerationHandler(FileSystemEventHandler):
f"Image never settled or vanished, skipping: {image_path}"
)
return
self.process_new_image(image_path, prompt_snapshot=prompt_snapshot)
self._queue_thumbnail(image_path)
try:
self.process_new_image(image_path, prompt_snapshot=prompt_snapshot)
self._queue_thumbnail(image_path)
finally:
# Linked or not, even if linking raised: the file exists.
self._report_output(image_path)
finally:
with self._pending_lock:
self._pending.pop(image_path, None)
@@ -285,6 +293,25 @@ class ImageGenerationHandler(FileSystemEventHandler):
return
self.logger.debug(f"Preview {outcome}: {os.path.basename(image_path)}")
def _report_output(self, image_path: str) -> None:
"""Tell the output listener a settled output exists, best effort.
The path is resolved here, on the worker thread, so the listener's
event loop does no file I/O. A failure is logged, never raised: the
output is then listed by the listing's 30 s refresh.
"""
listener = getattr(self, "output_listener", None)
if listener is None:
return
try:
listener(os.path.realpath(image_path))
except Exception:
self.logger.warning(
f"{os.path.basename(image_path)} could not be reported to the "
"Images listing",
exc_info=True,
)
def wait_until_settled(self, image_path: str, stop=None) -> bool:
"""True once the file's size is non-zero and unchanged across two reads.
@@ -848,6 +875,14 @@ class ImageMonitor:
self.logger = get_logger("prompt_manager.image_monitor")
self.review_recorder: Optional[Callable[[str], str]] = None
self.thumbnail_queue: Optional[Callable[[str], str]] = None
self.output_listener: Optional[Callable[[str], None]] = None
def set_output_listener(self, listener: Optional[Callable[[str], None]]) -> None:
"""Report every settled output's resolved path to ``listener`` (None
stops it), on the running handler and every later one (#348)."""
self.output_listener = listener
if self.handler is not None:
self.handler.output_listener = listener
def set_thumbnail_queue(
self, queue_preview: Optional[Callable[[str], str]]
@@ -925,6 +960,7 @@ class ImageMonitor:
handler = ImageGenerationHandler(self.db_manager, self.prompt_tracker)
handler.review_recorder = self.review_recorder
handler.thumbnail_queue = self.thumbnail_queue
handler.output_listener = self.output_listener
observer = Observer()
scheduled = []
for output_dir in valid_directories: