diff --git a/photomap/backend/embeddings.py b/photomap/backend/embeddings.py index 32eeb9c0..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. @@ -1891,6 +2006,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. """ @@ -2151,6 +2271,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. @@ -2215,7 +2353,21 @@ def remove_images_from_embeddings(self, indices: list[int]) -> None: **extras, ) - # 6. Re-prime the cache immediately to verify the write + # 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) except Exception as e: @@ -2288,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/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 098c3d7e..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 @@ -223,7 +247,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"] @@ -250,7 +274,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") @@ -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( @@ -676,7 +715,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 @@ -704,7 +743,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 642cda4f..725f4f3e 100644 --- a/photomap/backend/routers/umap.py +++ b/photomap/backend/routers/umap.py @@ -68,7 +68,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..288fc64d --- /dev/null +++ b/tests/backend/test_event_loop_blocking.py @@ -0,0 +1,287 @@ +"""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 +import threading +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. + # 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", + embeddings_module._NpzIndexCache(spy), + ) + embeddings_module._open_npz_file.cache_clear() + return threads + + +@pytest.mark.parametrize( + "request_for", + [ + 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", + ], +) +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, 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 + + +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 + ) + + +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