nomyo-ai/nomyo-router
3
2
Fork 0
Code Issues 2 Pull requests 7 Projects Releases 10 Packages 1 Wiki Activity Actions
nomyo-router/test/load/loadtest.py

803 lines
33 KiB
Python
Raw Permalink Blame History

#!/usr/bin/env python3
"""
NOMYO Router load test — asyncio + httpx driver with a built-in mock backend.
Three modes
-----------
1. Drive an already-running router (default)::
python test/load/loadtest.py --url http://127.0.0.1:12434 \
--concurrency 64 --duration 30 --stream
2. Fully self-contained "mock backend" mode — spins up a fast fake Ollama/OpenAI
backend AND the router (wired to that backend via a temp config), load-tests
them, then tears both down. This isolates the *router's* proxy overhead from
real GPU compute, so the numbers tell you how many concurrent connections the
router itself can sustain on this machine::
python test/load/loadtest.py --mock-backend \
--concurrency 128 --duration 30 --stream
3. Run just the mock backend (point your own router config at it)::
python test/load/loadtest.py --serve-mock --mock-port 11434
Both streaming and non-streaming are supported (--stream / --no-stream), against
either the Ollama API (--api ollama -> POST /api/chat) or the OpenAI-compatible
API (--api openai -> POST /v1/chat/completions).
Finding the concurrency ceiling
-------------------------------
Use --ramp to sweep concurrency levels and print a table; the "knee" is where
p99 latency climbs sharply or req/s stops increasing::
python test/load/loadtest.py --mock-backend --stream \
--ramp 16,32,64,128,256 --duration 15
The mock backend is a configurable "fake GPU": --mock-ttft-ms (prefill latency),
--mock-tokens (completion length) and --mock-tok-ms (per-token decode delay) let
you model anything from an instant echo (measure pure proxy overhead) to a slow,
long-streaming model (measure how many slow streams the box holds open).
"""
from __future__ import annotations
import argparse
import asyncio
import contextlib
import json
import math
import os
import signal
import socket
import statistics
import sys
import tempfile
import time
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional
import httpx
# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------
_WORDS = (
"the quick brown fox jumps over the lazy dog while a router proxies many "
"concurrent streaming completions across several ollama and openai backends "
"without dropping a single token under sustained synthetic load testing"
).split()
def _gen_text(n_tokens: int) -> str:
"""Deterministic pseudo-completion of roughly ``n_tokens`` space-separated tokens."""
return " ".join(_WORDS[i % len(_WORDS)] for i in range(max(0, n_tokens)))
def _rfc3339_now() -> str:
# Ollama-style timestamp, e.g. 2024-01-01T00:00:00.000000Z
return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%f") + "Z"
def _count_prompt_tokens(messages: list) -> int:
total = 0
for m in messages or []:
c = m.get("content")
if isinstance(c, str):
total += len(c.split())
elif isinstance(c, list):
for part in c:
if isinstance(part, dict) and isinstance(part.get("text"), str):
total += len(part["text"].split())
return max(1, total)
def _free_port() -> int:
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
s.bind(("127.0.0.1", 0))
port = s.getsockname()[1]
s.close()
return port
# ===========================================================================
# Mock backend (a fast, configurable fake Ollama + OpenAI-compatible server)
# ===========================================================================
def build_mock_app(models: list[str], ttft_ms: float, tokens: int, tok_ms: float):
"""Construct the aiohttp mock-backend application.
Serves the native-Ollama surface the router uses for discovery and the
`/api/chat` path (`/api/version`, `/api/tags`, `/api/ps`, `/api/chat`,
`/api/generate`) plus the OpenAI-compatible surface used by the
`/v1/chat/completions` path (`/v1/models`, `/v1/chat/completions`,
`/v1/completions`).
"""
from aiohttp import web # imported lazily so the driver has no hard aiohttp dep
ttft = ttft_ms / 1000.0
tok_delay = tok_ms / 1000.0
def _tag_entry(name: str) -> dict:
return {
"name": name,
"model": name,
"modified_at": _rfc3339_now(),
"size": 4_000_000_000,
"digest": "0" * 64,
"details": {
"parent_model": "",
"format": "gguf",
"family": "mock",
"families": ["mock"],
"parameter_size": "7B",
"quantization_level": "Q4_0",
},
}
async def version(_req):
return web.json_response({"version": "0.0.0-nomyo-mock"})
async def tags(_req):
return web.json_response({"models": [_tag_entry(m) for m in models]})
async def ps(_req):
# Report every advertised model as loaded with VRAM so choose_endpoint
# treats this endpoint as "loaded + free".
out = []
for m in models:
e = _tag_entry(m)
e["size_vram"] = e["size"]
e["expires_at"] = "2999-01-01T00:00:00Z"
out.append(e)
return web.json_response({"models": out})
async def v1_models(_req):
now = int(time.time())
return web.json_response({
"object": "list",
"data": [{"id": m, "object": "model", "created": now, "owned_by": "mock"} for m in models],
})
# ----- Ollama /api/chat -------------------------------------------------
async def api_chat(req):
payload = await req.json()
model = payload.get("model", models[0] if models else "mock")
stream = payload.get("stream", True)
prompt_tok = _count_prompt_tokens(payload.get("messages", []))
t0 = time.perf_counter()
if stream:
resp = web.StreamResponse(
status=200, headers={"Content-Type": "application/x-ndjson"}
)
await resp.prepare(req)
if ttft:
await asyncio.sleep(ttft)
for i in range(tokens):
if tok_delay and i:
await asyncio.sleep(tok_delay)
line = {
"model": model,
"created_at": _rfc3339_now(),
"message": {"role": "assistant", "content": _WORDS[i % len(_WORDS)] + " "},
"done": False,
}
await resp.write(json.dumps(line).encode() + b"\n")
dur_ns = int((time.perf_counter() - t0) * 1e9)
final = {
"model": model,
"created_at": _rfc3339_now(),
"message": {"role": "assistant", "content": ""},
"done": True,
"done_reason": "stop",
"total_duration": dur_ns,
"load_duration": 0,
"prompt_eval_count": prompt_tok,
"prompt_eval_duration": int(ttft * 1e9),
"eval_count": tokens,
"eval_duration": dur_ns,
}
await resp.write(json.dumps(final).encode() + b"\n")
await resp.write_eof()
return resp
# non-streaming: simulate the whole generation latency, then one object
await asyncio.sleep(ttft + tokens * tok_delay)
dur_ns = int((time.perf_counter() - t0) * 1e9)
return web.json_response({
"model": model,
"created_at": _rfc3339_now(),
"message": {"role": "assistant", "content": _gen_text(tokens)},
"done": True,
"done_reason": "stop",
"total_duration": dur_ns,
"load_duration": 0,
"prompt_eval_count": prompt_tok,
"prompt_eval_duration": int(ttft * 1e9),
"eval_count": tokens,
"eval_duration": dur_ns,
})
# ----- Ollama /api/generate --------------------------------------------
async def api_generate(req):
payload = await req.json()
model = payload.get("model", models[0] if models else "mock")
stream = payload.get("stream", True)
prompt_tok = max(1, len(str(payload.get("prompt", "")).split()))
t0 = time.perf_counter()
if stream:
resp = web.StreamResponse(status=200, headers={"Content-Type": "application/x-ndjson"})
await resp.prepare(req)
if ttft:
await asyncio.sleep(ttft)
for i in range(tokens):
if tok_delay and i:
await asyncio.sleep(tok_delay)
await resp.write(json.dumps({
"model": model, "created_at": _rfc3339_now(),
"response": _WORDS[i % len(_WORDS)] + " ", "done": False,
}).encode() + b"\n")
dur_ns = int((time.perf_counter() - t0) * 1e9)
await resp.write(json.dumps({
"model": model, "created_at": _rfc3339_now(), "response": "", "done": True,
"done_reason": "stop", "total_duration": dur_ns,
"prompt_eval_count": prompt_tok, "eval_count": tokens, "eval_duration": dur_ns,
}).encode() + b"\n")
await resp.write_eof()
return resp
await asyncio.sleep(ttft + tokens * tok_delay)
dur_ns = int((time.perf_counter() - t0) * 1e9)
return web.json_response({
"model": model, "created_at": _rfc3339_now(), "response": _gen_text(tokens),
"done": True, "done_reason": "stop", "total_duration": dur_ns,
"prompt_eval_count": prompt_tok, "eval_count": tokens, "eval_duration": dur_ns,
})
# ----- OpenAI /v1/chat/completions -------------------------------------
async def v1_chat(req):
payload = await req.json()
model = payload.get("model", models[0] if models else "mock")
stream = payload.get("stream", False)
want_usage = bool((payload.get("stream_options") or {}).get("include_usage"))
prompt_tok = _count_prompt_tokens(payload.get("messages", []))
created = int(time.time())
cid = "chatcmpl-mock"
if stream:
resp = web.StreamResponse(status=200, headers={"Content-Type": "text/event-stream"})
await resp.prepare(req)
if ttft:
await asyncio.sleep(ttft)
def _sse(obj: dict) -> bytes:
return b"data: " + json.dumps(obj).encode() + b"\n\n"
# first chunk carries the role
await resp.write(_sse({
"id": cid, "object": "chat.completion.chunk", "created": created, "model": model,
"choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}],
}))
for i in range(tokens):
if tok_delay and i:
await asyncio.sleep(tok_delay)
await resp.write(_sse({
"id": cid, "object": "chat.completion.chunk", "created": created, "model": model,
"choices": [{"index": 0, "delta": {"content": _WORDS[i % len(_WORDS)] + " "}, "finish_reason": None}],
}))
await resp.write(_sse({
"id": cid, "object": "chat.completion.chunk", "created": created, "model": model,
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
}))
if want_usage:
await resp.write(_sse({
"id": cid, "object": "chat.completion.chunk", "created": created, "model": model,
"choices": [],
"usage": {"prompt_tokens": prompt_tok, "completion_tokens": tokens,
"total_tokens": prompt_tok + tokens},
}))
await resp.write(b"data: [DONE]\n\n")
await resp.write_eof()
return resp
await asyncio.sleep(ttft + tokens * tok_delay)
return web.json_response({
"id": cid, "object": "chat.completion", "created": created, "model": model,
"choices": [{"index": 0, "message": {"role": "assistant", "content": _gen_text(tokens)},
"finish_reason": "stop", "logprobs": None}],
"usage": {"prompt_tokens": prompt_tok, "completion_tokens": tokens,
"total_tokens": prompt_tok + tokens},
})
# ----- OpenAI /v1/completions ------------------------------------------
async def v1_completions(req):
payload = await req.json()
model = payload.get("model", models[0] if models else "mock")
stream = payload.get("stream", False)
prompt_tok = max(1, len(str(payload.get("prompt", "")).split()))
created = int(time.time())
cid = "cmpl-mock"
if stream:
resp = web.StreamResponse(status=200, headers={"Content-Type": "text/event-stream"})
await resp.prepare(req)
if ttft:
await asyncio.sleep(ttft)
for i in range(tokens):
if tok_delay and i:
await asyncio.sleep(tok_delay)
await resp.write(b"data: " + json.dumps({
"id": cid, "object": "text_completion", "created": created, "model": model,
"choices": [{"index": 0, "text": _WORDS[i % len(_WORDS)] + " ", "finish_reason": None}],
}).encode() + b"\n\n")
await resp.write(b"data: [DONE]\n\n")
await resp.write_eof()
return resp
await asyncio.sleep(ttft + tokens * tok_delay)
return web.json_response({
"id": cid, "object": "text_completion", "created": created, "model": model,
"choices": [{"index": 0, "text": _gen_text(tokens), "finish_reason": "stop"}],
"usage": {"prompt_tokens": prompt_tok, "completion_tokens": tokens,
"total_tokens": prompt_tok + tokens},
})
app = web.Application(client_max_size=64 * 1024 * 1024)
app.add_routes([
web.get("/api/version", version),
web.get("/api/tags", tags),
web.get("/api/ps", ps),
web.post("/api/chat", api_chat),
web.post("/api/generate", api_generate),
web.get("/v1/models", v1_models),
web.post("/v1/chat/completions", v1_chat),
web.post("/v1/completions", v1_completions),
])
return app
def serve_mock(args) -> None:
from aiohttp import web
models = [m.strip() for m in args.mock_models.split(",") if m.strip()]
app = build_mock_app(models, args.mock_ttft_ms, args.mock_tokens, args.mock_tok_ms)
print(f"[mock] serving models={models} on http://{args.mock_host}:{args.mock_port} "
f"(ttft={args.mock_ttft_ms}ms tokens={args.mock_tokens} tok={args.mock_tok_ms}ms)",
flush=True)
web.run_app(app, host=args.mock_host, port=args.mock_port, print=None)
# ===========================================================================
# Load driver
# ===========================================================================
@dataclass
class Sample:
ok: bool
status: int
latency: float # full request wall time (s)
ttft: Optional[float] # time-to-first-byte for streaming (s), else None
err: Optional[str] = None
@dataclass
class Stats:
concurrency: int
wall: float = 0.0
samples: list = field(default_factory=list)
@property
def ok(self) -> list:
return [s for s in self.samples if s.ok]
@property
def n_total(self) -> int:
return len(self.samples)
@property
def n_ok(self) -> int:
return len(self.ok)
def _pct(values: list[float], p: float) -> float:
if not values:
return float("nan")
s = sorted(values)
if len(s) == 1:
return s[0]
k = (len(s) - 1) * (p / 100.0)
lo = math.floor(k)
hi = math.ceil(k)
if lo == hi:
return s[int(k)]
return s[lo] + (s[hi] - s[lo]) * (k - lo)
def _build_request(args):
"""Return (path, json_payload) for a single request."""
messages = [{"role": "user", "content": args.prompt}]
if args.api == "openai":
path = "/v1/chat/completions"
body = {"model": args.model, "messages": messages, "stream": args.stream}
else:
path = "/api/chat"
body = {"model": args.model, "messages": messages, "stream": args.stream}
return path, body
async def _one_request(client: httpx.AsyncClient, url: str, body: dict, stream: bool) -> Sample:
t0 = time.perf_counter()
try:
if stream:
ttft = None
async with client.stream("POST", url, json=body) as resp:
status = resp.status_code
async for _chunk in resp.aiter_bytes():
if ttft is None:
ttft = time.perf_counter() - t0
# drain complete
lat = time.perf_counter() - t0
ok = 200 <= status < 300
return Sample(ok=ok, status=status, latency=lat, ttft=ttft,
err=None if ok else f"HTTP {status}")
else:
resp = await client.post(url, json=body)
lat = time.perf_counter() - t0
ok = 200 <= resp.status_code < 300
# touch body so the full response is received
_ = resp.content
return Sample(ok=ok, status=resp.status_code, latency=lat, ttft=None,
err=None if ok else f"HTTP {resp.status_code}")
except Exception as e: # noqa: BLE001 — record any transport error as a failed sample
lat = time.perf_counter() - t0
return Sample(ok=False, status=0, latency=lat, ttft=None,
err=f"{type(e).__name__}: {str(e)[:120]}")
async def run_stage(args, concurrency: int) -> Stats:
path, body = _build_request(args)
url = args.url.rstrip("/") + path
stats = Stats(concurrency=concurrency)
limits = httpx.Limits(max_connections=concurrency + 50,
max_keepalive_connections=concurrency + 50)
timeout = httpx.Timeout(args.timeout, connect=15.0)
async with httpx.AsyncClient(limits=limits, timeout=timeout) as client:
# warmup (unmeasured): make a few requests so caches/connections are hot
if args.warmup > 0:
warm_deadline = time.perf_counter() + args.warmup
async def _warm():
while time.perf_counter() < warm_deadline:
await _one_request(client, url, body, args.stream)
await asyncio.gather(*[_warm() for _ in range(min(concurrency, 8))])
use_duration = args.requests is None
deadline = time.perf_counter() + args.duration if use_duration else None
remaining = args.requests if not use_duration else None
remaining_lock = asyncio.Lock()
async def worker():
nonlocal remaining
while True:
if use_duration:
if time.perf_counter() >= deadline:
return
else:
async with remaining_lock:
if remaining <= 0:
return
remaining -= 1
s = await _one_request(client, url, body, args.stream)
stats.samples.append(s)
wall0 = time.perf_counter()
await asyncio.gather(*[worker() for _ in range(concurrency)])
stats.wall = time.perf_counter() - wall0
return stats
def _print_stage(stats: Stats, args, header: bool) -> None:
lat = [s.latency * 1000 for s in stats.ok]
ttfts = [s.ttft * 1000 for s in stats.ok if s.ttft is not None]
rps = stats.n_ok / stats.wall if stats.wall else 0.0
errs = stats.n_total - stats.n_ok
if args.ramp:
if header:
cols = f"{'conc':>5} {'req':>7} {'ok':>7} {'err':>5} {'req/s':>9} " \
f"{'p50ms':>8} {'p90ms':>8} {'p99ms':>9} {'maxms':>9}"
if args.stream:
cols += f" {'ttftP50':>8} {'ttftP99':>8}"
print(cols)
print("-" * len(cols))
row = (f"{stats.concurrency:>5} {stats.n_total:>7} {stats.n_ok:>7} {errs:>5} "
f"{rps:>9.1f} {_pct(lat,50):>8.1f} {_pct(lat,90):>8.1f} "
f"{_pct(lat,99):>9.1f} {(max(lat) if lat else float('nan')):>9.1f}")
if args.stream:
row += f" {_pct(ttfts,50):>8.1f} {_pct(ttfts,99):>8.1f}"
print(row, flush=True)
return
# single-stage detailed report
print(f"\n=== Results (concurrency={stats.concurrency}, "
f"{'stream' if args.stream else 'non-stream'}, api={args.api}) ===")
print(f" wall time : {stats.wall:8.2f} s")
print(f" requests : {stats.n_total} total, {stats.n_ok} ok, {errs} failed")
print(f" throughput : {rps:8.1f} req/s")
if lat:
print(f" latency p50 : {_pct(lat,50):8.1f} ms")
print(f" p90 : {_pct(lat,90):8.1f} ms")
print(f" p95 : {_pct(lat,95):8.1f} ms")
print(f" p99 : {_pct(lat,99):8.1f} ms")
print(f" max : {max(lat):8.1f} ms")
print(f" mean : {statistics.mean(lat):8.1f} ms")
if ttfts:
print(f" TTFT p50 : {_pct(ttfts,50):8.1f} ms")
print(f" p90 : {_pct(ttfts,90):8.1f} ms")
print(f" p99 : {_pct(ttfts,99):8.1f} ms")
if errs:
by_err: dict[str, int] = {}
for s in stats.samples:
if not s.ok:
by_err[s.err or "unknown"] = by_err.get(s.err or "unknown", 0) + 1
print(" errors:")
for k, v in sorted(by_err.items(), key=lambda kv: -kv[1]):
print(f" {v:>6} {k}")
async def run_driver(args) -> list[Stats]:
stages = ([int(x) for x in args.ramp.split(",")] if args.ramp else [args.concurrency])
results: list[Stats] = []
for i, c in enumerate(stages):
stats = await run_stage(args, c)
_print_stage(stats, args, header=(i == 0))
results.append(stats)
return results
# ===========================================================================
# Orchestration: --mock-backend (spawn mock + router, run, tear down)
# ===========================================================================
PROJECT_ROOT = Path(__file__).resolve().parents[2]
async def _wait_http_ok(url: str, timeout: float, accept=(200,)) -> bool:
deadline = time.perf_counter() + timeout
async with httpx.AsyncClient(timeout=5.0) as client:
while time.perf_counter() < deadline:
try:
r = await client.get(url)
if r.status_code in accept:
return True
except Exception:
pass
await asyncio.sleep(0.25)
return False
def _write_temp_config(mock_url: str, models: list[str], max_conc: int) -> Path:
fd, path = tempfile.mkstemp(prefix="nomyo_loadtest_", suffix=".yaml")
os.close(fd)
cfg = (
"# Auto-generated by test/load/loadtest.py --mock-backend. Safe to delete.\n"
"endpoints:\n"
f" - {mock_url}\n"
"llama_server_endpoints: []\n"
f"max_concurrent_connections: {max_conc}\n"
"priority_routing: false\n"
"conversation_affinity: false\n"
"cache_enabled: false\n"
"nomyo-router-api-key: \"\"\n"
"api_keys:\n"
f" \"{mock_url}\": \"mock\"\n"
)
Path(path).write_text(cfg)
return Path(path)
async def run_with_mock_backend(args) -> list[Stats]:
mock_port = args.mock_port or _free_port()
router_port = args.router_port or _free_port()
mock_url = f"http://127.0.0.1:{mock_port}"
router_url = f"http://127.0.0.1:{router_port}"
models = [m.strip() for m in args.mock_models.split(",") if m.strip()]
# Size the router's per-endpoint admission limit so it does not artificially
# serialize the load (unless the user explicitly wants to measure that).
peak = max([int(x) for x in args.ramp.split(",")]) if args.ramp else args.concurrency
max_conc = args.router_max_conc if args.router_max_conc else max(peak, 1)
cfg_path = _write_temp_config(mock_url, models, max_conc)
db_path = Path(tempfile.gettempdir()) / f"nomyo_loadtest_{os.getpid()}.db"
env = dict(os.environ)
env["NOMYO_ROUTER_CONFIG_PATH"] = str(cfg_path)
env["NOMYO_ROUTER_DB_PATH"] = str(db_path)
mock_proc = None
router_proc = None
try:
# 1. mock backend first, so the router never caches it as "down"
mock_cmd = [
sys.executable, str(Path(__file__).resolve()), "--serve-mock",
"--mock-host", "127.0.0.1", "--mock-port", str(mock_port),
"--mock-models", args.mock_models,
"--mock-ttft-ms", str(args.mock_ttft_ms),
"--mock-tokens", str(args.mock_tokens),
"--mock-tok-ms", str(args.mock_tok_ms),
]
print(f"[orchestrator] starting mock backend: {mock_url}", flush=True)
mock_proc = await asyncio.create_subprocess_exec(*mock_cmd)
if not await _wait_http_ok(f"{mock_url}/api/version", timeout=15):
raise RuntimeError("mock backend did not become ready")
# 2. router
router_cmd = [
sys.executable, "-m", "uvicorn", "router:app",
"--host", "127.0.0.1", "--port", str(router_port),
# Per-request access logging is pure noise (and overhead) under load.
"--no-access-log",
]
if args.router_workers and args.router_workers > 1:
router_cmd += ["--workers", str(args.router_workers)]
print(f"[orchestrator] starting router: {router_url} "
f"(workers={args.router_workers}, max_concurrent_connections={max_conc})", flush=True)
router_proc = await asyncio.create_subprocess_exec(
*router_cmd, cwd=str(PROJECT_ROOT), env=env
)
# /health returns 200 only once it can reach the (healthy) mock backend
if not await _wait_http_ok(f"{router_url}/health", timeout=40, accept=(200,)):
raise RuntimeError("router did not become healthy")
print("[orchestrator] router healthy — starting load\n", flush=True)
# 3. drive load against the router
args.url = router_url
return await run_driver(args)
finally:
for name, proc in (("router", router_proc), ("mock", mock_proc)):
if proc and proc.returncode is None:
with contextlib.suppress(ProcessLookupError):
proc.send_signal(signal.SIGINT)
with contextlib.suppress(asyncio.TimeoutError):
await asyncio.wait_for(proc.wait(), timeout=8)
if proc.returncode is None:
with contextlib.suppress(ProcessLookupError):
proc.kill()
print(f"[orchestrator] stopped {name}", flush=True)
if args.keep_config:
print(f"[orchestrator] kept config: {cfg_path}", flush=True)
else:
with contextlib.suppress(FileNotFoundError):
cfg_path.unlink()
for suffix in ("", "-shm", "-wal"):
with contextlib.suppress(FileNotFoundError):
Path(str(db_path) + suffix).unlink()
# ===========================================================================
# CLI
# ===========================================================================
def build_parser() -> argparse.ArgumentParser:
p = argparse.ArgumentParser(
description="NOMYO Router load test (asyncio + httpx) with a built-in mock backend.",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
mode = p.add_argument_group("mode")
mode.add_argument("--serve-mock", action="store_true",
help="Run ONLY the mock backend (foreground) and exit on Ctrl-C.")
mode.add_argument("--mock-backend", action="store_true",
help="Spawn mock backend + router, load-test them, then tear down.")
tgt = p.add_argument_group("target (driver)")
tgt.add_argument("--url", default="http://127.0.0.1:12434",
help="Router base URL to drive (ignored with --mock-backend).")
tgt.add_argument("--api", choices=["ollama", "openai"], default="ollama",
help="ollama -> POST /api/chat ; openai -> POST /v1/chat/completions")
tgt.add_argument("--model", default="mock", help="Model name to request.")
tgt.add_argument("--prompt", default="Say hello and count to ten.",
help="User prompt sent in every request.")
stream_grp = tgt.add_mutually_exclusive_group()
stream_grp.add_argument("--stream", dest="stream", action="store_true",
help="Stream the response (default).")
stream_grp.add_argument("--no-stream", dest="stream", action="store_false",
help="Request a single non-streamed response.")
p.set_defaults(stream=True)
load = p.add_argument_group("load shape")
load.add_argument("--concurrency", type=int, default=32,
help="Number of concurrent virtual clients.")
load.add_argument("--duration", type=float, default=20.0,
help="Seconds to run each stage (ignored if --requests given).")
load.add_argument("--requests", type=int, default=None,
help="Send exactly N requests instead of running for --duration.")
load.add_argument("--ramp", default=None,
help="Comma-separated concurrency stages, e.g. 16,32,64,128 "
"(prints a table to find the knee).")
load.add_argument("--warmup", type=float, default=2.0,
help="Seconds of unmeasured warmup before each stage.")
load.add_argument("--timeout", type=float, default=120.0,
help="Per-request timeout (seconds).")
load.add_argument("--json", dest="json_out", default=None,
help="Also write the results as JSON to this path.")
mock = p.add_argument_group("mock backend tuning")
mock.add_argument("--mock-host", default="127.0.0.1")
mock.add_argument("--mock-port", type=int, default=0,
help="Mock backend port (0 = auto-pick a free port).")
mock.add_argument("--mock-models", default="mock",
help="Comma-separated model names the mock advertises.")
mock.add_argument("--mock-ttft-ms", type=float, default=0.0,
help="Simulated prefill latency before the first token (ms).")
mock.add_argument("--mock-tokens", type=int, default=64,
help="Completion length in tokens the mock emits.")
mock.add_argument("--mock-tok-ms", type=float, default=0.0,
help="Simulated per-token decode delay (ms) = inverse of tok/s.")
orch = p.add_argument_group("router orchestration (--mock-backend only)")
orch.add_argument("--router-port", type=int, default=0,
help="Router port (0 = auto-pick a free port).")
orch.add_argument("--router-workers", type=int, default=1,
help="uvicorn --workers for the spawned router.")
orch.add_argument("--router-max-conc", type=int, default=0,
help="max_concurrent_connections in the generated config "
"(0 = match peak concurrency so the router does not queue).")
orch.add_argument("--keep-config", action="store_true",
help="Do not delete the generated temp config on exit.")
return p
def _dump_json(path: str, args, results: list[Stats]) -> None:
out = {
"config": {k: getattr(args, k) for k in (
"api", "model", "stream", "duration", "requests", "warmup", "timeout",
"mock_tokens", "mock_ttft_ms", "mock_tok_ms")},
"stages": [],
}
for st in results:
lat = [s.latency * 1000 for s in st.ok]
ttfts = [s.ttft * 1000 for s in st.ok if s.ttft is not None]
out["stages"].append({
"concurrency": st.concurrency,
"wall_s": st.wall,
"requests": st.n_total,
"ok": st.n_ok,
"errors": st.n_total - st.n_ok,
"rps": (st.n_ok / st.wall) if st.wall else 0.0,
"latency_ms": {p: _pct(lat, p) for p in (50, 90, 95, 99)} | (
{"max": max(lat), "mean": statistics.mean(lat)} if lat else {}),
"ttft_ms": {p: _pct(ttfts, p) for p in (50, 90, 99)} if ttfts else {},
})
Path(path).write_text(json.dumps(out, indent=2))
print(f"\n[driver] wrote JSON results to {path}", flush=True)
def main() -> None:
args = build_parser().parse_args()
if args.serve_mock:
try:
serve_mock(args)
except KeyboardInterrupt:
pass
return
if args.requests is not None and args.requests <= 0:
print("--requests must be > 0", file=sys.stderr)
sys.exit(2)
if args.mock_backend:
results = asyncio.run(run_with_mock_backend(args)) or []
else:
results = asyncio.run(run_driver(args))
if args.json_out and results:
_dump_json(args.json_out, args, results)
if __name__ == "__main__":
main()
Reference in a new issue View git blame Copy permalink