| """Record one model answering one prompt through llama-server, with a timestamp for every streamed chunk. |
| |
| python3 scripts/race_record.py --server BIN --model GGUF --name "Dory-Beam-4B" --loops 2 \ |
| --problem math500-123 --seed 0 --out race/dory-beam.json [--server-args "..."] |
| |
| Starts llama-server, sends the problem as one user message (temperature 1.0, top_p 0.95, fixed seed), streams the |
| reply with raw <think> tags kept in the content, and writes {name, loops, problem, prompt, gold, events: [[t, text]], |
| completion_tokens, seconds, timings}. t is seconds since the request was sent. Stops the server afterwards. |
| """ |
|
|
| import argparse |
| import json |
| import shlex |
| import subprocess |
| import time |
| import urllib.request |
| from pathlib import Path |
|
|
| ROOT = Path(__file__).resolve().parent.parent |
| PORT = 8099 |
|
|
|
|
| def wait_ready(proc, timeout=600): |
| t0 = time.time() |
| while time.time() - t0 < timeout: |
| if proc.poll() is not None: |
| raise RuntimeError("llama-server exited during startup") |
| try: |
| with urllib.request.urlopen(f"http://127.0.0.1:{PORT}/health", timeout=2) as r: |
| if r.status == 200: |
| return |
| except Exception: |
| pass |
| time.sleep(1) |
| raise TimeoutError("llama-server did not become ready") |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--server", required=True) |
| ap.add_argument("--model", required=True) |
| ap.add_argument("--name", required=True) |
| ap.add_argument("--loops", type=int, default=2) |
| ap.add_argument("--problem", required=True) |
| ap.add_argument("--seed", type=int, default=0) |
| ap.add_argument("--max-tokens", type=int, default=16000) |
| ap.add_argument("--server-args", default="") |
| ap.add_argument("--effort", default="", help="reasoning_effort chat template argument (Dory-Beam-4B v2)") |
| ap.add_argument("--out", required=True) |
| args = ap.parse_args() |
| prob = next(json.loads(l) for l in open(ROOT / "data" / "evalsets.jsonl") if json.loads(l)["id"] == args.problem) |
|
|
| cmd = [args.server, "-m", args.model, "--host", "127.0.0.1", "--port", str(PORT), "-ngl", "999", |
| "--ctx-size", "20480", "--parallel", "1", "--no-context-shift", "--jinja", "--reasoning-format", "none", |
| "--override-kv", f"dory.num_loops=int:{args.loops},dory.recurrent_kv_cache_mode=str:per_loop", |
| *shlex.split(args.server_args)] |
| log = open(Path(args.out).with_suffix(".server.log"), "w") |
| proc = subprocess.Popen(cmd, stdout=log, stderr=subprocess.STDOUT) |
| try: |
| wait_ready(proc) |
| body = {"messages": [{"role": "user", "content": prob["prompt"]}], "stream": True, "temperature": 1.0, |
| "top_p": 0.95, "top_k": 0, "min_p": 0, "seed": args.seed, "max_tokens": args.max_tokens, |
| "cache_prompt": False, "timings_per_token": False} |
| if args.effort: |
| body["chat_template_kwargs"] = {"reasoning_effort": args.effort} |
| req = urllib.request.Request(f"http://127.0.0.1:{PORT}/v1/chat/completions", json.dumps(body).encode(), |
| {"Content-Type": "application/json"}) |
| events, timings, t0 = [], None, time.time() |
| with urllib.request.urlopen(req, timeout=7200) as resp: |
| for raw in resp: |
| line = raw.decode("utf-8", "replace").strip() |
| if not line.startswith("data: ") or line == "data: [DONE]": |
| continue |
| chunk = json.loads(line[6:]) |
| timings = chunk.get("timings", timings) |
| for ch in chunk.get("choices", []): |
| text = (ch.get("delta") or {}).get("content") |
| if text: |
| events.append([round(time.time() - t0, 3), text]) |
| seconds = time.time() - t0 |
| finally: |
| proc.terminate() |
| proc.wait(timeout=60) |
|
|
| text = "".join(e[1] for e in events) |
| out = {"name": args.name, "model_file": Path(args.model).name, "loops": args.loops, "effort": args.effort, |
| "problem": args.problem, |
| "prompt": prob["prompt"], "gold": prob["answer"], "seed": args.seed, "events": events, |
| "completion_tokens": (timings or {}).get("predicted_n"), "seconds": round(seconds, 2), "timings": timings, |
| "text": text} |
| Path(args.out).parent.mkdir(parents=True, exist_ok=True) |
| Path(args.out).write_text(json.dumps(out, ensure_ascii=False)) |
| print(f"{args.name}: {out['completion_tokens']} tokens in {seconds:.1f}s -> {args.out}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
