File size: 4,963 Bytes
f8a5189 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | #!/usr/bin/env python
"""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) # noqa: ASYNC251 - deliberately blocks the loop to show the contrast
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 # the in-process load would trip the per-IP rate limit
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())
|