#!/usr/bin/env python3 """GPQA-Diamond against llama-server with a shared (unified) KV pool. Sampling: temperature=1.0, top_p=0.95 (top_k and min_p off), reasoning_effort=high, max_tokens=65536, 1 sample. llama-server fails every running request when its shared KV pool (-kvu -c N) is full. To use many slots without that, this client keeps the sum of the live sequence lengths below the pool size: a question is admitted only when there is room, and when the pool is close to full the youngest request is aborted and re-queued (it restarts later as a fresh sample). Each finished question is appended to the output file, so an interrupted run continues where it stopped. Usage: pip install aiohttp datasets python gpqa.py --url http://127.0.0.1:8080 --pool 102400 --max-conc 10 --out gpqa_results.jsonl python gpqa.py --summary --out gpqa_results.jsonl --pool and --max-conc must match the server's -c and -np. """ import argparse, asyncio, json, os, re, time import aiohttp TEMPLATE = ("Answer the following multiple choice question. The last line of your response should be of the following " "format: 'Answer: $LETTER' (without quotes) where LETTER is one of ABCD. Think step by step before answering.\n\n{q}") def parse(text): m = re.findall(r"(?i)answer\s*[:：]\s*\**\s*\(?\$?([ABCD])\b", text or "") if m: return m[-1].upper() m = re.findall(r"\\boxed\{\s*(?:\\text\{)?([ABCD])\b", text or "") return m[-1] if m else None class Runner: def __init__(self, a): from datasets import load_dataset self.a = a # questions already contain the shuffled options "A. ... D."; "answer" is the letter ds = load_dataset(a.dataset, split="test") self.items = [{"idx": i, "q": x["question"], "gold": x["answer"].strip()} for i, x in enumerate(ds)] self.done = {} if os.path.exists(a.out): for line in open(a.out): r = json.loads(line) self.done[r["idx"]] = r self.pending = [it for it in self.items if it["idx"] not in self.done][: a.limit or None] self.active = {} # idx -> state self.attempts = {} self.preempted = 0 self.t0 = time.time() self.gen_tokens = 0 def usage(self): return sum(s["prompt"] + int(s["gen"] * 1.05) + 8 for s in self.active.values()) async def prompt_tokens(self, sess, q): body = {"messages": [{"role": "user", "content": TEMPLATE.format(q=q)}], "chat_template_kwargs": {"reasoning_effort": "high"}} async with sess.post(f"{self.a.url}/apply-template", json=body) as r: p = (await r.json())["prompt"] async with sess.post(f"{self.a.url}/tokenize", json={"content": p}) as r: return len((await r.json())["tokens"]) async def one(self, it, st): att = self.attempts.get(it["idx"], 0) body = {"messages": [{"role": "user", "content": TEMPLATE.format(q=it["q"])}], "max_tokens": self.a.max_tokens, "temperature": 1.0, "top_p": 0.95, "top_k": 0, "min_p": 0.0, "samplers": ["top_k", "top_p", "min_p", "temperature"], "seed": 1000 * it["idx"] + att, "chat_template_kwargs": {"reasoning_effort": "high"}, "stream": True, "stream_options": {"include_usage": True}} content, reasoning, finish, usage = [], [], None, None # a fresh connection per request: closing it makes the server stop the slot async with aiohttp.ClientSession(connector=aiohttp.TCPConnector(force_close=True), timeout=aiohttp.ClientTimeout(total=None, sock_read=900)) as s: async with s.post(f"{self.a.url}/v1/chat/completions", json=body) as r: if r.status != 200: raise RuntimeError(f"http {r.status}: {(await r.text())[:300]}") async for raw in r.content: line = raw.decode(errors="replace").strip() if not line.startswith("data:") or line == "data: [DONE]": continue j = json.loads(line[5:]) if "error" in j: raise RuntimeError(str(j["error"])[:300]) if j.get("usage"): usage = j["usage"] for ch in j.get("choices", []): d = ch.get("delta", {}) if d.get("content"): content.append(d["content"]) if d.get("reasoning_content"): reasoning.append(d["reasoning_content"]) if d.get("content") or d.get("reasoning_content"): st["gen"] += 1 self.gen_tokens += 1 if ch.get("finish_reason"): finish = ch["finish_reason"] content, reasoning = "".join(content), "".join(reasoning) pred = parse(content) return {"idx": it["idx"], "gold": it["gold"], "pred": pred, "correct": pred == it["gold"], "finish": finish, "completion_tokens": (usage or {}).get("completion_tokens", st["gen"]), "prompt_tokens": st["prompt"], "sec": time.time() - st["t"], "attempt": att, "content": content, "reasoning": reasoning, "fallback_pred_from_reasoning": parse(reasoning) if pred is None else None} def log(self, msg=""): n = len(self.done); c = sum(r["correct"] for r in self.done.values()) el = time.time() - self.t0 print(f"[{el / 60:6.1f} min] done {n}/{len(self.items)} acc {c}/{n} | active {len(self.active)} " f"pool {self.usage()}/{self.a.pool} | {self.gen_tokens / max(el, 1):.0f} tok/s | preempted {self.preempted} {msg}", flush=True) async def run(self): out = open(self.a.out, "a") async with aiohttp.ClientSession() as sess: for it in self.pending: it["ptok"] = await self.prompt_tokens(sess, it["q"]) print(f"pending {len(self.pending)}, prompt tokens max {max((i['ptok'] for i in self.pending), default=0)}", flush=True) queue = list(self.pending) last_log = 0 while queue or self.active: # admit while queue and len(self.active) < self.a.max_conc: it = queue[0] if self.active and self.usage() + it["ptok"] + self.a.reserve > self.a.pool - self.a.margin: break queue.pop(0) st = {"prompt": it["ptok"], "gen": 0, "t": time.time(), "it": it} st["task"] = asyncio.create_task(self.one(it, st)) self.active[it["idx"]] = st # preempt the youngest if the pool is about to overflow if len(self.active) > 1 and self.usage() > self.a.pool - self.a.margin: idx, st = max(self.active.items(), key=lambda kv: kv[1]["t"]) st["task"].cancel() try: await st["task"] except BaseException: pass del self.active[idx] self.attempts[idx] = self.attempts.get(idx, 0) + 1 self.preempted += 1 queue.insert(0, st["it"]) self.log(f"(preempted q{idx} at {st['gen']} tok)") # collect finished for idx, st in list(self.active.items()): if st["task"].done(): del self.active[idx] try: r = st["task"].result() except Exception as e: self.attempts[idx] = self.attempts.get(idx, 0) + 1 queue.append(st["it"]) self.log(f"(q{idx} error, re-queued: {str(e)[:120]})") await asyncio.sleep(5) continue self.done[idx] = r out.write(json.dumps(r, ensure_ascii=False) + "\n"); out.flush() self.log(f"(q{idx} {'OK ' if r['correct'] else 'BAD'} {r['completion_tokens']} tok {r['finish']})") if time.time() - last_log > 300: last_log = time.time(); self.log() await asyncio.sleep(0.5) self.log("FINISHED") def summary(path, n_total): rs = [json.loads(line) for line in open(path)] n = len(rs); c = sum(r["correct"] for r in rs) tok = sorted(r["completion_tokens"] for r in rs) print(f"answered {n}/{n_total}: accuracy {c}/{n} = {100 * c / max(n, 1):.1f}%") print(f" hit the token limit: {sum(r['finish'] == 'length' for r in rs)}, no answer parsed: {sum(r['pred'] is None for r in rs)}, " f"restarted after preemption: {sum(r['attempt'] > 0 for r in rs)}") if tok: print(f" completion tokens: mean {sum(tok) / n:.0f}, median {tok[n // 2]}, p90 {tok[int(n * 0.9)]}, max {tok[-1]}") if __name__ == "__main__": ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--url", default="http://127.0.0.1:8080", help="llama-server base URL") ap.add_argument("--out", default="gpqa_results.jsonl", help="results file (appended, resumable)") ap.add_argument("--dataset", default="fingertap/GPQA-Diamond", help="HF dataset with 'question' and 'answer' columns") ap.add_argument("--pool", type=int, default=102400, help="KV pool size in tokens, = server -c") ap.add_argument("--max-conc", type=int, default=10, help="max parallel requests, = server -np") ap.add_argument("--margin", type=int, default=3072, help="tokens kept free in the pool") ap.add_argument("--reserve", type=int, default=12288, help="free tokens required before a new question starts") ap.add_argument("--max-tokens", type=int, default=65536) ap.add_argument("--limit", type=int, default=0, help="only run the first N pending questions") ap.add_argument("--summary", action="store_true", help="print the score of --out and exit") a = ap.parse_args() if a.summary: summary(a.out, 198) else: asyncio.run(Runner(a).run())