#!/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("--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, ) -> 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"]) 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, ) ) 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", "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()