Local LLM Benchmark: Ollama vLLM llama cpp Compared

#!/usr/bin/env python3
"""
LLM Inference Benchmark Harness
Reproducible testing for Ollama, vLLM, and llama.cpp

Usage:
    python benchmark_harness.py --engine ollama --concurrency 8 --requests 100
"""

import argparse
import time
import statistics
import json
from openai import OpenAI
from concurrent.futures import ThreadPoolExecutor, as_completed
import psutil
import subprocess

class LLMBenchmark:
    def __init__(self, engine, base_url, model):
        self.engine = engine
        self.client = OpenAI(
            base_url=base_url,
            api_key="dummy"  # Not needed for local engines
        )
        self.model = model

    def get_vram_usage(self):
        """Query nvidia-smi for current VRAM usage in MB"""
        try:
            result = subprocess.run([
                "nvidia-smi",
                "--query-gpu=memory.used",
                "--format=csv,noheader,nounits"
            ], capture_output=True, text=True)
            return int(result.stdout.strip().split()[0])
        except:
            return 0

    def single_request(self, prompt, max_tokens=256):
        """Execute a single request and return timing metrics"""
        start_time = time.perf_counter()
        first_token_time = None
        tokens_received = 0

        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[{"role": "user", "content": prompt}],
                max_tokens=max_tokens,
                stream=True,
                temperature=0.7
            )

            full_response = ""
            for chunk in response:
                if chunk.choices[0].delta.content:
                    if first_token_time is None:
                        first_token_time = time.perf_counter()
                    tokens_received += 1
                    full_response += chunk.choices[0].delta.content

            end_time = time.perf_counter()

            ttft = (first_token_time - start_time) * 1000  # ms
            total_time = (end_time - start_time) * 1000  # ms
            tpot = (total_time - ttft) / max(tokens_received - 1, 1)
            throughput = tokens_received / (total_time / 1000)

            return {
                "ttft_ms": ttft,
                "tpot_ms": tpot,
                "total_time_ms": total_time,
                "tokens": tokens_received,
                "throughput_tps": throughput,
                "success": True
            }

        except Exception as e:
            return {
                "success": False,
                "error": str(e)
            }

    def run_concurrent_benchmark(self, prompts, concurrency, num_requests):
        """Run concurrent requests with specified parallelism"""
        results = []
        vram_start = self.get_vram_usage()

        with ThreadPoolExecutor(max_workers=concurrency) as executor:
            futures = []
            for i in range(num_requests):
                prompt = prompts[i % len(prompts)]
                futures.append(executor.submit(self.single_request, prompt))

            for future in as_completed(futures):
                results.append(future.result())

        vram_peak = self.get_vram_usage()

        # Calculate aggregate metrics
        successful = [r for r in results if r["success"]]
        failed = [r for r in results if not r["success"]]

        if not successful:
            return {"error": "All requests failed", "results": results}

        return {
            "total_requests": num_requests,
            "successful": len(successful),
            "failed": len(failed),
            "success_rate": len(successful) / num_requests * 100,
            "ttft_ms": statistics.median([r["ttft_ms"] for r in successful]),
            "tpot_ms": statistics.median([r["tpot_ms"] for r in successful]),
            "throughput_tps": statistics.mean([r["throughput_tps"] for r in successful]),
            "tokens_per_request": statistics.mean([r["tokens"] for r in successful]),
            "vram_start_mb": vram_start,
            "vram_peak_mb": vram_peak
        }

def load_prompts(filepath="prompts.json"):
    """Load diverse prompts for testing"""
    with open(filepath, 'r') as f:
        return json.load(f)

def main():
    parser = argparse.ArgumentParser(description="LLM Inference Benchmark")
    parser.add_argument("--engine", choices=["ollama", "vllm", "llama-cpp"], required=True)
    parser.add_argument("--concurrency", type=int, default=1)
    parser.add_argument("--requests", type=int, default=50)
    parser.add_argument("--output", default="results.json")

    args = parser.parse_args()

    # Engine configurations
    configs = {
        "ollama": {"url": "http://localhost:11434/v1", "model": "llama3:8b"},
        "vllm": {"url": "http://localhost:8000/v1", "model": "meta-llama/Meta-Llama-3-8B-Instruct"},
        "llama-cpp": {"url": "http://localhost:8080/v1", "model": "llama3-8b-q4_k_m"}
    }

    config = configs[args.engine]
    benchmark = LLMBenchmark(args.engine, config["url"], config["model"])

    print(f"Running benchmark: {args.engine} | Concurrency: {args.concurrency} | Requests: {args.requests}")

    prompts = load_prompts()
    results = benchmark.run_concurrent_benchmark(prompts, args.concurrency, args.requests)

    # Save results
    with open(args.output, 'w') as f:
        json.dump(results, f, indent=2)

    print(f"\nResults saved to {args.output}")
    print(f"Throughput: {results.get('throughput_tps', 0):.1f} tokens/sec")
    print(f"TTFT: {results.get('ttft_ms', 0):.1f} ms")
    print(f"Success Rate: {results.get('success_rate', 0):.1f}%")

if __name__ == "__main__":
    main()