| |
| """Warm streaming TTFT, decode-rate, and DSpark acceptance benchmark.""" |
| import argparse |
| import hashlib |
| import json |
| import os |
| import pathlib |
| import re |
| import statistics |
| import time |
| import urllib.request |
|
|
| CASES = [ |
| {"name": "count300", "prompt": "Count from 1 to 300, separated by commas. Return only the numbers and commas.", "max_tokens": 700}, |
| {"name": "dataclasses20", "prompt": "Write valid Python defining exactly 20 dataclasses named Record01 through Record20. Each has fields id: int, name: str, active: bool. Return only one Python code block and end it with the comment # END_DATACLASSES.", "max_tokens": 1400}, |
| {"name": "prose500", "prompt": "Write a coherent 500-word technical explanation of failure isolation in a two-node tensor-parallel inference service. Use plain prose, no headings or bullets, and finish with the exact marker END_PROSE.", "max_tokens": 1100}, |
| ] |
|
|
|
|
| def request_json(url, body=None, timeout=600): |
| data = json.dumps(body).encode() if body is not None else None |
| request = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"}) |
| with urllib.request.urlopen(request, timeout=timeout) as response: |
| return json.load(response) |
|
|
|
|
| def metrics(root): |
| text = urllib.request.urlopen(root + "/metrics", timeout=10).read().decode() |
| out = {} |
| for name in ("vllm:spec_decode_num_draft_tokens_total", "vllm:spec_decode_num_accepted_tokens_total"): |
| matches = re.findall(r"^" + re.escape(name) + r"\{[^\n]*\}\s+([0-9.eE+-]+)$", text, re.M) |
| out[name] = sum(float(value) for value in matches) |
| return out |
|
|
|
|
| def validate(name, content): |
| if name == "count300": |
| return content.strip() == ",".join(map(str, range(1, 301))) |
| if name == "dataclasses20": |
| return all(f"class Record{i:02d}:" in content for i in range(1, 21)) and content.rstrip("`\n ").endswith("# END_DATACLASSES") |
| if name == "prose500": |
| return content.rstrip().endswith("END_PROSE") |
| return False |
|
|
|
|
| def stream_one(root, model, case): |
| body = {"model": model, "messages": [{"role": "user", "content": case["prompt"]}], "temperature": 0, "max_tokens": case["max_tokens"], "stream": True, "stream_options": {"include_usage": True}} |
| request = urllib.request.Request(root + "/v1/chat/completions", data=json.dumps(body).encode(), headers={"Content-Type": "application/json"}) |
| started = time.perf_counter() |
| first = None |
| usage = None |
| pieces = [] |
| finish_reason = None |
| with urllib.request.urlopen(request, timeout=900) as response: |
| for raw in response: |
| line = raw.decode().strip() |
| if not line.startswith("data: ") or line == "data: [DONE]": |
| continue |
| event = json.loads(line[6:]) |
| choices = event.get("choices") or [] |
| if choices: |
| delta = choices[0].get("delta") or {} |
| piece = delta.get("content") or "" |
| if piece and first is None: |
| first = time.perf_counter() |
| pieces.append(piece) |
| finish_reason = choices[0].get("finish_reason") or finish_reason |
| if event.get("usage"): |
| usage = event["usage"] |
| finished = time.perf_counter() |
| content = "".join(pieces) |
| tokens = (usage or {}).get("completion_tokens", 0) |
| decode_elapsed = max(0.001, finished - (first or finished)) |
| return {"ttft_s": (first or finished) - started, "elapsed_s": finished - started, "completion_tokens": tokens, "decode_tok_s": tokens / decode_elapsed, "finish_reason": finish_reason, "content_sha256": hashlib.sha256(content.encode()).hexdigest(), "output_valid": validate(case["name"], content), "usage": usage} |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--base-url", default=os.environ.get("BASE_URL", "http://127.0.0.1:8000")) |
| parser.add_argument("--model", default=os.environ.get("MODEL", "deepseek-v4-flash-0731")) |
| parser.add_argument("--label", default="graph8-reproduction") |
| parser.add_argument("--output", default="benchmark.json") |
| parser.add_argument("--repetitions", type=int, default=3) |
| args = parser.parse_args() |
| root = args.base_url.rstrip("/") |
| report = {"schema_version": 1, "label": args.label, "model": args.model, "repetitions": args.repetitions, "cases": []} |
| for case in CASES: |
| stream_one(root, args.model, case) |
| before = metrics(root) |
| runs = [stream_one(root, args.model, case) for _ in range(args.repetitions)] |
| after = metrics(root) |
| draft = after["vllm:spec_decode_num_draft_tokens_total"] - before["vllm:spec_decode_num_draft_tokens_total"] |
| accepted = after["vllm:spec_decode_num_accepted_tokens_total"] - before["vllm:spec_decode_num_accepted_tokens_total"] |
| entry = {"name": case["name"], "median_decode_tok_s": statistics.median(r["decode_tok_s"] for r in runs), "median_ttft_s": statistics.median(r["ttft_s"] for r in runs), "draft_tokens_delta": draft, "accepted_tokens_delta": accepted, "acceptance_pct": 100 * accepted / draft if draft else None, "all_outputs_valid": all(r["output_valid"] for r in runs), "runs": runs} |
| report["cases"].append(entry) |
| pathlib.Path(args.output).write_text(json.dumps(report, indent=2, sort_keys=True) + "\n") |
| print(json.dumps({key: value for key, value in entry.items() if key != "runs"}, sort_keys=True), flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
