| """Speed of THX-01 on one GPU: single-request latency, batched throughput, and HTTP concurrency. |
| |
| python bench/speed.py --model ckpt/thx01 --out results/speed.json [--api http://localhost:8095 --key ...] |
| """ |
| import argparse, asyncio, json, sys, time |
| import numpy as np, torch |
| sys.path.insert(0, "./pkg") |
| import systemone |
|
|
| ap = argparse.ArgumentParser() |
| ap.add_argument("--model", default="ckpt/thx01") |
| ap.add_argument("--out", default="results/speed.json") |
| ap.add_argument("--api", default="") |
| ap.add_argument("--key", default="") |
| a = ap.parse_args() |
|
|
| ag = systemone.load(a.model, device="cuda") |
| STATES = ["Salam, mart ayı üçün hesabımızdan iki dəfə pul çıxılıb, zəhmət olmasa geri qaytarın.", |
| "Hi, we were billed twice for March. Please refund the duplicate today or we will cancel.", |
| "511 maddə nə tənzimləyir", "Сайт не работает уже два часа, срочно почините!", |
| "cox pis proqramdir islemir", "Can you write a SQL query that joins orders and customers by month?"] |
| Q = {"route": {"type": "choice", "instructions": "Which AI model should answer this?", |
| "criteria": {"legal": "laws and codes", "daily": "everyday tasks", "fast": "short facts", "reasoning": "hard problems"}}, |
| "urgent": {"type": "noul", "instructions": "Is this urgent?"}, |
| "sentiment": {"type": "score", "instructions": "How negative is the tone?", "criteria": ["calm", "annoyed", "angry"]}} |
| res = {"gpu": torch.cuda.get_device_name(0), "questions_per_request": len(Q)} |
|
|
|
|
| def timed(fn, n=30, warm=5): |
| for _ in range(warm): |
| fn() |
| ts = [] |
| for _ in range(n): |
| torch.cuda.synchronize(); t0 = time.perf_counter(); fn(); torch.cuda.synchronize() |
| ts.append((time.perf_counter() - t0) * 1000) |
| return float(np.percentile(ts, 50)), float(np.percentile(ts, 95)) |
|
|
|
|
| |
| p50, p95 = timed(lambda: ag.predict(STATES[0], Q)) |
| res["single_request"] = {"p50_ms": round(p50, 2), "p95_ms": round(p95, 2)} |
| q1 = {"route": Q["route"]} |
| p50, p95 = timed(lambda: ag.predict(STATES[0], q1)) |
| res["single_question"] = {"p50_ms": round(p50, 2), "p95_ms": round(p95, 2)} |
| print(res, flush=True) |
|
|
| |
| res["batched"] = {} |
| for n in (1, 10, 50, 150, 500, 1000): |
| sts = [STATES[i % len(STATES)] for i in range(n)] |
| p50, _ = timed(lambda: ag.predict_batch(sts, Q, batch_size=256), n=10, warm=2) |
| res["batched"][n] = {"total_ms": round(p50, 1), "ms_per_request": round(p50 / n, 3), |
| "requests_per_s": round(n / p50 * 1000, 1), "decisions_per_s": round(n * len(Q) / p50 * 1000, 1)} |
| print(n, res["batched"][n], flush=True) |
|
|
| |
| if a.api: |
| import httpx |
| H = {"Authorization": "Bearer " + a.key} |
|
|
| async def http(): |
| async with httpx.AsyncClient(timeout=60) as c: |
| body = lambda i: {"state": STATES[i % len(STATES)], "questions": Q, "model": "thx01"} |
| await c.post(a.api + "/v1/decide", json=body(0), headers=H) |
| t0 = time.perf_counter() |
| rs = await asyncio.gather(*(c.post(a.api + "/v1/decide", json=body(i), headers=H) for i in range(150))) |
| conc = (time.perf_counter() - t0) * 1000 |
| ok = sum(r.status_code == 200 for r in rs) |
| t0 = time.perf_counter() |
| r = await c.post(a.api + "/v1/decide/batch", json={"items": [{"state": STATES[i % 6], "questions": Q} for i in range(150)], |
| "model": "thx01"}, headers=H) |
| batch = (time.perf_counter() - t0) * 1000 |
| return conc, ok, batch, r.status_code |
| conc, ok, batch, code = asyncio.run(http()) |
| res["http"] = {"concurrent_150_total_ms": round(conc, 1), "concurrent_ok": ok, |
| "batch_150_total_ms": round(batch, 1), "batch_status": code} |
| print(res["http"]) |
| json.dump(res, open(a.out, "w"), indent=1) |
|
|
