Coverage for orchestrant/benchmark/lanes.py: 86%
207 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-18 00:48 +0000
1#!/usr/bin/env python3
2"""Concurrency benchmarks across one or more serving endpoints.
4--batching: does ONE server overlap concurrent requests?
5--lanes: do several servers ADD UP, or fight each other?
7Usage:
8 # does this server batch?
9 python3 bench_lanes.py --batching --backend ollama
11 # do these lanes add up? (names come from backends.json)
12 python3 bench_lanes.py --lanes geniex-npu geniex-cpu
14 # ...or spell an endpoint out in full
15 python3 bench_lanes.py --lanes \
16 npu=http://127.0.0.1:18181,model=qualcomm/Qwen3-4B-Instruct-2507:W4A16
17"""
19import argparse
20import json
21import random
22import threading
23import time
24import urllib.request
26from orchestrant.benchmark.answers import delta_pieces
29DEFAULT_PROMPT = (
30 "Write a Python function that merges two sorted lists. Explain briefly."
31)
34def stream_once(base_url, model, prompt, max_tokens=256, timeout=900, deadline=900):
35 """One streaming request; returns a timing dict and never raises.
37 `timeout` applies per socket read; `deadline` bounds the whole request (`gave_up`).
38 """
39 body = json.dumps(
40 {
41 "model": model,
42 "stream": True,
43 "max_tokens": max_tokens,
44 "temperature": 0,
45 "messages": [{"role": "user", "content": prompt}],
46 }
47 ).encode()
48 req = urllib.request.Request(
49 f"{base_url}/v1/chat/completions",
50 data=body,
51 headers={"Content-Type": "application/json"},
52 )
54 started = time.monotonic()
55 ttft = None
56 tokens = 0
57 gave_up = False
58 try:
59 with urllib.request.urlopen(req, timeout=timeout) as resp: # nosec B310 -- endpoint URL from config
60 for raw in resp:
61 if deadline and time.monotonic() - started > deadline:
62 gave_up = True
63 break
64 line = raw.decode("utf-8", "replace").strip()
65 # The space after "data:" is optional per the SSE spec.
66 if not line.startswith("data:"):
67 continue
68 payload = line[5:].lstrip()
69 if payload == "[DONE]":
70 break
71 try:
72 chunk = json.loads(payload)
73 except json.JSONDecodeError:
74 continue
75 choices = chunk.get("choices") or []
76 # Thinking is output too: uncounted, TTFT would read as the whole thought.
77 if choices and any(delta_pieces(choices[0].get("delta") or {})):
78 if ttft is None:
79 ttft = time.monotonic() - started
80 tokens += 1
81 except Exception as e:
82 return {
83 "error": f"{type(e).__name__}: {e}",
84 "wall_s": time.monotonic() - started,
85 }
87 wall = time.monotonic() - started
88 decode_window = wall - (ttft or 0)
89 return {
90 "ttft_s": round(ttft, 3) if ttft is not None else None,
91 "gave_up": gave_up,
92 "tokens": tokens,
93 "wall_s": round(wall, 2),
94 "decode_tok_per_sec": round((tokens - 1) / decode_window, 2)
95 if decode_window > 0 and tokens > 1
96 else 0.0,
97 }
100def _secs(value, width=6):
101 """Format a possibly-missing duration: a lane past its context returns 200 and no tokens."""
102 return (
103 f"{value:{width}.2f}"
104 if isinstance(value, (int, float))
105 else f"{'n/a':>{width}}"
106 )
109def _fresh(prompt):
110 """The prompt with a nonce: a repeated prompt is served from the runtime's cache."""
111 return f"{prompt} (request {random.SystemRandom().randrange(10**6)})"
114def run_parallel(jobs):
115 """Run (label, fn) jobs at the same time; return {label: result}."""
116 out = {}
117 lock = threading.Lock()
119 def work(label, fn):
120 result = fn()
121 with lock:
122 out[label] = result
124 threads = [threading.Thread(target=work, args=(lbl, fn)) for lbl, fn in jobs]
125 wall_start = time.monotonic()
126 for t in threads:
127 t.start()
128 for t in threads:
129 t.join()
130 return out, time.monotonic() - wall_start
133def probe_batching(base_url, model, prompt, max_tokens):
134 """Decide whether one endpoint overlaps two concurrent requests."""
135 print(f"\n Batching probe: two concurrent requests to {base_url}")
136 results, wall = run_parallel(
137 [
138 (f"req{i}", (lambda: stream_once(base_url, model, prompt, max_tokens)))
139 for i in range(2)
140 ]
141 )
143 bad = [k for k, v in results.items() if "error" in v]
144 if bad:
145 for k in bad:
146 print(f" {k}: ERROR {results[k]['error']}")
147 return None
149 if any(r.get("ttft_s") is None for r in results.values()):
150 print(
151 " VERDICT: inconclusive — a request returned no tokens at all "
152 "(past a context limit?), so there is no first-token time to compare"
153 )
154 return {
155 "verdict": "inconclusive",
156 "serialised": None,
157 "wall_s": round(wall, 2),
158 "requests": results,
159 }
160 ordered = sorted(results.values(), key=lambda r: r["ttft_s"])
161 first, second = ordered[0], ordered[1]
162 for name, r in sorted(results.items()):
163 print(
164 f" {name}: ttft={_secs(r['ttft_s'])}s "
165 f"{_secs(r['decode_tok_per_sec'])} tok/s wall={_secs(r['wall_s'])}s"
166 )
168 # A second TTFT near the first's wall time means the server ran them in turn.
169 serialised = (second["ttft_s"] or 0) >= first["wall_s"] * 0.8
170 verdict = "SERIALISED (no batching)" if serialised else "OVERLAPPED (batching)"
171 print(f" wall for both: {wall:.2f}s")
172 print(f" VERDICT: {verdict}")
173 if serialised:
174 print(" -> extra throughput needs MORE SERVERS, not more clients.")
175 return {
176 "verdict": verdict,
177 "serialised": serialised,
178 "wall_s": round(wall, 2),
179 "requests": results,
180 }
183def _serving(lanes):
184 """Each lane's build, serve flags and model files, before measuring; {name: runtime}."""
185 from orchestrant.benchmark.provenance import lane_runtimes, runtime_label
187 runtimes = lane_runtimes(lanes)
188 print("\n Serving:")
189 for name, runtime in runtimes.items():
190 print(f" {name:10s} {runtime_label(runtime)}")
191 return runtimes
194def run_lanes(lanes, prompt, max_tokens, sequential_baseline=True):
195 """Per-lane and aggregate throughput when lanes run together."""
196 report = {"lanes": {}, "baseline": {}, "runtimes": _serving(lanes)}
198 if sequential_baseline:
199 print("\n Baseline — each lane alone:")
200 for name, (url, model) in lanes.items():
201 r = stream_once(url, model, _fresh(prompt), max_tokens)
202 report["baseline"][name] = r
203 if "error" in r:
204 print(f" {name:10s} ERROR {r['error']}")
205 else:
206 print(
207 f" {name:10s} {_secs(r['decode_tok_per_sec'])} tok/s "
208 f"ttft={_secs(r['ttft_s'])}s"
209 )
211 print("\n Together — all lanes at once:")
212 results, wall = run_parallel(
213 [
214 (
215 name,
216 (lambda u=url, m=model: stream_once(u, m, _fresh(prompt), max_tokens)),
217 )
218 for name, (url, model) in lanes.items()
219 ]
220 )
221 report["lanes"] = results
223 total, tokens = 0.0, 0
224 for name in lanes:
225 r = results[name]
226 if "error" in r:
227 print(f" {name:10s} ERROR {r['error']}")
228 continue
229 total += r["decode_tok_per_sec"]
230 tokens += r["tokens"]
231 base = report["baseline"].get(name, {}).get("decode_tok_per_sec")
232 delta = ""
233 if base:
234 change = 100 * (r["decode_tok_per_sec"] - base) / base
235 delta = f" ({change:+.0f}% vs alone)"
236 print(
237 f" {name:10s} {_secs(r['decode_tok_per_sec'])} tok/s "
238 f"ttft={_secs(r['ttft_s'])}s{delta}"
239 )
241 best_alone = max(
242 (v.get("decode_tok_per_sec", 0) for v in report["baseline"].values()), default=0
243 )
244 # Tokens over the joint wall: a per-lane sum counts an early finisher's solo seconds.
245 delivered = tokens / wall if wall > 0 else 0.0
246 print(
247 f"\n AGGREGATE: {total:.1f} tok/s summed over {len(lanes)} lanes; "
248 f"delivered {delivered:.1f} tok/s ({tokens} tokens in the {wall:.1f}s wall)"
249 )
250 if best_alone:
251 print(
252 f" Best single lane alone: {best_alone:.1f} tok/s "
253 f"-> {delivered / best_alone:.2f}x delivered by running lanes together "
254 f"({total / best_alone:.2f}x by the sum of rates)"
255 )
256 print(
257 " NOTE: aggregate only materialises with that many CONCURRENT "
258 "requests.\n One agent waiting for one answer still sees a "
259 "single lane's speed."
260 )
261 report["aggregate_tok_per_sec"] = round(total, 2)
262 report["delivered_tok_per_sec"] = round(delivered, 2)
263 report["wall_s"] = round(wall, 2)
264 return report
267def sweep_nctx(base_url_template, model, values, prompt, max_tokens=256):
268 """Does --nctx change anything but memory? No CLI flag reaches this; call it from a script."""
269 print("\n nctx sweep (restart the lane between values):", flush=True)
270 rows = []
271 for value in values:
272 r = stream_once(base_url_template, model, prompt, max_tokens)
273 if "error" in r:
274 print(f" nctx={value:6d} ERROR {r['error'][:50]}", flush=True)
275 rows.append({"nctx": value, "error": r["error"][:120]})
276 continue
277 print(
278 f" nctx={value:6d} {r['decode_tok_per_sec']:6.2f} tok/s "
279 f"ttft={_secs(r['ttft_s'])}s",
280 flush=True,
281 )
282 rows.append({"nctx": value, **r})
283 return rows
286def parse_lane(spec):
287 """Parse 'name=URL,model=MODEL' into (name, url, model)."""
288 try:
289 head, model_part = spec.split(",model=", 1)
290 name, url = head.split("=", 1)
291 except ValueError:
292 raise argparse.ArgumentTypeError(
293 f"bad lane {spec!r}; expected name=URL,model=MODEL"
294 )
295 return name.strip(), url.strip().rstrip("/"), model_part.strip()
298def resolve_lane(spec, path=None):
299 """Accept a full 'name=URL,model=MODEL' spec or a bare backend name from `path`'s registry."""
300 if "=" in spec:
301 return parse_lane(spec)
303 from orchestrant.benchmark.openai_api import load_backends
305 backends, _ = load_backends(path)
306 if spec not in backends:
307 known = ", ".join(sorted(backends)) or "(none configured)"
308 raise argparse.ArgumentTypeError(
309 f"unknown backend {spec!r}. Known: {known}. "
310 "Or give a full spec: name=URL,model=MODEL"
311 )
312 entry = backends[spec]
313 model = entry.get("model")
314 if not model:
315 raise argparse.ArgumentTypeError(
316 f"backend {spec!r} has no default model in backends.json; "
317 f"use {spec}={entry['base_url']},model=<model> instead"
318 )
319 return spec, entry["base_url"].rstrip("/"), model
322def build_reports(batching=None, batching_endpoint=None, lane_run=None, lanes=None):
323 """Shape the probes' output as the shared envelope's reports[] rows; none is a score."""
324 reports = []
325 if batching_endpoint is not None:
326 url, model = batching_endpoint
327 row = {"label": "batching", "model": model, "base_url": url}
328 row.update(batching or {"verdict": "error", "serialised": None, "requests": {}})
329 reports.append(row)
330 if lane_run is not None:
331 for name, (url, model) in (lanes or {}).items():
332 together = lane_run["lanes"].get(name, {})
333 reports.append(
334 {
335 "label": name,
336 "model": model,
337 "base_url": url,
338 "alone": lane_run["baseline"].get(name),
339 "together": together,
340 "tok_per_sec": together.get("decode_tok_per_sec"),
341 "runtime": (lane_run.get("runtimes") or {}).get(name),
342 }
343 )
344 reports.append(
345 {
346 "label": "aggregate",
347 "model": None,
348 "base_url": None,
349 # Not sorted(lanes or {}): Cython 3.3's sorted() optimisation crashes on it.
350 "lanes": sorted(lanes) if lanes else [],
351 # bench_compare judges tokens over the joint wall; older reports carry the sum.
352 "tok_per_sec": lane_run.get(
353 "delivered_tok_per_sec", lane_run["aggregate_tok_per_sec"]
354 ),
355 "summed_tok_per_sec": lane_run["aggregate_tok_per_sec"],
356 "wall_s": lane_run["wall_s"],
357 }
358 )
359 return reports
362def main():
363 ap = argparse.ArgumentParser(
364 description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
365 )
366 ap.add_argument(
367 "--batching",
368 action="store_true",
369 help="LB5: does one endpoint overlap concurrent requests?",
370 )
371 ap.add_argument(
372 "--lanes",
373 nargs="+",
374 metavar="BACKEND|name=URL,model=MODEL",
375 help="LB4: drive these endpoints simultaneously. Either a "
376 "backend name from backends.json (e.g. geniex-npu) or "
377 "a full name=URL,model=MODEL spec.",
378 )
379 ap.add_argument(
380 "--base-url",
381 "--endpoint",
382 dest="base_url",
383 default=None,
384 help="Endpoint URL for --batching (overrides --backend). "
385 "--endpoint is kept as an alias: it was the original "
386 "spelling here while both sibling tools used "
387 "--base-url, and an operator scripting the suite had "
388 "to remember which tool wanted which.",
389 )
390 ap.add_argument(
391 "--backend",
392 default=None,
393 help="Named backend from backends.json for --batching",
394 )
395 ap.add_argument("--model", default=None, help="Model for --batching")
396 ap.add_argument("--prompt", default=DEFAULT_PROMPT)
397 ap.add_argument("--max-tokens", type=int, default=256)
398 ap.add_argument(
399 "--no-baseline",
400 action="store_true",
401 help="Skip the one-lane-at-a-time baseline",
402 )
403 ap.add_argument("--output", default=None, help="Write the report as JSON")
404 args = ap.parse_args()
406 if not args.batching and not args.lanes:
407 ap.error("nothing to do: pass --batching and/or --lanes")
409 from orchestrant.benchmark.client import run_start, write_report
411 batching = endpoint = lane_run = lanes = None
413 # Resolved first, so the start record can subtract the measured lane from the host load.
414 if args.batching:
415 from orchestrant.benchmark.openai_api import (
416 detect_model_via_api,
417 resolve_backend,
418 )
420 url, backend_model, source = resolve_backend(args.backend, args.base_url)
421 model = args.model or backend_model or detect_model_via_api(url)
422 print(f" Endpoint: {url} (from {source})")
423 endpoint = (url, model)
424 if args.lanes:
425 lanes = {
426 name: (url, model) for name, url, model in map(resolve_lane, args.lanes)
427 }
428 base_url = endpoint[0] if endpoint else next(iter((lanes or {}).values()))[0]
429 tool_files = ("lanes.py", "answers.py")
430 run = run_start(tool_files, base_url)
432 if endpoint:
433 batching = probe_batching(*endpoint, args.prompt, args.max_tokens)
434 if lanes:
435 lane_run = run_lanes(
436 lanes,
437 args.prompt,
438 args.max_tokens,
439 sequential_baseline=not args.no_baseline,
440 )
442 print()
443 if args.output:
444 # The shared envelope, so bench_report labels it and bench_compare can diff it.
445 write_report(
446 args.output,
447 "bench_lanes",
448 {
449 "prompt": args.prompt,
450 "max_tokens": args.max_tokens,
451 "baseline": not args.no_baseline,
452 },
453 build_reports(batching, endpoint, lane_run, lanes),
454 base_url,
455 tool_files,
456 run_start=run,
457 )
458 print(f" Report written to {args.output}")
461if __name__ == "__main__":
462 main()