49 lines
1.7 KiB
Python
Executable File
49 lines
1.7 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""bench-probe.py — RUNS ON DANKSERVER. The campaign's decode probe: one
|
|
/completion request, a prompt assembled from calib.jsonl to ~TARGET tokens,
|
|
temperature 0, max_tokens 512. Reports the engine's own timings: decode rate,
|
|
prefill rate, draft accept. Run 3x; the median is the arm's number."""
|
|
import json, sys, urllib.request
|
|
|
|
TARGET = 43600 # tokens
|
|
BYTES_PER_TOK = 3.8
|
|
URL = "http://localhost:7770/completion"
|
|
|
|
def build_prompt():
|
|
want = int(TARGET * BYTES_PER_TOK)
|
|
parts = []
|
|
n = 0
|
|
with open("/home/dank/corpora/out/calib.jsonl") as f:
|
|
for line in f:
|
|
d = json.loads(line)
|
|
p = d.get("prompt") or d.get("messages", [{}])[0].get("content", "")
|
|
parts.append(p)
|
|
n += len(p.encode())
|
|
if n >= want:
|
|
break
|
|
return "".join(parts)
|
|
|
|
def probe(salt):
|
|
prompt = f"[probe {salt}]\n" + build_prompt()
|
|
body = json.dumps({
|
|
"prompt": prompt, "temperature": 0.0, "max_tokens": 512,
|
|
"stream": False,
|
|
}).encode()
|
|
req = urllib.request.Request(URL, data=body,
|
|
headers={"Content-Type": "application/json"})
|
|
with urllib.request.urlopen(req, timeout=1800) as r:
|
|
d = json.load(r)
|
|
t = d.get("timings", {})
|
|
return {
|
|
"prompt_n": t.get("prompt_n"),
|
|
"prefill_tps": round(t.get("prompt_per_second", 0), 1),
|
|
"pred_n": t.get("predicted_n"),
|
|
"decode_tps": round(t.get("predicted_per_second", 0), 2),
|
|
"draft_n": t.get("draft_n"),
|
|
"accept": (round(t["draft_n_accepted"] / t["draft_n"], 4)
|
|
if t.get("draft_n") else None),
|
|
}
|
|
|
|
if __name__ == "__main__":
|
|
print(json.dumps(probe(sys.argv[1] if len(sys.argv) > 1 else "a")))
|