benchmark.py

#!/usr/bin/env python3
"""
M4 Pro 48GB Benchmark: LM Studio vs Ollama

Supports multiple models, large context windows, and auto backend detection.
Outputs per-model JSON results and a combined summary.
"""

import argparse
import json
import os
import random
import string
import time
import statistics
import sys
from dataclasses import dataclass, field
from pathlib import Path

import httpx

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

LMSTUDIO_BASE = "http://localhost:1234/v1"
OLLAMA_BASE = "http://localhost:11434/v1"
DEFAULT_TEMPERATURE = 0.0
DEFAULT_RUNS = 3


# ---------------------------------------------------------------------------
# Prompt generation
# ---------------------------------------------------------------------------


def generate_filler(token_target: int) -> str:
    """Generate roughly `token_target` tokens of realistic filler text.
    ~1 token ≈ 4 chars in English. We use a mix of code + prose."""
    chars_needed = token_target * 4
    blocks = [
        "The quick brown fox jumps over the lazy dog. ",
        "def fibonacci(n: int) -> int:\n    if n <= 1:\n        return n\n    return fibonacci(n-1) + fibonacci(n-2)\n\n",
        "Lorem ipsum dolor sit amet, consectetur adipiscing elit. ",
        "async function fetchData(url) {\n  const res = await fetch(url);\n  return res.json();\n}\n\n",
        "In computer science, a hash table is a data structure that implements an associative array. ",
        "SELECT u.id, u.name, COUNT(o.id) as order_count FROM users u LEFT JOIN orders o ON u.id = o.user_id GROUP BY u.id HAVING order_count > 5 ORDER BY order_count DESC;\n\n",
        "Machine learning models require careful feature engineering and hyperparameter tuning. ",
        "pub fn merge_sort<T: Ord + Clone>(arr: &mut [T]) {\n    let len = arr.len();\n    if len <= 1 { return; }\n    let mid = len / 2;\n}\n\n",
    ]
    result = []
    total = 0
    while total < chars_needed:
        block = random.choice(blocks)
        result.append(block)
        total += len(block)
    return "".join(result)[:chars_needed]


def build_prompts(context_sizes: list[int]) -> dict:
    """Build prompt suite including large context tests."""
    prompts = {
        "short": {
            "name": "Short gen (~100 tok)",
            "messages": [
                {
                    "role": "user",
                    "content": "Explain what a closure is in programming in 2-3 sentences.",
                }
            ],
            "max_tokens": 150,
        },
        "medium": {
            "name": "Medium gen (~500 tok)",
            "messages": [
                {
                    "role": "user",
                    "content": "Write a Python function that implements a binary search tree with insert, search, and delete operations. Include type hints.",
                }
            ],
            "max_tokens": 800,
        },
        "long": {
            "name": "Long gen (~2K tok)",
            "messages": [
                {
                    "role": "user",
                    "content": "Write a comprehensive guide on implementing a REST API in Go covering project structure, routing, middleware, error handling, database integration, and testing with code examples.",
                }
            ],
            "max_tokens": 2048,
        },
    }

    # Add context window tests
    for ctx_k in context_sizes:
        tok_count = ctx_k * 1024
        filler = generate_filler(tok_count)
        prompts[f"ctx_{ctx_k}k"] = {
            "name": f"Context {ctx_k}K input → 512 tok out",
            "messages": [
                {
                    "role": "user",
                    "content": f"Below is a large document. Read it carefully and write a 3-paragraph summary.\n\n---\n{filler}\n---\n\nSummarize the above document.",
                }
            ],
            "max_tokens": 512,
        }

    return prompts


# ---------------------------------------------------------------------------
# Data classes
# ---------------------------------------------------------------------------


@dataclass
class RunResult:
    ttft: float
    total_time: float
    prompt_tokens: int
    completion_tokens: int

    @property
    def tokens_per_sec(self) -> float:
        gen_time = self.total_time - self.ttft
        return self.completion_tokens / gen_time if gen_time > 0 else 0.0

    @property
    def prompt_tokens_per_sec(self) -> float:
        return self.prompt_tokens / self.ttft if self.ttft > 0 else 0.0


@dataclass
class BenchmarkResult:
    backend: str
    model: str
    prompt_name: str
    runs: list[RunResult] = field(default_factory=list)

    def _avg(self, attr: str) -> float:
        return statistics.mean(getattr(r, attr) for r in self.runs)

    @property
    def avg_ttft(self) -> float:
        return self._avg("ttft")

    @property
    def avg_tok_s(self) -> float:
        return self._avg("tokens_per_sec")

    @property
    def avg_prompt_tok_s(self) -> float:
        return self._avg("prompt_tokens_per_sec")

    @property
    def avg_total_time(self) -> float:
        return self._avg("total_time")

    @property
    def avg_completion_tokens(self) -> float:
        return self._avg("completion_tokens")


# ---------------------------------------------------------------------------
# API
# ---------------------------------------------------------------------------


def stream_chat(
    base_url: str,
    model: str,
    messages: list,
    max_tokens: int,
    temperature: float = DEFAULT_TEMPERATURE,
    no_think: bool = False,
) -> RunResult:
    url = f"{base_url}/chat/completions"

    all_messages = messages.copy()
    if no_think:
        all_messages = [{"role": "system", "content": "/no_think"}] + all_messages

    payload = {
        "model": model,
        "messages": all_messages,
        "max_tokens": max_tokens,
        "temperature": temperature,
        "stream": True,
        "stream_options": {"include_usage": True},
    }

    ttft = None
    start = time.perf_counter()
    prompt_tokens = 0
    completion_tokens = 0
    chunk_count = 0

    with httpx.Client(timeout=600.0) as client:
        with client.stream("POST", url, json=payload) as response:
            response.raise_for_status()
            for line in response.iter_lines():
                if not line.startswith("data: "):
                    continue
                data_str = line[6:]
                if data_str.strip() == "[DONE]":
                    break
                try:
                    chunk = json.loads(data_str)
                except json.JSONDecodeError:
                    continue

                if ttft is None:
                    choices = chunk.get("choices", [])
                    if choices:
                        delta = choices[0].get("delta", {})
                        if delta.get("content") or delta.get("reasoning"):
                            ttft = time.perf_counter() - start

                choices = chunk.get("choices", [])
                if choices:
                    delta = choices[0].get("delta", {})
                    if delta.get("content") or delta.get("reasoning"):
                        chunk_count += 1

                usage = chunk.get("usage")
                if usage:
                    prompt_tokens = usage.get("prompt_tokens", 0)
                    completion_tokens = usage.get("completion_tokens", 0)

    total_time = time.perf_counter() - start
    if ttft is None:
        ttft = total_time
    if completion_tokens == 0:
        completion_tokens = chunk_count

    return RunResult(
        ttft=ttft,
        total_time=total_time,
        prompt_tokens=prompt_tokens,
        completion_tokens=completion_tokens,
    )


def check_server(base_url: str) -> bool:
    try:
        with httpx.Client(timeout=5.0) as client:
            return client.get(f"{base_url}/models").status_code == 200
    except Exception:
        return False


def list_models(base_url: str) -> list[str]:
    try:
        with httpx.Client(timeout=10.0) as client:
            resp = client.get(f"{base_url}/models")
            resp.raise_for_status()
            return [m["id"] for m in resp.json().get("data", [])]
    except Exception:
        return []


# ---------------------------------------------------------------------------
# Runner
# ---------------------------------------------------------------------------


def run_benchmark(
    base_url: str,
    model: str,
    backend_name: str,
    prompts: dict,
    prompt_keys: list[str],
    runs: int = 3,
    warmup: bool = True,
    no_think: bool = False,
) -> list[BenchmarkResult]:
    results = []

    for pk in prompt_keys:
        prompt = prompts[pk]
        print(f"\n  [{pk}] {prompt['name']}")
        br = BenchmarkResult(backend=backend_name, model=model, prompt_name=pk)

        if warmup:
            print("    warmup...", end="", flush=True)
            try:
                stream_chat(
                    base_url,
                    model,
                    prompt["messages"],
                    prompt["max_tokens"],
                    no_think=no_think,
                )
                print(" ok")
            except Exception as e:
                print(f" FAIL: {e}")
                continue

        for i in range(runs):
            print(f"    run {i + 1}/{runs}...", end="", flush=True)
            try:
                r = stream_chat(
                    base_url,
                    model,
                    prompt["messages"],
                    prompt["max_tokens"],
                    no_think=no_think,
                )
                br.runs.append(r)
                print(
                    f" TTFT={r.ttft:.3f}s | {r.tokens_per_sec:.1f} tok/s | "
                    f"prompt={r.prompt_tokens} comp={r.completion_tokens} | {r.total_time:.2f}s"
                )
            except Exception as e:
                print(f" ERROR: {e}")

        if br.runs:
            results.append(br)

    return results


# ---------------------------------------------------------------------------
# Display
# ---------------------------------------------------------------------------


def print_comparison(
    lm_results: list[BenchmarkResult],
    ollama_results: list[BenchmarkResult],
    prompts: dict,
    model_label: str = "",
):
    lm_label = "LM Studio"
    ol_label = "Ollama"

    print("\n" + "=" * 90)
    title = f"RESULTS: {lm_label} vs {ol_label}"
    if model_label:
        title += f" — {model_label}"
    print(title)
    print(f"Hardware: Apple M4 Pro 14-core, 48GB | Memory BW: ~273 GB/s")
    print("=" * 90)

    lm_by_prompt = {r.prompt_name: r for r in lm_results}
    ol_by_prompt = {r.prompt_name: r for r in ollama_results}
    all_prompts = [pk for pk in prompts if pk in lm_by_prompt or pk in ol_by_prompt]

    for pk in all_prompts:
        prompt = prompts[pk]
        print(f"\n--- {prompt['name']} ---")
        print(f"{'Metric':<30} {lm_label:>18} {ol_label:>18} {'Winner':>16}")
        print("-" * 84)

        lm = lm_by_prompt.get(pk)
        ol = ol_by_prompt.get(pk)

        metrics = [
            ("TTFT (s)", "avg_ttft", True),
            ("Generation (tok/s)", "avg_tok_s", False),
            ("Prompt eval (tok/s)", "avg_prompt_tok_s", False),
            ("Total time (s)", "avg_total_time", True),
            ("Completion tokens", "avg_completion_tokens", None),
        ]

        for label, attr, lower_is_better in metrics:
            lm_val = getattr(lm, attr, None) if lm else None
            ol_val = getattr(ol, attr, None) if ol else None
            lm_str = f"{lm_val:.2f}" if lm_val is not None else "N/A"
            ol_str = f"{ol_val:.2f}" if ol_val is not None else "N/A"

            winner = ""
            if (
                lm_val is not None
                and ol_val is not None
                and lower_is_better is not None
            ):
                if lower_is_better:
                    name = "LM Studio" if lm_val < ol_val else "Ollama"
                    worse = max(lm_val, ol_val)
                else:
                    name = "LM Studio" if lm_val > ol_val else "Ollama"
                    worse = min(lm_val, ol_val)
                if worse > 0:
                    pct = abs(lm_val - ol_val) / worse * 100
                    winner = f"{name} +{pct:.0f}%"

            print(f"{label:<30} {lm_str:>18} {ol_str:>18} {winner:>16}")

    # Summary
    print("\n" + "=" * 90)
    print("SUMMARY")
    print("=" * 90)

    for label, results in [(lm_label, lm_results), (ol_label, ollama_results)]:
        if results:
            avg_toks = statistics.mean(r.avg_tok_s for r in results)
            avg_ttft = statistics.mean(r.avg_ttft for r in results)
            print(
                f"  {label:<20} avg gen: {avg_toks:>7.1f} tok/s | avg TTFT: {avg_ttft:.3f}s"
            )


def print_single(results: list[BenchmarkResult], prompts: dict):
    print("\n" + "=" * 60)
    print(f"RESULTS: {results[0].backend} — {results[0].model}")
    print("=" * 60)
    for br in results:
        p = prompts.get(br.prompt_name, {})
        print(f"\n  [{br.prompt_name}] {p.get('name', br.prompt_name)}")
        print(f"    TTFT:        {br.avg_ttft:.3f}s")
        print(f"    Gen speed:   {br.avg_tok_s:.1f} tok/s")
        if br.avg_prompt_tok_s > 0:
            print(f"    Prompt eval: {br.avg_prompt_tok_s:.1f} tok/s")
        print(f"    Total time:  {br.avg_total_time:.2f}s")
        print(f"    Avg tokens:  {br.avg_completion_tokens:.0f}")


def save_results(
    lm_results: list[BenchmarkResult],
    ollama_results: list[BenchmarkResult],
    output_path: Path,
    model_label: str = "",
):
    def serialize(results: list[BenchmarkResult]) -> list:
        return [
            {
                "backend": br.backend,
                "model": br.model,
                "prompt": br.prompt_name,
                "runs": [
                    {
                        "ttft": r.ttft,
                        "total_time": r.total_time,
                        "prompt_tokens": r.prompt_tokens,
                        "completion_tokens": r.completion_tokens,
                        "tokens_per_sec": r.tokens_per_sec,
                        "prompt_tokens_per_sec": r.prompt_tokens_per_sec,
                    }
                    for r in br.runs
                ],
                "averages": {
                    "ttft": br.avg_ttft,
                    "tok_s": br.avg_tok_s,
                    "prompt_tok_s": br.avg_prompt_tok_s,
                    "total_time": br.avg_total_time,
                },
            }
            for br in results
        ]

    data = {
        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
        "hardware": "Apple M4 Pro 14-core, 48GB",
        "model": model_label,
        "lm_studio": serialize(lm_results),
        "ollama": serialize(ollama_results),
    }
    output_path.write_text(json.dumps(data, indent=2, ensure_ascii=False))
    print(f"\nResults saved to: {output_path}")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def main():
    parser = argparse.ArgumentParser(
        description="Benchmark LM Studio vs Ollama on Apple Silicon"
    )
    parser.add_argument(
        "--lm-model", type=str, default=None, help="LM Studio model identifier"
    )
    parser.add_argument(
        "--ollama-model", type=str, default=None, help="Ollama model identifier"
    )
    parser.add_argument(
        "--prompts",
        type=str,
        default=None,
        help="Comma-separated prompt keys (default: all)",
    )
    parser.add_argument(
        "--context-sizes",
        type=str,
        default="4,8,16,32",
        help="Comma-separated context sizes in K to test (default: 4,8,16,32)",
    )
    parser.add_argument("--runs", type=int, default=DEFAULT_RUNS)
    parser.add_argument("--no-warmup", action="store_true")
    parser.add_argument(
        "--no-think",
        action="store_true",
        help="Disable thinking mode (for Qwen models)",
    )
    parser.add_argument("--only", type=str, choices=["lmstudio", "ollama"])
    parser.add_argument("--output", type=str, default="results.json")
    parser.add_argument(
        "--label",
        type=str,
        default="",
        help="Model label for output (e.g. 'Qwen3.5-35B-A3B')",
    )
    args = parser.parse_args()

    # Build prompts with context sizes
    ctx_sizes = [int(x) for x in args.context_sizes.split(",")]
    prompts = build_prompts(ctx_sizes)

    if args.prompts:
        prompt_keys = [p.strip() for p in args.prompts.split(",")]
        for pk in prompt_keys:
            if pk not in prompts:
                print(f"Unknown prompt: {pk}. Available: {list(prompts.keys())}")
                sys.exit(1)
    else:
        prompt_keys = list(prompts.keys())

    model_label = args.label or f"{args.lm_model or '?'} / {args.ollama_model or '?'}"

    print("=" * 60)
    print(f"M4 Pro 48GB Benchmark — {model_label}")
    print(f"Prompts: {', '.join(prompt_keys)}")
    print(f"Runs: {args.runs} | Warmup: {not args.no_warmup}")
    print("=" * 60)

    lm_results = []
    ollama_results = []

    if args.only != "ollama":
        print("\n[LM Studio] Checking...", end=" ")
        if check_server(LMSTUDIO_BASE):
            models = list_models(LMSTUDIO_BASE)
            lm_model = args.lm_model or (models[0] if models else None)
            if lm_model:
                print(f"OK — {lm_model}")
                lm_results = run_benchmark(
                    LMSTUDIO_BASE,
                    lm_model,
                    "LM Studio",
                    prompts,
                    prompt_keys,
                    args.runs,
                    warmup=not args.no_warmup,
                    no_think=args.no_think,
                )
            else:
                print("No model loaded.")
        else:
            print("Not running.")

    if args.only != "lmstudio":
        print("\n[Ollama] Checking...", end=" ")
        if check_server(OLLAMA_BASE):
            models = list_models(OLLAMA_BASE)
            ollama_model = args.ollama_model or (models[0] if models else None)
            if ollama_model:
                print(f"OK — {ollama_model}")
                ollama_results = run_benchmark(
                    OLLAMA_BASE,
                    ollama_model,
                    "Ollama",
                    prompts,
                    prompt_keys,
                    args.runs,
                    warmup=not args.no_warmup,
                    no_think=args.no_think,
                )
            else:
                print("No model available.")
        else:
            print("Not running.")

    if lm_results or ollama_results:
        if lm_results and ollama_results:
            print_comparison(lm_results, ollama_results, prompts, model_label)
        else:
            print_single(lm_results or ollama_results, prompts)
        save_results(lm_results, ollama_results, Path(args.output), model_label)
    else:
        print("\nNo results.")
        sys.exit(1)


if __name__ == "__main__":
    main()