| """Latency benchmark for Groot checkpoints (and, optionally, the Jev API). |
| |
| groot-bench --model runs/groot-1-s/final [--device cpu --threads 8] [--jev] |
| |
| Measures p50/p95 end-to-end latency of `GrootEngine.answer` for synthetic |
| requests at several state sizes and question counts. With --jev and a |
| TYPESAFE_API_KEY in the environment, the same requests are sent to |
| https://api.typesafe.ai/v1/systemone for a side-by-side table. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| import statistics |
| import time |
|
|
| import torch |
|
|
| from .schema import SystemOneRequest |
|
|
| LOREM = ( |
| "The customer reported that the invoice was charged twice on the corporate card and asked for a refund. " |
| "Support confirmed the duplicate transaction and escalated the case to billing with high priority. " |
| ) |
|
|
|
|
| def synthetic_request(state_words: int, n_questions: int) -> SystemOneRequest: |
| words = (LOREM * (state_words // len(LOREM.split()) + 1)).split()[:state_words] |
| state = " ".join(words) |
| questions = {} |
| kinds = ["choice", "score", "noul"] |
| for i in range(n_questions): |
| k = kinds[i % 3] |
| if k == "choice": |
| questions[f"q{i}"] = {"type": "choice", "instructions": "Which team should handle this?", "criteria": {"billing": "payments", "shipping": "delivery", "access": "login", "other": "anything else"}} |
| elif k == "score": |
| questions[f"q{i}"] = {"type": "score", "instructions": "How urgent is it?", "criteria": ["can wait", "this week", "today"]} |
| else: |
| questions[f"q{i}"] = {"type": "noul", "instructions": "Is the customer asking for a refund?"} |
| return SystemOneRequest(state=state, questions=questions) |
|
|
|
|
| def time_fn(fn, warmup: int = 2, runs: int = 10) -> dict: |
| for _ in range(warmup): |
| fn() |
| times = [] |
| for _ in range(runs): |
| t0 = time.perf_counter() |
| fn() |
| times.append((time.perf_counter() - t0) * 1000) |
| times.sort() |
| return {"p50_ms": round(statistics.median(times), 1), "p95_ms": round(times[int(0.95 * (len(times) - 1))], 1), "min_ms": round(times[0], 1)} |
|
|
|
|
| def main(argv: list[str] | None = None) -> None: |
| ap = argparse.ArgumentParser(prog="groot-bench") |
| ap.add_argument("--model", required=True) |
| ap.add_argument("--device", default=None) |
| ap.add_argument("--threads", type=int, default=None) |
| ap.add_argument("--runs", type=int, default=10) |
| ap.add_argument("--jev", action="store_true") |
| args = ap.parse_args(argv) |
| if args.threads: |
| torch.set_num_threads(args.threads) |
| from .serve import GrootEngine |
|
|
| engine = GrootEngine(args.model, args.device) |
| grid = [(64, 1), (256, 1), (256, 5), (1024, 5), (1024, 20), (4096, 5)] |
| rows = [] |
| for words, nq in grid: |
| req = synthetic_request(words, nq) |
| enc = (engine.members[0] if engine.members else engine).ser.encode_request(req.state, req.questions) |
| res = time_fn(lambda: engine.answer(req.state, req.questions), runs=args.runs) |
| row = {"state_tokens": enc.state.tokens, "chunks": len(enc.state.chunks), "questions": nq, **{f"groot_{k}": v for k, v in res.items()}} |
| if args.jev and os.environ.get("TYPESAFE_API_KEY"): |
| import httpx |
|
|
| body = {"model": "jev-latest", "state": req.state, "questions": {k: v.model_dump(exclude_none=True) for k, v in req.questions.items()}} |
| headers = {"Authorization": f"Bearer {os.environ['TYPESAFE_API_KEY']}"} |
| with httpx.Client(timeout=30) as client: |
| res_j = time_fn(lambda: client.post("https://api.typesafe.ai/v1/systemone", json=body, headers=headers).raise_for_status(), warmup=1, runs=max(3, args.runs // 2)) |
| row.update({f"jev_{k}": v for k, v in res_j.items()}) |
| rows.append(row) |
| print(json.dumps(row), flush=True) |
| print(json.dumps({"device": str(engine.device), "threads": torch.get_num_threads(), "model": engine.name, "rows": rows}, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
