Add bench.py performance benchmark script

bench.py

@@ -0,0 +1,391 @@
+#!/usr/bin/env python3
+"""
+bench.py — performance benchmarks for the vLLM server running Qwen3-Coder-Next-NVFP4-GB10.
+
+Measures:
+  - Time to first token (TTFT) via streaming
+  - Decode throughput (tok/s)
+  - Prefill throughput (prompt tok/s)
+  - Latency across prompt lengths: short / medium / long / max
+  - Concurrent request throughput (1, 4, 8, 16 parallel requests)
+  - Reasoning ON vs OFF overhead
+
+Usage:
+  python3 bench.py
+  python3 bench.py --host 192.168.1.50
+  python3 bench.py --host localhost --port 8000 --runs 3
+"""
+
+import argparse
+import json
+import statistics
+import sys
+import threading
+import time
+from dataclasses import dataclass, field
+from typing import Optional
+
+import urllib.request
+import urllib.error
+
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
+parser = argparse.ArgumentParser()
+parser.add_argument("--host", default="localhost")
+parser.add_argument("--port", type=int, default=8000)
+parser.add_argument("--runs", type=int, default=3, help="Runs per scenario (default: 3)")
+parser.add_argument("--no-color", action="store_true")
+args = parser.parse_args()
+
+BASE_URL = f"http://{args.host}:{args.port}/v1"
+
+if args.no_color or not sys.stdout.isatty():
+    GREEN = RED = YELLOW = CYAN = BOLD = NC = ""
+else:
+    GREEN  = "\033[0;32m"
+    RED    = "\033[0;31m"
+    YELLOW = "\033[0;33m"
+    CYAN   = "\033[0;36m"
+    BOLD   = "\033[1m"
+    NC     = "\033[0m"
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def get_model_id() -> str:
+    req = urllib.request.Request(f"{BASE_URL}/models")
+    with urllib.request.urlopen(req, timeout=10) as r:
+        data = json.loads(r.read())
+    return data["data"][0]["id"]
+
+
+def chat_stream(model: str, messages: list, max_tokens: int, enable_thinking: bool) -> tuple[float, float, int, int]:
+    """
+    Send a streaming chat completion request.
+    Returns: (ttft_s, total_s, prompt_tokens, completion_tokens)
+    """
+    payload = json.dumps({
+        "model": model,
+        "messages": messages,
+        "max_tokens": max_tokens,
+        "temperature": 0.1,
+        "stream": True,
+        "stream_options": {"include_usage": True},
+        "chat_template_kwargs": {"enable_thinking": enable_thinking},
+    }).encode()
+
+    req = urllib.request.Request(
+        f"{BASE_URL}/chat/completions",
+        data=payload,
+        headers={"Content-Type": "application/json"},
+        method="POST",
+    )
+
+    ttft = None
+    t0 = time.perf_counter()
+    prompt_tokens = 0
+    completion_tokens = 0
+
+    with urllib.request.urlopen(req, timeout=300) as resp:
+        for raw_line in resp:
+            line = raw_line.decode().strip()
+            if not line.startswith("data:"):
+                continue
+            chunk = line[5:].strip()
+            if chunk == "[DONE]":
+                break
+            try:
+                obj = json.loads(chunk)
+            except json.JSONDecodeError:
+                continue
+
+            # First token
+            if ttft is None:
+                choices = obj.get("choices", [])
+                if choices:
+                    delta = choices[0].get("delta", {})
+                    content = delta.get("content") or delta.get("reasoning_content")
+                    if content:
+                        ttft = time.perf_counter() - t0
+
+            # Usage (last chunk)
+            usage = obj.get("usage")
+            if usage:
+                prompt_tokens = usage.get("prompt_tokens", 0)
+                completion_tokens = usage.get("completion_tokens", 0)
+
+    total = time.perf_counter() - t0
+    if ttft is None:
+        ttft = total
+    return ttft, total, prompt_tokens, completion_tokens
+
+
+@dataclass
+class Result:
+    name: str
+    ttft_ms: list[float] = field(default_factory=list)
+    decode_tps: list[float] = field(default_factory=list)
+    prefill_tps: list[float] = field(default_factory=list)
+    total_s: list[float] = field(default_factory=list)
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+
+
+def run_scenario(name: str, model: str, messages: list, max_tokens: int,
+                 enable_thinking: bool = False, runs: int = 3) -> Result:
+    res = Result(name=name)
+    print(f"  {CYAN}{name}{NC}", end="", flush=True)
+    for i in range(runs):
+        try:
+            ttft, total, pt, ct = chat_stream(model, messages, max_tokens, enable_thinking)
+            decode_time = total - ttft
+            res.ttft_ms.append(ttft * 1000)
+            res.decode_tps.append(ct / decode_time if decode_time > 0.01 else 0)
+            res.prefill_tps.append(pt / ttft if ttft > 0.01 else 0)
+            res.total_s.append(total)
+            res.prompt_tokens = pt
+            res.completion_tokens = ct
+            print(f" {GREEN}·{NC}", end="", flush=True)
+        except Exception as e:
+            print(f" {RED}✗{NC}", end="", flush=True)
+    print()
+    return res
+
+
+def print_result(res: Result):
+    if not res.ttft_ms:
+        print(f"    {RED}all runs failed{NC}")
+        return
+    ttft_med = statistics.median(res.ttft_ms)
+    dtps_med = statistics.median(res.decode_tps)
+    ptps_med = statistics.median(res.prefill_tps)
+    total_med = statistics.median(res.total_s)
+    print(f"    prompt tokens   : {res.prompt_tokens}")
+    print(f"    completion tok  : {res.completion_tokens}")
+    print(f"    TTFT (median)   : {BOLD}{ttft_med:.0f} ms{NC}")
+    print(f"    decode (median) : {BOLD}{dtps_med:.1f} tok/s{NC}")
+    print(f"    prefill (median): {ptps_med:.0f} tok/s")
+    print(f"    total (median)  : {total_med:.1f} s")
+
+
+def run_concurrent(model: str, messages: list, max_tokens: int, concurrency: int,
+                   enable_thinking: bool = False) -> tuple[float, float]:
+    """Fire `concurrency` requests simultaneously, return (wall_s, aggregate_tps)."""
+    results = [None] * concurrency
+    errors = [None] * concurrency
+
+    def worker(idx):
+        try:
+            _, total, pt, ct = chat_stream(model, messages, max_tokens, enable_thinking)
+            results[idx] = (total, ct)
+        except Exception as e:
+            errors[idx] = e
+
+    threads = [threading.Thread(target=worker, args=(i,)) for i in range(concurrency)]
+    t0 = time.perf_counter()
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join()
+    wall = time.perf_counter() - t0
+
+    total_tokens = sum(r[1] for r in results if r)
+    agg_tps = total_tokens / wall if wall > 0 else 0
+    return wall, agg_tps
+
+
+# ---------------------------------------------------------------------------
+# Prompts
+# ---------------------------------------------------------------------------
+SHORT_PROMPT = "Write a Python one-liner that reverses a string."
+
+MEDIUM_PROMPT = (
+    "Write a Python class implementing a generic LRU cache with O(1) get and put. "
+    "Use OrderedDict. Include docstrings, type annotations, and a short usage example."
+)
+
+LONG_PROMPT = (
+    "You are a senior software engineer. Review the following Python code and provide "
+    "a detailed analysis covering: correctness, edge cases, performance, readability, "
+    "and security concerns. Suggest concrete improvements with code examples.\n\n"
+    "```python\n"
+    + "\n".join([
+        "import subprocess, os, json",
+        "from flask import Flask, request",
+        "",
+        "app = Flask(__name__)",
+        "",
+        "def run_query(user_input):",
+        "    result = subprocess.run(",
+        "        f'mysql -u root -ppassword mydb -e \"{user_input}\"',",
+        "        shell=True, capture_output=True, text=True",
+        "    )",
+        "    return result.stdout",
+        "",
+        "@app.route('/query')",
+        "def query():",
+        "    data = request.args.get('q', '')",
+        "    output = run_query(data)",
+        "    return json.dumps({'result': output, 'debug': os.environ})",
+        "",
+        "if __name__ == '__main__':",
+        "    app.run(debug=True, host='0.0.0.0')",
+    ])
+    + "\n```"
+)
+
+# ~2000 token prompt via repeated context
+CONTEXT_PROMPT = (
+    "You are given the following context about a distributed system architecture. "
+    "After reading it carefully, answer the questions at the end.\n\n"
+    + ("Context: " + "A microservices-based e-commerce platform consists of the following services: "
+       "UserService (authentication, profiles), ProductService (catalog, search), "
+       "OrderService (cart, checkout, order management), PaymentService (Stripe integration), "
+       "NotificationService (email/SMS), and AnalyticsService (event tracking). "
+       "All services communicate via gRPC internally and expose REST APIs externally. "
+       "A Redis cluster handles session data and caching. PostgreSQL with read replicas "
+       "serves as the primary database. Kafka handles async event streaming between services. "
+       "Kubernetes on AWS EKS manages deployment with HPA for auto-scaling. "
+       "A global CDN sits in front of the API gateway. ") * 12
+    + "\n\nQuestions:\n"
+    "1. What are the main single points of failure in this architecture?\n"
+    "2. How would you handle a PaymentService outage gracefully?\n"
+    "3. What observability stack would you recommend and why?\n"
+    "4. Suggest a strategy for zero-downtime database migrations.\n"
+)
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+def main():
+    print(f"\n{BOLD}{'='*60}{NC}")
+    print(f"{BOLD}  vLLM Performance Benchmark{NC}")
+    print(f"{BOLD}  {BASE_URL}{NC}")
+    print(f"{BOLD}{'='*60}{NC}\n")
+
+    # Check server
+    try:
+        model = get_model_id()
+        print(f"{GREEN}[OK]{NC} Server up. Model: {model}\n")
+    except Exception as e:
+        print(f"{RED}[FAIL]{NC} Cannot reach server: {e}")
+        sys.exit(1)
+
+    runs = args.runs
+    results = []
+
+    # -----------------------------------------------------------------------
+    # 1. Latency across prompt lengths
+    # -----------------------------------------------------------------------
+    print(f"{BOLD}1. Latency across prompt lengths (reasoning OFF){NC}")
+    print(f"   ({runs} runs each, streaming, median reported)\n")
+
+    for name, prompt, max_tok in [
+        ("short  (~10 prompt tok,  200 output)", SHORT_PROMPT,   200),
+        ("medium (~80 prompt tok,  500 output)", MEDIUM_PROMPT,  500),
+        ("long   (~400 prompt tok, 800 output)", LONG_PROMPT,    800),
+        ("ctx    (~2K prompt tok,  600 output)", CONTEXT_PROMPT, 600),
+    ]:
+        messages = [{"role": "user", "content": prompt}]
+        res = run_scenario(name, model, messages, max_tok, enable_thinking=False, runs=runs)
+        print_result(res)
+        results.append(res)
+        print()
+
+    # -----------------------------------------------------------------------
+    # 2. Reasoning ON vs OFF
+    # -----------------------------------------------------------------------
+    print(f"{BOLD}2. Reasoning ON vs OFF (medium prompt, 800 output tokens){NC}\n")
+
+    messages = [{"role": "user", "content": MEDIUM_PROMPT}]
+    for label, thinking in [("reasoning OFF", False), ("reasoning ON ", True)]:
+        res = run_scenario(label, model, messages, 800, enable_thinking=thinking, runs=runs)
+        print_result(res)
+        print()
+
+    # -----------------------------------------------------------------------
+    # 3. Concurrent requests
+    # -----------------------------------------------------------------------
+    print(f"{BOLD}3. Concurrent requests throughput (short prompt, 300 output tok){NC}\n")
+    messages = [{"role": "user", "content": SHORT_PROMPT}]
+
+    print(f"  {'concurrency':<14} {'wall_s':>8} {'agg tok/s':>12}")
+    print(f"  {'-'*36}")
+    for c in [1, 2, 4, 8, 16]:
+        wall, agg = run_concurrent(model, messages, 300, c, enable_thinking=False)
+        print(f"  {c:<14} {wall:>8.1f} {agg:>12.1f}")
+    print()
+
+    # -----------------------------------------------------------------------
+    # 4. Tool calling smoke
+    # -----------------------------------------------------------------------
+    print(f"{BOLD}4. Tool calling latency{NC}\n")
+
+    tool_messages = [{"role": "user", "content": "What is the weather in Warsaw? Use the get_weather tool."}]
+    tool_payload = json.dumps({
+        "model": model,
+        "messages": tool_messages,
+        "max_tokens": 200,
+        "temperature": 0.1,
+        "tools": [{
+            "type": "function",
+            "function": {
+                "name": "get_weather",
+                "description": "Get current weather for a city",
+                "parameters": {
+                    "type": "object",
+                    "properties": {"city": {"type": "string"}},
+                    "required": ["city"],
+                }
+            }
+        }],
+        "tool_choice": "auto",
+        "chat_template_kwargs": {"enable_thinking": False},
+    }).encode()
+
+    times = []
+    print(f"  tool_call latency", end="", flush=True)
+    for _ in range(runs):
+        try:
+            req = urllib.request.Request(
+                f"{BASE_URL}/chat/completions",
+                data=tool_payload,
+                headers={"Content-Type": "application/json"},
+                method="POST",
+            )
+            t0 = time.perf_counter()
+            with urllib.request.urlopen(req, timeout=60) as resp:
+                data = json.loads(resp.read())
+            elapsed = time.perf_counter() - t0
+            tool_calls = data["choices"][0]["message"].get("tool_calls")
+            if tool_calls:
+                times.append(elapsed * 1000)
+                print(f" {GREEN}·{NC}", end="", flush=True)
+            else:
+                print(f" {YELLOW}?{NC}", end="", flush=True)
+        except Exception:
+            print(f" {RED}✗{NC}", end="", flush=True)
+    print()
+    if times:
+        print(f"    latency (median): {BOLD}{statistics.median(times):.0f} ms{NC}")
+    print()
+
+    # -----------------------------------------------------------------------
+    # Summary
+    # -----------------------------------------------------------------------
+    print(f"{BOLD}{'='*60}{NC}")
+    print(f"{BOLD}  Summary{NC}")
+    print(f"{BOLD}{'='*60}{NC}")
+    print(f"  {'scenario':<40} {'TTFT ms':>8} {'tok/s':>8}")
+    print(f"  {'-'*58}")
+    for res in results:
+        if res.ttft_ms:
+            print(f"  {res.name:<40} {statistics.median(res.ttft_ms):>8.0f} {statistics.median(res.decode_tps):>8.1f}")
+    print()
+
+
+if __name__ == "__main__":
+    main()