perf(images): share and warm the Images listing (#348, PR 1) (#349)

* perf(images): share and warm one listing per output root (#348)

The Images page, its sidebar counts, the subfolder list and the legacy
gallery each rebuilt the output listing when the 30 s cache expired, all
at once, and the subfolder list computed relative paths on the event
loop. Thumbnail writes also threw the listing away.

Keep one RootListing per output root (py/api/gallery_listing.py):
- one walk per root at a time; concurrent readers wait for it (shielded)
- a listing older than 30 s is served at once and refreshed in the
  background
- the walk, the comparison with the previous listing and both folder
  sets run in the executor; an unchanged walk keeps the same list object
- os.scandir probe: a missing root fails the walk and keeps the previous
  listing (one warning per failure streak)
- Rescan waits for a walk started after the click and is a 500 when it
  fails, even with a previous listing
- deletes keep the blocking invalidation; a generation counter that
  survives drops discards walks that raced it
- thumbnail jobs, conversion and the legacy generate/clear routes no
  longer drop the listing
- warm-up on the host's on_startup, cancelled on on_cleanup

Adds tools/vendor/perf/images_load_probe.mjs (maintainer tool, not CI)
and the TESTING.md section "Images listing cache (#348)".

* fix(images): drop listings on the loop when invalidated from a worker thread (#348)

A settings save (gallery roots change) and the browser-test seeds call
invalidate_gallery_cache() from the executor. Listing records are only
touched on the event loop, so an off-loop call now hands the drop to the
loop and waits for it, keeping the call synchronous for its callers.

The shared-Rescan test waits for the later Rescans to take their walk
number instead of sleeping.

* test(images): bound the shared-listing test waits (#348)

A macOS py3.12 run stopped inside the shared-listing tests and sat until
the 6 h job limit with no traceback. It does not reproduce locally (about
300 runs on 3.12 and 3.14, pinned to one core). Requests and background
settles in TestSharedListing now give up after 30 s, so a recurrence
fails with the stuck await instead of hanging the run.

* test(images): let the settle loops yield so a finished walk leaves the set (#348)

On CI (ubuntu and macOS py3.12) the runner's watchdog dumped
test_a_root_that_recovers_warns_again_on_its_next_failure spinning in
_settle: asyncio.gather over tasks that are already done returns without
yielding, so the set's discard callback, queued when the walk finished,
never ran and the loop never ended (wait_for could not fire either).
Delaying discard by one loop turn reproduces it locally; both settle
loops now yield one turn after each gather.

* test(browser): poll the import review's rows instead of reading them once

gridOf reads the dialog's rows once; right after the click the review may
not have rendered them yet (firefox CI: received undefined).

* perf(images): walk the output with scandir entries and build folders from strings (#348)

The live test showed a read during a background refresh taking 2.2 s instead
of 0.07 s: the walk held the GIL for about 2 s in the server. py-spy split
that walk into two halves: Path.relative_to per file in folder_sets, and
Path.stat + Path.is_symlink + Path joins per file in the scan.

- _walk_media visits folders in os.walk's order with os.scandir entries:
  the file type comes from the listing (no lstat per file) and the mtime
  from entry.stat (one call, none on Windows). _iter_media_files is now a
  thin wrapper, so its other caller is unchanged.
- folder_sets slices the path strings (in the path flavour's separator)
  instead of calling relative_to, falling back to it for a path that does
  not start with the root string.

On the maintainer's 19,935-file output folder: scan 0.386 s -> 0.112 s,
folder sets 0.452 s -> 0.007 s, both results identical to the previous code.
A reference copy of the os.walk version checks order, filters, the depth
cap, symlinks and the file cap.

* perf(images): build the workspace caches behind the warm-up and changed walks (#348)

The live test after the faster walk still took 2.9-5.2 s for the first
Images open after a restart. py-spy showed the walk was now small; the time
went to the workspace's own caches, cold after a restart and built by three
parallel first requests at once: the favorites listing (Path.relative_to per
file), the per-file values and newest order, and the prompt-link read.

- Once the host's startup warm-up has run, a stored walk that changed a
  listing schedules one executor job that builds what a first page reads,
  through the request path itself (_select_workspace_page with the default
  filters, then _workspace_favorite_paths). A walk stored during a build
  asks for one more build; an unchanged walk (same list kept) builds
  nothing; a failed build only warns; on_cleanup stops it. Servers that
  never ran the warm-up (the test servers) start no build.
- _favorite_listing slices the path strings as _listed does, falling back
  to relative_to for a path outside the root string.

* perf(images): start a stale listing's refresh walk a second after the read (#348)

After the faster walk, a read that found the listing stale still took
0.5-1.0 s against 0.07 s: it was served from the old listing but ran beside
the refresh walk it had just started, sharing the GIL. The refresh now starts
REFRESH_DELAY (1 s) after that read through a per-root rewalk_handle, so the
page's own requests finish first. One pending refresh per root and none while
a walk runs; any walk start (a cold read, Rescan, the warm-up) cancels a
pending one; on_cleanup cancels them before its first await.

* test(browser): accept KB or MB for the picked safety backup's size

The restore preview showed '2.5 MB' for the test server's database copy in a
CI coverage shard: its size depends on what earlier tests on that server
wrote, and the formatter switches to MB at 1 MiB.

* fix(images): keep the 3.10-3.13 suffix rule in the scandir walk and harden the warm-up (#348)

Review of the walk speedup:
- os.path.splitext drops names like ..png that Path.suffix listed on Python
  3.10-3.13 (3.14 changed Path.suffix, which hid it from the parity test on
  a 3.14 venv). _name_suffix spells out the 3.10-3.13 rule; the test now
  asserts ..png and ...webp are listed and g. is not.
- _warm_workspace reads the roots inside its try, so a failing
  _get_all_output_dirs warns instead of leaving an unretrieved task error.
- The folder_sets fallback test now takes the fallback (a Windows case
  mismatch); the coalescing test pins the walk count.
This commit is contained in:
Vito
2026-10-26 11:46:54 -07:00
committed by GitHub
parent 5772de8d27
commit 9ac471eb72
16 changed files with 1661 additions and 72 deletions
+47
View File
@@ -294,6 +294,53 @@ change the budget in the spec with a reason) in a follow-up PR. Never rerun it t
green. It is not a required check, and it cannot be one while it skips pull
requests (`.github/branch-protection.json`).
## Images listing cache (#348)
`/prompt_manager/images/output`, `/prompt_manager/gallery/subfolders` and
`POST /prompt_manager/gallery/rescan` read one shared listing per output root
(`RootListing` in `py/api/gallery_listing.py`, `_listing` in
`py/api/images.py`). The tests are `tests/test_gallery_listing.py` (pure
helpers) and `TestSharedListing` in `tests/test_api_images.py` (routes); they
hold a walk open with a gated scan instead of sleeping.
- **Single flight.** At most one walk per root runs at a time; every reader
that needs that root while it runs waits for it (shielded, so a client that
disconnects never cancels it). The walk, the comparison with the previous
listing and both folder sets run in the executor (`walk_root`). A walk that
finds the same `(path, mtime)` list keeps the previous list object, so the
favorites lookup and the sort orders are not rebuilt.
- **30 s background refresh.** A listing older than `REFRESH_AFTER` (30 s) is
served at once and one background walk replaces it. Files added or removed
by hand therefore show up one open later than before; Rescan shows them at
once.
- **Blocking invalidation after deletes.** Duplicate cleanup and the legacy
duplicate delete still call `invalidate_gallery_cache()`: the listing is
dropped and the next read waits for a walk, so a deleted file is never
listed. A walk running during the drop is discarded (the record's
`generation` outlives every drop); a cold reader gives up with a 500 after 3
discarded walks in a row. Thumbnail writes (jobs, conversion, the legacy
generate and clear routes) no longer drop the listing: thumbnail URLs are
computed per request.
- **Startup warm-up.** The root `__init__.py` registers
`api.warm_gallery_listings` on the host's `on_startup`, which starts one
background walk per root, and `api.cleanup_gallery_listings` on
`on_cleanup`, which cancels running walks. A warm-up failure never stops
startup; the first read retries.
- **Failures and Rescan.** `walk_root` opens the root with `os.scandir` first,
so a deleted or unreadable root fails the walk instead of listing nothing.
A failed walk keeps the previous listing, logs one warning per streak of
failures and is retried at most once per `REFRESH_AFTER`; a reader with no
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.
- **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
`node tools/vendor/perf/images_load_probe.mjs http://127.0.0.1:8188` against
a running ComfyUI. It prints the time to the first decoded thumbnail, the
time to network idle and the 10 slowest requests, from a fresh browser
context each run.
## Windows portability preflight
Before changing paths, SQLite, streams, restore or cleanup:
+5
View File
@@ -100,6 +100,11 @@ try:
if host_app is not None and hasattr(host_app, "on_cleanup"):
host_app.on_cleanup.append(api.cleanup_history_reads)
host_app.on_cleanup.append(api.cleanup_annotation_captures)
host_app.on_cleanup.append(api.cleanup_gallery_listings)
if host_app is not None and hasattr(host_app, "on_startup"):
# Build each output root's Images listing once the server's loop runs,
# so the first open does not wait for the walk (#348).
host_app.on_startup.append(api.warm_gallery_listings)
if host_app is not None and hasattr(host_app, "on_shutdown"):
# Open Needs review change streams end before the server waits on them.
host_app.on_shutdown.append(api.close_review_events)
+101 -11
View File
@@ -31,6 +31,7 @@ import math
import ntpath
import os
import re
import threading
from pathlib import Path
from aiohttp import web
@@ -47,6 +48,7 @@ from .review_events import ReviewEventsMixin
from .saved_views import SavedViewRoutesMixin
from .history import HistoryRoutesMixin
from .images import ImageRoutesMixin
from .gallery_listing import WALK_STOPPED, RootListing, WalkOutcome
from .admin import AdminRoutesMixin
from .settings_routes import SettingsRoutesMixin
from .job_routes import JobRoutesMixin
@@ -87,6 +89,18 @@ def _get_project_root():
return os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
# How long an executor thread waits for the loop to drop the listings.
_OFF_LOOP_DROP_WAIT = 30
def _running_loop():
"""The event loop running in this thread, or None."""
try:
return asyncio.get_running_loop()
except RuntimeError:
return None
def _public_path(path) -> str:
"""Render a server path for an API response without revealing the layout.
@@ -344,8 +358,17 @@ class PromptManagerAPI(
self._captures_closed = False
self._cached_output_dir = None # Lazy-cached by _find_comfyui_output_dir()
self._html_cache = {} # Cached HTML file contents keyed by path
self._gallery_cache = {} # dict: path_str -> (files, timestamp)
self._gallery_cache_ttl = 30 # Cache TTL in seconds
# One shared listing per output root, keyed by str(root), and the
# 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.
self._listing_loop = None
# 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
self._workspace_warm_task = None
self._workspace_warm_again = False
# Workspace listing values and sort orders derived from a cached scan
# (py/api/images.py _workspace_listed/_workspace_orders).
self._workspace_listed_cache = {}
@@ -465,20 +488,87 @@ class PromptManagerAPI(
return await loop.run_in_executor(None, call)
return await loop.run_in_executor(None, func, *args)
def invalidate_gallery_cache(self):
"""Invalidate the cached gallery file listing.
def _root_listing(self, root) -> RootListing:
"""The shared listing record of one output root (created empty)."""
return self._root_listings.setdefault(str(root), RootListing())
Called by the image monitor when new files are detected. The workspace
keeps its per-file values and card facts, so the next request rebuilds
only new or changed files (see ImageRoutesMixin._workspace_listed and
_card_facts); its sort orders
are dropped, as the rescan replaces them anyway, and so is the
favorites lookup, which holds the old scan list.
def invalidate_gallery_cache(self):
"""Drop every root's listing, so the next read waits for a new walk.
Called after files are deleted through the app (duplicate cleanup and
the legacy duplicate delete), never for thumbnail writes. A walk
running now is not stored (its generation moved on). The workspace
keeps its per-file values and card facts, so the next request
rebuilds only new or changed files (see ImageRoutesMixin.
_workspace_listed and _card_facts); its sort orders are dropped, as
the walk replaces them anyway, and so is the favorites lookup, which
holds the old list.
"""
self._gallery_cache = {}
loop = self._listing_loop
if loop is not None and loop.is_running() and _running_loop() is not loop:
# Called from an executor thread (a settings save, a test seed):
# listing records are touched only on the loop, and the caller
# expects the drop done when this returns.
dropped = threading.Event()
def drop():
try:
self._drop_listings()
finally:
dropped.set()
loop.call_soon_threadsafe(drop)
if not dropped.wait(_OFF_LOOP_DROP_WAIT):
self.logger.warning("The Images listing drop is still queued")
return
self._drop_listings()
def _drop_listings(self):
for record in list(self._root_listings.values()):
record.drop()
self._workspace_order_cache = {}
self._favorite_listing_cache = {}
async def warm_gallery_listings(self, app=None):
"""Start one background walk per output root (aiohttp on_startup).
Registered by the root ``__init__.py``, so the first Images open after
ComfyUI starts finds the listing built. It never fails startup: a
root that cannot be walked is retried by the first read.
"""
# 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
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}")
async def cleanup_gallery_listings(self, app=None):
"""Stop the background listing walks (aiohttp on_cleanup)."""
self._workspace_warm_enabled = False
# 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
if handle is not None:
handle.cancel()
warm, self._workspace_warm_task = self._workspace_warm_task, None
if warm is not None:
warm.cancel()
await asyncio.gather(warm, return_exceptions=True)
tasks = list(self._listing_tasks)
for task in tasks:
task.cancel()
if tasks:
await asyncio.gather(*tasks, return_exceptions=True)
self._listing_tasks.clear()
# A walk cancelled before its first step never ran its own cleanup.
for record in self._root_listings.values():
flight, record.flight = record.flight, None
if flight is not None and not flight.done():
flight.set_result(WalkOutcome.failed(RuntimeError(WALK_STOPPED)))
def reset_gallery_roots_cache(self):
"""Drop every cached listing and derived workspace value (roots changed)."""
self.invalidate_gallery_cache()
+146
View File
@@ -0,0 +1,146 @@
"""One shared listing per output root, and the pure helpers behind it (#348).
The API keeps one ``RootListing`` per output root (``_root_listings`` in
``py/api/__init__.py``). The record is read and written only on the event
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.
"""
from __future__ import annotations
import asyncio
import os
from dataclasses import dataclass
from pathlib import PurePath
from typing import Callable, NamedTuple
# A listing older than this (seconds) is served as it is, and one background
# walk replaces it. Read as ``gallery_listing.REFRESH_AFTER`` so tests can
# patch it.
REFRESH_AFTER = 30
# The stale read is served at once; its refresh walk starts this many seconds
# later, so the walk does not share the GIL with the page's own requests
# (#348 live test: reads during the walk took 2.2 s instead of 0.07 s).
REFRESH_DELAY = 1.0
# 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
WALK_STOPPED = "The listing walk was stopped"
STORED = "stored"
DISCARDED = "discarded"
FAILED = "failed"
class WalkOutcome(NamedTuple):
"""What one walk ended with; a walk's flight future resolves to this."""
kind: str
error: BaseException | None = None
@classmethod
def stored(cls) -> "WalkOutcome":
return cls(STORED)
@classmethod
def discarded(cls) -> "WalkOutcome":
return cls(DISCARDED)
@classmethod
def failed(cls, error: BaseException) -> "WalkOutcome":
return cls(FAILED, error)
@dataclass
class RootListing:
"""The shared listing of one output root.
A record is never deleted: ``drop()`` clears the listing and bumps
``generation``, so a walk that started before a drop can never be stored
after it.
"""
# (path, mtime) newest first, as _scan_gallery_files_sync returns it.
files: list | None = None
folders: frozenset | None = None
folders_with_ancestors: frozenset | None = None
# Monotonic time the last walk finished, stored or failed.
walked_at: float | None = None
# Number of the last walk started; a running walk is always that one.
walk_seq: int = 0
# Number of the walk whose files are stored (Rescan waits on it).
stored_seq: int = 0
generation: int = 0
# The running walk's future (resolves to a WalkOutcome), or None.
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.
rewalk_handle: asyncio.TimerHandle | None = None
def drop(self) -> None:
self.files = None
self.folders = None
self.folders_with_ancestors = None
self.walked_at = None
self.generation += 1
def folder_sets(root, files) -> tuple[frozenset, frozenset]:
"""``(folders, folders_with_ancestors)`` of the files listed under *root*.
Folders keep the spelling ``/gallery/subfolders`` has always answered:
``str(relative.parent)`` (native separators) for the folders holding
files, plus every ancestor joined by ``/``.
"""
# From the path strings: relative_to per file was half of a large walk.
# The separator is the path flavour's own (tests use PureWindowsPath).
sep = str(type(root)("a", "b"))[1]
prefix = str(root)
if not prefix.endswith(sep):
prefix += sep
folders = set()
for path, _mtime in files:
text = str(path)
if text.startswith(prefix):
rel_dir = text[len(prefix) :].rpartition(sep)[0]
else:
rel_dir = str(path.relative_to(root).parent)
if rel_dir and rel_dir != ".":
folders.add(rel_dir)
with_ancestors = set(folders)
for folder in folders:
parts = folder.replace("\\", "/").split("/")
for i in range(1, len(parts)):
with_ancestors.add("/".join(parts[:i]))
return frozenset(folders), frozenset(with_ancestors)
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:
return False
return old is new or old == new
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)``.
Opens the root first, so a missing or unreadable root raises ``OSError``
instead of listing nothing. Returns *previous_files* itself when nothing
changed, which keeps the identity-keyed caches built on it warm.
"""
with os.scandir(root):
pass
files = scan_fn(root)
if same_listing(previous_files, files):
files = previous_files
folders, with_ancestors = folder_sets(root, files)
return files, folders, with_ancestors
+258 -47
View File
@@ -22,6 +22,15 @@ from aiohttp import web
from PIL import Image, UnidentifiedImageError
from .jobs import MAX_OUTCOME_ROWS, THUMBNAILS_JOB_KIND
from . import gallery_listing
from .gallery_listing import (
DISCARDED,
FAILED,
MAX_DISCARDED_WALKS,
WALK_STOPPED,
WalkOutcome,
walk_root,
)
from . import thumbnail_settings
from .thumbnail_encode import (
PREVIEW_HELD,
@@ -763,32 +772,65 @@ def _favorite_root_key(root) -> str:
return _favorite_relative_key(os.path.realpath(str(root)))
def _iter_media_files(output_path, max_files):
"""Yield media files under *output_path* with bounded effort.
def _name_suffix(name):
"""``Path(name).suffix`` as Python 3.10-3.13 compute it.
The last ``.``-part unless the dot is the first or last character, so
``..png`` is ``.png`` and ``.png`` has none. Spelled out because 3.14
changed ``Path.suffix`` and ``os.path.splitext`` differs for leading dots.
"""
dot = name.rfind(".")
return name[dot:] if 0 < dot < len(name) - 1 else ""
def _walk_media(output_path, max_files):
"""Yield the ``os.DirEntry`` of each media file under *output_path*.
Symlinks are never followed, ``thumbnails/`` is skipped, descent stops
at MAX_SCAN_DEPTH and at most *max_files* paths are produced.
at MAX_SCAN_DEPTH and at most *max_files* entries are produced. The order
is ``os.walk``'s (top-down, subfolders in ``os.scandir`` order, files by
name), which the gallery's ties in mtime keep. Each entry's type comes
from the directory listing itself, so a walk builds no ``Path`` and makes
no extra ``lstat`` per file (#348: a 20,000-file walk spent half its time
there).
"""
output_path = Path(output_path)
root_depth = len(output_path.parts)
stack = [(os.fspath(output_path), ())]
produced = 0
for root, dirs, files in os.walk(output_path):
root_path = Path(root)
rel_parts = root_path.parts[root_depth:]
while stack:
top, rel_parts = stack.pop()
if "thumbnails" in rel_parts:
dirs[:] = []
continue
if len(rel_parts) >= MAX_SCAN_DEPTH:
dirs[:] = []
for name in sorted(files):
if Path(name).suffix.lower() not in MEDIA_EXTENSIONS:
dirs, files = [], []
try:
with os.scandir(top) as entries:
for entry in entries:
try:
is_dir = entry.is_dir()
except OSError:
is_dir = False
(dirs if is_dir else files).append(entry)
except OSError:
continue # as os.walk: an unreadable folder is skipped
for entry in sorted(files, key=lambda item: item.name):
if _name_suffix(entry.name).lower() not in MEDIA_EXTENSIONS:
continue
if (root_path / name).is_symlink():
if entry.is_symlink():
continue # never read through a link out of the output tree
if produced >= max_files:
return
produced += 1
yield root_path / name
yield entry
if len(rel_parts) >= MAX_SCAN_DEPTH:
continue
for entry in reversed(dirs):
if not entry.is_symlink():
stack.append((entry.path, rel_parts + (entry.name,)))
def _iter_media_files(output_path, max_files):
"""Yield the ``Path`` of each media file under *output_path* (see _walk_media)."""
for entry in _walk_media(output_path, max_files):
yield Path(entry.path)
def media_thumbnail_rel(rel_path, settings=None):
@@ -1278,6 +1320,60 @@ class ImageRoutesMixin:
{"success": False, "error": _safe_error(e)}, status=500
)
def _schedule_workspace_warm(self):
"""Rebuild the workspace caches behind a changed listing (on the loop).
Only after the host's startup warm-up (warm_gallery_listings), so the
first Images open after a restart, and the first read after a refresh
found new files, pay no per-file rebuild. A walk stored while a build
runs asks for one more build instead of a second one at once.
"""
if not self._workspace_warm_enabled:
return
task = self._workspace_warm_task
if task is not None and not task.done():
self._workspace_warm_again = True
return
self._workspace_warm_task = asyncio.get_running_loop().create_task(
self._warm_workspace()
)
async def _warm_workspace(self):
while True:
self._workspace_warm_again = False
try:
listings = self._stored_listings()
if listings is None:
return # a root has no listing yet; its walk asks again
await self._run_in_executor(self._build_workspace_caches, listings)
except Exception as e:
self.logger.warning(f"Could not prepare the Images page: {e}")
return
if not self._workspace_warm_again:
return
def _stored_listings(self):
"""``[(index, root, files)]`` of every output root, or None if one has none."""
listings = []
for index, root in enumerate(self._get_all_output_dirs()):
record = self._root_listings.get(str(root))
if record is None or record.files is None:
return None
listings.append((index, root, record.files))
return listings
def _build_workspace_caches(self, listings):
"""Build what a first workspace page reads, as that page builds it (blocking).
The default page's selection builds the per-file values, the newest
order, the wanted paths and reads the prompt links (warming SQLite's
cache); the favorites filter builds each root's favorite listing.
"""
self._select_workspace_page(
listings, _workspace_filters({}), None, 1, "workspace-warm-up"
)
self._workspace_favorite_paths(listings)
def _scan_gallery_files_sync(self, output_path):
"""Scan output directory for media files (blocking I/O, run in executor).
@@ -1285,27 +1381,141 @@ class ImageRoutesMixin:
walk is bounded by MAX_SCAN_DEPTH and MAX_GALLERY_FILES.
"""
found = []
for media_path in _iter_media_files(output_path, MAX_GALLERY_FILES):
for entry in _walk_media(output_path, MAX_GALLERY_FILES):
try:
found.append((media_path, media_path.stat().st_mtime))
# Never a symlink (_walk_media skips them), so this is the
# file's own mtime, read from the listing where it can be.
mtime = entry.stat(follow_symlinks=False).st_mtime
except OSError:
continue
found.append((Path(entry.path), mtime))
if len(found) >= MAX_GALLERY_FILES:
self.logger.warning(f"Gallery scan capped at {MAX_GALLERY_FILES} files")
found.sort(key=lambda x: x[1], reverse=True)
return found
async def _get_gallery_files(self, output_path):
"""Get gallery files with per-directory TTL cache."""
now = _time.monotonic()
key = str(output_path)
cached = self._gallery_cache.get(key)
if cached and (now - cached[1]) < self._gallery_cache_ttl:
return cached[0]
"""The shared listing of one output root: ``[(path, mtime)]``, newest first."""
return (await self._listing(output_path)).files
files = await self._run_in_executor(self._scan_gallery_files_sync, output_path)
self._gallery_cache[key] = (files, now)
return files
async def _listing(self, output_path, *, fresh=False):
"""The shared ``RootListing`` of *output_path*, with a listing (#348).
A listing that exists is returned at once; when it is older than
``REFRESH_AFTER`` one background walk replaces it. With no listing,
the caller waits for the root's single walk. ``fresh`` (Rescan) waits
for a walk that started after this call, and raises when it fails
even if a previous listing exists.
"""
record = self._root_listing(output_path)
needed = record.walk_seq + 1 if fresh else 0
discarded = 0
while True:
if record.files is not None and record.stored_seq >= needed:
if not fresh:
self._refresh_if_stale(output_path, record)
return record
flight = record.flight
if flight is None:
flight = self._start_walk(output_path, record)
joined = record.walk_seq
# Shielded: a client that goes away never cancels a shared walk.
outcome = await asyncio.shield(flight)
if joined < needed:
continue # an older walk ended; Rescan needs a newer one
if outcome.kind == FAILED:
raise outcome.error
if outcome.kind == DISCARDED:
discarded += 1
if discarded >= MAX_DISCARDED_WALKS:
raise RuntimeError(
"The output listing kept changing while it was read"
)
def _refresh_if_stale(self, output_path, record):
"""Walk a stale listing again, REFRESH_DELAY after this read (on the loop).
At most one pending refresh per root, none while a walk runs.
"""
age = _time.monotonic() - (record.walked_at or 0.0)
if age < gallery_listing.REFRESH_AFTER:
return
if record.flight is not None or record.rewalk_handle is not None:
return
record.rewalk_handle = asyncio.get_running_loop().call_later(
gallery_listing.REFRESH_DELAY, self._delayed_walk, output_path, record
)
def _delayed_walk(self, output_path, record):
record.rewalk_handle = None
self._schedule_walk(output_path)
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)
def _start_walk(self, output_path, record):
"""Start the root's one walk; returns its flight future (on the loop)."""
loop = asyncio.get_running_loop()
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
record.walk_seq += 1
flight = loop.create_future()
record.flight = flight
task = loop.create_task(
self._run_walk(
output_path, record, flight, record.generation, record.walk_seq
)
)
self._listing_tasks.add(task)
task.add_done_callback(self._listing_tasks.discard)
return flight
async def _run_walk(self, output_path, record, flight, generation, seq):
"""Walk in the executor, then store, discard or fail (on the loop).
The flight is always resolved with a WalkOutcome, never failed, so an
error never reaches a future nobody awaits.
"""
outcome = None
try:
files, folders, with_ancestors = await self._run_in_executor(
walk_root, self._scan_gallery_files_sync, output_path, record.files
)
except asyncio.CancelledError:
outcome = WalkOutcome.failed(RuntimeError(WALK_STOPPED))
raise
except Exception as exc:
record.walked_at = _time.monotonic()
if not record.failing:
self.logger.warning(
f"Could not list output folder {output_path}: {exc}"
)
record.failing = True
outcome = WalkOutcome.failed(exc)
else:
record.failing = False
if record.generation == generation:
changed = files is not record.files
record.files = files
record.folders = folders
record.folders_with_ancestors = with_ancestors
record.walked_at = _time.monotonic()
record.stored_seq = seq
outcome = WalkOutcome.stored()
if changed:
self._schedule_workspace_warm()
else:
outcome = WalkOutcome.discarded()
finally:
if record.flight is flight:
record.flight = None
if not flight.done():
flight.set_result(outcome)
async def get_output_images(self, request):
"""Get all images from ComfyUI output folder(s)."""
@@ -1832,17 +2042,18 @@ class ImageRoutesMixin:
)
async def rescan_gallery(self, request):
"""Forget the cached listing and rewalk every output root now.
"""Rewalk every output root now and answer with the new listing.
The listing is cached for ``_gallery_cache_ttl`` seconds; the gallery's
Rescan button needs files added or removed on disk to show up at once.
Readers are served the shared listing and refresh it in the
background (``gallery_listing.REFRESH_AFTER``); the gallery's Rescan
button needs files added or removed on disk to show up at once, so it
waits for a walk that started after the click and fails when it does.
"""
try:
self.invalidate_gallery_cache()
output_dirs = list(self._get_all_output_dirs())
total = 0
for output_path in output_dirs:
total += len(await self._get_gallery_files(output_path))
total += len((await self._listing(output_path, fresh=True)).files)
return web.json_response(
{"success": True, "total": total, "roots": len(output_dirs)}
)
@@ -1853,20 +2064,14 @@ class ImageRoutesMixin:
)
async def _root_subfolders(self, output_path, include_ancestors):
"""The distinct folders holding listed files under one output root."""
subfolders = set()
for f, _mtime in await self._get_gallery_files(output_path):
rel_dir = str(f.relative_to(output_path).parent)
if rel_dir and rel_dir != ".":
subfolders.add(rel_dir)
"""The distinct folders holding listed files under one output root.
Built with the walk, in the executor (gallery_listing.folder_sets).
"""
listing = await self._listing(output_path)
if include_ancestors:
ancestors = set()
for folder in subfolders:
parts = folder.replace("\\", "/").split("/")
for i in range(1, len(parts)):
ancestors.add("/".join(parts[:i]))
subfolders.update(ancestors)
return subfolders
return listing.folders_with_ancestors
return listing.folders
async def get_gallery_subfolders(self, request):
"""Get distinct subfolders from gallery output directories.
@@ -1940,7 +2145,6 @@ class ImageRoutesMixin:
result = await self._run_in_executor(
self._generate_gallery_thumbnails_sync, output_dirs
)
self.invalidate_gallery_cache()
return web.json_response(result)
except Exception as e:
@@ -2416,7 +2620,6 @@ class ImageRoutesMixin:
cleared_count, cleared_size = await self._run_in_executor(
self._clear_gallery_thumbnails_sync, output_dirs
)
self.invalidate_gallery_cache()
def format_size(bytes_size: int) -> str:
size = float(bytes_size)
@@ -2764,9 +2967,17 @@ class ImageRoutesMixin:
cached = self._favorite_listing_cache.get(cache_key)
if cached is not None and cached[0] is files:
return cached[1]
# From the path strings, as _listed does: relative_to per file was the
# largest cost of a cold first open (#348).
prefix = os.path.join(str(root), "")
listing = {}
for path, _mtime in files:
posix = path.relative_to(root).as_posix()
text = str(path)
if text.startswith(prefix):
relative = text[len(prefix) :]
posix = relative if os.sep == "/" else relative.replace(os.sep, "/")
else:
posix = path.relative_to(root).as_posix()
listing.setdefault(_favorite_relative_key(posix), posix)
self._favorite_listing_cache[cache_key] = (files, listing)
return listing
-1
View File
@@ -175,7 +175,6 @@ class ThumbnailConvertMixin:
job.log_line(f"Conversion finished · {NOTHING_OUTDATED.lower()}")
return "finished"
totals, left = await self._convert_in_batches(job, options, plan, output_dirs)
self.invalidate_gallery_cache()
if totals["stopped"]:
job.update(status_line=FORMAT_CHANGED.format(converted=totals["converted"]))
job.log_line(
-1
View File
@@ -261,7 +261,6 @@ class ThumbnailJobsMixin:
stats, start, left, attempted = await self._generate_in_batches(
job, options, plan, roots, output_dirs
)
self.invalidate_gallery_cache()
generated = _plural(stats["generated"], "preview")
if job.cancelled:
return self._end_cancelled(
@@ -485,5 +485,6 @@ test('a safety backup file (no .db extension) can be picked and reviewed (M3)',
const answer = await chooseBackup(page, bytes, 'prompts.db.backup_20261025_093000_000000');
expect((await answer).status()).toBe(200);
await expect(dialog(page).locator('#modal-title')).toHaveText('Restore preview');
expect((await gridOf(page))[0][1]).toMatch(/^prompts\.db\.backup_20261025_093000_000000 · \d+ KB · .+$/);
// The copy's size depends on what earlier tests on this server wrote: KB or MB.
expect((await gridOf(page))[0][1]).toMatch(/^prompts\.db\.backup_20261025_093000_000000 · [\d,.]+ (KB|MB) · .+$/);
});
@@ -381,7 +381,8 @@ test('saved pairs block the import: the review explains it and offers only Close
await openIntegrations(page);
await enable(page);
await layoutOf(page).getByRole('button', { name: 'Review import', exact: true }).click();
expect((await gridOf(page))[2]).toEqual(['Protected', '1 imported prompt used by saved pairs']);
// gridOf reads once; poll until the review has rendered its rows.
await expect.poll(async () => (await gridOf(page))[2]).toEqual(['Protected', '1 imported prompt used by saved pairs']);
await expect(dialog(page).locator('p.small')).toHaveText(
'Delete or change the saved pairs that use these prompts, then review the import again. Nothing was changed.');
expect(await dialog(page).locator('.dialog-actions button').allTextContents()).toEqual(['Close']);
+16 -2
View File
@@ -687,9 +687,23 @@ class TestInstanceHelpers(FolderPathsFixture):
self.assertEqual(self._run(self.api._run_in_executor(add, 1, 2)), 3)
def test_invalidate_gallery_cache(self):
self.api._gallery_cache = {"x": 1}
record = self.api._root_listing(self.output_dir)
record.files = [(self.output_dir / "a.png", 1.0)]
record.folders = record.folders_with_ancestors = frozenset()
record.walked_at = 1.0
self.api._favorite_listing_cache = {"x": 1}
self.api._workspace_order_cache = {"x": 1}
self.api.invalidate_gallery_cache()
self.assertEqual(self.api._gallery_cache, {})
# The record stays: its generation outlives every drop (#348).
self.assertIs(self.api._root_listings[str(self.output_dir)], record)
self.assertIsNone(record.files)
self.assertIsNone(record.folders)
self.assertIsNone(record.walked_at)
self.assertEqual(record.generation, 1)
self.assertEqual(self.api._favorite_listing_cache, {})
self.assertEqual(self.api._workspace_order_cache, {})
self.api.reset_gallery_roots_cache()
self.assertEqual(record.generation, 2)
def test_enrich_images_adds_urls_and_thumbnails(self):
GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)]
+758 -8
View File
@@ -14,7 +14,7 @@ import tempfile
import threading
import unittest
from contextlib import ExitStack
from pathlib import Path
from pathlib import Path, PurePath
from unittest.mock import AsyncMock, MagicMock, patch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
@@ -34,6 +34,7 @@ from PIL import Image # noqa: E402
from database.operations import PromptDatabase # noqa: E402
import py.api.images as images_module
from py.api import gallery_listing # noqa: E402
from py.api.thumbnail_settings import ThumbnailSettings # noqa: E402
from py.api import PromptManagerAPI # noqa: E402
from py.api.prompts import MAX_PAGE_LIMIT # noqa: E402
@@ -1454,12 +1455,15 @@ class TestGalleryListing(ImageRouteCoverageCase):
self.assertEqual(data["images"], [])
def test_gallery_scan_skips_unreadable_entries(self):
ghost = self.output_dir / "ghost.png"
make_png(self.output_dir / "real.png")
listed = list(images_module._walk_media(self.output_dir, 10))
ghost = MagicMock(path=str(self.output_dir / "ghost.png"))
ghost.stat.side_effect = FileNotFoundError("gone")
with patch.object(
images_module, "_iter_media_files", return_value=iter([ghost])
images_module, "_walk_media", return_value=iter([ghost, *listed])
):
found = self.api._scan_gallery_files_sync(self.output_dir)
self.assertEqual(found, [])
self.assertEqual([path.name for path, _mtime in found], ["real.png"])
class TestThumbnailUnits(ImageRouteCoverageCase):
@@ -1805,10 +1809,13 @@ class TestClearThumbnails(ImageRouteCoverageCase):
self.assertFalse(data["success"])
self.assertTrue(original.is_file())
async def test_cleanup_runs_off_loop_and_invalidates_gallery_cache(self):
async def test_cleanup_runs_off_loop_and_keeps_the_listing(self):
import threading
make_png(self.output_dir / "thumbnails" / "a_thumb.png")
make_png(self.output_dir / "a.png")
await self._json("GET", "/prompt_manager/images/output")
listing = self.api._root_listings[str(self.output_dir)].files
loop_thread = threading.get_ident()
deletion_threads = []
original_unlink = Path.unlink
@@ -1817,14 +1824,14 @@ class TestClearThumbnails(ImageRouteCoverageCase):
deletion_threads.append(threading.get_ident())
return original_unlink(path, *args, **kwargs)
self.api._gallery_cache = {"previous": object()}
with patch.object(Path, "unlink", unlink):
status, data = await self._json("POST", self.URL)
self.assertEqual(status, 200)
self.assertEqual(data["cleared_files"], 1)
self.assertTrue(deletion_threads)
self.assertNotIn(loop_thread, deletion_threads)
self.assertEqual(self.api._gallery_cache, {})
# Previews are not part of the listing (#348): it stays as it was.
self.assertIs(self.api._root_listings[str(self.output_dir)].files, listing)
class TestImagePromptLookup(ImageRouteCoverageCase):
@@ -1979,7 +1986,10 @@ class TestGalleryRescan(ImageAPITestCase):
return resp.status, await resp.json()
async def test_rescan_makes_new_files_visible_before_the_cache_expires(self):
self.api._gallery_cache_ttl = 3600 # the cache alone would hide the new file
# The listing alone would hide the new file for an hour.
refresh = patch.object(gallery_listing, "REFRESH_AFTER", 3600)
refresh.start()
self.addCleanup(refresh.stop)
make_png(self.output_dir / "a.png")
_, first = await self._json("GET", self.URL)
self.assertEqual(first["total"], 1)
@@ -2026,6 +2036,78 @@ class TestOutputScanHelpers(ImageAPITestCase):
expected = hashlib.sha1(str(media).encode("utf-8")).hexdigest()[:16]
self.assertEqual(entry["id"], expected)
def test_media_walk_matches_the_os_walk_reference(self):
# The scandir walk (#348) must list exactly what the os.walk version
# listed, in the same order: ties in mtime keep this order.
output = Path(self.output_dir)
for rel in (
"z.png",
"a.PNG",
"notes.txt",
"b/2.png",
"b/1.webp",
"b/c/deep.mp4",
"b/c/d/too_deep.png",
"thumbnails/t.png",
"b/thumbnails/u.png",
"e/.hidden.png",
"e/f.jpeg",
"e/..png",
"e/...webp",
"e/g.",
):
make_png(output / rel)
try:
os.symlink(output / "z.png", output / "b" / "link.png")
os.symlink(output / "b", output / "dirlink", target_is_directory=True)
except (OSError, NotImplementedError):
pass # the rest of the tree still checks order and filters
with patch.object(images_module, "MAX_SCAN_DEPTH", 2):
listed = list(images_module._iter_media_files(output, 100))
expected = list(_reference_iter_media_files(output, 100))
capped = list(images_module._iter_media_files(output, 3))
self.assertEqual(listed, expected)
self.assertEqual(capped, expected[:3])
# Path.suffix of Python 3.10-3.13, which CI and hosts run (3.14 changed
# it): "..png" has the suffix ".png", ".png" and "g." have none.
names = {path.name for path in listed}
self.assertTrue({"..png", "...webp", ".hidden.png"} <= names, names)
self.assertFalse({"g.", "notes.txt"} & names, names)
self.assertTrue(all(isinstance(path, Path) for path in listed))
def test_gallery_scan_reads_each_file_type_and_mtime_from_scandir(self):
# Path.stat and Path.is_symlink per file were the other half of the
# live walk's cost (#348); the scandir entries already carry both.
output = Path(self.output_dir)
make_png(output / "a.png")
make_png(output / "sub" / "b.png")
with (
patch.object(Path, "stat", side_effect=AssertionError),
patch.object(Path, "is_symlink", side_effect=AssertionError),
):
found = self.api._scan_gallery_files_sync(output)
self.assertEqual(
sorted(path.name for path, _mtime in found), ["a.png", "b.png"]
)
for path, mtime in found:
self.assertEqual(mtime, os.stat(path).st_mtime)
def test_favorite_listing_slices_the_path_strings(self):
# relative_to per file was the largest cold cost of the first open.
output = Path(self.output_dir)
files = [
(output / "a.png", 1.0),
(output / "Sub" / "B.png", 2.0),
(output / "sub" / "deep" / "c.webp", 3.0),
]
expected = {}
for path, _mtime in files:
posix = path.relative_to(output).as_posix()
expected.setdefault(images_module._favorite_relative_key(posix), posix)
with patch.object(PurePath, "relative_to", side_effect=AssertionError):
listing = self.api._favorite_listing(output, files)
self.assertEqual(listing, expected)
def test_iter_media_files_skips_symlinked_files(self):
output = Path(self.output_dir)
make_png(output / "real.png")
@@ -2342,3 +2424,671 @@ class TestRootScopeServing(ImageAPITestCase):
self._make_symlink(roots[1], link)
for path in ("image.png", "thumbnails/image_thumb.png"):
await self.assert_no_access(path, scope=token)
# Seconds any single wait in the shared-listing tests may take.
_HANG_LIMIT = 30
def _reference_iter_media_files(output_path, max_files):
"""``_iter_media_files`` as it was before the scandir walk (#348)."""
output_path = Path(output_path)
root_depth = len(output_path.parts)
produced = 0
for root, dirs, files in os.walk(output_path):
root_path = Path(root)
rel_parts = root_path.parts[root_depth:]
if "thumbnails" in rel_parts:
dirs[:] = []
continue
if len(rel_parts) >= images_module.MAX_SCAN_DEPTH:
dirs[:] = []
for name in sorted(files):
# Path.suffix as Python 3.10-3.13 computed it (3.14 changed it).
dot = name.rfind(".")
suffix = name[dot:] if 0 < dot < len(name) - 1 else ""
if suffix.lower() not in images_module.MEDIA_EXTENSIONS:
continue
if (root_path / name).is_symlink():
continue
if produced >= max_files:
return
produced += 1
yield root_path / name
class _GatedScan:
"""Wraps the real scan, counts calls and can hold one call after it listed.
``hold_next()`` makes the next call list the files, then wait on the
returned gate, so a test can change the disk or the cache while a walk
is in flight and decide when it ends.
"""
def __init__(self, real):
self.real = real
self.calls = 0
self._gate = None
self.entered = threading.Event()
self._gates = []
def hold_next(self):
self._gate = threading.Event()
self._gates.append(self._gate)
self.entered.clear()
return self._gate
def release_all(self):
for gate in self._gates:
gate.set()
def __call__(self, root):
self.calls += 1
files = self.real(root)
gate, self._gate = self._gate, None
if gate is not None:
self.entered.set()
gate.wait(30)
return files
async def wait_entered(self):
loop = asyncio.get_running_loop()
entered = await loop.run_in_executor(None, self.entered.wait, 30)
if not entered:
raise AssertionError("the walk never reached the gate")
class TestSharedListing(ImageAPITestCase):
"""One listing per output root, shared and refreshed in the background (#348)."""
URL = "/prompt_manager/images/output"
async def asyncSetUp(self):
await super().asyncSetUp()
self.scan = _GatedScan(self.api._scan_gallery_files_sync)
self.api._scan_gallery_files_sync = self.scan
self.addCleanup(self.scan.release_all)
# The stale refresh waits REFRESH_DELAY after the read (#348); these
# tests start it on the next loop turn and _settle waits for it.
delay = patch.object(gallery_listing, "REFRESH_DELAY", 0)
delay.start()
self.addCleanup(delay.stop)
async def _json(self, method, path, **kwargs):
# Bounded, like every wait in this class: a stuck walk fails the test
# with a traceback instead of hanging the whole run.
async def call():
resp = await self.client.request(method, path, **kwargs)
return resp.status, await resp.json()
return await asyncio.wait_for(call(), _HANG_LIMIT)
def _record(self):
return self.api._root_listings[str(self.output_dir)]
def _pending(self):
return any(
record.rewalk_handle is not None
for record in self.api._root_listings.values()
)
def _warming(self):
task = self.api._workspace_warm_task
return task is not None and not task.done()
def _age(self, seconds=None):
from py.api import gallery_listing
record = self._record()
record.walked_at -= (seconds or gallery_listing.REFRESH_AFTER) + 1
async def _settle(self):
"""Wait for every background walk this API started."""
async def settle():
while self.api._listing_tasks or self._warming() or self._pending():
if self._pending():
await asyncio.sleep(0.01) # a delayed refresh not started yet
await asyncio.gather(*list(self.api._listing_tasks))
# A finished walk leaves the set in a done callback, which needs
# a loop turn; gather over finished tasks does not yield one.
await asyncio.sleep(0)
if self._warming():
await asyncio.wait([self.api._workspace_warm_task])
await asyncio.wait_for(settle(), _HANG_LIMIT)
async def test_concurrent_cold_readers_share_one_walk(self):
make_png(self.output_dir / "sub" / "a.png")
results = await asyncio.gather(
self._json("GET", self.URL),
self._json("GET", self.URL + "?limit=1"),
self._json("GET", self.URL + "?mode=workspace&limit=1"),
self._json("GET", "/prompt_manager/gallery/subfolders"),
)
self.assertEqual([status for status, _ in results], [200] * 4)
self.assertEqual(self.scan.calls, 1)
self.assertEqual(results[0][1]["total"], 1)
self.assertEqual(results[3][1]["subfolders"], ["sub"])
async def test_a_stale_listing_is_served_at_once_and_refreshed_behind(self):
make_png(self.output_dir / "a.png")
_, first = await self._json("GET", self.URL)
make_png(self.output_dir / "b.png")
self._age()
gate = self.scan.hold_next()
_, stale = await self._json("GET", self.URL)
# Served from the old listing while the walk is still held.
self.assertEqual((first["total"], stale["total"]), (1, 1))
await self.scan.wait_entered()
# A second stale read joins the running walk instead of starting one.
_, again = await self._json("GET", self.URL)
self.assertEqual(again["total"], 1)
gate.set()
await self._settle()
_, fresh = await self._json("GET", self.URL)
self.assertEqual(fresh["total"], 2)
self.assertEqual(self.scan.calls, 2)
async def test_a_fresh_listing_is_not_walked_again(self):
make_png(self.output_dir / "a.png")
for _ in range(3):
await self._json("GET", self.URL)
await self._json("GET", "/prompt_manager/gallery/subfolders")
self.assertEqual(self.scan.calls, 1)
self.assertEqual(self.api._listing_tasks, set())
async def test_rescan_waits_for_a_walk_started_after_the_click(self):
make_png(self.output_dir / "a.png")
await self._json("GET", self.URL)
self._age()
gate = self.scan.hold_next()
await self._json("GET", self.URL) # starts a background walk
await self.scan.wait_entered() # it has listed: it cannot see c.png
make_png(self.output_dir / "c.png")
rescan = asyncio.ensure_future(
self._json("POST", "/prompt_manager/gallery/rescan")
)
await asyncio.sleep(0)
self.assertFalse(rescan.done())
gate.set()
status, data = await rescan
self.assertEqual(status, 200)
self.assertEqual((data["total"], data["roots"]), (2, 1))
self.assertEqual(self.scan.calls, 3)
async def test_rescans_clicked_during_a_walk_share_the_next_one(self):
make_png(self.output_dir / "a.png")
await self._json("GET", self.URL)
original = self.api._listing
rescans = []
clicked = asyncio.Event()
async def listing(root, *, fresh=False):
if fresh:
rescans.append(root)
if len(rescans) == 3:
clicked.set()
# Runs on in this step up to its first wait, so the Rescan's
# walk number is taken before the test resumes.
return await original(root, fresh=fresh)
self.api._listing = listing
gate = self.scan.hold_next()
first = asyncio.ensure_future(
self._json("POST", "/prompt_manager/gallery/rescan")
)
await self.scan.wait_entered()
# Clicked while the first Rescan's walk runs: neither may use it,
# and both wait for the same next walk.
later = [
asyncio.ensure_future(self._json("POST", "/prompt_manager/gallery/rescan"))
for _ in range(2)
]
await asyncio.wait_for(clicked.wait(), 30)
gate.set()
results = await asyncio.gather(first, *later)
self.assertEqual([status for status, _ in results], [200] * 3)
self.assertEqual(self.scan.calls, 3)
async def test_rescan_failure_with_a_previous_listing_is_500(self):
make_png(self.output_dir / "a.png")
await self._json("GET", self.URL)
self.api._scan_gallery_files_sync = _raise
with self.assertLogs(self.api.logger, "WARNING"):
status, data = await self._json("POST", "/prompt_manager/gallery/rescan")
self.assertEqual(status, 500)
self.assertFalse(data["success"])
self.assertNotIn("boom", json.dumps(data))
# Plain readers keep the previous listing, and do not retry at once.
status, listed = await self._json("GET", self.URL)
self.assertEqual((status, listed["total"]), (200, 1))
self.assertEqual(self.api._listing_tasks, set())
async def test_a_cold_read_whose_walk_fails_is_500(self):
self.api._scan_gallery_files_sync = _raise
with self.assertLogs(self.api.logger, "WARNING"):
status, data = await self._json("GET", self.URL)
self.assertEqual(status, 500)
self.assertFalse(data["success"])
async def test_a_deleted_root_keeps_the_previous_listing_and_warns_once(self):
import shutil
make_png(self.output_dir / "a.png")
await self._json("GET", self.URL)
shutil.rmtree(self.output_dir)
with self.assertLogs(self.api.logger, "WARNING") as logs:
for _ in range(2):
self._age()
status, data = await self._json("GET", self.URL)
self.assertEqual((status, data["total"]), (200, 1))
await self._settle()
warnings = [r for r in logs.records if r.levelname == "WARNING"]
self.assertEqual(len(warnings), 1, [r.getMessage() for r in warnings])
self.assertEqual(self.scan.calls, 1) # the probe failed before the scan
async def test_a_root_that_recovers_warns_again_on_its_next_failure(self):
make_png(self.output_dir / "a.png")
await self._json("GET", self.URL)
real = self.api._scan_gallery_files_sync
with self.assertLogs(self.api.logger, "WARNING") as logs:
for scan in (_raise, real, _raise):
self.api._scan_gallery_files_sync = scan
self._age()
await self._json("GET", self.URL)
await self._settle()
warnings = [r for r in logs.records if r.levelname == "WARNING"]
self.assertEqual(len(warnings), 2)
async def test_an_invalidation_during_a_walk_discards_it(self):
make_png(self.output_dir / "a.png")
gate = self.scan.hold_next()
reader = asyncio.ensure_future(self._json("GET", self.URL))
await self.scan.wait_entered()
make_png(self.output_dir / "b.png")
self.api.invalidate_gallery_cache()
gate.set()
status, data = await reader
self.assertEqual((status, data["total"]), (200, 2))
self.assertEqual(self.scan.calls, 2)
async def test_an_invalidation_from_a_worker_thread_drops_on_the_loop(self):
make_png(self.output_dir / "a.png")
await self._json("GET", self.URL)
record = self._record()
loop_thread = threading.get_ident()
threads = []
original = record.drop
def drop():
threads.append(threading.get_ident())
original()
record.drop = drop
loop = asyncio.get_running_loop()
await loop.run_in_executor(None, self.api.invalidate_gallery_cache)
# Dropped on the loop, before invalidate_gallery_cache returned.
self.assertEqual(threads, [loop_thread])
self.assertIsNone(record.files)
async def test_three_discarded_walks_are_a_500(self):
loop = asyncio.get_running_loop()
real = self.scan.real
async def invalidate():
self.api.invalidate_gallery_cache()
def scan(root):
asyncio.run_coroutine_threadsafe(invalidate(), loop).result(30)
return real(root)
self.scan.real = scan
status, data = await self._json("GET", self.URL)
self.assertEqual(status, 500)
self.assertFalse(data["success"])
self.assertEqual(self.scan.calls, 3)
async def test_a_disconnected_waiter_does_not_cancel_the_walk(self):
make_png(self.output_dir / "a.png")
gate = self.scan.hold_next()
waiter = asyncio.ensure_future(self.api._listing(self.output_dir))
await self.scan.wait_entered()
waiter.cancel()
with self.assertRaises(asyncio.CancelledError):
await waiter
flight = self._record().flight
self.assertIsNotNone(flight)
gate.set()
await asyncio.shield(flight)
self.assertEqual(len(self._record().files), 1)
self.assertIsNone(self._record().flight)
async def test_a_duplicate_delete_is_gone_on_the_very_next_read(self):
make_png(self.output_dir / "keep.png")
make_png(self.output_dir / "dupe.png")
_, before = await self._json("GET", self.URL)
self.assertEqual(before["total"], 2)
status, data = await self._json(
"POST",
"/prompt_manager/delete_duplicate_images",
json={"image_paths": ["dupe.png"]},
)
self.assertEqual((status, data["deleted_count"]), (200, 1))
_, after = await self._json("GET", self.URL)
self.assertEqual([image["filename"] for image in after["images"]], ["keep.png"])
async def test_response_keys_are_unchanged(self):
make_png(self.output_dir / "sub" / "a.png")
_, legacy = await self._json("GET", self.URL)
self.assertEqual(
set(legacy),
{"success", "images", "total", "offset", "limit", "has_more"},
)
_, workspace = await self._json("GET", self.URL + "?mode=workspace")
self.assertEqual(
set(workspace),
{
"success",
"images",
"total",
"positive_count",
"limit",
"has_more",
"next_cursor",
},
)
_, folders = await self._json("GET", "/prompt_manager/gallery/subfolders")
self.assertEqual(set(folders), {"success", "subfolders"})
_, scoped = await self._json(
"GET", "/prompt_manager/gallery/subfolders?mode=workspace"
)
self.assertEqual(set(scoped), {"success", "subfolders", "roots"})
self.assertEqual(set(scoped["roots"][0]), {"index", "label", "subfolders"})
_, rescan = await self._json("POST", "/prompt_manager/gallery/rescan")
self.assertEqual(set(rescan), {"success", "total", "roots"})
async def test_subfolders_with_ancestors_come_from_the_listing(self):
make_png(self.output_dir / "a" / "b" / "c.png")
_, plain = await self._json("GET", "/prompt_manager/gallery/subfolders")
_, nested = await self._json(
"GET", "/prompt_manager/gallery/subfolders?include_ancestors=true"
)
self.assertEqual(plain["subfolders"], [str(Path("a") / "b")])
self.assertEqual(
sorted(nested["subfolders"]), sorted(["a", str(Path("a") / "b")])
)
self.assertEqual(self.scan.calls, 1)
async def test_warm_up_walks_every_root_before_the_first_read(self):
second = self.comfy_dir / "second"
make_png(second / "x.png")
make_png(self.output_dir / "a.png")
self.api._get_all_output_dirs = lambda: [self.output_dir, second]
await self.api.warm_gallery_listings(web.Application())
self.assertEqual(len(self.api._listing_tasks), 2)
await self._settle()
self.assertEqual(self.scan.calls, 2)
_, data = await self._json("GET", self.URL)
self.assertEqual(data["total"], 2)
self.assertEqual(self.scan.calls, 2)
async def test_warm_up_never_fails_startup(self):
def broken():
raise RuntimeError("no roots")
self.api._get_all_output_dirs = broken
with self.assertLogs(self.api.logger, "WARNING"):
await self.api.warm_gallery_listings(web.Application())
self.assertEqual(self.api._listing_tasks, set())
async def test_warm_up_failure_is_retried_by_the_first_read(self):
make_png(self.output_dir / "a.png")
self.api._scan_gallery_files_sync = _raise
with self.assertLogs(self.api.logger, "WARNING"):
await self.api.warm_gallery_listings(web.Application())
await self._settle()
self.api._scan_gallery_files_sync = self.scan
status, data = await self._json("GET", self.URL)
self.assertEqual((status, data["total"]), (200, 1))
async def test_cleanup_cancels_running_walks(self):
make_png(self.output_dir / "a.png")
gate = self.scan.hold_next()
await self.api.warm_gallery_listings(web.Application())
await self.scan.wait_entered()
tasks = set(self.api._listing_tasks)
flight = self._record().flight
await self.api.cleanup_gallery_listings(web.Application())
gate.set()
self.assertTrue(all(task.cancelled() for task in tasks))
self.assertEqual(self.api._listing_tasks, set())
self.assertTrue(flight.done())
self.assertIsNone(self._record().flight)
self.assertIsNone(self._record().files)
async def test_cleanup_resolves_a_walk_cancelled_before_it_ran(self):
await self.api.warm_gallery_listings(web.Application())
flight = self._record().flight
waiter = asyncio.ensure_future(asyncio.shield(flight))
# No await between the warm-up and the cleanup: the task never ran.
await self.api.cleanup_gallery_listings(web.Application())
outcome = await waiter
self.assertEqual(outcome.kind, gallery_listing.FAILED)
self.assertIsNone(self._record().flight)
self.assertEqual(self.scan.calls, 0)
async def test_a_stale_read_starts_its_refresh_after_a_delay(self):
# The read is served from the old listing; the walk starts
# REFRESH_DELAY later so it does not share the GIL with the page's
# own requests (#348 live test: 2.2 s reads during the walk).
make_png(self.output_dir / "a.png")
await self._json("GET", self.URL)
self._age()
with patch.object(gallery_listing, "REFRESH_DELAY", 30):
status, data = await self._json("GET", self.URL)
await self._json("GET", self.URL) # still one pending refresh
record = self._record()
handle = record.rewalk_handle
self.assertEqual((status, data["total"]), (200, 1))
self.assertIsNotNone(handle)
self.assertIsNone(record.flight)
self.assertEqual(self.scan.calls, 1)
self.assertAlmostEqual(
handle.when() - asyncio.get_running_loop().time(), 30, delta=2
)
handle.cancel()
record.rewalk_handle = None
async def test_the_delayed_refresh_walks_once_when_it_fires(self):
make_png(self.output_dir / "a.png")
await self._json("GET", self.URL)
make_png(self.output_dir / "b.png")
self._age()
await self._json("GET", self.URL)
await self._json("GET", self.URL)
await self._settle()
self.assertEqual(self.scan.calls, 2)
self.assertIsNone(self._record().rewalk_handle)
_, data = await self._json("GET", self.URL)
self.assertEqual(data["total"], 2)
async def test_a_walk_started_meanwhile_cancels_the_delayed_refresh(self):
make_png(self.output_dir / "a.png")
await self._json("GET", self.URL)
self._age()
with patch.object(gallery_listing, "REFRESH_DELAY", 30):
await self._json("GET", self.URL)
handle = self._record().rewalk_handle
status, _ = await self._json("POST", "/prompt_manager/gallery/rescan")
self.assertEqual(status, 200)
self.assertTrue(handle.cancelled())
self.assertIsNone(self._record().rewalk_handle)
self.assertEqual(self.scan.calls, 2)
async def test_cleanup_cancels_a_delayed_refresh(self):
make_png(self.output_dir / "a.png")
await self._json("GET", self.URL)
self._age()
with patch.object(gallery_listing, "REFRESH_DELAY", 30):
await self._json("GET", self.URL)
handle = self._record().rewalk_handle
await self.api.cleanup_gallery_listings(web.Application())
self.assertTrue(handle.cancelled())
self.assertIsNone(self._record().rewalk_handle)
self.assertEqual(self.scan.calls, 1)
WORKSPACE_URL = "/prompt_manager/images/output?mode=workspace&limit=1"
async def _warmed(self):
await self.api.warm_gallery_listings(web.Application())
await self._settle()
def _workspace_ready(self):
"""True when a first workspace page needs no per-file rebuild."""
files = self._record().files
favorites = self.api._favorite_listing_cache.get(str(self.output_dir))
listed = self.api._workspace_listed_cache.get((0, str(self.output_dir)))
return (
favorites is not None
and favorites[0] is files
and listed is not None
and listed[0] is files
and bool(self.api._workspace_order_cache)
)
async def test_warm_up_also_builds_the_workspace_caches(self):
# The first open after a restart paid ~3 s building these (#348 live).
make_png(self.output_dir / "a.png")
make_png(self.output_dir / "sub" / "b.png")
await self._warmed()
self.assertTrue(self._workspace_ready())
with patch.object(images_module, "_listed", side_effect=AssertionError):
status, data = await self._json("GET", self.WORKSPACE_URL)
self.assertEqual((status, data["total"]), (200, 2))
async def test_a_changed_refresh_rebuilds_the_workspace_caches_behind(self):
make_png(self.output_dir / "a.png")
await self._warmed()
make_png(self.output_dir / "b.png")
self._age()
await self._json("GET", self.URL) # stale: one background walk
await self._settle()
self.assertEqual(len(self._record().files), 2)
self.assertTrue(self._workspace_ready())
async def test_an_unchanged_refresh_does_not_rebuild_them(self):
make_png(self.output_dir / "a.png")
await self._warmed()
with patch.object(
self.api, "_build_workspace_caches", side_effect=AssertionError
):
self._age()
await self._json("GET", self.URL)
await self._settle()
self.assertTrue(self._workspace_ready())
async def test_without_the_warm_up_reads_start_no_workspace_build(self):
make_png(self.output_dir / "a.png")
await self._json("GET", self.URL)
await self._settle()
self.assertIsNone(self.api._workspace_warm_task)
self.assertEqual(self.api._favorite_listing_cache, {})
async def test_a_failed_workspace_build_warns_and_reads_still_work(self):
make_png(self.output_dir / "a.png")
with (
patch.object(
self.api, "_select_workspace_page", side_effect=RuntimeError("boom")
),
self.assertLogs(self.api.logger, "WARNING") as logs,
):
await self._warmed()
self.assertTrue(any("Images page" in line for line in logs.output))
status, data = await self._json("GET", self.URL)
self.assertEqual((status, data["total"]), (200, 1))
async def test_walks_during_a_workspace_build_coalesce_into_one_more(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)
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"): # two changed walks while it builds
make_png(self.output_dir / name)
self._age()
await self._json("GET", self.URL)
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, 3])
self.assertEqual(self.scan.calls, 3) # two stored walks, one more build
self.assertTrue(self._workspace_ready())
async def test_a_build_that_cannot_list_the_roots_only_warns(self):
make_png(self.output_dir / "a.png")
await self._warmed()
def broken():
raise RuntimeError("no roots")
self.api._get_all_output_dirs = broken
with self.assertLogs(self.api.logger, "WARNING") as logs:
self.api._schedule_workspace_warm()
task = self.api._workspace_warm_task
await asyncio.wait([task])
self.assertIsNone(task.exception())
self.assertTrue(any("Images page" in line for line in logs.output))
async def test_cleanup_stops_a_running_workspace_build(self):
make_png(self.output_dir / "a.png")
release, entered = threading.Event(), threading.Event()
def build(_listings):
entered.set()
release.wait(_HANG_LIMIT)
self.addCleanup(release.set)
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)
task = self.api._workspace_warm_task
await self.api.cleanup_gallery_listings(web.Application())
release.set()
self.assertTrue(task.done())
self.assertIsNone(self.api._workspace_warm_task)
class TestGalleryListingHostHooks(unittest.IsolatedAsyncioTestCase):
async def test_root_registers_the_warm_up_and_its_cleanup(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_hooks_pkg",
str(Path(directory) / "db.sqlite"),
directory,
server,
) 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)
+38
View File
@@ -171,6 +171,44 @@ class TestWorkspaceImages(ImageAPITestCase):
self.assertIsNone(second["next_cursor"])
self.assertEqual(second["total"], 4)
async def test_an_unchanged_refresh_keeps_the_listing_and_derived_caches(self):
from py.api import gallery_listing
self.media("a.png", 300)
self.media("b.png", 200)
_, first = await self.get()
response = await self.client.put(
"/prompt_manager/images/favorite",
json={"root_index": 0, "path": "a.png"},
)
self.assertEqual(response.status, 200)
record = self.api._root_listings[str(self.output_dir)]
files = record.files
key = ((0, str(self.output_dir)),)
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
with patch.object(
self.api,
"_scan_gallery_files_sync",
wraps=self.api._scan_gallery_files_sync,
) as scan:
with patch.object(gallery_listing, "REFRESH_DELAY", 0):
await self.get() # served at once; one refresh walk follows
while record.rewalk_handle is not None:
await asyncio.sleep(0.01)
while self.api._listing_tasks:
await asyncio.gather(*list(self.api._listing_tasks))
# A finished walk leaves the set in a done callback, which needs
# a loop turn; gather over finished tasks does not yield one.
await asyncio.sleep(0)
_, again = await self.get()
self.assertEqual(scan.call_count, 1)
self.assertIs(record.files, files)
self.assertIs(self.api._workspace_order_cache[key][1]["newest"], newest)
self.assertIs(self.api._favorite_listing_cache[str(self.output_dir)], favorites)
self.assertEqual(again["total"], first["total"])
async def test_filters_apply_before_limit_and_sorts_are_complete(self):
self.media("folder/Zebra.png", 200)
self.media("folder/apple.png", 300)
+185
View File
@@ -0,0 +1,185 @@
"""Pure helpers of the shared Images listing (py/api/gallery_listing.py, #348)."""
import os
import sys
import tempfile
import unittest
from pathlib import Path, PurePath, PureWindowsPath
from unittest.mock import patch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from py.api import gallery_listing # noqa: E402
from py.api.gallery_listing import ( # noqa: E402
DISCARDED,
FAILED,
STORED,
RootListing,
WalkOutcome,
folder_sets,
same_listing,
walk_root,
)
def _legacy_subfolders(root, files, include_ancestors):
"""``_root_subfolders`` as it was before #348, the reference for folder_sets."""
subfolders = set()
for f, _mtime in files:
rel_dir = str(f.relative_to(root).parent)
if rel_dir and rel_dir != ".":
subfolders.add(rel_dir)
if include_ancestors:
ancestors = set()
for folder in subfolders:
parts = folder.replace("\\", "/").split("/")
for i in range(1, len(parts)):
ancestors.add("/".join(parts[:i]))
subfolders.update(ancestors)
return subfolders
class FolderSetsTests(unittest.TestCase):
def test_nested_folders_match_the_previous_subfolder_listing(self):
root = Path("/out")
files = [
(root / "top.png", 1.0),
(root / "a" / "one.png", 2.0),
(root / "a" / "b" / "c" / "two.png", 3.0),
(root / "d" / "three.png", 4.0),
]
folders, with_ancestors = folder_sets(root, files)
self.assertEqual(folders, _legacy_subfolders(root, files, False))
self.assertEqual(with_ancestors, _legacy_subfolders(root, files, True))
self.assertIn("a/b", with_ancestors)
self.assertNotIn("a/b", folders)
self.assertIsInstance(folders, frozenset)
self.assertIsInstance(with_ancestors, frozenset)
def test_windows_separators_keep_the_previous_spelling(self):
root = PureWindowsPath("C:/out")
files = [
(root / "x" / "y" / "z.png", 1.0),
(root / "w.png", 2.0),
]
folders, with_ancestors = folder_sets(root, files)
self.assertEqual(folders, {"x\\y"})
self.assertEqual(with_ancestors, {"x\\y", "x"})
self.assertEqual(folders, _legacy_subfolders(root, files, False))
self.assertEqual(with_ancestors, _legacy_subfolders(root, files, True))
def test_an_empty_listing_has_no_folders(self):
self.assertEqual(folder_sets(Path("/out"), []), (frozenset(), frozenset()))
def test_folders_come_from_the_path_strings_not_relative_to(self):
# relative_to per file was half of a 20,000-file walk (#348 live test).
root = Path("/out")
files = [(root / "a" / "b" / "one.png", 1.0), (root / "two.png", 2.0)]
with patch.object(PurePath, "relative_to", side_effect=AssertionError):
folders, with_ancestors = folder_sets(root, files)
self.assertEqual(folders, {os.path.join("a", "b")})
self.assertEqual(with_ancestors, {os.path.join("a", "b"), "a"})
def test_a_path_not_spelled_like_the_root_falls_back_to_relative_to(self):
# Windows folds case: the string prefix misses, relative_to matches.
root = PureWindowsPath("C:/Out")
files = [(PureWindowsPath("c:/out/x/y.png"), 1.0)]
folders, with_ancestors = folder_sets(root, files)
self.assertEqual(folders, {"x"})
self.assertEqual(folders, _legacy_subfolders(root, files, False))
self.assertEqual(with_ancestors, _legacy_subfolders(root, files, True))
class SameListingTests(unittest.TestCase):
def test_equal_path_and_mtime_lists_are_the_same(self):
old = [(Path("/o/a.png"), 1.0), (Path("/o/b.png"), 2.0)]
self.assertTrue(same_listing(old, list(old)))
self.assertTrue(same_listing(old, old))
def test_a_changed_mtime_path_or_length_differs(self):
old = [(Path("/o/a.png"), 1.0), (Path("/o/b.png"), 2.0)]
self.assertFalse(same_listing(old, [(Path("/o/a.png"), 1.5), old[1]]))
self.assertFalse(same_listing(old, [(Path("/o/c.png"), 1.0), old[1]]))
self.assertFalse(same_listing(old, old[:1]))
self.assertFalse(same_listing(None, old))
class WalkRootTests(unittest.TestCase):
def setUp(self):
self._tmp = tempfile.TemporaryDirectory()
self.addCleanup(self._tmp.cleanup)
self.root = Path(self._tmp.name) / "output"
(self.root / "sub").mkdir(parents=True)
def _scan(self, files):
calls = []
def scan(root):
calls.append(root)
return list(files)
return scan, calls
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)
self.assertEqual(calls, [self.root])
self.assertEqual(listed, files)
self.assertEqual(folders, frozenset({"sub"}))
self.assertEqual(with_ancestors, frozenset({"sub"}))
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)
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)
self.assertIsNot(listed, previous)
self.assertEqual(listed, files)
def test_a_deleted_root_raises_instead_of_listing_nothing(self):
scan, calls = self._scan([])
missing = self.root / "gone"
with self.assertRaises(OSError):
walk_root(scan, missing, None)
self.assertEqual(calls, [])
class RecordTests(unittest.TestCase):
def test_drop_clears_the_listing_and_bumps_the_generation(self):
record = RootListing()
record.files = [(Path("/o/a.png"), 1.0)]
record.folders = frozenset({"x"})
record.folders_with_ancestors = frozenset({"x"})
record.walked_at = 5.0
record.drop()
record.drop()
self.assertIsNone(record.files)
self.assertIsNone(record.folders)
self.assertIsNone(record.folders_with_ancestors)
self.assertIsNone(record.walked_at)
self.assertEqual(record.generation, 2)
def test_outcomes(self):
error = OSError("gone")
self.assertEqual(WalkOutcome.stored().kind, STORED)
self.assertEqual(WalkOutcome.discarded().kind, DISCARDED)
failed = WalkOutcome.failed(error)
self.assertEqual(failed.kind, FAILED)
self.assertIs(failed.error, error)
def test_refresh_constants(self):
self.assertEqual(gallery_listing.REFRESH_AFTER, 30)
self.assertEqual(gallery_listing.REFRESH_DELAY, 1.0)
self.assertIsNone(RootListing().rewalk_handle)
if __name__ == "__main__":
unittest.main()
+8
View File
@@ -85,6 +85,14 @@ class TestConvertOutdated(ConvertCase):
)
self.assertEqual(self.api._jobs.list(), [])
async def test_conversion_keeps_the_listing(self):
self._legacy("a")
await self.client.get("/prompt_manager/images/output?mode=workspace")
files = self.api._root_listings[str(self.output_dir)].files
_summary, job = await self._convert()
self.assertEqual(job.state, "finished")
self.assertIs(self.api._root_listings[str(self.output_dir)].files, files)
async def test_converts_from_the_source_and_removes_the_old_preview(self):
old = self._legacy("a")
original = (self.output_dir / "a.png").read_bytes()
+14
View File
@@ -345,6 +345,20 @@ class TestThumbnailJob(ThumbnailJobCase):
sizes[quality] = thumb.stat().st_size
self.assertLess(sizes["low"], sizes["high"])
async def test_a_job_keeps_the_listing_and_the_next_page_shows_the_preview(self):
make_png(self.output_dir / "a.png", size=(64, 64))
url = "/prompt_manager/images/output?mode=workspace"
before = (await (await self.client.get(url)).json())["images"][0]
files = self.api._root_listings[str(self.output_dir)].files
status, body = await self._start(mode="missing")
self.assertEqual(status, 202)
await self._finished(body["job"]["id"])
# Previews are not part of the listing (#348): no rewalk is forced.
self.assertIs(self.api._root_listings[str(self.output_dir)].files, files)
after = (await (await self.client.get(url)).json())["images"][0]
self.assertNotIn("v=", before["thumbnail_url"])
self.assertIn("v=", after["thumbnail_url"])
async def test_no_output_folder_fails_the_job(self):
self.api._get_all_output_dirs = lambda: []
_, body = await self._start()
+81
View File
@@ -0,0 +1,81 @@
// Maintainer tool, run by hand (never in CI): how long the workspace Images
// page takes to show its first thumbnail (#348, TESTING.md "Images listing
// cache (#348)").
//
// Once: npm ci --prefix tools/vendor
// npx --prefix tools/vendor playwright install chromium
// Then: node tools/vendor/perf/images_load_probe.mjs http://127.0.0.1:8188
//
// Prints the time to the first loaded thumbnail, the time to network idle and
// the 10 slowest requests. Each run opens a fresh browser context (no cache).
import { chromium } from '@playwright/test';
const SLOWEST = 10;
const TIMEOUT_MS = 120_000;
function usage() {
console.error('usage: node tools/vendor/perf/images_load_probe.mjs <base-url>');
console.error('example: node tools/vendor/perf/images_load_probe.mjs http://127.0.0.1:8188');
process.exit(2);
}
function parseBase(argument) {
if (!argument) usage();
let url;
try {
url = new URL(argument);
} catch {
usage();
}
if (url.protocol !== 'http:' && url.protocol !== 'https:') usage();
return url.origin + url.pathname.replace(/\/+$/, '');
}
const ms = value => `${Math.round(value)} ms`;
async function probe(base) {
const browser = await chromium.launch();
try {
const context = await browser.newContext({ viewport: { width: 1440, height: 900 } });
const page = await context.newPage();
const requests = [];
page.on('requestfinished', request => {
const timing = request.timing();
requests.push({
url: request.url().replace(base, ''),
status: null,
ms: timing.responseEnd >= 0 ? timing.responseEnd : NaN,
response: request.response(),
});
});
const started = Date.now();
await page.goto(`${base}/prompt_manager/workspace#/images`, { timeout: TIMEOUT_MS });
// The first Images thumbnail the browser has decoded.
await page.waitForFunction(() => [...document.images].some(image =>
image.complete && image.naturalWidth > 0
&& image.src.includes('/prompt_manager/images/serve/')), null,
{ timeout: TIMEOUT_MS, polling: 'raf' });
const firstThumbnail = Date.now() - started;
await page.waitForLoadState('networkidle', { timeout: TIMEOUT_MS });
const networkIdle = Date.now() - started;
for (const entry of requests) {
entry.status = (await entry.response)?.status() ?? 'n/a';
}
const slowest = requests.filter(entry => Number.isFinite(entry.ms))
.sort((a, b) => b.ms - a.ms).slice(0, SLOWEST);
console.log(`first thumbnail: ${ms(firstThumbnail)}`);
console.log(`network idle: ${ms(networkIdle)}`);
console.log(`requests: ${requests.length}`);
console.log(`slowest ${slowest.length}:`);
for (const entry of slowest) {
console.log(` ${ms(entry.ms).padStart(9)} ${String(entry.status).padEnd(3)} ${entry.url}`);
}
} finally {
await browser.close();
}
}
probe(parseBase(process.argv[2])).catch(error => {
console.error(`probe failed: ${error.message}`);
process.exit(1);
});