262 lines
9.2 KiB
Python
262 lines
9.2 KiB
Python
#!/usr/bin/env python3
|
|
"""Replay a private exact-trace anchor without emitting prompt text."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import hashlib
|
|
import json
|
|
import math
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
TARGET_PASS_RATE = 0.95
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--host", default="127.0.0.1")
|
|
parser.add_argument("--port", type=int, required=True)
|
|
parser.add_argument("--requests-file", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument(
|
|
"--served-model",
|
|
default=None,
|
|
help="Optional server-side model alias used only for HTTP routing.",
|
|
)
|
|
parser.add_argument("--tpot-slo-ms", type=float, default=150.0)
|
|
parser.add_argument("--timeout-seconds", type=float, default=1800.0)
|
|
return parser.parse_args()
|
|
|
|
|
|
def percentile(values: list[float], fraction: float) -> float | None:
|
|
if not values:
|
|
return None
|
|
ordered = sorted(values)
|
|
return ordered[math.ceil(fraction * len(ordered)) - 1]
|
|
|
|
|
|
def ttft_slo_ms(input_tokens: int) -> float:
|
|
return 1000.0 + 1000.0 * input_tokens / 8000.0
|
|
|
|
|
|
def row_vector_sha256(rows: list[dict[str, Any]]) -> str:
|
|
digest = hashlib.sha256()
|
|
for row in rows:
|
|
digest.update(
|
|
json.dumps(
|
|
[
|
|
row["source_index"],
|
|
row["arrived_at"],
|
|
row["input_length"],
|
|
row["output_length"],
|
|
row["session_id"],
|
|
row["runtime_block_ids"],
|
|
],
|
|
separators=(",", ":"),
|
|
).encode()
|
|
)
|
|
digest.update(b"\n")
|
|
return digest.hexdigest()
|
|
|
|
|
|
async def request_one(
|
|
session: aiohttp.ClientSession,
|
|
row: dict[str, Any],
|
|
*,
|
|
scheduled_at: float,
|
|
benchmark_start: float,
|
|
tpot_slo_ms: float,
|
|
served_model: str | None,
|
|
) -> dict[str, Any]:
|
|
loop = asyncio.get_running_loop()
|
|
delay = scheduled_at - loop.time()
|
|
if delay > 0:
|
|
await asyncio.sleep(delay)
|
|
admitted_at = loop.time()
|
|
record: dict[str, Any] = {
|
|
"source_index": int(row["source_index"]),
|
|
"session_id": int(row["session_id"]),
|
|
"scheduled_s": scheduled_at - benchmark_start,
|
|
"admitted_s": admitted_at - benchmark_start,
|
|
"admission_lag_ms": (admitted_at - scheduled_at) * 1000.0,
|
|
"input_tokens": int(row["input_length"]),
|
|
"requested_output_tokens": int(row["output_length"]),
|
|
"success": False,
|
|
}
|
|
body = dict(row["body"])
|
|
if served_model is not None:
|
|
body["model"] = served_model
|
|
body.update(
|
|
{
|
|
"temperature": 0,
|
|
"stream": True,
|
|
"stream_options": {"include_usage": True},
|
|
"return_token_ids": True,
|
|
}
|
|
)
|
|
try:
|
|
started = loop.time()
|
|
async with session.post("/v1/completions", json=body) as response:
|
|
if response.status != 200:
|
|
detail = (await response.text())[:1000]
|
|
raise RuntimeError(f"HTTP {response.status}: {detail}")
|
|
first_token_at = None
|
|
last_token_at = None
|
|
streamed_tokens = 0
|
|
usage = None
|
|
while True:
|
|
raw = await response.content.readline()
|
|
if not raw:
|
|
break
|
|
line = raw.decode(errors="replace").strip()
|
|
if not line.startswith("data:"):
|
|
continue
|
|
data = line[5:].strip()
|
|
if data == "[DONE]":
|
|
break
|
|
payload = json.loads(data)
|
|
if payload.get("usage"):
|
|
usage = payload["usage"]
|
|
emitted = 0
|
|
for choice in payload.get("choices") or []:
|
|
token_ids = choice.get("token_ids") or []
|
|
emitted += len(token_ids) if token_ids else int(bool(choice.get("text")))
|
|
if emitted:
|
|
now = loop.time()
|
|
first_token_at = first_token_at or now
|
|
last_token_at = now
|
|
streamed_tokens += emitted
|
|
finished = loop.time()
|
|
if first_token_at is None or last_token_at is None or usage is None:
|
|
raise RuntimeError("missing streaming token or usage")
|
|
actual_input = int(usage["prompt_tokens"])
|
|
actual_output = int(usage["completion_tokens"])
|
|
if actual_input != int(row["input_length"]) or actual_output != int(
|
|
row["output_length"]
|
|
):
|
|
raise RuntimeError(f"usage mismatch: {actual_input}+{actual_output}")
|
|
ttft = (first_token_at - started) * 1000.0
|
|
tpot = (
|
|
(last_token_at - first_token_at) * 1000.0 / (actual_output - 1)
|
|
if actual_output > 1
|
|
else None
|
|
)
|
|
record.update(
|
|
{
|
|
"success": True,
|
|
"actual_input_tokens": actual_input,
|
|
"actual_output_tokens": actual_output,
|
|
"streamed_token_count": streamed_tokens,
|
|
"ttft_ms": ttft,
|
|
"tpot_ms": tpot,
|
|
"e2e_ms": (finished - started) * 1000.0,
|
|
"slo_pass": ttft <= ttft_slo_ms(actual_input)
|
|
and (tpot is None or tpot <= tpot_slo_ms),
|
|
}
|
|
)
|
|
except Exception as error:
|
|
record.update(
|
|
{
|
|
"error": f"{type(error).__name__}: {error}",
|
|
"slo_pass": False,
|
|
}
|
|
)
|
|
return record
|
|
|
|
|
|
async def replay(args: argparse.Namespace, rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
|
import aiohttp
|
|
|
|
timeout = aiohttp.ClientTimeout(total=args.timeout_seconds)
|
|
connector = aiohttp.TCPConnector(limit=0, ttl_dns_cache=300)
|
|
benchmark_start = asyncio.get_running_loop().time() + 2.0
|
|
async with aiohttp.ClientSession(
|
|
base_url=f"http://{args.host}:{args.port}",
|
|
timeout=timeout,
|
|
connector=connector,
|
|
) as session:
|
|
tasks = [
|
|
asyncio.create_task(
|
|
request_one(
|
|
session,
|
|
row,
|
|
scheduled_at=benchmark_start + float(row["arrived_at"]),
|
|
benchmark_start=benchmark_start,
|
|
tpot_slo_ms=args.tpot_slo_ms,
|
|
served_model=args.served_model,
|
|
)
|
|
)
|
|
for row in rows
|
|
]
|
|
return await asyncio.gather(*tasks)
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
if args.tpot_slo_ms <= 0 or args.timeout_seconds <= 0:
|
|
raise ValueError("SLO and timeout must be positive")
|
|
rows = [json.loads(line) for line in args.requests_file.open() if line.strip()]
|
|
if not rows:
|
|
raise ValueError("requests file is empty")
|
|
arrivals = [float(row["arrived_at"]) for row in rows]
|
|
if any(right < left for left, right in zip(arrivals, arrivals[1:])):
|
|
raise ValueError("request arrival order drift")
|
|
requests = asyncio.run(replay(args, rows))
|
|
requests.sort(key=lambda row: int(row["source_index"]))
|
|
completed = [row for row in requests if row["success"]]
|
|
passed = sum(bool(row["slo_pass"]) for row in requests)
|
|
ttfts = [float(row["ttft_ms"]) for row in completed]
|
|
tpots = [
|
|
float(row["tpot_ms"])
|
|
for row in completed
|
|
if row["tpot_ms"] is not None
|
|
]
|
|
pass_rate = passed / len(requests)
|
|
payload = {
|
|
"schema": "qwen30-exact-trace-anchor-v1",
|
|
"contract": {
|
|
"requests_file": str(args.requests_file.resolve()),
|
|
"requests_file_sha256": hashlib.sha256(
|
|
args.requests_file.read_bytes()
|
|
).hexdigest(),
|
|
"row_vector_sha256": row_vector_sha256(rows),
|
|
"requests": len(rows),
|
|
"first_arrival_s": arrivals[0],
|
|
"last_arrival_s": arrivals[-1],
|
|
"arrival": "original_trace_timestamp_and_order",
|
|
"input_output_prompt": "exact_source_values",
|
|
"served_model_alias": args.served_model,
|
|
"ttft_slo": "1000ms + 1000ms * input_tokens / 8000",
|
|
"tpot_slo_ms": args.tpot_slo_ms,
|
|
"target_pass_rate": TARGET_PASS_RATE,
|
|
"prompt_text_emitted": False,
|
|
},
|
|
"summary": {
|
|
"completed": len(completed),
|
|
"failed": len(requests) - len(completed),
|
|
"passed": passed,
|
|
"pass_rate": pass_rate,
|
|
"feasible": pass_rate >= TARGET_PASS_RATE,
|
|
"ttft_p50_ms": percentile(ttfts, 0.50),
|
|
"ttft_p95_ms": percentile(ttfts, 0.95),
|
|
"tpot_p50_ms": percentile(tpots, 0.50),
|
|
"tpot_p95_ms": percentile(tpots, 0.95),
|
|
"admission_lag_p95_ms": percentile(
|
|
[float(row["admission_lag_ms"]) for row in requests], 0.95
|
|
),
|
|
},
|
|
"requests": requests,
|
|
}
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
|
|
print(json.dumps(payload["summary"], sort_keys=True), flush=True)
|
|
if len(completed) != len(requests):
|
|
raise SystemExit(2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|