diff --git a/CHANGELOG.md b/CHANGELOG.md index bc95f09..554d45f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- `make bench-serving` — serving-layer benchmark replaying a scripted + conversation through the pipeline; measures per-turn wall time, TTFT + (`--stream`), token usage and VRAM peak, with an optional interleaved + second tenant. Baseline instrument for the serving migration. + ### Changed - `make setup` now ends with a `pytest --collect-only` pass so a broken environment diff --git a/Makefile b/Makefile index 7429b40..9e9c0ea 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -.PHONY: help setup run test test-unit test-integration test-contracts lint typecheck clean +.PHONY: help setup run test test-unit test-integration test-contracts bench-serving lint typecheck clean VENV := .venv PYTHON := $(VENV)/bin/python @@ -50,6 +50,9 @@ test-integration: ## Run integration tests (needs Claude/Ollama) test-contracts: ## Wire-level contract tests against live service boundaries $(PYTEST) tests/contracts -v --no-cov +bench-serving: ## Serving benchmark vs a live backend (BENCH_ARGS="--base-url ... --stream") + $(PYTHON) scripts/bench_serving.py $(BENCH_ARGS) + lint: ## Run ruff linter and formatter check $(RUFF) check src tests $(RUFF) format --check src tests diff --git a/scripts/bench_serving.py b/scripts/bench_serving.py new file mode 100644 index 0000000..a1128ea --- /dev/null +++ b/scripts/bench_serving.py @@ -0,0 +1,355 @@ +"""Serving-layer benchmark: a scripted conversation through the real pipeline. + +Measures what the serving rework is judged on (workspace serving plan, +2026-09-11): per-turn wall time, time-to-first-token in --stream mode, token +usage per turn, and VRAM peak while the conversation runs. Run it against a +backend to produce a baseline, run it again after a serving change, and +compare the JSON artifacts side by side. + +The conversation replays full history each turn, exactly like a real client. +Planted facts are prefixed BENCH-MARKER so anything the biographer learns +from a bench run is recognizable in memory stores later. + +Usage: + .venv/bin/python scripts/bench_serving.py # dev server + .venv/bin/python scripts/bench_serving.py --base-url http://localhost:8000 + .venv/bin/python scripts/bench_serving.py --stream # TTFT mode + .venv/bin/python scripts/bench_serving.py --interleave-url http://localhost:11434/v1 --interleave-every 20 + +Exit codes: 0 ok, 69 backend unreachable. +""" + +from __future__ import annotations + +import argparse +import json +import subprocess +import sys +import threading +import time +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +import httpx + +TURN_TIMEOUT_S = 300.0 + +# Labels name the pipeline path each turn is expected to exercise. +CONVERSATION: list[tuple[str, str]] = [ + ( + "plant-fact", + "For the record: BENCH-MARKER my favorite tea is lapsang souchong. " + "Please remember that.", + ), + ("tool-arithmetic", "What is 61 plus 12?"), + ("small-talk", "Good evening Tatlock, how are the household systems?"), + ("memory-recall", "What did I tell you my favorite tea was?"), + ( + "librarian-routed", + "What do you know about the desklock device in the living room?", + ), + ("tool-arithmetic-2", "Multiply 17 by 23, then add 100."), + ("synthesis-light", "Suggest a two-line dinner idea using rice."), + ( + "context-summary", + "Summarize what we discussed in this conversation in one sentence.", + ), +] + + +class VramPoller: + """Samples nvidia-smi memory.used until stopped; keeps peak and samples.""" + + def __init__(self, interval_s: float = 0.5): + self.interval_s = interval_s + self.samples_mib: list[int] = [] + self._stop = threading.Event() + self._thread = threading.Thread(target=self._run, daemon=True) + + def _read_mib(self) -> int | None: + try: + out = subprocess.run( + [ + "nvidia-smi", + "--query-gpu=memory.used", + "--format=csv,noheader,nounits", + ], + capture_output=True, + text=True, + timeout=5, + ) + return int(out.stdout.strip().splitlines()[0]) + except (OSError, ValueError, IndexError, subprocess.TimeoutExpired): + return None + + def _run(self) -> None: + while not self._stop.is_set(): + mib = self._read_mib() + if mib is not None: + self.samples_mib.append(mib) + self._stop.wait(self.interval_s) + + def __enter__(self) -> VramPoller: + self._thread.start() + return self + + def __exit__(self, *exc: object) -> None: + self._stop.set() + self._thread.join(timeout=5) + + @property + def peak_mib(self) -> int | None: + return max(self.samples_mib) if self.samples_mib else None + + +class Interleaver: + """Fires a plain completion at a second backend on a fixed cadence. + + Simulates an Open-WebUI-style tenant competing for the same card while + the main conversation runs — the slot-eviction case the serving plan + calls out. Each request's wall time is recorded so contention shows on + both sides of the measurement. + """ + + def __init__(self, url: str, model: str, every_s: float): + self.url = url.rstrip("/") + self.model = model + self.every_s = every_s + self.results: list[dict[str, Any]] = [] + self._stop = threading.Event() + self._thread = threading.Thread(target=self._run, daemon=True) + + def _run(self) -> None: + n = 0 + while not self._stop.is_set(): + n += 1 + started = time.monotonic() + try: + r = httpx.post( + f"{self.url}/chat/completions", + json={ + "model": self.model, + "messages": [ + { + "role": "user", + "content": f"Interleave probe {n}: name one European capital.", + } + ], + "max_tokens": 32, + }, + timeout=TURN_TIMEOUT_S, + ) + self.results.append( + { + "n": n, + "wall_s": round(time.monotonic() - started, 3), + "status": r.status_code, + } + ) + except httpx.HTTPError as e: + self.results.append( + { + "n": n, + "wall_s": round(time.monotonic() - started, 3), + "error": type(e).__name__, + } + ) + self._stop.wait(self.every_s) + + def __enter__(self) -> Interleaver: + self._thread.start() + return self + + def __exit__(self, *exc: object) -> None: + self._stop.set() + self._thread.join(timeout=10) + + +def run_turn( + client: httpx.Client, base_url: str, model: str, messages: list[dict[str, str]] +) -> dict[str, Any]: + """One non-streamed turn: wall time and usage.""" + started = time.monotonic() + r = client.post( + f"{base_url}/v1/chat/completions", + json={"model": model, "messages": messages}, + timeout=TURN_TIMEOUT_S, + ) + wall = time.monotonic() - started + r.raise_for_status() + body = r.json() + reply = body["choices"][0]["message"]["content"] or "" + return { + "wall_s": round(wall, 3), + "usage": body.get("usage"), + "reply_chars": len(reply), + "reply": reply, + } + + +def run_turn_streamed( + client: httpx.Client, base_url: str, model: str, messages: list[dict[str, str]] +) -> dict[str, Any]: + """One streamed turn: wall time and time-to-first-token.""" + started = time.monotonic() + ttft: float | None = None + chunks: list[str] = [] + with client.stream( + "POST", + f"{base_url}/v1/chat/completions", + json={"model": model, "messages": messages, "stream": True}, + timeout=TURN_TIMEOUT_S, + ) as r: + r.raise_for_status() + for line in r.iter_lines(): + if not line.startswith("data: "): + continue + payload = line[len("data: ") :] + if payload.strip() == "[DONE]": + break + try: + delta = json.loads(payload)["choices"][0]["delta"] + except (json.JSONDecodeError, KeyError, IndexError): + continue + content = delta.get("content") + if content: + if ttft is None: + ttft = time.monotonic() - started + chunks.append(content) + reply = "".join(chunks) + return { + "wall_s": round(time.monotonic() - started, 3), + "ttft_s": round(ttft, 3) if ttft is not None else None, + "reply_chars": len(reply), + "reply": reply, + } + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", default="http://localhost:8777") + parser.add_argument( + "--model", + default="Tatlock", + help="model id as the target serves it (tatlock registry: 'Tatlock'; raw backend: e.g. 'gemma4:e2b')", + ) + parser.add_argument("--turns", type=int, default=len(CONVERSATION)) + parser.add_argument( + "--stream", + action="store_true", + help="measure time-to-first-token instead of usage", + ) + parser.add_argument( + "--interleave-url", + default=None, + help="OpenAI-compat base (e.g. http://localhost:11434/v1) for a competing tenant", + ) + parser.add_argument("--interleave-every", type=float, default=20.0) + parser.add_argument( + "--interleave-model", + default="gemma4:e2b", + help="model id at the interleave backend (raw Ollama/llama-server tag)", + ) + parser.add_argument("--label", default=None, help="run label stored in the artifact") + parser.add_argument("--out-dir", default="build/bench") + args = parser.parse_args() + + base_url = args.base_url.rstrip("/") + client = httpx.Client() + + try: + health = client.get(f"{base_url}/health", timeout=10).json() + except httpx.HTTPError as e: + print(f"backend unreachable at {base_url}: {type(e).__name__}", file=sys.stderr) + return 69 + + started_at = datetime.now(UTC).isoformat() + turns: list[dict[str, Any]] = [] + messages: list[dict[str, str]] = [] + + interleaver = ( + Interleaver(args.interleave_url, args.interleave_model, args.interleave_every) + if args.interleave_url + else None + ) + + with VramPoller() as vram: + ctx = interleaver if interleaver else _NullCtx() + with ctx: + for label, text in CONVERSATION[: args.turns]: + messages.append({"role": "user", "content": text}) + runner = run_turn_streamed if args.stream else run_turn + try: + result = runner(client, base_url, args.model, messages) + except httpx.HTTPStatusError as e: + detail = e.response.text[:200] + print( + f" {label}: FAILED {e.response.status_code} {detail}", + file=sys.stderr, + ) + turns.append( + { + "label": label, + "error": type(e).__name__, + "status": e.response.status_code, + "detail": detail, + } + ) + messages.pop() + continue + except httpx.HTTPError as e: + print(f" {label}: FAILED ({type(e).__name__})", file=sys.stderr) + turns.append({"label": label, "error": type(e).__name__}) + messages.pop() + continue + messages.append({"role": "assistant", "content": result.pop("reply")}) + turns.append({"label": label, **result}) + shown = result.get("ttft_s") if args.stream else result.get("wall_s") + print(f" {label}: {shown}s ({result})") + + artifact = { + "label": args.label, + "started_at": started_at, + "base_url": base_url, + "model": args.model, + "mode": "stream" if args.stream else "usage", + "backend_health": health, + "interleave": { + "url": args.interleave_url, + "every_s": args.interleave_every, + "results": interleaver.results if interleaver else [], + }, + "vram": {"peak_mib": vram.peak_mib, "samples": len(vram.samples_mib)}, + "turns": turns, + } + + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + stamp = started_at.replace(":", "").replace("-", "")[:15] + out_path = out_dir / f"bench-{stamp}-{args.label or 'run'}.json" + out_path.write_text(json.dumps(artifact, indent=2)) + + ok_walls = [t["wall_s"] for t in turns if "wall_s" in t] + print(f"\nturns: {len(ok_walls)}/{len(turns)} ok") + if ok_walls: + print(f"wall total {sum(ok_walls):.1f}s mean {sum(ok_walls) / len(ok_walls):.1f}s") + if args.stream: + ttfts = [t["ttft_s"] for t in turns if t.get("ttft_s") is not None] + if ttfts: + print(f"ttft mean {sum(ttfts) / len(ttfts):.2f}s max {max(ttfts):.2f}s") + print(f"vram peak {vram.peak_mib} MiB over {len(vram.samples_mib)} samples") + print(f"artifact: {out_path}") + return 0 + + +class _NullCtx: + def __enter__(self) -> _NullCtx: + return self + + def __exit__(self, *exc: object) -> None: + return None + + +if __name__ == "__main__": + sys.exit(main())