llamacpp_benchmark.py

#!/usr/bin/env python3
"""
Benchmark script for llama.cpp server using HTTP API.

This script connects to a running llama.cpp server and benchmarks its performance
across different context sizes.

Usage:
    # Start llama.cpp server first:
    # ./llama-server -m model.gguf --host 0.0.0.0 --port 8080
    
    # Then run benchmark (default localhost:8080):
    python llamacpp_benchmark.py gpt-oss:20b
    
    # Or specify custom host and port:
    python llamacpp_benchmark.py gpt-oss:20b --host localhost --port 9000
"""

import argparse
import sys
import time
from pathlib import Path
from typing import Dict, Optional

import requests

from benchmark_common import (
    add_throughput_metrics,
    create_output_directory,
    find_context_files,
    format_hardware_string,
    get_hardware_info,
    make_cache_buster,
    print_benchmark_summary,
    run_benchmark_peak,
    run_benchmark_peak_per_run,
    save_all_outputs,
    save_generated_text,
    setup_common_args,
)


def test_server_connection(server_url: str) -> bool:
    """Test if the llama.cpp server is running and accessible."""
    try:
        response = requests.get(f"{server_url}/health", timeout=5)
        if response.status_code == 200:
            return True
    except requests.exceptions.RequestException as e:
        print(f"Error connecting to server: {e}")
    return False


def get_server_info(server_url: str) -> Dict:
    """Get model and server information from llama.cpp server."""
    try:
        # Try to get model metadata
        response = requests.get(f"{server_url}/props", timeout=5)
        if response.status_code == 200:
            return response.json()
    except (requests.exceptions.RequestException, Exception):
        pass

    # Return empty dict if no info available
    return {}


def _hardware_folder_label(label: str) -> str:
    """Return a compact filesystem-safe hardware label."""
    return "".join(ch for ch in label.replace("Apple ", "") if ch.isalnum())


def benchmark_llamacpp(
    server_url: str,
    context_file: Path,
    max_tokens: int = 128,
    timeout: int = 3600,
    cold_prefill: bool = True,
    _run_idx: Optional[int] = None,
) -> Optional[Dict]:
    """Benchmark llama.cpp server with a given context file.

    Uses the OpenAI-compatible /v1/chat/completions endpoint so the server
    applies the model's chat template. Hitting the raw /completion endpoint
    with a big chunk of document text causes the model to occasionally sample
    EOS on its first step (there is no "assistant is responding" framing), in
    which case predicted_n=1 and predicted_ms≈0, producing a bogus
    ~1,000,000 tokens/sec reading.

    Args:
        server_url: URL of the llama.cpp server
        context_file: Path to the context file
        max_tokens: Maximum number of tokens to generate
        timeout: Request timeout in seconds
        cold_prefill: Prevent server-side KV cache reuse (default: True)

    Returns:
        Dictionary with benchmark results or None if failed
    """
    with open(context_file, "r", encoding="utf-8") as f:
        prompt = f.read()

    if cold_prefill:
        prompt = make_cache_buster() + prompt
    elif _run_idx is not None:
        prompt = make_cache_buster(run_idx=_run_idx) + prompt

    # OpenAI-style chat payload. llama.cpp accepts `cache_prompt` as an
    # extension field on this endpoint too.
    payload = {
        "model": "default",
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": max_tokens,
        "temperature": 0.7,
        "top_p": 0.95,
        "stream": True,
        "cache_prompt": not cold_prefill,
    }

    start_time = time.time()
    first_token_time = None
    generated_text = ""
    finish_reason: Optional[str] = None
    timings: Dict = {}

    try:
        import json

        response = requests.post(
            f"{server_url}/v1/chat/completions",
            json=payload,
            timeout=timeout,
            stream=True,
        )
        response.raise_for_status()

        for line in response.iter_lines():
            if not line:
                continue
            line = line.decode("utf-8")
            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

            choices = chunk.get("choices") or []
            if choices:
                delta = choices[0].get("delta") or {}
                content = delta.get("content") or ""
                if content:
                    if first_token_time is None:
                        first_token_time = time.time()
                    generated_text += content
                if choices[0].get("finish_reason"):
                    finish_reason = choices[0]["finish_reason"]

            # llama.cpp-specific: final chunk includes a `timings` object.
            if chunk.get("timings"):
                timings = chunk["timings"]
    except requests.exceptions.RequestException as e:
        print(f"Error during benchmark: {e}")
        return None

    total_time = time.time() - start_time
    time_to_first_token = (first_token_time - start_time) if first_token_time else 0.0

    prompt_tokens = timings.get("prompt_n", 0)
    generation_tokens = timings.get("predicted_n", 0)

    prompt_ms = timings.get("prompt_ms", 0)
    predict_ms = timings.get("predicted_ms", 0)

    prompt_time = prompt_ms / 1000.0
    predict_time = predict_ms / 1000.0

    # Guard against pathological results (e.g. model sampled EOS on first
    # step so predicted_n=1 and predicted_ms≈0). Require at least 2 decoded
    # tokens and a non-trivial decode duration before trusting the TPS ratio.
    if generation_tokens < 2 or predict_time < 0.01:
        print(
            f"    Warning: degenerate run (predicted_n={generation_tokens}, "
            f"predicted_ms={predict_ms}, finish_reason={finish_reason}); discarding"
        )
        return None

    prompt_tps = prompt_tokens / prompt_time if prompt_time > 0 else 0
    generation_tps = generation_tokens / predict_time if predict_time > 0 else 0

    result = {
        "context_size": context_file.stem,
        "prompt_tokens": prompt_tokens,
        "generation_tokens": generation_tokens,
        "prompt_time": prompt_time,
        "prompt_eval_duration": prompt_time,
        "time_to_first_token": time_to_first_token,
        "eval_duration": predict_time,
        "total_time": total_time,
        "prompt_tps": prompt_tps,
        "generation_tps": generation_tps,
        "generated_text": generated_text,
        "wall_time": total_time,
    }
    return add_throughput_metrics(result, prompt_text=prompt)


def main() -> int:
    """Main function to run the benchmark."""
    parser = argparse.ArgumentParser(description="Benchmark llama.cpp server across different context sizes")
    parser.add_argument(
        "model",
        help="Model name or identifier (used for display purposes)",
    )
    parser.add_argument(
        "--host",
        default="localhost",
        help="Host of the llama.cpp server (default: localhost)",
    )
    parser.add_argument(
        "--port",
        type=int,
        default=8080,
        help="Port of the llama.cpp server (default: 8080)",
    )
    parser.add_argument(
        "--cold-prefill",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Prevent server-side KV cache reuse for cold-prefill numbers (default: enabled; "
        "use --no-cold-prefill for cached/warm-reuse numbers)",
    )
    parser.add_argument(
        "--no-kl-capture",
        action="store_true",
        help="Skip top-K logprob capture used by compare_benchmarks --kl-baseline",
    )
    parser.add_argument(
        "--server-hardware",
        help="Remote server hardware label to store in results (e.g. 'GB10' or 'NVIDIA RTX 5090')",
    )
    parser.add_argument(
        "--server-memory-gb",
        type=float,
        help="Remote server memory in GB to store in results",
    )
    parser.add_argument(
        "--server-cores",
        type=int,
        help="Remote server CPU/core count to store in results",
    )

    # Add common arguments
    setup_common_args(parser)

    args = parser.parse_args()

    # Construct server URL from host and port
    server_url = f"http://{args.host}:{args.port}"

    # Test server connection
    print(f"Testing connection to llama.cpp server at {server_url}...")
    if not test_server_connection(server_url):
        print(f"Error: Cannot connect to llama.cpp server at {server_url}")
        print("Please ensure the server is running:")
        print("  ./llama-server -m model.gguf --host 0.0.0.0 --port 8080")
        return 1

    print("Successfully connected to llama.cpp server")

    # Get server info
    server_info = get_server_info(server_url)
    # Use the model name from args, fallback to server info if available
    model_name = args.model or server_info.get("default_generation_settings", {}).get("model", "llama.cpp model")

    # Get hardware info
    hardware_info = get_hardware_info()
    if args.server_hardware:
        hardware_info["client_chip"] = hardware_info.get("chip")
        hardware_info["client_model"] = hardware_info.get("model")
        hardware_info["chip"] = args.server_hardware
        hardware_info["model"] = "Remote llama.cpp server"
        hardware_info["api_endpoint"] = server_url
        if args.server_memory_gb is not None:
            hardware_info["memory_gb"] = args.server_memory_gb
        if args.server_cores is not None:
            hardware_info["total_cores"] = args.server_cores
    hardware_str = format_hardware_string(hardware_info)

    print(f"\nHardware: {hardware_str}")
    print(f"Model: {model_name}")
    print(f"Max tokens: {args.max_tokens}")
    print(f"Cold prefill: {'enabled (cache_prompt=false)' if args.cold_prefill else 'disabled (cache reuse allowed)'}")

    # Find context files
    context_files = find_context_files(args.contexts)
    if not context_files:
        return 1

    # Create output directory using common function
    machine_name = _hardware_folder_label(args.server_hardware) if args.server_hardware else None
    output_dir = create_output_directory("llamacpp", args.model, cold_prefill=args.cold_prefill, machine_name=machine_name)

    # Capture top-K logprobs on a fixed reference text for later KL comparison
    if not args.no_kl_capture:
        from kl_capture import DEFAULT_REF_FILE, capture_logprobs_llamacpp, save_logprobs

        ref_path = Path(DEFAULT_REF_FILE)
        if not ref_path.exists():
            print(f"\nKL capture skipped: reference file {DEFAULT_REF_FILE} not found")
        else:
            print(f"\nCapturing top-K logprobs from {DEFAULT_REF_FILE} for KL comparison...")
            try:
                kl_data = capture_logprobs_llamacpp(server_url, ref_path, timeout=args.timeout)
                kl_data["engine"] = "llamacpp"
                kl_data["model"] = args.model
                save_logprobs(kl_data, output_dir / "logprobs.json")
                print(
                    f"  Captured {kl_data['num_positions']} positions × top-{kl_data['top_k']} "
                    f"→ {output_dir / 'logprobs.json'}"
                )
            except Exception as e:
                print(f"  KL capture failed (continuing): {e}")

    results = []

    # Run benchmarks
    import time

    start_time = time.time()
    if args.cold_prefill:
        for context_file in context_files:
            print(f"\n{'=' * 50}")
            print(f"Benchmarking {context_file.name}...")
            print(f"{'=' * 50}")

            result = run_benchmark_peak(
                benchmark_llamacpp,
                server_url,
                context_file,
                args.max_tokens,
                args.timeout,
                cold_prefill=args.cold_prefill,
                n_runs=args.runs,
            )

            if result:
                results.append(result)

                # Print results
                print(f"\nResults for {context_file.name}:")
                print(f"  Prompt tokens: {result['prompt_tokens']}")
                print(f"  Generated tokens: {result['generation_tokens']}")
                print(f"  Time to first token: {result['time_to_first_token']:.2f}s")
                print(f"  Prompt time: {result['prompt_time']:.2f}s")
                print(f"  Generation time: {result['eval_duration']:.2f}s")
                print(f"  Total time: {result['total_time']:.2f}s")
                print(f"  Prompt TPS: {result['prompt_tps']:.1f} tokens/sec")
                print(f"  Generation TPS: {result['generation_tps']:.1f} tokens/sec")

                # Save response if requested
                if args.save_responses:
                    response_file = output_dir / f"response_{context_file.stem}.txt"
                    save_generated_text(result, model_name, response_file, framework="llama.cpp")
    else:
        results = run_benchmark_peak_per_run(
            benchmark_llamacpp,
            context_files=context_files,
            n_runs=args.runs,
            server_url=server_url,
            max_tokens=args.max_tokens,
            timeout=args.timeout,
            cold_prefill=args.cold_prefill,
        )
        if args.save_responses:
            for result in results:
                ctx_name = result.get("context_size", "unknown")
                response_file = output_dir / f"response_{ctx_name}.txt"
                save_generated_text(result, model_name, response_file, framework="llama.cpp")

    total_benchmark_time = time.time() - start_time

    if not results:
        print("\nNo successful benchmark results")
        return 1

    # Save all outputs using common function
    save_all_outputs(results, output_dir, model_name, "llama.cpp", hardware_info, args)

    # Print summary using common function
    print_benchmark_summary(results, model_name, "llama.cpp", hardware_info, output_dir, total_benchmark_time)

    return 0


if __name__ == "__main__":
    sys.exit(main())