From b78152744acedd1588084e06f194c3c199ab0041 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Wed, 19 Aug 2026 20:01:28 -0400 Subject: [PATCH 1/2] fix: load an album's index off the event loop MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `Embeddings.indexes` and `open_cached_embeddings` are a full `np.load` of the index whenever the three-entry `lru_cache` misses — which is any request touching a fourth album, and the first request after the index is rewritten. The load unpickles the per-image metadata object array, copies four arrays and sorts one. Measured at 0.34s for a 50,000-image index with modest metadata and a warm page cache; a real library with InvokeAI generation metadata and a cold cache is several times that. Six async endpoints did that inside the coroutine, so the whole server stopped for the duration — the slideshow, thumbnails, and the indexing-progress polling the user is watching while they wait: umap.py /umap_data umap_embeddings + open_cached_embeddings search.py /image_info .indexes /get_metadata .indexes /lookup_image_indices, /get_image_by_name invoke.py /recall_parameters, /use_ref_image (via the sync helpers) index.py /index_metadata open_cached_embeddings `load_indexes()` and `load_cached_embeddings()` do it in a thread; the two `invoke.py` helpers stay synchronous and hop at their call sites. Moving reads into threads makes one existing race reachable, so it is fixed here rather than left behind. The delete path clears the index cache, writes, then re-primes "to verify the write" — but a reader that missed the cache before the clear can finish loading at any point after it and store the pre-delete snapshot, in which case the re-prime is a cache hit that verifies nothing and every later request keeps serving a deleted image. Clearing again immediately before the re-prime makes it a real load, matching what `update_image_path` already does. Not included: the delete endpoints themselves still run their load-and-rewrite on the loop. That one cannot simply be threaded — the event loop is currently what serializes concurrent deletes, and moving it without a per-index lock would let two deletes race and lose one. Co-Authored-By: Claude Opus 5 (1M context) --- photomap/backend/embeddings.py | 31 +++++- photomap/backend/routers/index.py | 5 +- photomap/backend/routers/invoke.py | 13 ++- photomap/backend/routers/search.py | 8 +- photomap/backend/routers/umap.py | 4 +- tests/backend/test_event_loop_blocking.py | 122 ++++++++++++++++++++++ 6 files changed, 173 insertions(+), 10 deletions(-) create mode 100644 tests/backend/test_event_loop_blocking.py diff --git a/photomap/backend/embeddings.py b/photomap/backend/embeddings.py index bad78669..7444e59a 100644 --- a/photomap/backend/embeddings.py +++ b/photomap/backend/embeddings.py @@ -1890,6 +1890,11 @@ def indexes(self) -> dict[str, np.ndarray]: """ Load all indexes from the embeddings file. + Blocking, and not always cheaply: on a cache miss this is a full + ``np.load`` of the index — including unpickling the per-image + metadata object array — plus copies and a sort. Async callers must + use :meth:`load_indexes`. + Returns: Dict[str, np.ndarray]: Dictionary containing all indexes. """ @@ -2144,6 +2149,24 @@ def retrieve_image( len(sorted_filenames), ) + async def load_indexes(self) -> dict[str, np.ndarray]: + """:attr:`indexes` off the event loop. + + The cache behind it holds three entries, so any request touching a + fourth album — or the first request after an index is rewritten — + pays the full load. Measured at 0.34s for a 50,000-image index with + modest metadata, page cache warm; a large library with real + generation metadata and a cold cache is several times that, and the + whole server is stopped for the duration. + """ + return await asyncio.to_thread(lambda: self.indexes) + + async def load_cached_embeddings(self) -> dict[str, Any]: + """:meth:`open_cached_embeddings` for this index, off the event loop.""" + return await asyncio.to_thread( + self.open_cached_embeddings, self.embeddings_path + ) + def remove_image_from_embeddings(self, index: int) -> None: """ Remove an image from the embeddings file. @@ -2208,7 +2231,13 @@ def remove_images_from_embeddings(self, indices: list[int]) -> None: **extras, ) - # 6. Re-prime the cache immediately to verify the write + # 6. Re-prime the cache to verify the write. Cleared again first: + # a reader that missed the cache before step 4 can finish its load + # at any point after it and store the pre-delete snapshot, and + # then this would be a cache hit that verifies nothing and goes on + # serving images that are no longer in the index. Readers run in + # worker threads, so this is not hypothetical. + _open_npz_file.cache_clear() _open_npz_file(self.embeddings_path) except Exception as e: diff --git a/photomap/backend/routers/index.py b/photomap/backend/routers/index.py index a7fee279..a825363b 100644 --- a/photomap/backend/routers/index.py +++ b/photomap/backend/routers/index.py @@ -4,6 +4,7 @@ It allows creating, deleting, and checking the existence of embeddings indices for albums. """ +import asyncio import logging import os import shutil @@ -285,7 +286,9 @@ async def index_metadata(album_config: AlbumDep) -> EmbeddingsIndexMetadata: marker = index_path.parent / LAST_UPDATED_FILENAME if marker.exists(): last_modified = max(last_modified, marker.stat().st_mtime) - filenames = Embeddings.open_cached_embeddings(index_path)["filenames"] + filenames = ( + await asyncio.to_thread(Embeddings.open_cached_embeddings, index_path) + )["filenames"] filename_count = len(filenames) video_count = sum(1 for f in filenames if is_video(Path(str(f)))) diff --git a/photomap/backend/routers/invoke.py b/photomap/backend/routers/invoke.py index feeb1d30..6559ae82 100644 --- a/photomap/backend/routers/invoke.py +++ b/photomap/backend/routers/invoke.py @@ -24,6 +24,7 @@ from __future__ import annotations +import asyncio import logging import mimetypes import re @@ -544,7 +545,9 @@ async def recall_parameters(request: RecallRequest) -> dict: ), ) - raw_metadata = _load_raw_metadata(request.album_key, request.index) + raw_metadata = await asyncio.to_thread( + _load_raw_metadata, request.album_key, request.index + ) payload = _build_recall_payload(raw_metadata, include_seed=request.include_seed) if not payload: raise HTTPException( @@ -699,7 +702,9 @@ async def use_ref_image(request: UseRefImageRequest) -> dict: ), ) - image_path = _load_image_path(request.album_key, request.index) + image_path = await asyncio.to_thread( + _load_image_path, request.album_key, request.index + ) if not image_path.is_file(): raise HTTPException( status_code=404, detail=f"Image file not found on disk: {image_path.name}" @@ -710,7 +715,9 @@ async def use_ref_image(request: UseRefImageRequest) -> dict: # still matches InvokeAI's ``{uuid}.{ext}`` convention, or the PNG # carries Invoke generation metadata. Loading the metadata here is the # same lookup used by /invokeai/recall, so it's cheap and local. - raw_metadata = _load_raw_metadata(request.album_key, request.index) + raw_metadata = await asyncio.to_thread( + _load_raw_metadata, request.album_key, request.index + ) filename_matches = _looks_like_invoke_filename(image_path.name) metadata_matches = _has_invoke_metadata(raw_metadata) should_probe = filename_matches or metadata_matches diff --git a/photomap/backend/routers/search.py b/photomap/backend/routers/search.py index eab84526..9072db55 100644 --- a/photomap/backend/routers/search.py +++ b/photomap/backend/routers/search.py @@ -211,7 +211,7 @@ async def image_info( embeddings: EmbeddingsDep, ) -> ImageData: """Retrieve basic metadata on an image.""" - data = embeddings.indexes + data = await embeddings.load_indexes() sorted_filenames = data["sorted_filenames"] filename_map = data["filename_map"] modification_times = data["sorted_modification_times"] @@ -238,7 +238,7 @@ async def get_metadata(album_key: str, index: int, embeddings: EmbeddingsDep): """ Download the JSON-formatted metadata for an image by album key and index. """ - indexes = embeddings.indexes + indexes = await embeddings.load_indexes() metadata = indexes["sorted_metadata"] if index < 0 or index >= len(metadata): raise HTTPException(status_code=404, detail="Index out of range") @@ -664,7 +664,7 @@ async def lookup_image_indices( as clickable thumbnails. Filenames not found in the album map to ``null``. Duplicate basenames in the album resolve to the first matching index. """ - sorted_filenames = embeddings.indexes["sorted_filenames"] + sorted_filenames = (await embeddings.load_indexes())["sorted_filenames"] basename_to_index: dict[str, int] = {} for idx, full_path in enumerate(sorted_filenames): basename = Path(full_path).name @@ -692,7 +692,7 @@ async def get_image_by_name( if Path(filename).suffix.lower() not in SUPPORTED_EXTENSIONS: raise HTTPException(status_code=403, detail="Unsupported image type") - indexes = embeddings.indexes + indexes = await embeddings.load_indexes() # inefficient linear search for the filename, but still pretty quick! absolute_paths = [ x for x in indexes["sorted_filenames"] if Path(x).name == filename diff --git a/photomap/backend/routers/umap.py b/photomap/backend/routers/umap.py index 98df84ad..70cfad1a 100644 --- a/photomap/backend/routers/umap.py +++ b/photomap/backend/routers/umap.py @@ -62,7 +62,9 @@ async def get_umap_data( cluster_min_samples, ) - embeddings = embeddings.open_cached_embeddings(embeddings.embeddings_path) + # Threaded for the same reason as the coordinates above: on a cache miss + # this is a full np.load of the index, metadata unpickling included. + embeddings = await embeddings.load_cached_embeddings() filenames = embeddings["filenames"] filename_map = embeddings["filename_map"] diff --git a/tests/backend/test_event_loop_blocking.py b/tests/backend/test_event_loop_blocking.py new file mode 100644 index 00000000..7aef762f --- /dev/null +++ b/tests/backend/test_event_loop_blocking.py @@ -0,0 +1,122 @@ +"""Endpoints must not load an album's index on the event loop. + +``Embeddings.indexes`` and ``open_cached_embeddings`` are backed by an +``lru_cache`` holding three entries, so any request touching a fourth album — +or the first request after the index is rewritten — pays a full ``np.load`` +of the index, unpickling the per-image metadata array along the way. Measured +at 0.34s for 50,000 images with modest metadata and a warm page cache; a real +library with generation metadata and a cold cache is several times that. + +Doing that inside a coroutine stops every other request for the duration: +the slideshow, thumbnail fetches, and the indexing-progress polling that the +user is watching while they wait. + +``asyncio.get_running_loop()`` succeeds only on the thread running the loop, +which is what makes "did this happen on the event loop" checkable without +depending on thread names. +""" + +import asyncio +from pathlib import Path + +import pytest +from fixtures import build_index + +from photomap.backend import embeddings as embeddings_module + + +@pytest.fixture +def loop_thread_loads(monkeypatch): + """Record every index load, tagged with the kind of thread it ran on.""" + threads = [] + real = embeddings_module._open_npz_file.__wrapped__ + + def spy(path): + try: + asyncio.get_running_loop() + except RuntimeError: + threads.append("worker") + else: + threads.append("event-loop") + return real(path) + + # Patch under the cache so every miss is seen, and start from empty. + monkeypatch.setattr( + embeddings_module, + "_open_npz_file", + __import__("functools").lru_cache(maxsize=3)(spy), + ) + embeddings_module._open_npz_file.cache_clear() + return threads + + +@pytest.mark.parametrize( + "request_for", + [ + lambda key: ("GET", f"/umap_data/{key}"), + lambda key: ("GET", f"/image_info/{key}/0"), + lambda key: ("GET", f"/get_metadata/{key}/0"), + lambda key: ("GET", f"/index_metadata/{key}"), + ], + ids=["umap_data", "image_info", "get_metadata", "index_metadata"], +) +def test_endpoints_do_not_load_the_index_on_the_event_loop( + client, new_album, loop_thread_loads, request_for +): + build_index(client, new_album) + embeddings_module._open_npz_file.cache_clear() + loop_thread_loads.clear() + + method, url = request_for(new_album["key"]) + assert client.request(method, url).status_code == 200 + + assert loop_thread_loads, f"{url} never loaded the index; the test proves nothing" + assert "event-loop" not in loop_thread_loads + + +def test_a_concurrent_reader_cannot_restore_a_deleted_image( + client, new_album, monkeypatch +): + """The delete path re-primes the index cache to verify its own write. + + A reader that missed the cache before the delete cleared it can finish + loading at any point afterwards and store the pre-delete snapshot — and + then the re-prime is a cache hit that verifies nothing, so every later + request goes on serving an image that is no longer in the index. Now + that readers run in worker threads, that interleaving is reachable. + + Driven deterministically by priming the cache from the still-unmodified + file during the write itself, which is precisely what a reader that + loaded just before the rename would have left behind. + """ + build_index(client, new_album) + + from photomap.backend.routers.album import get_embeddings_for_album + + embeddings = get_embeddings_for_album(new_album["key"]) + path = embeddings.embeddings_path + before = len(embeddings.indexes["sorted_filenames"]) + + real_savez = embeddings_module.atomic_savez + raced = {"done": False} + + def savez_with_a_racing_reader(target, **arrays): + if not raced["done"] and Path(target) == Path(path): + raced["done"] = True + # The file on disk is still the pre-delete one right now. + embeddings_module._open_npz_file(path) + return real_savez(target, **arrays) + + monkeypatch.setattr(embeddings_module, "atomic_savez", savez_with_a_racing_reader) + + embeddings.remove_image_from_embeddings(0) + assert raced["done"], "the racing reader never ran; the test proves nothing" + + monkeypatch.undo() + assert ( + len(get_embeddings_for_album(new_album["key"]).indexes["sorted_filenames"]) + == before - 1 + ) + assert client.get(f"/index_metadata/{new_album['key']}").json()["image_count"] == ( + before - 1 + ) From e49b94f4075e45ea3e563aec87704a2e16b31416 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Fri, 21 Aug 2026 19:47:31 -0400 Subject: [PATCH 2/2] fix: dedupe and properly invalidate the index cache, and finish the sweep Follow-up to the previous commit, from an adversarial review of it. Three problems, all consequences of the same change: once index reads run in worker threads, the lru_cache underneath them is no longer protected by the event loop's serialization. 1. The delete race was not actually closed. Clearing the cache again before the re-prime narrows the window but does not remove it, because the re-prime is itself a full load: a reader that lands *during* it installs the pre-delete snapshot, and CPython's lru_cache then refuses to overwrite a key that appeared while the wrapped call ran, so the re-prime's fresh result is silently discarded. The cache goes on serving a deleted image indefinitely while the delete reports success. Reproduced against the real code before fixing. 2. Concurrent misses were no longer deduped. lru_cache does not collapse in-flight calls; the event loop used to, by never letting two overlap. Eight concurrent requests for an uncached 50,000-image index measured eight full loads and 1.7 GB peak RSS, against one load and 0.26 GB for a single request. Both are fixed by replacing the lru_cache with _NpzIndexCache: one load per path at a time, and a generation counter so a load stores its result only if the generation it started under is still current. Measured after: exactly one load per burst at 1, 2, 4, 8 and 16 concurrent readers, with peak RSS flat. 3. The sweep missed the endpoints that matter most. /retrieve_image/ is the slideshow's own endpoint and fires once per slide, so it is the one most likely to *take* the cold miss the previous commit set out to move off the loop; /thumbnails/ fires in a burst per grid. Those two, plus /image_path/, /video_frame/, /download_images_zip/, /search_with_text_and_image/ and /curate_sync, all still ran the same blocking load inside their coroutine. Search goes behind a semaphore as well as a thread, so that threading it does not newly allow two concurrent CLIP encodes on one GPU. Regression tests, each confirmed to fail without its fix: the narrow delete race, one-load-per-burst under eight concurrent readers, the six newly threaded endpoints, and an empty-index zip download (the priming call added to /download_images_zip/ would otherwise turn a missing index into a 500 where both loops below already tolerate it). 741 backend tests pass; ruff clean. Co-Authored-By: Claude Opus 5 (1M context) --- photomap/backend/embeddings.py | 155 ++++++++++++++++-- photomap/backend/routers/curation.py | 6 +- photomap/backend/routers/search.py | 85 +++++++--- tests/backend/test_event_loop_blocking.py | 181 +++++++++++++++++++++- 4 files changed, 380 insertions(+), 47 deletions(-) diff --git a/photomap/backend/embeddings.py b/photomap/backend/embeddings.py index 427743ae..a6d8b49a 100644 --- a/photomap/backend/embeddings.py +++ b/photomap/backend/embeddings.py @@ -8,13 +8,12 @@ """ import asyncio -import functools import gc import logging import os import threading import warnings -from collections import deque +from collections import OrderedDict, deque from collections.abc import Callable, Generator from collections.abc import Set as AbstractSet from concurrent.futures import ThreadPoolExecutor @@ -381,11 +380,12 @@ def _copy_non_per_image_keys(data: Any) -> dict[str, Any]: return {key: data[key].copy() for key in data.files if key not in _PER_IMAGE_KEYS} -@functools.lru_cache(maxsize=3) -def _open_npz_file(embeddings_path: Path) -> dict[str, Any]: - """ - Global helper to open .npz files with caching. - Uses context manager to ensure file handles are released. +def _load_npz_file(embeddings_path: Path) -> dict[str, Any]: + """Read one .npz index off disk and derive its sorted views. + + The uncached half of :data:`_open_npz_file` — call that instead unless + you specifically want to bypass the cache. Uses a context manager so + the file handle is released before we return. """ embeddings_path = Path(embeddings_path).resolve() @@ -434,6 +434,121 @@ def _open_npz_file(embeddings_path: Path) -> dict[str, Any]: } +class _NpzCacheInfo(NamedTuple): + """``functools.lru_cache``-shaped stats, so existing callers still read.""" + + hits: int + misses: int + maxsize: int + currsize: int + + +class _NpzIndexCache: + """Path-keyed cache of loaded ``.npz`` indexes. + + This was a ``functools.lru_cache`` until index reads moved onto worker + threads (see :meth:`Embeddings.load_indexes`). Two things that + ``lru_cache`` cannot express became load-bearing at that point: + + **One load per path at a time.** ``lru_cache`` does not dedupe in-flight + calls, so N concurrent requests for an uncached album each ran a full + ``np.load``. The event loop used to serialize them down to one; worker + threads do not. Measured on a 50,000-image / 111 MB index: eight + concurrent requests meant eight loads and 1.7 GB peak RSS, against one + load and 0.26 GB for a single request. Latecomers now wait on the + in-flight load and take its result. + + **Invalidation that outruns an in-flight load.** Whether a loaded + snapshot may be cached depends on when the load *started*, not on when + it finished: a reader that began before a rewrite is holding pre-rewrite + data no matter how late it returns. Every :meth:`cache_clear` bumps a + generation counter and a load stores its result only if the generation + it started under is still current. ``lru_cache`` got this backwards -- + it *refuses* to overwrite a key that appeared while the wrapped call was + running, so a stale reader that landed mid-reload won and the fresh + value was silently dropped, leaving deleted images served forever. + """ + + def __init__( + self, + loader: Callable[[Path], dict[str, Any]], + maxsize: int = 3, + ) -> None: + # Named ``__wrapped__`` for parity with the lru_cache this replaced: + # tests reach through it to instrument the real load. + self.__wrapped__ = loader + self._maxsize = maxsize + self._guard = threading.Lock() + self._entries: OrderedDict[Path, dict[str, Any]] = OrderedDict() + # One lock per index path. Never evicted, for the same reason as + # _UMAP_BUILD_LOCKS: one lock per album per process is nothing, and + # dropping one while a thread still held it would let the next + # caller build a second, unrelated lock and defeat the exclusion. + self._load_locks: dict[Path, threading.Lock] = {} + self._generation = 0 + self._hits = 0 + self._misses = 0 + + def __call__(self, embeddings_path: Path) -> dict[str, Any]: + # Keyed on the path as given, like the lru_cache this replaced. + # Resolving here would put a filesystem round-trip on every cache + # hit -- /retrieve_image/ takes one per slide, and album paths can + # be network mounts -- to dedupe spellings that no call site + # actually produces. Invalidation does not need it either: + # cache_clear() drops every entry, not one key. + key = embeddings_path + + with self._guard: + cached = self._entries.get(key) + if cached is not None: + self._entries.move_to_end(key) + self._hits += 1 + return cached + load_lock = self._load_locks.setdefault(key, threading.Lock()) + + # Held across the load so concurrent callers queue behind it rather + # than each running their own. cache_clear() deliberately does not + # take this lock: a writer must never wait on a reader. + with load_lock: + with self._guard: + cached = self._entries.get(key) + if cached is not None: + # Filled by whoever held the lock before us. + self._entries.move_to_end(key) + self._hits += 1 + return cached + generation = self._generation + self._misses += 1 + + data = self.__wrapped__(embeddings_path) + + with self._guard: + if generation == self._generation: + self._entries[key] = data + self._entries.move_to_end(key) + while len(self._entries) > self._maxsize: + self._entries.popitem(last=False) + # Otherwise the index was rewritten while we were reading it. + # Our caller still gets this snapshot -- it was valid when + # asked for -- but it must not outlive the request. + return data + + def cache_clear(self) -> None: + """Drop every entry, and disqualify every load already in flight.""" + with self._guard: + self._entries.clear() + self._generation += 1 + + def cache_info(self) -> _NpzCacheInfo: + with self._guard: + return _NpzCacheInfo( + self._hits, self._misses, self._maxsize, len(self._entries) + ) + + +_open_npz_file = _NpzIndexCache(_load_npz_file) + + class IndexResult(BaseModel): """ Result of an indexing operation. @@ -2238,12 +2353,20 @@ def remove_images_from_embeddings(self, indices: list[int]) -> None: **extras, ) - # 6. Re-prime the cache to verify the write. Cleared again first: - # a reader that missed the cache before step 4 can finish its load - # at any point after it and store the pre-delete snapshot, and - # then this would be a cache hit that verifies nothing and goes on - # serving images that are no longer in the index. Readers run in - # worker threads, so this is not hypothetical. + # 6. Clear once more, then re-prime to verify the write. + # + # The second clear is the load-bearing one, and it has to come + # *after* the rename: a reader that missed the cache before step 4 + # is holding the pre-delete snapshot, and bumping the generation + # here is what stops it from caching that snapshot whenever it + # happens to finish. Readers run in worker threads now, so that + # interleaving is reachable; without this the index goes on + # serving images that are no longer in it. + # + # Clearing alone is not enough, which is why _NpzIndexCache tracks + # generations rather than just emptying a dict: the stale reader + # can finish *during* the re-prime below, and a plain cache would + # take its value and discard the fresh one. _open_npz_file.cache_clear() _open_npz_file(self.embeddings_path) @@ -2317,8 +2440,10 @@ def update_image_path(self, index: int, new_path: Path) -> None: logger.error(f"Failed to update image path in embeddings: {e}") raise - # Re-clear after the write so any reader that primed the cache mid-flight - # is also invalidated. + # Re-clear after the write. As in remove_images_from_embeddings, this + # is what disqualifies a reader still loading the pre-update file -- + # it bumps the cache generation, so that reader's result is dropped + # rather than cached, whenever it lands. _open_npz_file.cache_clear() # This is not used in the current implementation, but can be useful for testing. diff --git a/photomap/backend/routers/curation.py b/photomap/backend/routers/curation.py index 3ade9d43..111e348a 100644 --- a/photomap/backend/routers/curation.py +++ b/photomap/backend/routers/curation.py @@ -1,3 +1,4 @@ +import asyncio import logging import os import random @@ -251,7 +252,10 @@ async def run_curation_sync(request: CurationRequest): try: _validate_curation_request(request) logger.info(f"Curation: Running {request.method.upper()} x{request.iterations}...") - return _compute_curation(request) + # Off the loop: this is N rounds of kmeans/FPS over the whole index + # plus the index load itself, all of which used to run inside the + # coroutine and freeze every other request for the duration. + return await asyncio.to_thread(_compute_curation, request) except HTTPException: raise diff --git a/photomap/backend/routers/search.py b/photomap/backend/routers/search.py index 90ddd76f..73a4f244 100644 --- a/photomap/backend/routers/search.py +++ b/photomap/backend/routers/search.py @@ -146,25 +146,33 @@ async def search_with_text_and_image( f"Search request: {req.min_search_score=}, {req.max_search_results=}" ) try: - results, scores = embeddings.search_images_by_text_and_image( - query_image_data=query_image_data, - positive_query=req.positive_query, - negative_query=req.negative_query, - image_weight=req.image_weight, - positive_weight=req.positive_weight, - negative_weight=req.negative_weight, - # Omitted means "this album's floor" — the album knows one, - # resolved from its encoder when it was created. Falling - # straight through to the encoder default would ignore a - # value the user tuned. - minimum_score=( - req.min_search_score - if req.min_search_score is not None - else album_config.min_search_score - ), - top_k=req.max_search_results, - use_query_optimization=req.use_query_optimization, - ) + # Threaded like the other index readers, but behind a semaphore: + # the loop is what serializes searches today, and a search is a + # CLIP encode plus a full-index matmul. Letting two run at once + # would be a new way to OOM the GPU (the handler below already + # treats that as a live outcome), so the gate keeps concurrency + # exactly where it was and only the blocking goes away. + async with _search_gate(): + results, scores = await asyncio.to_thread( + embeddings.search_images_by_text_and_image, + query_image_data=query_image_data, + positive_query=req.positive_query, + negative_query=req.negative_query, + image_weight=req.image_weight, + positive_weight=req.positive_weight, + negative_weight=req.negative_weight, + # Omitted means "this album's floor" — the album knows one, + # resolved from its encoder when it was created. Falling + # straight through to the encoder default would ignore a + # value the user tuned. + minimum_score=( + req.min_search_score + if req.min_search_score is not None + else album_config.min_search_score + ), + top_k=req.max_search_results, + use_query_optimization=req.use_query_optimization, + ) except HTTPException: # Pass-through (e.g. AlbumDep / EmbeddingsDep already raised # a useful HTTPException; don't bury it under a generic one). @@ -194,6 +202,19 @@ async def search_with_text_and_image( temp_path.unlink(missing_ok=True) +# Process-wide gate around search. Created lazily on first use because +# asyncio.Semaphore wants a running event loop, matching the indexing and +# scan semaphores in embeddings.py. +_search_semaphore: asyncio.Semaphore | None = None + + +def _search_gate() -> asyncio.Semaphore: + global _search_semaphore + if _search_semaphore is None: + _search_semaphore = asyncio.Semaphore(1) + return _search_semaphore + + # Image Retrieval Routes @search_router.get( "/retrieve_image/{album_key}/{index}", @@ -206,7 +227,10 @@ async def retrieve_image( embeddings: EmbeddingsDep, ) -> SlideSummary: """Retrieve metadata for a specific image.""" - slide_metadata = embeddings.retrieve_image(index) + # Threaded for the same reason as /image_info/ below. This one matters + # most: the slideshow calls it per slide, so it is usually the endpoint + # that *takes* the cold miss after an album switch or an index rewrite. + slide_metadata = await asyncio.to_thread(embeddings.retrieve_image, index) create_slide_url(slide_metadata, album_key) return slide_metadata @@ -345,7 +369,9 @@ async def serve_thumbnail( raise HTTPException(status_code=400, detail="Invalid color parameter") try: - image_path = embeddings.get_image_path(index) + # A thumbnail grid fires this many times at once; on a cold index + # every one of them would otherwise be a full np.load on the loop. + image_path = await asyncio.to_thread(embeddings.get_image_path, index) except Exception as e: raise HTTPException( status_code=404, detail=f"Image not found for index {index}: {e}" @@ -453,7 +479,7 @@ async def serve_video_frame( instead of leaving a broken image on screen. """ try: - video_path = embeddings.get_image_path(index) + video_path = await asyncio.to_thread(embeddings.get_image_path, index) except Exception as e: raise HTTPException( status_code=404, detail=f"Image not found for index {index}: {e}" @@ -574,6 +600,19 @@ async def download_images_zip( """ Download multiple images as a ZIP file. """ + # Prime the index off the loop once. Both loops below resolve paths + # through get_image_path, which reads the same cached index -- with it + # warm they are dictionary lookups, so threading each of them + # individually would buy nothing and cost a hop per file. + # + # Best-effort: both loops already treat an unreadable index as "no + # files matched" and return an empty archive. Letting a missing index + # escape from here would turn that into a 500. + try: + await embeddings.load_indexes() + except Exception as e: # noqa: BLE001 - priming only; the loops re-raise + logger.debug(f"Could not prime the index for {album_key}: {e}") + # The archive is assembled entirely in memory, which was fine for photos # but is not for video: twenty bookmarked 200 MB clips would be several # gigabytes resident. Refuse above a ceiling rather than exhausting the @@ -642,7 +681,7 @@ async def get_image_path(album_key: str, index: int, embeddings: EmbeddingsDep) Return the image path for a given index in the album. """ try: - image_path = embeddings.get_image_path(index) + image_path = await asyncio.to_thread(embeddings.get_image_path, index) return image_path.as_posix() except Exception as e: raise HTTPException( diff --git a/tests/backend/test_event_loop_blocking.py b/tests/backend/test_event_loop_blocking.py index 7aef762f..288fc64d 100644 --- a/tests/backend/test_event_loop_blocking.py +++ b/tests/backend/test_event_loop_blocking.py @@ -17,6 +17,7 @@ """ import asyncio +import threading from pathlib import Path import pytest @@ -41,10 +42,12 @@ def spy(path): return real(path) # Patch under the cache so every miss is seen, and start from empty. + # Wrapping the real cache type rather than a bare lru_cache keeps the + # dedupe and generation behaviour these endpoints depend on. monkeypatch.setattr( embeddings_module, "_open_npz_file", - __import__("functools").lru_cache(maxsize=3)(spy), + embeddings_module._NpzIndexCache(spy), ) embeddings_module._open_npz_file.cache_clear() return threads @@ -53,12 +56,35 @@ def spy(path): @pytest.mark.parametrize( "request_for", [ - lambda key: ("GET", f"/umap_data/{key}"), - lambda key: ("GET", f"/image_info/{key}/0"), - lambda key: ("GET", f"/get_metadata/{key}/0"), - lambda key: ("GET", f"/index_metadata/{key}"), + lambda key: ("GET", f"/umap_data/{key}", None), + lambda key: ("GET", f"/image_info/{key}/0", None), + lambda key: ("GET", f"/get_metadata/{key}/0", None), + lambda key: ("GET", f"/index_metadata/{key}", None), + # The slideshow's own endpoint, and the one most likely to take + # the cold miss: it fires once per slide. + lambda key: ("GET", f"/retrieve_image/{key}/0", None), + # A thumbnail grid fires these in a burst. + lambda key: ("GET", f"/thumbnails/{key}/0", None), + lambda key: ("GET", f"/image_path/{key}/0", None), + lambda key: ("POST", f"/search_with_text_and_image/{key}", + {"positive_query": "a photo"}), + lambda key: ("POST", f"/download_images_zip/{key}", {"indices": [0]}), + lambda key: ("POST", "/api/curation/curate_sync", + {"target_count": 2, "iterations": 1, "album": key, + "method": "fps", "excluded_indices": []}), + ], + ids=[ + "umap_data", + "image_info", + "get_metadata", + "index_metadata", + "retrieve_image", + "thumbnails", + "image_path", + "search", + "download_images_zip", + "curate_sync", ], - ids=["umap_data", "image_info", "get_metadata", "index_metadata"], ) def test_endpoints_do_not_load_the_index_on_the_event_loop( client, new_album, loop_thread_loads, request_for @@ -67,8 +93,8 @@ def test_endpoints_do_not_load_the_index_on_the_event_loop( embeddings_module._open_npz_file.cache_clear() loop_thread_loads.clear() - method, url = request_for(new_album["key"]) - assert client.request(method, url).status_code == 200 + method, url, body = request_for(new_album["key"]) + assert client.request(method, url, json=body).status_code == 200 assert loop_thread_loads, f"{url} never loaded the index; the test proves nothing" assert "event-loop" not in loop_thread_loads @@ -120,3 +146,142 @@ def savez_with_a_racing_reader(target, **arrays): assert client.get(f"/index_metadata/{new_album['key']}").json()["image_count"] == ( before - 1 ) + + +def test_a_reader_landing_mid_reprime_cannot_restore_a_deleted_image( + client, new_album, monkeypatch +): + """The narrow version of the race above: the reader lands *during* step 6. + + The test before this one drives a reader that finishes while the rewrite + is in flight, which clearing the cache afterwards is enough to handle. It + is not enough one step later. The re-prime is itself a full load, and a + plain cache takes whichever value arrives first and refuses to overwrite + it -- so a reader finishing inside that window installs the *pre*-delete + snapshot and the re-prime's own fresh result is silently discarded. The + result is a cache that serves a deleted image indefinitely while the + delete path reports success. + + Driven deterministically: the reader loads the still-unmodified file + during the write, then is held until the re-prime has started its own + load, which is exactly the interleaving described above. + """ + build_index(client, new_album) + + from photomap.backend.routers.album import get_embeddings_for_album + + embeddings = get_embeddings_for_album(new_album["key"]) + path = embeddings.embeddings_path + before = len(embeddings.indexes["sorted_filenames"]) + + real_load = embeddings_module._open_npz_file.__wrapped__ + reader_loaded = threading.Event() + reprime_started = threading.Event() + reader_finished = threading.Event() + write_done = threading.Event() + + def spy(target): + if threading.current_thread().name == "racing-reader": + data = real_load(target) # the pre-delete file is still on disk + reader_loaded.set() + reprime_started.wait(10) # hold until step 6's load is underway + return data + if write_done.is_set(): # this call is step 6's re-prime + reprime_started.set() + reader_finished.wait(10) + return real_load(target) + return real_load(target) + + monkeypatch.setattr( + embeddings_module, "_open_npz_file", embeddings_module._NpzIndexCache(spy) + ) + embeddings_module._open_npz_file.cache_clear() + + real_savez = embeddings_module.atomic_savez + + def savez_with_a_racing_reader(target, **arrays): + reader = threading.Thread( + target=lambda: ( + embeddings_module._open_npz_file(path), + reader_finished.set(), + ), + name="racing-reader", + ) + reader.start() + assert reader_loaded.wait(10), "the reader never loaded the pre-delete file" + result = real_savez(target, **arrays) + write_done.set() + return result + + monkeypatch.setattr(embeddings_module, "atomic_savez", savez_with_a_racing_reader) + + embeddings.remove_image_from_embeddings(0) + assert reader_finished.wait(10), "the racing reader never ran; the test proves nothing" + + served = len(embeddings_module._open_npz_file(path)["sorted_filenames"]) + assert served == before - 1, ( + f"the cache serves {served} filenames after a delete left {before - 1} " + "on disk: the racing reader's pre-delete snapshot won" + ) + + +def test_concurrent_readers_share_one_load(client, new_album, monkeypatch): + """A burst of requests for an uncached index must not each load it. + + Nothing deduped in-flight loads while readers ran on the event loop -- + the loop did it, by never letting two overlap. Worker threads removed + that, and a full load is hundreds of megabytes: eight concurrent requests + for a 50,000-image index measured 1.7 GB peak RSS against 0.26 GB for + one. Latecomers have to wait on the load already running. + """ + build_index(client, new_album) + + from photomap.backend.routers.album import get_embeddings_for_album + + embeddings = get_embeddings_for_album(new_album["key"]) + real_load = embeddings_module._open_npz_file.__wrapped__ + loads = [] + all_readers_waiting = threading.Event() + + def spy(target): + loads.append(target) + # Hold the first load open until every reader has had a chance to + # queue up behind it. Without the wait the first could plausibly + # finish before the others start, and the test would pass vacuously. + all_readers_waiting.wait(5) + return real_load(target) + + monkeypatch.setattr( + embeddings_module, "_open_npz_file", embeddings_module._NpzIndexCache(spy) + ) + embeddings_module._open_npz_file.cache_clear() + + async def eight_at_once(): + readers = [embeddings.load_indexes() for _ in range(8)] + gathered = asyncio.gather(*readers) + await asyncio.sleep(0.5) # let them all reach the cache + all_readers_waiting.set() + return await gathered + + results = asyncio.run(eight_at_once()) + + assert len(results) == 8 + assert all(r is results[0] for r in results), "readers got different snapshots" + assert len(loads) == 1, f"{len(loads)} concurrent loads of the same index, expected 1" + + +def test_zip_download_without_an_index_still_answers(client, new_album): + """Priming the index for the zip must not turn a missing index into a 500. + + Both loops in ``/download_images_zip/`` already treat an unresolvable + index as "no files matched" and hand back an empty archive. The priming + call added in front of them reads the same index and would otherwise + raise ``FileNotFoundError`` straight out of the handler. + """ + embeddings_module._open_npz_file.cache_clear() + + response = client.post( + f"/download_images_zip/{new_album['key']}", json={"indices": [0, 1]} + ) + + assert response.status_code == 200