load_test.py

"""Benchmark v3 — concurrent load: how throughput & tail latency shift under load."""
import asyncio
import time
import httpx

OLLAMA_URL = "http://localhost:11434/v1/chat/completions"
MODEL = "qwen2.5:1.5b"
REQUESTS_PER_LEVEL = 20
MAX_TOKENS = 64
PROMPT = "Explain what a GPU is, in 2 sentences."
CONCURRENCY_LEVELS = [1, 4, 8]          # how many requests run at the same time


def percentile(sorted_data, p):
    k = (len(sorted_data) - 1) * (p / 100)
    lo = int(k)
    hi = min(lo + 1, len(sorted_data) - 1)
    return sorted_data[lo] + (sorted_data[hi] - sorted_data[lo]) * (k - lo)


async def one_request(client, sem):
    payload = {
        "model": MODEL,
        "messages": [{"role": "user", "content": PROMPT}],
        "max_tokens": MAX_TOKENS,
    }
    async with sem:                       # the semaphore caps how many run at once
        start = time.perf_counter()
        resp = await client.post(OLLAMA_URL, json=payload, timeout=120.0)
        elapsed = time.perf_counter() - start
    tokens = resp.json()["usage"]["completion_tokens"]
    return elapsed, tokens


async def run_level(concurrency):
    sem = asyncio.Semaphore(concurrency)
    async with httpx.AsyncClient() as client:
        wall_start = time.perf_counter()
        results = await asyncio.gather(
            *[one_request(client, sem) for _ in range(REQUESTS_PER_LEVEL)]
        )
        wall = time.perf_counter() - wall_start
    latencies = sorted(r[0] for r in results)
    total_tokens = sum(r[1] for r in results)
    return {
        "c": concurrency,
        "throughput": total_tokens / wall,
        "p50": percentile(latencies, 50),
        "p95": percentile(latencies, 95),
        "p99": percentile(latencies, 99),
    }


async def main():
    print(f"Load test: {MODEL}  ·  {REQUESTS_PER_LEVEL} requests per level\n")
    print(f"{'concurrency':>11} | {'throughput':>11} | {'p50':>7} | {'p95':>7} | {'p99':>7}")
    print("-" * 58)
    for c in CONCURRENCY_LEVELS:
        r = await run_level(c)
        print(f"{r['c']:>11} | {r['throughput']:>8.1f}t/s | {r['p50']:>6.2f}s | "
              f"{r['p95']:>6.2f}s | {r['p99']:>6.2f}s")


if __name__ == "__main__":
    asyncio.run(main())