| |
| """Load-test the API's /query endpoint against a *simulated* model latency. |
| |
| What it measures: how the API behaves when each answer waits on a slow language |
| model. The model is simulated (a fixed sleep) and retrieval runs offline on the |
| bundled corpus, so the numbers say nothing about real model speed. They show how |
| the server copes while waiting, which is what the async work changes. |
| |
| Two variants run the same requests against the real FastAPI app: |
| |
| non-blocking the model call is awaited (what the async client does) |
| blocking the model call blocks the event loop (what a synchronous SDK call |
| inside async code does); shown to make the cost visible |
| |
| For each variant it reports throughput, request latency, and the event-loop lag: |
| the longest stretch in which a 10 ms heartbeat task could not run. Lag is what turns |
| one slow model call into a stall for every other client, including /health. |
| |
| Usage: |
| python scripts/loadtest.py |
| python scripts/loadtest.py --latency 0.2 --requests 100 --concurrency 25 |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import asyncio |
| import importlib |
| import statistics |
| import sys |
| import time |
| from pathlib import Path |
|
|
| ROOT = Path(__file__).parent.parent |
| sys.path.insert(0, str(ROOT)) |
|
|
| import httpx |
|
|
| from lumen_rag.engine import RagEngine |
| from lumen_rag.ingestion.loaders import load_file |
| from lumen_rag.llm import Answer |
|
|
| api = importlib.import_module("lumen_rag.api.app") |
|
|
|
|
| def percentile(values: list[float], pct: float) -> float: |
| ordered = sorted(values) |
| return ordered[min(len(ordered) - 1, round(pct / 100 * (len(ordered) - 1)))] |
|
|
|
|
| async def run_variant(name: str, blocking: bool, latency: float, requests: int, concurrency: int) -> dict: |
| async def simulated_model(question, chunks): |
| if blocking: |
| time.sleep(latency) |
| else: |
| await asyncio.sleep(latency) |
| return Answer("simulated answer", []) |
|
|
| import lumen_rag.engine as engine_module |
|
|
| engine_module.answer_async = simulated_model |
|
|
| lags: list[float] = [] |
| stop = False |
|
|
| async def heartbeat(): |
| loop = asyncio.get_running_loop() |
| last = loop.time() |
| while not stop: |
| await asyncio.sleep(0.01) |
| now = loop.time() |
| lags.append(max(0.0, now - last - 0.01)) |
| last = now |
|
|
| latencies: list[float] = [] |
| gate = asyncio.Semaphore(concurrency) |
| transport = httpx.ASGITransport(app=api.app) |
| async with httpx.AsyncClient(transport=transport, base_url="http://load") as client: |
|
|
| async def one(): |
| async with gate: |
| t0 = time.perf_counter() |
| r = await client.post("/query", json={"question": "how many days of annual leave", "k": 3}) |
| r.raise_for_status() |
| latencies.append(time.perf_counter() - t0) |
|
|
| beat = asyncio.create_task(heartbeat()) |
| await asyncio.sleep(0.05) |
| started = time.perf_counter() |
| await asyncio.gather(*[one() for _ in range(requests)]) |
| elapsed = time.perf_counter() - started |
| stop = True |
| await beat |
|
|
| return { |
| "variant": name, |
| "throughput": requests / elapsed, |
| "p50": statistics.median(latencies), |
| "p95": percentile(latencies, 95), |
| "lag_max": max(lags), |
| "lag_p95": percentile(lags, 95), |
| } |
|
|
|
|
| async def main() -> None: |
| parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| parser.add_argument("--latency", type=float, default=0.1, help="simulated model latency in seconds") |
| parser.add_argument("--requests", type=int, default=60) |
| parser.add_argument("--concurrency", type=int, default=20) |
| args = parser.parse_args() |
|
|
| engine = RagEngine() |
| engine.add_documents([load_file(p) for p in sorted((ROOT / "data" / "docs").glob("*.md"))]) |
| api.engine = engine |
| api.enforce = lambda *a, **k: None |
|
|
| print( |
| f"{args.requests} requests, {args.concurrency} in flight, simulated model latency " |
| f"{args.latency * 1000:.0f} ms, offline retrieval on {len(engine.store)} chunks\n" |
| ) |
| rows = [ |
| await run_variant("non-blocking model call", False, args.latency, args.requests, args.concurrency), |
| await run_variant("blocking model call in the loop", True, args.latency, args.requests, args.concurrency), |
| ] |
| print("| Variant | requests/s | p50 latency | p95 latency | loop lag max | loop lag p95 |") |
| print("|---|---|---|---|---|---|") |
| for r in rows: |
| print( |
| f"| {r['variant']} | {r['throughput']:.1f} | {r['p50'] * 1000:.0f} ms | {r['p95'] * 1000:.0f} ms | " |
| f"{r['lag_max'] * 1000:.0f} ms | {r['lag_p95'] * 1000:.0f} ms |" |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| asyncio.run(main()) |
|
|
