Track simulator fidelity experiment artifacts
This commit is contained in:
200
runs/frontier-multicase-sufficiency-v1/t0_smoke_client.py
Normal file
200
runs/frontier-multicase-sufficiency-v1/t0_smoke_client.py
Normal file
@@ -0,0 +1,200 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Issue exact fixed-shape completion requests and record streaming latency."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import concurrent.futures
|
||||
import http.client
|
||||
import json
|
||||
import statistics
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
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("--input-tokens", type=int, default=2048)
|
||||
parser.add_argument("--output-tokens", type=int, default=128)
|
||||
parser.add_argument("--concurrency", type=int, required=True)
|
||||
parser.add_argument("--requests", type=int, required=True)
|
||||
parser.add_argument("--timeout-seconds", type=float, default=600.0)
|
||||
parser.add_argument("--output", type=Path, required=True)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def run_request(
|
||||
*,
|
||||
request_index: int,
|
||||
args: argparse.Namespace,
|
||||
prompt_token_id: int,
|
||||
start_barrier: threading.Barrier,
|
||||
) -> dict[str, Any]:
|
||||
body = {
|
||||
"model": args.served_model,
|
||||
"prompt": [prompt_token_id] * args.input_tokens,
|
||||
"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,
|
||||
}
|
||||
encoded = json.dumps(body, separators=(",", ":")).encode()
|
||||
connection = http.client.HTTPConnection(
|
||||
args.host, args.port, timeout=args.timeout_seconds
|
||||
)
|
||||
start_barrier.wait()
|
||||
started = time.perf_counter()
|
||||
connection.request(
|
||||
"POST",
|
||||
"/v1/completions",
|
||||
body=encoded,
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
response = connection.getresponse()
|
||||
if response.status != 200:
|
||||
detail = response.read().decode(errors="replace")
|
||||
raise RuntimeError(f"request {request_index} failed: HTTP {response.status}: {detail}")
|
||||
|
||||
first_token_at: float | None = None
|
||||
last_token_at: float | None = None
|
||||
streamed_token_count = 0
|
||||
usage: dict[str, Any] | None = 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 []
|
||||
if token_ids:
|
||||
emitted += len(token_ids)
|
||||
elif choice.get("text"):
|
||||
emitted += 1
|
||||
if emitted:
|
||||
now = time.perf_counter()
|
||||
if first_token_at is None:
|
||||
first_token_at = now
|
||||
last_token_at = now
|
||||
streamed_token_count += emitted
|
||||
finished = time.perf_counter()
|
||||
connection.close()
|
||||
|
||||
if first_token_at is None or last_token_at is None or usage is None:
|
||||
raise RuntimeError(
|
||||
f"request {request_index} missing streaming token or usage metadata"
|
||||
)
|
||||
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"request {request_index} usage mismatch: prompt={prompt_tokens}, "
|
||||
f"completion={completion_tokens}"
|
||||
)
|
||||
|
||||
ttft_ms = (first_token_at - started) * 1000.0
|
||||
tpot_ms = (
|
||||
(last_token_at - first_token_at) * 1000.0 / (completion_tokens - 1)
|
||||
if completion_tokens > 1
|
||||
else 0.0
|
||||
)
|
||||
return {
|
||||
"request_index": request_index,
|
||||
"prompt_tokens": prompt_tokens,
|
||||
"completion_tokens": completion_tokens,
|
||||
"streamed_token_count": streamed_token_count,
|
||||
"ttft_ms": ttft_ms,
|
||||
"tpot_ms": tpot_ms,
|
||||
"e2e_ms": (finished - started) * 1000.0,
|
||||
"ttft_slo_ms": 1000.0 + args.input_tokens / 8.0,
|
||||
"tpot_slo_ms": 40.0,
|
||||
"joint_slo_pass": ttft_ms <= 1000.0 + args.input_tokens / 8.0
|
||||
and tpot_ms <= 40.0,
|
||||
}
|
||||
|
||||
|
||||
def percentile(values: list[float], fraction: float) -> float:
|
||||
ordered = sorted(values)
|
||||
index = min(len(ordered) - 1, max(0, round((len(ordered) - 1) * fraction)))
|
||||
return ordered[index]
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
if args.concurrency <= 0 or args.requests < args.concurrency:
|
||||
raise ValueError("requests must be at least concurrency, and both must be positive")
|
||||
if args.input_tokens <= 0 or args.output_tokens <= 0:
|
||||
raise ValueError("token lengths must be positive")
|
||||
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
|
||||
candidate_ids = tokenizer.encode(" hello", add_special_tokens=False)
|
||||
if not candidate_ids:
|
||||
raise RuntimeError("tokenizer returned no prompt token id")
|
||||
prompt_token_id = int(candidate_ids[0])
|
||||
|
||||
results: list[dict[str, Any]] = []
|
||||
for batch_start in range(0, args.requests, args.concurrency):
|
||||
batch_count = min(args.concurrency, args.requests - batch_start)
|
||||
barrier = threading.Barrier(batch_count)
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=batch_count) as pool:
|
||||
futures = [
|
||||
pool.submit(
|
||||
run_request,
|
||||
request_index=batch_start + offset,
|
||||
args=args,
|
||||
prompt_token_id=prompt_token_id,
|
||||
start_barrier=barrier,
|
||||
)
|
||||
for offset in range(batch_count)
|
||||
]
|
||||
results.extend(future.result() for future in futures)
|
||||
|
||||
ttfts = [float(row["ttft_ms"]) for row in results]
|
||||
tpots = [float(row["tpot_ms"]) for row in results]
|
||||
payload = {
|
||||
"schema": "qwen235b-t0-smoke-v1",
|
||||
"workload": {
|
||||
"input_tokens": args.input_tokens,
|
||||
"output_tokens": args.output_tokens,
|
||||
"uniform_qps": None,
|
||||
"prefix_caching": False,
|
||||
"concurrency": args.concurrency,
|
||||
"request_count": args.requests,
|
||||
"prompt_token_id": prompt_token_id,
|
||||
},
|
||||
"summary": {
|
||||
"completed_requests": len(results),
|
||||
"joint_slo_pass_count": sum(bool(row["joint_slo_pass"]) for row in results),
|
||||
"ttft_mean_ms": statistics.fmean(ttfts),
|
||||
"ttft_p95_ms": percentile(ttfts, 0.95),
|
||||
"tpot_mean_ms": statistics.fmean(tpots),
|
||||
"tpot_p95_ms": percentile(tpots, 0.95),
|
||||
},
|
||||
"requests": sorted(results, key=lambda row: int(row["request_index"])),
|
||||
}
|
||||
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 __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user