Add a server benchmark for continuous batching (ml-explore#728)

Author: Awni Hannun

benchmarks/server_benchmark.py

@@ -0,0 +1,348 @@
+"""
+Spin up the local server:
+
+    mlx_lm.server
+
+Then run the benchmark:
+
+    python server_benchmark.py --concurrency 4
+"""
+
+import argparse
+import asyncio
+import json
+import math
+import time
+from collections import defaultdict
+from itertools import cycle
+from typing import Any, Dict, List, Optional, Tuple
+
+import aiohttp
+from tqdm import tqdm
+
+# Default prompts if no file is provided
+DEFAULT_PROMPTS = [
+    "Explain quantum computing in simple terms.",
+    "What are the main differences between Python and JavaScript?",
+    "Describe the process of photosynthesis in plants.",
+    "How does a neural network learn from data?",
+    "What is the significance of the Turing test in AI?",
+    "Explain the concept of blockchain technology.",
+    "What causes seasons on Earth?",
+    "How do vaccines work in the human body?",
+    "Describe the water cycle and its importance.",
+    "What is the theory of relativity proposed by Einstein?",
+    "How do electric cars help reduce carbon emissions?",
+    "What are the key features of a market economy?",
+    "Explain how DNA replication works in cells.",
+    "What is machine learning and its real-world applications?",
+    "Describe the structure and function of the human heart.",
+]
+
+
+def tokens_per_second(tokens):
+    start = math.floor(tokens[0])
+    stop = math.ceil(tokens[-1])
+    n_bins = int(stop - start) * 10
+    bins = [0] * n_bins
+    for t in tokens:
+        bins[int(n_bins * (t - start) / (stop - start))] += 1
+
+    result = []
+
+    ms = 0
+    cnt = 0
+    for i, b in enumerate(bins):
+        ms += b
+        if cnt == 10:
+            ms -= bins[i - 10]
+        else:
+            cnt += 1
+
+        result.append(10 * ms / cnt)
+
+    times = [start]
+    while times[-1] < stop:
+        times.append(times[-1] + 0.1)
+
+    return times, result
+
+
+def plot_generation(times, tokens_per_sec, start=None, interval=1.0, width=50):
+    c = "█"
+    start = start or times[0]
+    stop = times[-1]
+
+    bar_times = [start]
+    while bar_times[-1] < stop:
+        bar_times.append(bar_times[-1] + interval)
+
+    bar_values = [[] for _ in bar_times]
+    bar_idx = 0
+
+    for t, v in zip(times, tokens_per_sec):
+        while t > bar_times[bar_idx] + interval:
+            bar_idx += 1
+        bar_values[bar_idx].append(v)
+
+    bar_values = [sum(v) / len(v) if v else 0 for v in bar_values]
+    m = max(bar_values)
+
+    for t, v in zip(bar_times, bar_values):
+        t = t - start
+        b = c * int(v * width / m)
+        print(f"{t:3.2f} {b} ({v})")
+
+
+def percentile(data, percent):
+    if not data:
+        return 0
+    data = sorted(data)
+    k = (len(data) - 1) * percent / 100
+    f = math.floor(k)
+    c = math.ceil(k)
+    return (
+        data[int(f)]
+        if f == c
+        else data[int(f)] + (data[int(c)] - data[int(f)]) * (k - f)
+    )
+
+
+def median(data):
+    return percentile(data, 50)
+
+
+async def make_request(
+    session: aiohttp.ClientSession,
+    url: str,
+    api_key: str,
+    model: str,
+    prompt: str,
+    max_tokens: int,
+) -> Tuple[bool, float, list]:
+    """
+    Make a single streaming API request and return
+
+        - whether the request succeeded
+        - the request start time
+        - the time of every generated token
+    """
+    payload = {
+        "model": model,
+        "messages": [{"role": "user", "content": prompt}],
+        "max_tokens": max_tokens,
+        "stream": True,
+    }
+    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
+
+    start_time = time.perf_counter()
+    tokens = []
+
+    try:
+        async with session.post(url, json=payload, headers=headers) as response:
+            if response.status != 200:
+                error_body = await response.text()
+                print(f"Error {response.status}: {error_body}")
+                return (False, 0, [])
+
+            # Process streaming response
+            async for chunk in response.content:
+                if chunk:
+                    chunk_str = chunk.decode("utf-8").strip()
+                    if chunk_str.startswith("data:"):
+                        data_str = chunk_str[5:].strip()
+                        if data_str == "[DONE]":
+                            break
+
+                        try:
+                            data = json.loads(data_str)
+                            if choices := data.get("choices", False):
+                                if choices[0].get("finish_reason") != "length":
+                                    tokens.append(time.perf_counter())
+                        except json.JSONDecodeError:
+                            continue
+
+            return (bool(tokens), start_time, tokens)
+
+    except Exception as e:
+        print(f"Request failed: {str(e)}")
+        return (False, 0, [])
+
+
+async def run_benchmark(
+    url: str,
+    api_key: str,
+    model: str,
+    max_tokens: int,
+    concurrency: int,
+    total_requests: int,
+    prompts: List[str],
+) -> Dict[str, Any]:
+    prompt_cycle = cycle(prompts)
+    semaphore = asyncio.Semaphore(concurrency)
+    results = []
+    request_times = []
+    bar = tqdm(total=total_requests)
+
+    async def worker():
+        async with semaphore:
+            prompt = next(prompt_cycle)
+            result = await make_request(
+                session, url, api_key, model, prompt, max_tokens
+            )
+            bar.update(1)
+            return result
+
+    async with aiohttp.ClientSession() as session:
+        tasks = []
+        for _ in range(total_requests):
+            task = asyncio.create_task(worker())
+            tasks.append(task)
+            await asyncio.sleep(0.01)  # Stagger requests slightly
+
+        for task in tasks:
+            result = await task
+            results.append(result)
+        bar.close()
+
+    successful_requests = [r for r in results if r[0]]
+    total_tokens = sum(len(r[2]) for r in successful_requests)
+
+    # Gather all the tokens generated with their corresponding timestamps
+    all_tokens = []
+    for r in successful_requests:
+        all_tokens.extend(r[2])
+    all_tokens.sort()
+    full_generation = tokens_per_second(all_tokens)
+    start = min(r[1] for r in successful_requests)
+
+    # Aggregate metrics
+    metrics = {
+        "total_requests": total_requests,
+        "successful_requests": len(successful_requests),
+        "failed_requests": total_requests - len(successful_requests),
+        "total_tokens": total_tokens,
+        "total_time": all_tokens[-1] - start,
+        "aggregate_tokens_per_sec": median(full_generation[1]),
+        "per_request": [],
+        "start": start,
+        "full_generation": full_generation,
+    }
+
+    # Per-request metrics
+    for i, (_, start, tokens) in enumerate(successful_requests):
+        metrics["per_request"].append(
+            {
+                "request_id": i + 1,
+                "time_to_first_token": tokens[0] - start,
+                "total_time": tokens[-1] - start,
+                "tokens_received": len(tokens),
+                "tokens_per_sec": median(tokens_per_second(tokens)[1]),
+            }
+        )
+
+    # Calculate percentiles
+    ttft_values = [m["time_to_first_token"] for m in metrics["per_request"]]
+    tps_values = [m["tokens_per_sec"] for m in metrics["per_request"]]
+
+    metrics["aggregate_metrics"] = {
+        "time_to_first_token": {
+            "min": min(ttft_values) if ttft_values else 0,
+            "max": max(ttft_values) if ttft_values else 0,
+            "avg": sum(ttft_values) / len(ttft_values) if ttft_values else 0,
+            "p95": percentile(ttft_values, 95) if ttft_values else 0,
+        },
+        "tokens_per_sec": {
+            "min": min(tps_values) if tps_values else 0,
+            "max": max(tps_values) if tps_values else 0,
+            "avg": sum(tps_values) / len(tps_values) if tps_values else 0,
+            "p95": percentile(tps_values, 95) if tps_values else 0,
+        },
+    }
+
+    return metrics
+
+
+def main():
+    parser = argparse.ArgumentParser(description="LLM API Benchmark Tool")
+    parser.add_argument(
+        "--url",
+        default="http://localhost:8080/v1/chat/completions",
+        help="Chat completions API endpoint URL",
+    )
+    parser.add_argument("--api-key", default="none", help="API key")
+    parser.add_argument("--model", default="default_model", help="Model name")
+    parser.add_argument(
+        "--max-tokens", type=int, default=100, help="Max tokens to generate"
+    )
+    parser.add_argument(
+        "--concurrency", type=int, default=1, help="Number of concurrent requests"
+    )
+    parser.add_argument(
+        "--total-requests", type=int, default=10, help="Total requests to make"
+    )
+    parser.add_argument("--prompt-file", help="File containing prompts (one per line)")
+    parser.add_argument("--output", help="Output file for results (JSON format)")
+
+    args = parser.parse_args()
+
+    # Load prompts
+    if args.prompt_file:
+        with open(args.prompt_file, "r") as f:
+            prompts = [line.strip() for line in f if line.strip()]
+    else:
+        prompts = DEFAULT_PROMPTS
+
+    print(
+        f"Starting benchmark with {args.concurrency} concurrency and {args.total_requests} total requests..."
+    )
+    start_time = time.perf_counter()
+
+    # Run benchmark
+    results = asyncio.run(
+        run_benchmark(
+            url=args.url,
+            api_key=args.api_key,
+            model=args.model,
+            max_tokens=args.max_tokens,
+            concurrency=args.concurrency,
+            total_requests=args.total_requests,
+            prompts=prompts,
+        )
+    )
+
+    duration = time.perf_counter() - start_time
+    print(f"\nBenchmark completed in {duration:.2f} seconds")
+    print(
+        f"Successful requests: {results['successful_requests']}/{args.total_requests}"
+    )
+    print(f"Total tokens generated: {results['total_tokens']}")
+    print(f"Aggregate tokens/sec: {results['aggregate_tokens_per_sec']:.2f}")
+
+    # Print summary
+    if results["successful_requests"] > 0:
+        ttft = results["aggregate_metrics"]["time_to_first_token"]
+        tps = results["aggregate_metrics"]["tokens_per_sec"]
+
+        print("\nTime to First Token (seconds):")
+        print(
+            f"  Min: {ttft['min']:.4f} | Max: {ttft['max']:.4f} | Avg: {ttft['avg']:.4f} | P95: {ttft['p95']:.4f}"
+        )
+
+        print("\nTokens per Second (per request):")
+        print(
+            f"  Min: {tps['min']:.2f} | Max: {tps['max']:.2f} | Avg: {tps['avg']:.2f} | P95: {tps['p95']:.2f}"
+        )
+
+        print()
+        plot_generation(*results["full_generation"], results["start"])
+
+    # Save results
+    if args.output:
+        with open(args.output, "w") as f:
+            json.dump(results, f, indent=2)
+        print(f"\nResults saved to {args.output}")
+
+
+if __name__ == "__main__":
+    main()

mlx_lm/server.py

@@ -1548,7 +1548,11 @@ def run(
         "it only implements basic security checks."
     )
     logging.info(f"Starting httpd at {host} on port {port}...")
-    httpd.serve_forever()
+    try:
+        httpd.serve_forever()
+    except KeyboardInterrupt:
+        httpd.shutdown()
+        response_generator.stop_and_join()
 
 
 def main():