#!/usr/bin/env python3 """Open-loop fixed-shape prefill-only workload for one real offered-load anchor.""" from __future__ import annotations import argparse import concurrent.futures import hashlib import http.client import json import math import time from pathlib import Path from typing import Any TTFT_SLO_MS = 1256.0 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("--served-model", required=True) parser.add_argument("--model-path", type=Path, required=True) parser.add_argument("--rate", type=float, required=True) parser.add_argument("--requests", type=int, default=64) parser.add_argument("--input-tokens", type=int, default=2048) parser.add_argument("--timeout-seconds", type=float, default=900.0) parser.add_argument("--output", type=Path, required=True) return parser.parse_args() def percentile(values: list[float], fraction: float) -> float | None: if not values: return None ordered = sorted(values) index = min(len(ordered) - 1, max(0, math.ceil(fraction * len(ordered)) - 1)) return ordered[index] def run_request( *, request_index: int, scheduled_at: float, benchmark_start: float, args: argparse.Namespace, prompt_ids: list[int], ) -> dict[str, Any]: delay = scheduled_at - time.perf_counter() if delay > 0: time.sleep(delay) admitted = time.perf_counter() record: dict[str, Any] = { "request_index": request_index, "scheduled_s": scheduled_at - benchmark_start, "admitted_s": admitted - benchmark_start, "admission_lag_ms": (admitted - scheduled_at) * 1000.0, "success": False, } connection = http.client.HTTPConnection(args.host, args.port, timeout=args.timeout_seconds) body = { "model": args.served_model, "prompt": prompt_ids, "min_tokens": 1, "max_tokens": 1, "ignore_eos": True, "temperature": 0, "stream": True, "stream_options": {"include_usage": True}, "return_token_ids": True, } try: started = time.perf_counter() connection.request( "POST", "/v1/completions", body=json.dumps(body, separators=(",", ":")).encode(), headers={"Content-Type": "application/json"}, ) response = connection.getresponse() if response.status != 200: raise RuntimeError( f"HTTP {response.status}: {response.read().decode(errors='replace')}" ) first_token_at = None streamed_tokens = 0 usage = None while True: raw = response.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: first_token_at = first_token_at or time.perf_counter() streamed_tokens += emitted finished = time.perf_counter() if first_token_at is None or usage is None: raise RuntimeError("missing streaming token or usage") prompt_tokens = int(usage["prompt_tokens"]) completion_tokens = int(usage["completion_tokens"]) if prompt_tokens != args.input_tokens or completion_tokens != 1: raise RuntimeError(f"usage mismatch: {prompt_tokens}+{completion_tokens}") ttft = (first_token_at - started) * 1000.0 record.update( { "success": True, "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, "streamed_token_count": streamed_tokens, "ttft_ms": ttft, "e2e_ms": (finished - started) * 1000.0, "slo_pass": ttft <= TTFT_SLO_MS, } ) except Exception as error: # Failed requests remain in the SLO denominator. record["error"] = f"{type(error).__name__}: {error}" record["slo_pass"] = False finally: connection.close() return record def main() -> None: args = parse_args() if args.rate <= 0 or args.requests <= 0 or args.input_tokens <= 0: raise ValueError("rate, requests, and input tokens must be positive") from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True) excluded = set(tokenizer.all_special_ids) candidates = [ token_id for token_id in range(tokenizer.vocab_size) if token_id not in excluded ] if len(candidates) < args.requests + 1: raise RuntimeError("tokenizer has too few non-special token IDs") base_id = candidates[0] prompts = [ [candidates[index + 1], *([base_id] * (args.input_tokens - 1))] for index in range(args.requests) ] prompt_hash = hashlib.sha256( "\n".join(",".join(map(str, prompt)) for prompt in prompts).encode() ).hexdigest() benchmark_start = time.perf_counter() + 2.0 with concurrent.futures.ThreadPoolExecutor(max_workers=args.requests) as pool: futures = [ pool.submit( run_request, request_index=index, scheduled_at=benchmark_start + index / args.rate, benchmark_start=benchmark_start, args=args, prompt_ids=prompts[index], ) for index in range(args.requests) ] requests = [future.result() for future in futures] requests.sort(key=lambda row: int(row["request_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] pass_rate = passed / len(requests) payload = { "schema": "qwen30-prefill-rate-anchor-v1", "workload": { "offered_request_rate": args.rate, "request_count": args.requests, "input_tokens": args.input_tokens, "output_tokens": 1, "prefix_caching": False, "arrival": "open_loop_uniform", "last_scheduled_arrival_s": (args.requests - 1) / args.rate, "prompt_vector_sha256": prompt_hash, }, "summary": { "completed": len(completed), "failed": len(requests) - len(completed), "ttft_p50_ms": percentile(ttfts, 0.50), "ttft_p95_ms": percentile(ttfts, 0.95), "ttft_max_ms": max(ttfts) if ttfts else None, "admission_lag_max_ms": max( float(row["admission_lag_ms"]) for row in requests ), "slo": { "ttft_threshold_ms": TTFT_SLO_MS, "passed": passed, "pass_rate": pass_rate, "feasible": pass_rate >= TARGET_PASS_RATE, }, }, "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) != args.requests: raise SystemExit(2) if __name__ == "__main__": main()