| |
| """harvest pod benchmark — throughput + GPU-utilization sweep. |
| |
| Runs advisory-style chat completions at increasing concurrency against a local |
| llama-server, with reasoning ON or OFF, while sampling nvidia-smi. Reports |
| aggregate output tok/s (the budget number), latency, TTFT, and GPU util/VRAM. |
| |
| python3 bench-pod.py --base-url http://127.0.0.1:18080/v1 --model qwen3.6-35b \ |
| --concurrencies 1,8,32,64 --thinking off --max-tokens 256 --out r.json |
| """ |
| import argparse, json, statistics, subprocess, sys, threading, time |
| from concurrent.futures import ThreadPoolExecutor, as_completed |
|
|
| try: |
| import requests |
| except Exception: |
| sys.exit("needs `requests` (pip install requests)") |
|
|
| SYSTEM = ("You are one advisory agent in a multi-agent system serving a smallholder " |
| "farmer across a full cropping season. Give concise, decision-focused advice.") |
| USER = ("Farmer profile: 0.8 ha rainfed plot, sandy-loam soil, monsoon onset delayed ~2 weeks, " |
| "cotton, moderate bollworm pressure reported nearby, limited cash for inputs. " |
| "It is the sowing window. Recommend: sow now vs wait, variety duration, and the single " |
| "most important early-season action. Explain your reasoning in a few sentences.") |
|
|
|
|
| class GpuSampler(threading.Thread): |
| def __init__(self): |
| super().__init__(daemon=True) |
| self.samples, self.on = [], True |
| def run(self): |
| while self.on: |
| try: |
| out = subprocess.run( |
| ["nvidia-smi", "--query-gpu=utilization.gpu,memory.used", |
| "--format=csv,noheader,nounits"], |
| capture_output=True, text=True, timeout=5).stdout.strip().splitlines()[0] |
| u, m = [int(x.strip()) for x in out.split(",")] |
| self.samples.append((u, m)) |
| except Exception: |
| pass |
| time.sleep(1) |
| def window(self, start_idx): |
| w = self.samples[start_idx:] |
| if not w: return {} |
| return {"gpu_util_mean": round(statistics.mean(s[0] for s in w), 1), |
| "gpu_util_max": max(s[0] for s in w), |
| "vram_mb_max": max(s[1] for s in w)} |
|
|
|
|
| def one_request(base_url, model, key, max_tokens, thinking): |
| body = {"model": model, |
| "messages": [{"role": "system", "content": SYSTEM}, |
| {"role": "user", "content": USER}], |
| "max_tokens": max_tokens, "temperature": 0.7, "stream": True, |
| "stream_options": {"include_usage": True}, |
| "chat_template_kwargs": {"enable_thinking": thinking}} |
| t0 = time.perf_counter(); ttft = None; ctoks = 0 |
| try: |
| with requests.post(base_url.rstrip("/") + "/chat/completions", |
| headers={"Authorization": f"Bearer {key}"}, |
| json=body, stream=True, timeout=900) as r: |
| r.raise_for_status() |
| for line in r.iter_lines(): |
| if not line: continue |
| s = line.decode("utf-8", "ignore") |
| if not s.startswith("data: "): continue |
| s = s[6:] |
| if s.strip() == "[DONE]": break |
| try: ch = json.loads(s) |
| except json.JSONDecodeError: continue |
| if ttft is None and ch.get("choices"): |
| d = ch["choices"][0].get("delta", {}) |
| if d.get("content") or d.get("reasoning_content"): |
| ttft = time.perf_counter() - t0 |
| if ch.get("usage"): |
| ctoks = ch["usage"].get("completion_tokens", 0) |
| return ttft, time.perf_counter() - t0, ctoks, True, None |
| except Exception as e: |
| return ttft, time.perf_counter() - t0, ctoks, False, f"{type(e).__name__}: {e}" |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--base-url", required=True) |
| ap.add_argument("--model", required=True) |
| ap.add_argument("--api-key", default="sk-harvest-local") |
| ap.add_argument("--concurrencies", default="1,8,32,64") |
| ap.add_argument("--rounds", type=int, default=2) |
| ap.add_argument("--max-tokens", type=int, default=256) |
| ap.add_argument("--thinking", choices=["on", "off"], default="off") |
| ap.add_argument("--out", default=None) |
| a = ap.parse_args() |
| thinking = a.thinking == "on" |
|
|
| sampler = GpuSampler(); sampler.start() |
| print(f"model={a.model} thinking={a.thinking} max_tokens={a.max_tokens} rounds={a.rounds}") |
| print(f"{'conc':>5} {'ok':>7} {'agg tok/s':>10} {'lat p50':>8} {'ttft p50':>9} " |
| f"{'util avg':>9} {'util max':>9} {'vram MB':>8}") |
| results = [] |
| for c in [int(x) for x in a.concurrencies.split(",") if x.strip()]: |
| idx = len(sampler.samples) |
| n = c * a.rounds |
| lat, ttfts, toks, oks = [], [], [], 0 |
| w0 = time.perf_counter() |
| with ThreadPoolExecutor(max_workers=c) as ex: |
| futs = [ex.submit(one_request, a.base_url, a.model, a.api_key, |
| a.max_tokens, thinking) for _ in range(n)] |
| for f in as_completed(futs): |
| ttft, total, ct, ok, err = f.result() |
| if ok: |
| oks += 1; lat.append(total); toks.append(ct) |
| if ttft is not None: ttfts.append(ttft) |
| elif err: |
| print(f" [err] {err[:90]}", file=sys.stderr) |
| wall = time.perf_counter() - w0 |
| gpu = sampler.window(idx) |
| agg = sum(toks) / wall if wall else 0 |
| row = {"concurrency": c, "ok": oks, "n": n, "agg_tok_s": round(agg, 1), |
| "lat_p50_s": round(statistics.median(lat), 2) if lat else None, |
| "ttft_p50_s": round(statistics.median(ttfts), 3) if ttfts else None, |
| "mean_completion_tokens": round(statistics.mean(toks), 1) if toks else 0, |
| "wall_s": round(wall, 1), **gpu} |
| results.append(row) |
| print(f"{c:>5} {oks:>3}/{n:<3} {row['agg_tok_s']:>10} {row['lat_p50_s'] or 0:>8} " |
| f"{row['ttft_p50_s'] or 0:>9} {row.get('gpu_util_mean',0):>9} " |
| f"{row.get('gpu_util_max',0):>9} {row.get('vram_mb_max',0):>8}") |
| sampler.on = False |
| peak = max((r["agg_tok_s"] for r in results), default=0) |
| print(f"\npeak aggregate: {peak:.0f} tok/s (thinking={a.thinking})") |
| if a.out: |
| json.dump({"model": a.model, "thinking": a.thinking, |
| "max_tokens": a.max_tokens, "results": results, |
| "peak_agg_tok_s": peak}, open(a.out, "w"), indent=2) |
| print("wrote", a.out) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
