Spaces:
Paused
Paused
File size: 2,306 Bytes
9f6ffb8 83dc8bf 9f6ffb8 | 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 | """Measure end-to-end retrieval latency (embed + FAISS search) against the
50ms budget defined in app/config.py.
Usage:
python -m app.benchmark [n_queries]
"""
import statistics
import sys
from pathlib import Path
# Ensure root workspace directory is in sys.path
_ROOT_DIR = Path(__file__).resolve().parent.parent
if str(_ROOT_DIR) not in sys.path:
sys.path.insert(0, str(_ROOT_DIR))
from rag_service.config import LATENCY_BUDGET_MS
from rag_service.retriever import search, warmup
QUERIES = [
"What is FAISS used for?",
"How does HNSW indexing work?",
"What is retrieval augmented generation?",
"Which embedding model is fast on CPU?",
"How do you reduce RAG latency?",
"What does efSearch control?",
"Why normalize embeddings before indexing?",
"What are the stages of a RAG pipeline?",
]
def percentile(values: list[float], pct: float) -> float:
values = sorted(values)
k = (len(values) - 1) * (pct / 100)
f, c = int(k), min(int(k) + 1, len(values) - 1)
if f == c:
return values[f]
return values[f] + (k - f) * (values[c] - values[f])
def main():
n = int(sys.argv[1]) if len(sys.argv) > 1 else 50
print("Warming up (model load + first inference)...")
warmup()
total_ms, embed_ms, search_ms = [], [], []
for i in range(n):
query = QUERIES[i % len(QUERIES)]
resp = search(query, top_k=5)
total_ms.append(resp.total_ms)
embed_ms.append(resp.embed_ms)
search_ms.append(resp.search_ms)
print(f"\nRan {n} queries\n")
print(f"{'stage':<12}{'avg':>8}{'p50':>8}{'p95':>8}{'p99':>8} (ms)")
for name, values in [("embed", embed_ms), ("search", search_ms), ("total", total_ms)]:
print(
f"{name:<12}"
f"{statistics.mean(values):>8.2f}"
f"{percentile(values, 50):>8.2f}"
f"{percentile(values, 95):>8.2f}"
f"{percentile(values, 99):>8.2f}"
)
p95_total = percentile(total_ms, 95)
print(f"\nLatency budget: {LATENCY_BUDGET_MS}ms | p95 total: {p95_total:.2f}ms")
if p95_total <= LATENCY_BUDGET_MS:
print("PASS: within budget")
else:
print("FAIL: over budget -- see README 'Tuning latency' section")
sys.exit(1)
if __name__ == "__main__":
main()
|