256 lines
9.6 KiB
Python
256 lines
9.6 KiB
Python
|
import argparse
|
||
|
|
import concurrent.futures
|
||
|
|
import json
|
||
|
|
import math
|
||
|
|
import statistics
|
||
|
|
import time
|
||
|
|
import urllib.error
|
||
|
|
import urllib.request
|
||
|
|
from dataclasses import dataclass, asdict
|
||
|
|
from datetime import datetime
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class RequestResult:
|
||
|
|
ok: bool
|
||
|
|
elapsed_sec: float
|
||
|
|
ttft_sec: float | None
|
||
|
|
prompt_tokens: int
|
||
|
|
cached_tokens: int
|
||
|
|
completion_tokens: int
|
||
|
|
reasoning_tokens: int
|
||
|
|
output_tps: float | None
|
||
|
|
error: str | None = None
|
||
|
|
|
||
|
|
|
||
|
|
def percentile(values, pct):
|
||
|
|
if not values:
|
||
|
|
return None
|
||
|
|
values = sorted(values)
|
||
|
|
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 load_dataset(path):
|
||
|
|
text = Path(path).read_text(encoding="utf-8")
|
||
|
|
if path.endswith(".jsonl"):
|
||
|
|
data = [json.loads(line) for line in text.splitlines() if line.strip()]
|
||
|
|
else:
|
||
|
|
data = json.loads(text)
|
||
|
|
|
||
|
|
# Official-like format: each item is already an OpenAI chat completion
|
||
|
|
# request with messages/tools/tool_choice/stream/etc.
|
||
|
|
if data and isinstance(data[0], dict) and "messages" in data[0]:
|
||
|
|
return data
|
||
|
|
|
||
|
|
requests = []
|
||
|
|
for item in data:
|
||
|
|
system_prompt = item.get("system_prompt") or "You are a helpful assistant."
|
||
|
|
history = [{"role": "system", "content": system_prompt}]
|
||
|
|
for question in item.get("user_questions", []):
|
||
|
|
messages = history + [{"role": "user", "content": question}]
|
||
|
|
requests.append({"messages": messages})
|
||
|
|
# Synthetic assistant placeholder keeps later prompts multi-turn.
|
||
|
|
history = messages + [{"role": "assistant", "content": "好的，我们继续。"}]
|
||
|
|
return requests
|
||
|
|
|
||
|
|
|
||
|
|
def post_stream(url, model, request_item, max_tokens, timeout):
|
||
|
|
allowed_fields = {
|
||
|
|
"messages",
|
||
|
|
"tools",
|
||
|
|
"tool_choice",
|
||
|
|
"model",
|
||
|
|
"max_tokens",
|
||
|
|
"temperature",
|
||
|
|
"top_p",
|
||
|
|
"stop",
|
||
|
|
"presence_penalty",
|
||
|
|
"frequency_penalty",
|
||
|
|
"n",
|
||
|
|
"response_format",
|
||
|
|
}
|
||
|
|
payload = {k: v for k, v in dict(request_item).items() if k in allowed_fields}
|
||
|
|
payload["model"] = payload.get("model") or model
|
||
|
|
# The benchmark CLI controls output length, even if the synthetic dataset
|
||
|
|
# stores a larger official-like max_tokens value.
|
||
|
|
payload["max_tokens"] = max_tokens
|
||
|
|
payload["temperature"] = payload.get("temperature", 0)
|
||
|
|
payload["stream"] = True
|
||
|
|
payload["stream_options"] = {"include_usage": True}
|
||
|
|
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_token_at = None
|
||
|
|
usage = {}
|
||
|
|
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 {}
|
||
|
|
text = delta.get("content") or delta.get("reasoning_content") or ""
|
||
|
|
if text and first_token_at is None:
|
||
|
|
first_token_at = time.perf_counter()
|
||
|
|
elapsed = time.perf_counter() - start
|
||
|
|
ttft = first_token_at - start if first_token_at else None
|
||
|
|
completion_tokens = int(usage.get("completion_tokens") or 0)
|
||
|
|
prompt_tokens = int(usage.get("prompt_tokens") or 0)
|
||
|
|
reasoning_tokens = int(usage.get("reasoning_tokens") or 0)
|
||
|
|
details = usage.get("prompt_tokens_details") or {}
|
||
|
|
cached_tokens = int(details.get("cached_tokens") or 0)
|
||
|
|
decode_sec = elapsed - ttft if ttft is not None else elapsed
|
||
|
|
output_tps = completion_tokens / decode_sec if completion_tokens and decode_sec > 0 else None
|
||
|
|
return RequestResult(
|
||
|
|
ok=True,
|
||
|
|
elapsed_sec=elapsed,
|
||
|
|
ttft_sec=ttft,
|
||
|
|
prompt_tokens=prompt_tokens,
|
||
|
|
cached_tokens=cached_tokens,
|
||
|
|
completion_tokens=completion_tokens,
|
||
|
|
reasoning_tokens=reasoning_tokens,
|
||
|
|
output_tps=output_tps,
|
||
|
|
)
|
||
|
|
except urllib.error.HTTPError as exc:
|
||
|
|
body = ""
|
||
|
|
try:
|
||
|
|
body = exc.read().decode("utf-8", errors="replace")
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
elapsed = time.perf_counter() - start
|
||
|
|
return RequestResult(
|
||
|
|
ok=False,
|
||
|
|
elapsed_sec=elapsed,
|
||
|
|
ttft_sec=None,
|
||
|
|
prompt_tokens=0,
|
||
|
|
cached_tokens=0,
|
||
|
|
completion_tokens=0,
|
||
|
|
reasoning_tokens=0,
|
||
|
|
output_tps=None,
|
||
|
|
error=f"HTTPError {exc.code}: {body[:1000]}",
|
||
|
|
)
|
||
|
|
except Exception as exc:
|
||
|
|
elapsed = time.perf_counter() - start
|
||
|
|
return RequestResult(
|
||
|
|
ok=False,
|
||
|
|
elapsed_sec=elapsed,
|
||
|
|
ttft_sec=None,
|
||
|
|
prompt_tokens=0,
|
||
|
|
cached_tokens=0,
|
||
|
|
completion_tokens=0,
|
||
|
|
reasoning_tokens=0,
|
||
|
|
output_tps=None,
|
||
|
|
error=f"{type(exc).__name__}: {exc}",
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
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("--dataset", required=True)
|
||
|
|
parser.add_argument("--concurrency", type=int, default=1)
|
||
|
|
parser.add_argument("--max-requests", type=int, default=20)
|
||
|
|
parser.add_argument("--max-tokens", type=int, default=128)
|
||
|
|
parser.add_argument("--timeout", type=int, default=600)
|
||
|
|
parser.add_argument("--out", default="/root/work/logs/formal_perf_result.json")
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
reqs = load_dataset(args.dataset)
|
||
|
|
if not reqs:
|
||
|
|
raise SystemExit("empty dataset")
|
||
|
|
scheduled = [reqs[i % len(reqs)] for i in range(args.max_requests)]
|
||
|
|
|
||
|
|
wall_start = time.perf_counter()
|
||
|
|
results = []
|
||
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=args.concurrency) as pool:
|
||
|
|
futures = [
|
||
|
|
pool.submit(post_stream, args.url, args.model, request_item, args.max_tokens, args.timeout)
|
||
|
|
for request_item in scheduled
|
||
|
|
]
|
||
|
|
for fut in concurrent.futures.as_completed(futures):
|
||
|
|
results.append(fut.result())
|
||
|
|
print(f"done {len(results)}/{len(scheduled)} ok={results[-1].ok}", flush=True)
|
||
|
|
wall_sec = time.perf_counter() - wall_start
|
||
|
|
|
||
|
|
ok_results = [r for r in results if r.ok]
|
||
|
|
success_rate = len(ok_results) / len(results) if results else 0.0
|
||
|
|
ttfts = [r.ttft_sec for r in ok_results if r.ttft_sec is not None]
|
||
|
|
output_tps_values = [r.output_tps for r in ok_results if r.output_tps is not None]
|
||
|
|
|
||
|
|
prompt_tokens = sum(r.prompt_tokens for r in ok_results)
|
||
|
|
cached_tokens = sum(r.cached_tokens for r in ok_results)
|
||
|
|
input_tokens_uncached = max(prompt_tokens - cached_tokens, 0)
|
||
|
|
output_tokens = sum(r.completion_tokens for r in ok_results)
|
||
|
|
|
||
|
|
aggregate_output_tps = output_tokens / wall_sec if wall_sec > 0 else 0.0
|
||
|
|
aggregate_input_tps = input_tokens_uncached / wall_sec if wall_sec > 0 else 0.0
|
||
|
|
aggregate_cache_tps = cached_tokens / wall_sec if wall_sec > 0 else 0.0
|
||
|
|
weighted = aggregate_output_tps * 16.796 + aggregate_input_tps * 2.799 + aggregate_cache_tps * 0.56
|
||
|
|
cache_hit_rate = cached_tokens / prompt_tokens if prompt_tokens else 0.0
|
||
|
|
|
||
|
|
summary = {
|
||
|
|
"created_at": datetime.now().isoformat(timespec="seconds"),
|
||
|
|
"url": args.url,
|
||
|
|
"model": args.model,
|
||
|
|
"dataset": args.dataset,
|
||
|
|
"concurrency": args.concurrency,
|
||
|
|
"max_requests": args.max_requests,
|
||
|
|
"max_tokens": args.max_tokens,
|
||
|
|
"wall_sec": wall_sec,
|
||
|
|
"success_rate": success_rate,
|
||
|
|
"ttft_p90_sec": percentile(ttfts, 90),
|
||
|
|
"output_tps_p10_per_request": percentile(output_tps_values, 10),
|
||
|
|
"aggregate_output_tps": aggregate_output_tps,
|
||
|
|
"aggregate_input_tps_uncached": aggregate_input_tps,
|
||
|
|
"aggregate_cache_tps": aggregate_cache_tps,
|
||
|
|
"cache_hit_rate": cache_hit_rate,
|
||
|
|
"weighted_token_throughput": weighted,
|
||
|
|
"totals": {
|
||
|
|
"requests": len(results),
|
||
|
|
"success": len(ok_results),
|
||
|
|
"prompt_tokens": prompt_tokens,
|
||
|
|
"input_tokens_uncached": input_tokens_uncached,
|
||
|
|
"cached_tokens": cached_tokens,
|
||
|
|
"completion_tokens": output_tokens,
|
||
|
|
"reasoning_tokens": sum(r.reasoning_tokens for r in ok_results),
|
||
|
|
},
|
||
|
|
"targets": {
|
||
|
|
"output_tps_p10_per_request_gte_20": (percentile(output_tps_values, 10) or 0) >= 20,
|
||
|
|
"ttft_p90_lte_5": (percentile(ttfts, 90) or 999) <= 5,
|
||
|
|
"cache_hit_rate_gte_50pct": cache_hit_rate >= 0.5,
|
||
|
|
"success_rate_gte_99pct": success_rate >= 0.99,
|
||
|
|
"weighted_token_throughput_gte_8000": weighted >= 8000,
|
||
|
|
},
|
||
|
|
"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()
|