diff --git a/serve/src/pixelrag_serve/api.py b/serve/src/pixelrag_serve/api.py index 7b7f003..01c4716 100644 --- a/serve/src/pixelrag_serve/api.py +++ b/serve/src/pixelrag_serve/api.py @@ -279,6 +279,10 @@ class StatusResponse(BaseModel): _MAX_IMAGE_B64_LEN = 10 * 1024 * 1024 _MAX_IMAGE_PIXELS = 25_000_000 +# Most on-demand renders one /search request may trigger (see the search handler). +# A real caller asks for n_docs=10 with include_images, so this is generous. +_MAX_ONDEMAND_RENDERS = 16 + def _parse_queries( queries: list[Query], instruction: str | None = None @@ -540,6 +544,18 @@ async def search(req: SearchRequest, request: Request): # Build results tiles_dir = _state.get("tiles_dir", "") + # Budget for on-demand renders across the whole request. A render is a + # Chrome page load under a process-global lock, bounded only by + # PIXELRAG_RENDER_TIMEOUT (120s by default), and nothing else caps how many + # one request can ask for: at this endpoint's own limits (32 queries, + # n_docs <= 1000) a single include_images request could queue tens of + # thousands of them, stalling every other caller's renders behind the same + # lock. Past the budget a hit is returned without its image — what a failed + # render or a missing tile already yields — and the client can still fetch + # it from /tile. Cache hits are cheap and don't spend the budget. + ondemand = _state.get("ondemand") + renders_left = _MAX_ONDEMAND_RENDERS + results = [] for qi in range(len(req.queries)): hits = [] @@ -559,11 +575,17 @@ async def search(req: SearchRequest, request: Request): if req.include_images and tile_path and os.path.exists(tile_path): with open(tile_path, "rb") as fp: img_b64 = base64.b64encode(fp.read()).decode() - elif req.include_images and _state.get("ondemand") is not None: - # Render off the event loop: _ondemand_chunk_b64 -> render_url uses - # asyncio.run(), which raises "cannot be called from a running event - # loop" if invoked directly here. Offload to a worker thread. - img_b64 = await asyncio.to_thread(_ondemand_chunk_b64, aid, ti, ci, th) + elif req.include_images and ondemand is not None: + cached = ondemand.cached_chunk_path(aid, ti, ci) is not None + if cached or renders_left > 0: + if not cached: + renders_left -= 1 + # Render off the event loop: _ondemand_chunk_b64 -> render_url uses + # asyncio.run(), which raises "cannot be called from a running event + # loop" if invoked directly here. Offload to a worker thread. + img_b64 = await asyncio.to_thread( + _ondemand_chunk_b64, aid, ti, ci, th + ) # Expose a relative tile path, not the absolute server filesystem # path (avoids leaking the host's directory layout; clients fetch # tiles via /tile/{article_id}/{tile_index}/{chunk_index}). diff --git a/serve/src/pixelrag_serve/render_ondemand.py b/serve/src/pixelrag_serve/render_ondemand.py index 316decc..198838b 100644 --- a/serve/src/pixelrag_serve/render_ondemand.py +++ b/serve/src/pixelrag_serve/render_ondemand.py @@ -112,14 +112,27 @@ def _ensure_chrome(self) -> str: def _article_dir(self, article_id: int) -> str: return os.path.join(self.cache_dir, f"{article_id}.png.tiles") + def _chunk_file(self, article_id: int, tile_index: int, chunk_index: int) -> str: + chunk_name = f"chunk_{tile_index:04d}_{chunk_index:02d}.png" + return os.path.join(self._article_dir(article_id), chunk_name) + + def cached_chunk_path(self, article_id: int, tile_index: int, chunk_index: int): + """Path to an already-rendered chunk, or None. Never renders. + + Lets a caller tell a cheap cache hit from a page render before + committing to one. + """ + cpath = self._chunk_file(article_id, tile_index, chunk_index) + return cpath if os.path.exists(cpath) else None + def chunk_path( self, article_id: int, title: str, tile_index: int, chunk_index: int ): """Path to chunk_{ti}_{ci}.png, rendering+chunking the page on a cache miss.""" - chunk_name = f"chunk_{tile_index:04d}_{chunk_index:02d}.png" - cpath = os.path.join(self._article_dir(article_id), chunk_name) - if os.path.exists(cpath): - return cpath + cached = self.cached_chunk_path(article_id, tile_index, chunk_index) + if cached: + return cached + cpath = self._chunk_file(article_id, tile_index, chunk_index) if not title: return None with _render_lock: diff --git a/tests/test_serve_ondemand_budget.py b/tests/test_serve_ondemand_budget.py new file mode 100644 index 0000000..21a5d77 --- /dev/null +++ b/tests/test_serve_ondemand_budget.py @@ -0,0 +1,124 @@ +"""Cap on on-demand renders per /search request (DoS hardening). + +An on-demand render is a Chrome page load under a process-global lock with a +120s timeout. Without a per-request budget, one `include_images` request could +queue one per hit — at the endpoint's own limits, tens of thousands — and every +other caller's renders stall behind the same lock. + +Driven through the real handler with pre-computed embeddings, so no model or +index is needed. +""" + +import numpy as np +import pytest + +pytest.importorskip("fastapi", reason="serve extra not installed") +from fastapi.testclient import TestClient +from pixelrag_serve import api + +DIM = 4 +N_HITS = 40 # more hits than the budget allows renders + + +class _CountingOnDemand: + """Stands in for OnDemandTiles: counts renders, never touches Chrome.""" + + def __init__(self, cached: set[int] | None = None): + self.cached = cached or set() + self.renders = 0 + + def cached_chunk_path(self, article_id, tile_index, chunk_index): + return f"/cached/{article_id}.png" if article_id in self.cached else None + + def chunk_path(self, article_id, title, tile_index, chunk_index): + if article_id not in self.cached: + self.renders += 1 + + +class _StubBackend: + dimension = DIM + + def set_nprobe(self, n): + pass + + def reset_nprobe(self): + pass + + def raw_search(self, query_vectors, k, **kw): + return [ + [ + { + "score": 1.0, + "vector_id": i, + "article_id": i, + "tile_index": 0, + "chunk_index": 0, + "y_offset": 0, + "tile_height": 100, + } + for i in range(min(k, N_HITS)) + ] + for _ in query_vectors + ] + + +@pytest.fixture +def client(tmp_path): + def _make(ondemand): + api._state.clear() + api._state.update( + { + "backend": _StubBackend(), + "articles": [f"Article_{i}" for i in range(N_HITS)], + "dimension": DIM, + "tiles_dir": str(tmp_path), # empty: nothing is on disk + "ondemand": ondemand, + } + ) + api._article_pages.cache_clear() + return TestClient(api.app) + + yield _make + api._state.clear() + + +def _search(client, n_docs): + body = { + "queries": [{"embedding": np.ones(DIM).tolist()}], + "n_docs": n_docs, + "include_images": True, + } + resp = client.post("/search", json=body) + assert resp.status_code == 200, resp.text + return resp.json() + + +def test_renders_are_capped_per_request(client): + od = _CountingOnDemand() + body = _search(client(od), N_HITS) + # Every hit still comes back — only the images past the budget are dropped. + assert len(body["results"][0]["hits"]) == N_HITS + assert od.renders == api._MAX_ONDEMAND_RENDERS + + +def test_budget_spans_all_queries_in_one_request(client): + od = _CountingOnDemand() + c = client(od) + resp = c.post( + "/search", + json={ + "queries": [{"embedding": np.ones(DIM).tolist()} for _ in range(4)], + "n_docs": N_HITS, + "include_images": True, + }, + ) + assert resp.status_code == 200, resp.text + # One budget for the request, not one per query. + assert od.renders == api._MAX_ONDEMAND_RENDERS + + +def test_cache_hits_do_not_spend_the_budget(client): + # Every article already rendered: serving them must cost no renders at all. + od = _CountingOnDemand(cached=set(range(N_HITS))) + _search(client(od), N_HITS) + assert od.renders == 0