File size: 5,854 Bytes
df41178 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 | #!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""同じ条件で llama-server の速度を測る小さな再利用可能ハーネス。"""
from __future__ import annotations
import argparse
import datetime as dt
import json
import statistics
import sys
import time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent))
from public_bench import ( # noqa: E402
HARNESS_VERSION,
LocalClient,
atomic_write,
collect_runtime_metadata,
collect_server_metadata,
redact_url,
)
PROMPTS = [
"Return only the next prime number after 100.",
"Write one short sentence explaining why the sky appears blue.",
"Compute 37 * 29. Return only the integer.",
"Return a Python function named square(x) that returns x*x, with no explanation.",
]
def percentile(values: list[float], q: float) -> float | None:
if not values:
return None
xs = sorted(values)
index = (len(xs) - 1) * q
low = int(index)
high = min(low + 1, len(xs) - 1)
return xs[low] + (xs[high] - xs[low]) * (index - low)
def write_json(path: Path, obj: dict) -> None:
atomic_write(path, json.dumps(obj, ensure_ascii=False, indent=2) + "\n")
def main() -> int:
ap = argparse.ArgumentParser()
ap.add_argument("--url", default="http://127.0.0.1:8080/v1/chat/completions")
ap.add_argument("--out-dir", default="")
ap.add_argument("--repeat", type=int, default=1)
ap.add_argument("--warmup", type=int, default=1)
ap.add_argument("--max-tokens", type=int, default=128)
ap.add_argument("--timeout", type=int, default=900)
ap.add_argument("--retries", type=int, default=2)
ap.add_argument("--retry-backoff", type=float, default=0.25)
ap.add_argument("--seed", type=int, default=20260909)
ap.add_argument("--think", action="store_true")
args = ap.parse_args()
if args.repeat <= 0 or args.warmup < 0:
ap.error("--repeat は正数、--warmup は0以上が必要です")
stamp = dt.datetime.now().strftime("%Y%m%d-%H%M%S")
out = Path(args.out_dir or Path(__file__).resolve().parent / "bench_results" /
f"speed-{stamp}")
out.mkdir(parents=True, exist_ok=True)
client = LocalClient(args.url, args.think, args.max_tokens, args.timeout,
retries=args.retries, retry_backoff=args.retry_backoff,
seed=args.seed)
for i in range(args.warmup):
result = client.call(PROMPTS[i % len(PROMPTS)], f"warmup/{i}")
if result.get("error"):
print(f"warmup error: {result['error']}", file=sys.stderr)
rows = []
started = time.monotonic()
for repeat in range(args.repeat):
for prompt_index, prompt in enumerate(PROMPTS):
rid = f"p{prompt_index}/r{repeat}"
result = client.call(prompt, rid)
rows.append({
"id": rid,
"prompt_index": prompt_index,
"repeat": repeat,
"elapsed_sec": result.get("elapsed_sec", 0),
"timings": result.get("timings") or {},
"usage": result.get("usage"),
"model": result.get("model"),
"error": result.get("error"),
"error_type": result.get("error_type"),
"http_status": result.get("http_status"),
"attempts": result.get("attempts") or [],
"retry_count": result.get("retry_count", 0),
})
print(f" {rid}: {rows[-1]['elapsed_sec']:.3f}s "
f"{rows[-1]['error'] or 'ok'}", flush=True)
successful = [r for r in rows if not r.get("error")]
elapsed = [float(r.get("elapsed_sec", 0)) for r in successful]
tps = [float((r.get("timings") or {}).get("predicted_per_second", 0) or 0)
for r in successful if (r.get("timings") or {}).get("predicted_per_second")]
prompt_tps = [float((r.get("timings") or {}).get("prompt_per_second", 0) or 0)
for r in successful if (r.get("timings") or {}).get("prompt_per_second")]
summary = {
"harness_version": HARNESS_VERSION,
"created_at": dt.datetime.now(dt.timezone.utc).isoformat(),
"url": redact_url(args.url),
"repeat": args.repeat,
"warmup": args.warmup,
"n": len(rows),
"successful_n": len(successful),
"failed_n": len(rows) - len(successful),
"elapsed_total_sec": time.monotonic() - started,
"mean_sec": statistics.mean(elapsed) if elapsed else None,
"median_sec": statistics.median(elapsed) if elapsed else None,
"p95_sec": percentile(elapsed, 0.95),
"mean_predicted_tps": statistics.mean(tps) if tps else None,
"median_predicted_tps": statistics.median(tps) if tps else None,
"mean_prompt_tps": statistics.mean(prompt_tps) if prompt_tps else None,
"error_types": {
str(k): sum(1 for r in rows if (r.get("error_type") or "error") == k)
for k in sorted({r.get("error_type") or "error" for r in rows
if r.get("error")})
},
"runtime": collect_runtime_metadata(),
"server": collect_server_metadata(args.url),
}
jsonl = out / "speed.jsonl"
atomic_write(jsonl, "".join(json.dumps(r, ensure_ascii=False) + "\n" for r in rows))
write_json(out / "speed.summary.json", summary)
write_json(out / "config.json", {
"harness_version": HARNESS_VERSION,
"argv": sys.argv,
"url": redact_url(args.url),
"max_tokens": args.max_tokens,
"think": args.think,
"seed": args.seed,
"repeat": args.repeat,
"warmup": args.warmup,
})
print(json.dumps(summary, ensure_ascii=False, indent=2))
print(f"結果: {out}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|