104 lines
3.5 KiB
Python
Executable File
104 lines
3.5 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""pp-bench.py — drive the PP head API and report server-side timings.
|
|
|
|
Runs on the HEAD machine (Dankclaw). Targets the running PP head
|
|
(serve-pp-head.sh, port from argument). Builds a ~Nk-token prose prompt,
|
|
sends one warmup completion, then REPS measured completions, and reads the
|
|
authoritative `slot print_timing` lines from the head log.
|
|
|
|
Usage:
|
|
python3 pp-bench.py <prompt_k> <out_tokens> <reps> <port> <head_log>
|
|
"""
|
|
import json
|
|
import re
|
|
import sys
|
|
import time
|
|
import urllib.request
|
|
|
|
prompt_k, out_tokens, reps, port, head_log = sys.argv[1:6]
|
|
out_tokens = int(out_tokens)
|
|
reps = int(reps)
|
|
target = int(prompt_k) * 1000
|
|
|
|
# ~38 tokens per 160-char prose unit; size the filler to hit the target.
|
|
unit = ("The careful study of natural language reveals patterns of meaning "
|
|
"that emerge from the deliberate arrangement of words and phrases. ")
|
|
approx = max(1, target // 38)
|
|
prompt = unit * approx + (
|
|
"\n\nNow, in a detailed and flowing paragraph, "
|
|
"continue this essay about the philosophy of "
|
|
"language and its influence on modern computation. "
|
|
)
|
|
|
|
|
|
def req(n):
|
|
body = json.dumps({"prompt": prompt, "max_tokens": n,
|
|
"temperature": 0.0, "stream": False}).encode()
|
|
r = urllib.request.Request(
|
|
f"http://127.0.0.1:{port}/v1/completions",
|
|
data=body, headers={"Content-Type": "application/json"})
|
|
with urllib.request.urlopen(r, timeout=7200) as resp:
|
|
return json.load(resp)
|
|
|
|
|
|
def log_timings():
|
|
try:
|
|
lines = open(head_log, encoding="utf-8", errors="ignore").read().splitlines()
|
|
except OSError:
|
|
return []
|
|
out = []
|
|
for l in lines:
|
|
if re.search(r"prompt eval time|eval time|draft acceptance", l):
|
|
out.append(l.strip())
|
|
return out
|
|
|
|
|
|
GEN_RE = re.compile(
|
|
r"eval time =\s+([\d.]+) ms /\s+(\d+) tokens.*?([\d.]+) tokens per second")
|
|
PRE_RE = re.compile(
|
|
r"prompt eval time =\s+([\d.]+) ms /\s+(\d+) tokens.*?([\d.]+) tokens per second")
|
|
ACC_RE = re.compile(r"draft acceptance = ([\d.]+)")
|
|
|
|
# warmup (first run pays graph-compile / JIT costs)
|
|
warm = req(16)
|
|
print(f"warmup: prompt_tokens={warm['usage']['prompt_tokens']} "
|
|
f"completion_tokens={warm['usage']['completion_tokens']}")
|
|
|
|
gen_rates, prefill_rates, accepts = [], [], []
|
|
for i in range(reps):
|
|
t0 = time.time()
|
|
out = req(out_tokens)
|
|
dt = time.time() - t0
|
|
u = out["usage"]
|
|
pt, ct = u["prompt_tokens"], u["completion_tokens"]
|
|
gen = pre = acc = None
|
|
for l in log_timings():
|
|
if "prompt eval" not in l:
|
|
m = GEN_RE.search(l)
|
|
if m:
|
|
gen = float(m.group(3))
|
|
m = PRE_RE.search(l)
|
|
if m:
|
|
pre = float(m.group(3))
|
|
m = ACC_RE.search(l)
|
|
if m:
|
|
acc = float(m.group(1))
|
|
gen_rates.append(gen if gen is not None else ct / dt)
|
|
if pre is not None:
|
|
prefill_rates.append(pre)
|
|
if acc is not None:
|
|
accepts.append(acc)
|
|
print(f"rep {i + 1}: prompt={pt} gen={ct} wall={dt:.1f}s "
|
|
f"gen={gen_rates[-1]:.2f} t/s "
|
|
f"prefill={f'{pre:.1f}' if pre is not None else 'n/a'} t/s "
|
|
f"acceptance={f'{acc:.3f}' if acc is not None else 'n/a'}")
|
|
|
|
print("== RESULT pp ==")
|
|
print(f"avg gen tok/s = {sum(gen_rates) / len(gen_rates):.2f}")
|
|
if prefill_rates:
|
|
print(f"avg prefill t/s = {sum(prefill_rates) / len(prefill_rates):.1f}")
|
|
if accepts:
|
|
print(f"avg acceptance = {sum(accepts) / len(accepts):.3f}")
|
|
for l in log_timings()[-3:]:
|
|
print(" " + l)
|