224 lines
8.1 KiB
Python
224 lines
8.1 KiB
Python
#!/usr/bin/env python3
|
|
"""Open-loop fixed-shape completion workload for one T0 offered-load anchor."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import concurrent.futures
|
|
import hashlib
|
|
import http.client
|
|
import json
|
|
import math
|
|
import statistics
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
TPOT_SLOS_MS = (40.0, 120.0, 150.0, 180.0)
|
|
TTFT_SLO_MS = 1256.0
|
|
|
|
|
|
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("--output-tokens", type=int, default=128)
|
|
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": args.output_tokens,
|
|
"max_tokens": args.output_tokens,
|
|
"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
|
|
last_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:
|
|
now = time.perf_counter()
|
|
first_token_at = first_token_at or now
|
|
last_token_at = now
|
|
streamed_tokens += emitted
|
|
finished = time.perf_counter()
|
|
if first_token_at is None or last_token_at is None or usage is None:
|
|
raise RuntimeError("missing streaming tokens or usage")
|
|
prompt_tokens = int(usage["prompt_tokens"])
|
|
completion_tokens = int(usage["completion_tokens"])
|
|
if prompt_tokens != args.input_tokens or completion_tokens != args.output_tokens:
|
|
raise RuntimeError(f"usage mismatch: {prompt_tokens}+{completion_tokens}")
|
|
ttft = (first_token_at - started) * 1000.0
|
|
tpot = (last_token_at - first_token_at) * 1000.0 / (completion_tokens - 1)
|
|
record.update(
|
|
{
|
|
"success": True,
|
|
"prompt_tokens": prompt_tokens,
|
|
"completion_tokens": completion_tokens,
|
|
"streamed_token_count": streamed_tokens,
|
|
"ttft_ms": ttft,
|
|
"tpot_ms": tpot,
|
|
"e2e_ms": (finished - started) * 1000.0,
|
|
}
|
|
)
|
|
except Exception as error: # Preserve failed requests as SLO failures.
|
|
record["error"] = f"{type(error).__name__}: {error}"
|
|
finally:
|
|
connection.close()
|
|
return record
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
if args.rate <= 0 or args.requests <= 0:
|
|
raise ValueError("rate and requests 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"]]
|
|
slos = {}
|
|
for limit in TPOT_SLOS_MS:
|
|
passed = sum(
|
|
row["success"]
|
|
and float(row["ttft_ms"]) <= TTFT_SLO_MS
|
|
and float(row["tpot_ms"]) <= limit
|
|
for row in requests
|
|
)
|
|
slos[f"tpot_{int(limit)}ms"] = {
|
|
"passed": passed,
|
|
"pass_rate": passed / len(requests),
|
|
"feasible": passed / len(requests) >= 0.95,
|
|
}
|
|
ttfts = [float(row["ttft_ms"]) for row in completed]
|
|
tpots = [float(row["tpot_ms"]) for row in completed]
|
|
payload = {
|
|
"schema": "qwen235b-t0-rate-anchor-v1",
|
|
"workload": {
|
|
"offered_request_rate": args.rate,
|
|
"request_count": args.requests,
|
|
"input_tokens": args.input_tokens,
|
|
"output_tokens": args.output_tokens,
|
|
"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),
|
|
"tpot_p50_ms": percentile(tpots, 0.50),
|
|
"tpot_p95_ms": percentile(tpots, 0.95),
|
|
"admission_lag_max_ms": max(float(row["admission_lag_ms"]) for row in requests),
|
|
"slos": slos,
|
|
},
|
|
"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()
|