#!/usr/bin/env python3
"""四种负载下的解码吞吐与 MTP 接受率。prompt 与 2026-09-08 那轮逐字相同，可直接对比。

用法: python3 bench-workloads.py <api-key> [标签]
接受率取自 /metrics 的窗口差值，而不是服务端日志的滚动均值。
"""
import json, re, sys, time, urllib.request

KEY = sys.argv[1]
TAG = sys.argv[2] if len(sys.argv) > 2 else "run"
URL = "http://127.0.0.1:8000"

CASES = [
    ("中文散文", "请写一段 400 字的短文，主题：分布式推理系统里的 KV cache 管理。"),
    ("Python代码", "Write a complete Python class ThreadSafeLRUCache with get/put/delete, an "
                   "OrderedDict backing store, a threading.RLock, TTL support, and full type hints. "
                   "Include docstrings."),
    ("重复结构", "List the numbers 1 through 60. For each, write one line in exactly this format: "
                 "'Item N: value=N, squared=N*N, status=active'."),
    ("JSON输出", "Generate a JSON array of 25 user records. Each object has fields: id (int), "
                 "name (string), email (string), age (int), city (string), active (bool). "
                 "Output only valid JSON."),
]


def once(prompt):
    """返回 (TTFT, 解码 tok/s, completion_tokens)。解码速率不含 TTFT。"""
    req = urllib.request.Request(
        URL + "/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": 0, "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=300) 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


def metrics():
    """抓取 spec_decode 相关计数器，返回 {名字(含标签): 值}。"""
    req = urllib.request.Request(URL + "/metrics",
                                 headers={"Authorization": "Bearer " + KEY})
    out = {}
    with urllib.request.urlopen(req, timeout=30) as r:
        for line in r.read().decode().splitlines():
            if not line.startswith("vllm:spec_decode"):
                continue
            m = re.match(r'^(\S+)\s+([0-9.eE+-]+)$', line)
            if not m:
                continue
            name, val = m.group(1), float(m.group(2))
            # 丢掉 model_name 之类的高基数标签，只保留 position
            name = re.sub(r'model_name="[^"]*",?', '', name)
            name = re.sub(r'engine="[^"]*",?', '', name)
            name = name.replace('{,', '{').replace(',}', '}').replace('{}', '')
            out[name] = out.get(name, 0.0) + val
    return out


def delta(a, b):
    return {k: b.get(k, 0.0) - a.get(k, 0.0) for k in set(a) | set(b)}


def pick(d, key):
    """只取 *_total 计数器：同名的 *_created 是 Unix 时间戳，窗口差值为 0。"""
    pre = "vllm:spec_decode_" + key + "_total"
    for k, v in d.items():
        if k.startswith(pre):
            return v
    return 0.0


def per_pos(d):
    """按位置的接受计数，按 position 标签排序。"""
    pre = "vllm:spec_decode_num_accepted_tokens_per_pos_total"
    out = []
    for k, v in d.items():
        if not k.startswith(pre):
            continue
        m = re.search(r'position="(\d+)"', k)
        out.append((int(m.group(1)) if m else 0, v))
    return [v for _, v in sorted(out)]


print("### %s  %s" % (TAG, time.strftime("%Y-%m-%dT%H:%M:%S")))
rows = []
for name, prompt in CASES:
    once(prompt)                                  # 预热，不计入
    m0 = metrics()
    rates = sorted(once(prompt)[1] for _ in range(3))
    ttft, _, ct = once(prompt)
    m1 = metrics()
    d = delta(m0, m1)
    drafted = pick(d, "num_draft_tokens")
    accepted = pick(d, "num_accepted_tokens")
    drafts = pick(d, "num_drafts")
    acc = accepted / drafted * 100 if drafted else float("nan")
    per_len = accepted / drafts if drafts else float("nan")
    pos = per_pos(d)
    pos_txt = ""
    if pos and drafts:
        pos_txt = "  按位置 " + str(["%.0f%%" % (v / drafts * 100) for v in pos])
    rows.append((name, rates[1], ttft, ct, acc, per_len, pos_txt))
    print("  %-10s 解码 %5.1f tok/s   TTFT %.2fs  (%d token)   接受率 %5.1f%%  每步接受 %.2f%s"
          % (name, rates[1], ttft, ct, acc, per_len, pos_txt))

print("\n  最快负载: %s  %.1f tok/s" % max(((r[0], r[1]) for r in rows), key=lambda x: x[1]))