point_system/llm_benchmark.py

#!/usr/bin/env python3
import time
import requests
import json
import argparse
import sys
from concurrent.futures import ThreadPoolExecutor, wait, FIRST_COMPLETED


def _run_single_request(url, headers, payload, show_output):
    start_time = time.perf_counter()
    first_token_time = None
    last_token_time = None

    content = ""
    chunk_count = 0
    exact_completion_tokens = None

    try:
        response = requests.post(url, headers=headers, json=payload, stream=True)
        response.raise_for_status()

        for line in response.iter_lines():
            if line:
                line = line.decode('utf-8')
                if line.startswith("data: "):
                    data_str = line[6:]

                    if data_str == "[DONE]":
                        break

                    try:
                        data = json.loads(data_str)

                        if "usage" in data and data["usage"] is not None:
                            exact_completion_tokens = data["usage"].get("completion_tokens")

                        choices = data.get("choices", [])
                        if choices:
                            delta = choices[0].get("delta", {})
                            if "content" in delta:
                                chunk = delta["content"]
                                if chunk:
                                    if first_token_time is None:
                                        first_token_time = time.perf_counter()

                                    content += chunk
                                    chunk_count += 1
                                    last_token_time = time.perf_counter()

                                    if show_output:
                                        sys.stdout.write(chunk)
                                        sys.stdout.flush()

                    except json.JSONDecodeError:
                        continue

    except requests.exceptions.RequestException as e:
        return {"success": False, "error": str(e), "url": url, "model": payload.get("model")}

    if show_output:
        print("\n")

    if first_token_time is None:
        return {"success": False, "error": "No tokens were generated", "url": url, "model": payload.get("model")}

    total_time = last_token_time - start_time
    generation_time = last_token_time - first_token_time
    final_token_count = exact_completion_tokens if exact_completion_tokens is not None else chunk_count

    if final_token_count > 1 and generation_time > 0:
        tps = (final_token_count - 1) / generation_time
    elif final_token_count == 1 and total_time > 0:
        tps = 1 / total_time
    else:
        tps = 0

    return {
        "success": True,
        "url": url,
        "model": payload.get("model"),
        "ttft": first_token_time - start_time,
        "generation_time": generation_time,
        "total_time": total_time,
        "token_count": final_token_count,
        "tps": tps,
        "exact": exact_completion_tokens is not None,
    }


def _print_summary(results, wall_time, timed_out=False):
    success = [r for r in results if r.get("success")]
    failures = [r for r in results if not r.get("success")]

    print("\n" + "-" * 50)
    print("🎯 BENCHMARK RESULTS")
    print("-" * 50)
    if results:
        print(f"URL:                    {results[0].get('url')}")
        print(f"Model:                  {results[0].get('model')}")
    print(f"Completed Requests:     {len(results)}")
    print(f"Successful Requests:    {len(success)}")
    print(f"Failed Requests:        {len(failures)}")
    print(f"Total Wall Time:        {wall_time:.3f} s")
    if timed_out:
        print(f"Timed out:            stopped at {wall_time:.1f} s limit")

    if not success:
        print("No successful requests. Check the model name, API key, and endpoint.")
        return

    total_tokens = sum(r["token_count"] for r in success)
    tokens_exact = all(r["exact"] for r in success) and any(r["exact"] for r in success)
    avg_tps = sum(r["tps"] for r in success) / len(success)
    min_tps = min(r["tps"] for r in success)
    max_tps = max(r["tps"] for r in success)
    avg_ttft = sum(r["ttft"] for r in success) / len(success)
    avg_total = sum(r["total_time"] for r in success) / len(success)
    overall_tps = total_tokens / wall_time
    rps = len(success) / wall_time

    print(f"Total Tokens:           {total_tokens}")
    print(f"Overall Throughput:     {overall_tps:.2f} tokens/sec")
    print(f"Request Throughput:     {rps:.2f} req/sec")
    print(f"Average TPS:            {avg_tps:.2f} tokens/sec")
    print(f"Min / Max TPS:          {min_tps:.2f} / {max_tps:.2f} tokens/sec")
    print(f"Average TTFT:           {avg_ttft:.3f} s")
    print(f"Average Request Time:   {avg_total:.3f} s")
    if len(success) == 1:
        only = success[0]
        print(f"Tokens Generated:       {only['token_count']} {'(Exact)' if only['exact'] else '(Approx chunk count)'}")

    if failures:
        print(f"\nFirst errors:")
        for r in failures[:5]:
            print(f"  - {r.get('error')}")


def benchmark_llm(url, model, api_key, prompt, max_tokens, temperature, show_output,
                  concurrency, max_time):
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }

    payload = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": max_tokens,
        "temperature": temperature,
        "stream": True,
        "stream_options": {"include_usage": True}
    }

    print(f"Benchmarking model: '{model}' at '{url}'...")
    print(f"Prompt: {prompt[:50]}...")
    print(f"Concurrency: {concurrency}, Max Time: {max_time:.0f}s (run until time limit)\n")
    print("-" * 50)

    results = []
    completed = 0
    timed_out = False
    wall_start = time.perf_counter()
    deadline = wall_start + max_time

    executor = ThreadPoolExecutor(max_workers=concurrency)
    pending = [
        executor.submit(_run_single_request, url, headers, payload, show_output)
        for _ in range(concurrency)
    ]
    try:
        while pending:
            remaining = deadline - time.perf_counter()
            if remaining <= 0:
                timed_out = True
                break
            done, pending = wait(pending, return_when=FIRST_COMPLETED, timeout=remaining)
            pending = list(pending)
            if not done:
                timed_out = True
                break
            for future in done:
                result = future.result()
                results.append(result)
                completed += 1
                elapsed = time.perf_counter() - wall_start
                if result.get("success"):
                    print(f"[{completed}] ok    tps={result['tps']:.2f}  "
                          f"ttft={result['ttft']:.3f}s  tokens={result['token_count']}  @{elapsed:.1f}s")
                else:
                    print(f"[{completed}] fail  {result.get('error')}  @{elapsed:.1f}s")
            while len(pending) < concurrency and time.perf_counter() < deadline:
                pending.append(executor.submit(_run_single_request, url, headers, payload, show_output))
    finally:
        for f in pending:
            f.cancel()
        executor.shutdown(wait=False)

    wall_time = time.perf_counter() - wall_start
    _print_summary(results, wall_time, timed_out)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Benchmark an OpenAI-compatible LLM endpoint.")

    parser.add_argument("--url", type=str, default="http://192.168.221.15/v1/chat/completions",
                        help="API endpoint URL (default: http://192.168.221.15/v1/chat/completions)")
    parser.add_argument("--model", type=str, default="qwen3.8-27b",
                        help="The name of the model to benchmark (default: qwen3.8-27b)")
    parser.add_argument("--api-key", type=str, default="EMPTY",
                        help="API Key (default: EMPTY for local engines)")
    parser.add_argument("--prompt", type=str, default="Explain the theory of relativity and quantum mechanics in great detail. Write at least 4 paragraphs.",
                        help="The prompt to send to the model.")
    parser.add_argument("--max-tokens", type=int, default=512,
                        help="Maximum number of tokens to generate.")
    parser.add_argument("--temperature", type=float, default=0.0,
                        help="Sampling temperature (default 0.0 for deterministic output).")
    parser.add_argument("--show-output", action="store_true",
                        help="Print the model's text response as it generates (not for high concurrency).")
    parser.add_argument("--concurrency", type=int, default=10,
                        help="Number of concurrent requests to run at once (default: 10).")
    parser.add_argument("--max-time", type=float, default=300,
                        help="Stop after this many seconds (default: 300 = 5 min).")

    args = parser.parse_args()

    benchmark_llm(args.url, args.model, args.api_key, args.prompt, args.max_tokens,
                  args.temperature, args.show_output, args.concurrency, args.max_time)