From b7b99004c7a0af6cc5f0896dd4b426ad6a6cae45 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Mon, 26 Oct 2026 10:14:31 -0700 Subject: [PATCH] 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. --- py/api/__init__.py | 13 ++++ py/api/images.py | 67 +++++++++++++++++- tests/test_api_images.py | 145 ++++++++++++++++++++++++++++++++++++++- 3 files changed, 222 insertions(+), 3 deletions(-) diff --git a/py/api/__init__.py b/py/api/__init__.py index 2d137763..c9b95ced 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -364,6 +364,11 @@ class PromptManagerAPI( 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 = {} @@ -531,6 +536,9 @@ class PromptManagerAPI( 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) @@ -539,6 +547,11 @@ class PromptManagerAPI( async def cleanup_gallery_listings(self, app=None): """Stop the background listing walks (aiohttp on_cleanup).""" + self._workspace_warm_enabled = False + 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() diff --git a/py/api/images.py b/py/api/images.py index 87017bc7..0229a3d6 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -1309,6 +1309,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 + listings = self._stored_listings() + if listings is None: + return # a root has no listing yet; its walk asks again + try: + 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). @@ -1419,12 +1473,15 @@ class ImageRoutesMixin: 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: @@ -2883,9 +2940,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 diff --git a/tests/test_api_images.py b/tests/test_api_images.py index f4a22743..e3b6af69 100644 --- a/tests/test_api_images.py +++ b/tests/test_api_images.py @@ -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__)))) @@ -2084,6 +2084,22 @@ class TestOutputScanHelpers(ImageAPITestCase): 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") @@ -2494,6 +2510,10 @@ class TestSharedListing(ImageAPITestCase): def _record(self): return self.api._root_listings[str(self.output_dir)] + 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 @@ -2504,11 +2524,13 @@ class TestSharedListing(ImageAPITestCase): """Wait for every background walk this API started.""" async def settle(): - while self.api._listing_tasks: + while self.api._listing_tasks or self._warming(): 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) @@ -2832,6 +2854,125 @@ class TestSharedListing(ImageAPITestCase): self.assertIsNone(self._record().flight) self.assertEqual(self.scan.calls, 0) + 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.assertTrue(self._workspace_ready()) + + 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):