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

@@ -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()

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

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):