#!/usr/bin/env bash
# Archived benchmark: run against a separately configured local endpoint.
#   1) 流式单流解码 tok/s（T=0 与 T=0.7），按首 token 到末 token 计时，不含 TTFT
#   2) 并发 1/2/4/8 的聚合 tok/s（非流式，200 token，wall time）
#   3) MTP 草稿接受率（总量 + 按位置）
set -uo pipefail
KEY="${VLLM_API_KEY:?Set VLLM_API_KEY}"; URL="http://127.0.0.1:8000"
TAG="${1:-run}"; OUT=~/bench-results/$(date +%Y%m%d-%H%M%S)-$TAG.txt; mkdir -p ~/bench-results
PROMPT="请写一段 400 字的短文，主题：分布式推理系统里的 KV cache 管理。"
{
echo "### $TAG $(date -Is)"
echo "config: $(curl -s $URL/metrics >/dev/null && docker exec ray-head bash -c 'grep -oE "non-default args: \{.*" /tmp/vllm.log | tail -1 | cut -c1-600')"
echo "--- 1) streaming single-stream decode ---"
for T in 0 0.7; do
python3 - "$KEY" "$T" "$PROMPT" <<'PY'
import json, sys, time, urllib.request
key, T, prompt = sys.argv[1], float(sys.argv[2]), sys.argv[3]
def run():
    req = urllib.request.Request("http://127.0.0.1:8000/v1/chat/completions", method="POST",
        headers={"Authorization": "Bearer "+key, "Content-Type": "application/json"},
        data=json.dumps({"model":"qwen3.8-flash-next","messages":[{"role":"user","content":prompt}],
                         "max_tokens":300,"temperature":T,"stream":True,"stream_options":{"include_usage":True},"chat_template_kwargs":{"enable_thinking":False}}).encode())
    t0=time.time(); first=last=None; usage=None
    with urllib.request.urlopen(req, timeout=600) as r:
        for line in r:
            line=line.decode().strip()
            if not line.startswith("data:") or line=="data: [DONE]": continue
            d=json.loads(line[5:])
            if d.get("usage"): usage=d["usage"]
            ch=d.get("choices") or []
            if ch and (ch[0].get("delta") or {}).get("content"):
                now=time.time(); last=now; first=first or now
    ct=usage["completion_tokens"]; return first-t0, (ct-1)/(last-first), ct
run()                       # 预热一次（页缓存 / JIT）
res=[run() for _ in range(3)]
dec=sorted(r[1] for r in res)[1]
print(f"T={T}: TTFT={min(r[0] for r in res):.2f}s  decode(median of 3)={dec:.1f} tok/s  tokens={[r[2] for r in res]}")
PY
done
echo "--- 2) concurrency sweep (200 tokens each, T=0.7) ---"
bench() {
  N=$1; t0=$(date +%s.%N)
  for i in $(seq 1 $N); do
    curl -s -m 600 $URL/v1/chat/completions -H "Authorization: Bearer $KEY" -H 'Content-Type: application/json' \
      -d "{\"model\":\"qwen3.8-flash-next\",\"messages\":[{\"role\":\"user\",\"content\":\"$PROMPT 编号 $i。\"}],\"max_tokens\":200,\"temperature\":0.7,\"chat_template_kwargs\":{\"enable_thinking\":false}}" \
      | python3 -c 'import json,sys; d=json.load(sys.stdin); print(d["usage"]["completion_tokens"])' > /tmp/bench.$i.out 2>/dev/null &
  done; wait
  t1=$(date +%s.%N); tot=$(cat /tmp/bench.*.out 2>/dev/null | paste -sd+ | bc); rm -f /tmp/bench.*.out
  python3 -c "t=$t1-$t0; print(f'N={$N:>2}  wall={t:6.1f}s  tokens={$tot:>5}  aggregate={$tot/t:6.1f} tok/s  per-stream={$tot/t/$N:5.1f} tok/s')"
}
for n in 1 2 4 8; do bench $n; done
echo "--- 3) MTP acceptance ---"
curl -s $URL/metrics | python3 -c '
import re,sys
m={}
for l in sys.stdin:
    r=re.match(r"vllm:spec_decode_num_(draft_tokens|accepted_tokens|drafts)_total\{[^}]*\} ([0-9.e+]+)", l)
    if r: m[r.group(1)]=float(r.group(2))
    r=re.match(r"vllm:spec_decode_num_accepted_tokens_per_pos_total\{[^}]*position=\"(\d+)\"[^}]*\} ([0-9.e+]+)", l)
    if r: m["pos"+r.group(1)]=float(r.group(2))
d,a,n=m.get("draft_tokens",0),m.get("accepted_tokens",0),m.get("drafts",1)
pos=[m[k] for k in sorted(k for k in m if k.startswith("pos"))]
pp = ["%.0f%%" % (100*p/max(n,1)) for p in pos]
print("drafts=%.0f draft_tokens=%.0f accepted=%.0f acceptance=%.1f%%  mean_accepted_per_step=%.2f  per-position=%s" % (n, d, a, 100*a/max(d,1), a/max(n,1), pp))'
} | tee "$OUT"
echo "saved: $OUT"