Inference benchmark harness

#!/usr/bin/env python3
"""Concurrent latency/load harness for an OpenAI-compatible endpoint."""
import argparse, concurrent.futures, json, math, statistics, time, urllib.request

def percentile(xs, q):
    if not xs: return float("nan")
    ys=sorted(xs); pos=(len(ys)-1)*q; lo=math.floor(pos); hi=math.ceil(pos)
    return ys[lo] if lo==hi else ys[lo]*(hi-pos)+ys[hi]*(pos-lo)

def one(url, prompt, timeout):
    body=json.dumps({"model":"course-model","messages":[{"role":"user","content":prompt}],"max_tokens":64}).encode()
    req=urllib.request.Request(url,data=body,method="POST",headers={"content-type":"application/json"})
    start=time.perf_counter()
    try:
        with urllib.request.urlopen(req,timeout=timeout) as response:
            payload=response.read()
        elapsed=(time.perf_counter()-start)*1000
        data=json.loads(payload); tokens=data.get("usage",{}).get("completion_tokens",0)
        return {"ok":True,"total_ms":elapsed,"tokens":tokens}
    except Exception as exc:
        return {"ok":False,"error":type(exc).__name__}

def run(url, concurrency, requests, timeout):
    prompt="Explain why goodput is more useful than peak throughput for an inference service."
    started=time.perf_counter()
    with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as pool:
        results=list(pool.map(lambda _: one(url,prompt,timeout),range(requests)))
    wall=time.perf_counter()-started; ok=[r for r in results if r["ok"]]; lat=[r["total_ms"] for r in ok]
    return {"concurrency":concurrency,"requests":requests,"success":len(ok),"errors":requests-len(ok),
        "success_rps":round(len(ok)/wall,3),"p50_ms":round(percentile(lat,.5),2),
        "p95_ms":round(percentile(lat,.95),2),"p99_ms":round(percentile(lat,.99),2),
        "goodput_rps_at_250ms":round(sum(x<=250 for x in lat)/wall,3)}

if __name__=="__main__":
    p=argparse.ArgumentParser(); p.add_argument("--url",required=True); p.add_argument("--concurrency",default="1,4,8")
    p.add_argument("--requests",type=int,default=40); p.add_argument("--timeout",type=float,default=30)
    a=p.parse_args()
    for c in [int(x) for x in a.concurrency.split(",")]:
        print(json.dumps(run(a.url,c,a.requests,a.timeout)))