add decode throughput analysis
This commit is contained in:
314
worklogs/decode_microbench.py
Normal file
314
worklogs/decode_microbench.py
Normal file
@@ -0,0 +1,314 @@
|
||||
import argparse
|
||||
import concurrent.futures
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from dataclasses import asdict, dataclass
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
@dataclass
|
||||
class RequestResult:
|
||||
ok: bool
|
||||
elapsed_sec: float
|
||||
ttft_sec: float | None
|
||||
completion_tokens: int
|
||||
prompt_tokens: int
|
||||
cached_tokens: int
|
||||
reasoning_tokens: int
|
||||
output_tps: float | None
|
||||
chars: int
|
||||
error: str | None = None
|
||||
|
||||
|
||||
def percentile(values, pct):
|
||||
values = sorted(v for v in values if v is not None)
|
||||
if not values:
|
||||
return None
|
||||
if len(values) == 1:
|
||||
return values[0]
|
||||
pos = (len(values) - 1) * pct / 100.0
|
||||
lo = math.floor(pos)
|
||||
hi = math.ceil(pos)
|
||||
if lo == hi:
|
||||
return values[lo]
|
||||
return values[lo] * (hi - pos) + values[hi] * (pos - lo)
|
||||
|
||||
|
||||
def make_tools(count):
|
||||
tools = []
|
||||
for i in range(count):
|
||||
tools.append({
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": f"tool_{i:02d}_exec",
|
||||
"description": "Run a deterministic diagnostic action.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"command": {"type": "string"},
|
||||
"path": {"type": "string"},
|
||||
},
|
||||
"required": ["command"],
|
||||
},
|
||||
},
|
||||
})
|
||||
return tools
|
||||
|
||||
|
||||
def make_request(args, idx):
|
||||
if args.prompt_mode == "short":
|
||||
user = (
|
||||
"Do not explain. Output a comma-separated sequence of four digit "
|
||||
"numbers starting at 0001. Continue until the token limit stops you."
|
||||
)
|
||||
messages = [{"role": "user", "content": user}]
|
||||
elif args.prompt_mode == "tool":
|
||||
messages = [
|
||||
{"role": "system", "content": "You are a coding agent. Return concise tool-call-like JSON text."},
|
||||
{"role": "user", "content": "Create a shell command to list Python files and print the answer as JSON."},
|
||||
]
|
||||
else:
|
||||
raise ValueError(f"unknown prompt_mode: {args.prompt_mode}")
|
||||
|
||||
payload = {
|
||||
"model": args.model,
|
||||
"messages": messages,
|
||||
"max_tokens": args.max_tokens,
|
||||
"temperature": 0,
|
||||
"stream": True,
|
||||
"stream_options": {"include_usage": True},
|
||||
}
|
||||
if args.with_tools:
|
||||
payload["tools"] = make_tools(args.tool_count)
|
||||
payload["tool_choice"] = "auto"
|
||||
return payload
|
||||
|
||||
|
||||
def post_stream(url, payload, timeout):
|
||||
req = urllib.request.Request(
|
||||
url.rstrip("/") + "/v1/chat/completions",
|
||||
data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
start = time.perf_counter()
|
||||
first_at = None
|
||||
usage = {}
|
||||
chars = 0
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||||
for raw in resp:
|
||||
line = raw.decode("utf-8", errors="replace").strip()
|
||||
if not line or not line.startswith("data:"):
|
||||
continue
|
||||
data = line[5:].strip()
|
||||
if data == "[DONE]":
|
||||
break
|
||||
obj = json.loads(data)
|
||||
if obj.get("usage"):
|
||||
usage = obj["usage"]
|
||||
for choice in obj.get("choices") or []:
|
||||
delta = choice.get("delta") or {}
|
||||
parts = [
|
||||
delta.get("content") or "",
|
||||
delta.get("reasoning_content") or "",
|
||||
]
|
||||
for tool_call in delta.get("tool_calls") or []:
|
||||
fn = tool_call.get("function") or {}
|
||||
parts.append(fn.get("name") or "")
|
||||
parts.append(fn.get("arguments") or "")
|
||||
added = sum(len(p) for p in parts)
|
||||
if added and first_at is None:
|
||||
first_at = time.perf_counter()
|
||||
chars += added
|
||||
elapsed = time.perf_counter() - start
|
||||
ttft = first_at - start if first_at is not None else None
|
||||
details = usage.get("prompt_tokens_details") or {}
|
||||
completion_tokens = int(usage.get("completion_tokens") or 0)
|
||||
decode_sec = elapsed - ttft if ttft is not None else elapsed
|
||||
return RequestResult(
|
||||
ok=True,
|
||||
elapsed_sec=elapsed,
|
||||
ttft_sec=ttft,
|
||||
completion_tokens=completion_tokens,
|
||||
prompt_tokens=int(usage.get("prompt_tokens") or 0),
|
||||
cached_tokens=int(details.get("cached_tokens") or 0),
|
||||
reasoning_tokens=int(usage.get("reasoning_tokens") or 0),
|
||||
output_tps=completion_tokens / decode_sec if completion_tokens and decode_sec > 0 else None,
|
||||
chars=chars,
|
||||
)
|
||||
except urllib.error.HTTPError as exc:
|
||||
body = exc.read().decode("utf-8", errors="replace")[:1000]
|
||||
return RequestResult(False, time.perf_counter() - start, None, 0, 0, 0, 0, None, chars, f"HTTP {exc.code}: {body}")
|
||||
except Exception as exc:
|
||||
return RequestResult(False, time.perf_counter() - start, None, 0, 0, 0, 0, None, chars, f"{type(exc).__name__}: {exc}")
|
||||
|
||||
|
||||
def parse_ixsmi(raw):
|
||||
samples = []
|
||||
for line in raw.splitlines():
|
||||
m = re.search(r"\|\s*\d+%\s+\d+C\s+\S+\s+(\d+)W\s*/\s*(\d+)W\s*\|\s*(\d+)MiB\s*/\s*(\d+)MiB\s*\|\s*(\d+)%", line)
|
||||
if m:
|
||||
samples.append({
|
||||
"power_w": int(m.group(1)),
|
||||
"power_cap_w": int(m.group(2)),
|
||||
"mem_mib": int(m.group(3)),
|
||||
"mem_total_mib": int(m.group(4)),
|
||||
"util_pct": int(m.group(5)),
|
||||
})
|
||||
return samples
|
||||
|
||||
|
||||
def monitor_ixsmi(stop_event, out_path, interval):
|
||||
records = []
|
||||
env = dict(os.environ)
|
||||
env["LD_LIBRARY_PATH"] = ":".join([
|
||||
"/usr/local/corex/lib64",
|
||||
"/usr/local/corex/lib",
|
||||
"/usr/local/iluvatar/lib64",
|
||||
env.get("LD_LIBRARY_PATH", ""),
|
||||
])
|
||||
candidates = [
|
||||
"/usr/local/corex/bin/ixsmi",
|
||||
"/usr/local/iluvatar/bin/ixsmi",
|
||||
"ixsmi",
|
||||
]
|
||||
while not stop_event.is_set():
|
||||
ts = datetime.now().isoformat(timespec="seconds")
|
||||
try:
|
||||
last_error = None
|
||||
raw = None
|
||||
for binary in candidates:
|
||||
try:
|
||||
raw = subprocess.check_output([binary], text=True, stderr=subprocess.STDOUT, timeout=10, env=env)
|
||||
break
|
||||
except Exception as exc:
|
||||
last_error = exc
|
||||
if raw is None:
|
||||
raise last_error or RuntimeError("ixsmi not found")
|
||||
records.append({"ts": ts, "ok": True, "raw": raw, "parsed": parse_ixsmi(raw)})
|
||||
except Exception as exc:
|
||||
records.append({"ts": ts, "ok": False, "error": repr(exc)})
|
||||
stop_event.wait(interval)
|
||||
Path(out_path).write_text(json.dumps(records, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
|
||||
|
||||
def summarize_monitor(path):
|
||||
if not path or not Path(path).exists():
|
||||
return None
|
||||
records = json.loads(Path(path).read_text(encoding="utf-8"))
|
||||
util = []
|
||||
mem = []
|
||||
power = []
|
||||
for rec in records:
|
||||
for gpu in rec.get("parsed") or []:
|
||||
util.append(gpu["util_pct"])
|
||||
mem.append(gpu["mem_mib"])
|
||||
power.append(gpu["power_w"])
|
||||
if not util:
|
||||
return {"records": len(records), "parsed_samples": 0}
|
||||
return {
|
||||
"records": len(records),
|
||||
"parsed_samples": len(util),
|
||||
"avg_gpu_util_pct": sum(util) / len(util),
|
||||
"max_gpu_util_pct": max(util),
|
||||
"avg_mem_mib": sum(mem) / len(mem),
|
||||
"max_mem_mib": max(mem),
|
||||
"avg_power_w": sum(power) / len(power),
|
||||
"max_power_w": max(power),
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--url", default="http://127.0.0.1:1111")
|
||||
parser.add_argument("--model", default="llm")
|
||||
parser.add_argument("--label", required=True)
|
||||
parser.add_argument("--prompt-mode", choices=["short", "tool"], default="short")
|
||||
parser.add_argument("--with-tools", action="store_true")
|
||||
parser.add_argument("--tool-count", type=int, default=16)
|
||||
parser.add_argument("--concurrency", type=int, default=1)
|
||||
parser.add_argument("--requests", type=int, default=4)
|
||||
parser.add_argument("--max-tokens", type=int, default=256)
|
||||
parser.add_argument("--timeout", type=int, default=900)
|
||||
parser.add_argument("--monitor-out")
|
||||
parser.add_argument("--monitor-interval", type=float, default=1.0)
|
||||
parser.add_argument("--out", required=True)
|
||||
args = parser.parse_args()
|
||||
|
||||
stop_event = threading.Event()
|
||||
monitor_thread = None
|
||||
if args.monitor_out:
|
||||
monitor_thread = threading.Thread(
|
||||
target=monitor_ixsmi,
|
||||
args=(stop_event, args.monitor_out, args.monitor_interval),
|
||||
daemon=True,
|
||||
)
|
||||
monitor_thread.start()
|
||||
|
||||
started = time.perf_counter()
|
||||
results = []
|
||||
try:
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=args.concurrency) as pool:
|
||||
futures = [
|
||||
pool.submit(post_stream, args.url, make_request(args, i), args.timeout)
|
||||
for i in range(args.requests)
|
||||
]
|
||||
for fut in concurrent.futures.as_completed(futures):
|
||||
result = fut.result()
|
||||
results.append(result)
|
||||
print(f"done {len(results)}/{args.requests} ok={result.ok} tps={result.output_tps}", flush=True)
|
||||
finally:
|
||||
stop_event.set()
|
||||
if monitor_thread:
|
||||
monitor_thread.join(timeout=15)
|
||||
|
||||
wall = time.perf_counter() - started
|
||||
ok = [r for r in results if r.ok]
|
||||
tps_values = [r.output_tps for r in ok if r.output_tps is not None]
|
||||
ttft_values = [r.ttft_sec for r in ok if r.ttft_sec is not None]
|
||||
completion = sum(r.completion_tokens for r in ok)
|
||||
prompt = sum(r.prompt_tokens for r in ok)
|
||||
cached = sum(r.cached_tokens for r in ok)
|
||||
summary = {
|
||||
"created_at": datetime.now().isoformat(timespec="seconds"),
|
||||
"label": args.label,
|
||||
"url": args.url,
|
||||
"model": args.model,
|
||||
"prompt_mode": args.prompt_mode,
|
||||
"with_tools": args.with_tools,
|
||||
"tool_count": args.tool_count if args.with_tools else 0,
|
||||
"concurrency": args.concurrency,
|
||||
"requests": args.requests,
|
||||
"max_tokens": args.max_tokens,
|
||||
"wall_sec": wall,
|
||||
"success_rate": len(ok) / len(results) if results else 0,
|
||||
"ttft_p50_sec": percentile(ttft_values, 50),
|
||||
"ttft_p90_sec": percentile(ttft_values, 90),
|
||||
"output_tps_p10_per_request": percentile(tps_values, 10),
|
||||
"output_tps_p50_per_request": percentile(tps_values, 50),
|
||||
"aggregate_output_tps": completion / wall if wall > 0 else 0,
|
||||
"prompt_tokens": prompt,
|
||||
"cached_tokens": cached,
|
||||
"completion_tokens": completion,
|
||||
"reasoning_tokens": sum(r.reasoning_tokens for r in ok),
|
||||
"chars": sum(r.chars for r in ok),
|
||||
"monitor": summarize_monitor(args.monitor_out) if args.monitor_out else None,
|
||||
"results": [asdict(r) for r in results],
|
||||
}
|
||||
Path(args.out).parent.mkdir(parents=True, exist_ok=True)
|
||||
Path(args.out).write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
print(json.dumps({k: v for k, v in summary.items() if k != "results"}, ensure_ascii=False, indent=2))
|
||||
print("RESULT_FILE", args.out)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user