Gahow Wang
d5b276180d
Workload-conditioned operator profiling on patched vLLM 0.24.0 + Qwen3-30B-A3B/H20. H1b PASS (irregular patterns carry +23-45pp R64 raggedness, 8-45% token-efficiency loss vs rectangular controls); mechanism decomposition kills the padding narrative and finds the arrival-uniformization artifact (-12.9%); cross-version churn surface shows TP2/MNS64 -29.4% across vLLM 0.20->0.24 while the argmax held. Raw Layer-1 JSONL streams (507 MB) stay on disk, git-ignored; footer sidecars and metrics are tracked. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
785 lines
29 KiB
Python
785 lines
29 KiB
Python
#!/usr/bin/env python3
|
|
"""Token-exact fixed-duration client for the OpProf Phase-3 protocol."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import gzip
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
import time
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import aiohttp
|
|
|
|
SCHEMA = 1
|
|
TOKEN_BASE = 1000
|
|
TOKEN_SPAN = 100000
|
|
|
|
|
|
class ManifestExhausted(RuntimeError):
|
|
pass
|
|
|
|
|
|
def sha256_file(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as f:
|
|
for chunk in iter(lambda: f.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def atomic_json(path: Path, value: Any, mode: int = 0o640) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
tmp = path.with_name(path.name + f".tmp.{os.getpid()}")
|
|
fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_EXCL, mode)
|
|
with os.fdopen(fd, "w", encoding="utf-8") as f:
|
|
json.dump(value, f, sort_keys=True, indent=2)
|
|
f.write("\n")
|
|
f.flush()
|
|
os.fsync(f.fileno())
|
|
os.replace(tmp, path)
|
|
|
|
|
|
def atomic_jsonl(path: Path, rows: list[dict[str, Any]], mode: int = 0o640) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
tmp = path.with_name(path.name + f".tmp.{os.getpid()}")
|
|
fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_EXCL, mode)
|
|
with os.fdopen(fd, "w", encoding="utf-8") as f:
|
|
for row in rows:
|
|
f.write(json.dumps(row, sort_keys=True, separators=(",", ":")) + "\n")
|
|
f.flush()
|
|
os.fsync(f.fileno())
|
|
os.replace(tmp, path)
|
|
|
|
|
|
def parse_range(value: str) -> tuple[int, int]:
|
|
lo_text, hi_text = value.split(":", 1)
|
|
lo, hi = int(lo_text), int(hi_text)
|
|
if lo <= 0 or hi < lo:
|
|
raise argparse.ArgumentTypeError(f"invalid positive range: {value}")
|
|
return lo, hi
|
|
|
|
|
|
def _integer_counts(weights: list[float], total: int) -> list[int]:
|
|
raw = [w * total for w in weights]
|
|
counts = [math.floor(x) for x in raw]
|
|
order = sorted(
|
|
range(len(raw)), key=lambda i: raw[i] - counts[i], reverse=True
|
|
)
|
|
for idx in order[: total - sum(counts)]:
|
|
counts[idx] += 1
|
|
return counts
|
|
|
|
|
|
def numeric_sanity(values: list[float | int]) -> dict[str, Any]:
|
|
finite = [float(x) for x in values if math.isfinite(float(x))]
|
|
return {
|
|
"n": len(values),
|
|
"finite_n": len(finite),
|
|
"missing_n": len(values) - len(finite),
|
|
"min": min(finite) if finite else None,
|
|
"max": max(finite) if finite else None,
|
|
"distinct_n": len(set(finite)),
|
|
}
|
|
|
|
|
|
def manifest_summary(rows: list[dict[str, Any]]) -> dict[str, Any]:
|
|
return {
|
|
"schema": SCHEMA,
|
|
"rows": len(rows),
|
|
"input_tokens": numeric_sanity([int(r["input_tokens"]) for r in rows]),
|
|
"output_tokens": numeric_sanity([int(r["output_tokens"]) for r in rows]),
|
|
"arrival_values": sorted({str(r["arrival"]) for r in rows}),
|
|
"pattern_values": sorted({str(r["pattern_id"]) for r in rows}),
|
|
}
|
|
|
|
|
|
def materialize(args: argparse.Namespace) -> dict[str, Any]:
|
|
import numpy as np
|
|
|
|
rng = np.random.default_rng(args.workload_seed)
|
|
n = args.num_requests
|
|
if args.kind == "prefix-pool":
|
|
if args.num_prefixes <= 0 or args.prefix_len <= 0 or args.suffix_fixed <= 0:
|
|
raise ValueError("prefix-pool requires positive pool/prefix/suffix")
|
|
lengths = np.full(n, args.prefix_len + args.suffix_fixed, dtype=np.int64)
|
|
prefix_ids = np.arange(n, dtype=np.int64) % args.num_prefixes
|
|
rng.shuffle(prefix_ids)
|
|
else:
|
|
prefix_ids = np.full(n, -1, dtype=np.int64)
|
|
if args.input_uniform:
|
|
lo, hi = parse_range(args.input_uniform)
|
|
lengths = rng.integers(lo, hi + 1, n, dtype=np.int64)
|
|
elif args.input_fixed:
|
|
lengths = np.full(n, args.input_fixed, dtype=np.int64)
|
|
elif args.input_mixture:
|
|
spec = json.loads(args.input_mixture)
|
|
if not isinstance(spec, dict) or not spec:
|
|
raise ValueError("input mixture must be a non-empty JSON object")
|
|
keys = list(spec)
|
|
weights = [float(spec[key]) for key in keys]
|
|
if any(w < 0 for w in weights) or not math.isclose(sum(weights), 1.0):
|
|
raise ValueError("mixture weights must be non-negative and sum to 1")
|
|
pieces = []
|
|
for key, count in zip(
|
|
keys, _integer_counts(weights, n), strict=True
|
|
):
|
|
kind, lo_text, hi_text = key.split(":")
|
|
if kind != "uniform":
|
|
raise ValueError(f"unsupported mixture component: {key}")
|
|
pieces.append(
|
|
rng.integers(
|
|
int(lo_text), int(hi_text) + 1, count, dtype=np.int64
|
|
)
|
|
)
|
|
lengths = np.concatenate(pieces)
|
|
rng.shuffle(lengths)
|
|
else:
|
|
raise ValueError("exactly one input distribution is required")
|
|
if args.output_fixed <= 0 or args.arrival not in {"steady", "burst:8"}:
|
|
raise ValueError("invalid output length or arrival class")
|
|
|
|
rows = []
|
|
for i in range(n):
|
|
row = {
|
|
"schema": SCHEMA,
|
|
"request_id": f"{args.id}-{i:05d}",
|
|
"pattern_id": args.id,
|
|
"kind": args.kind,
|
|
"input_tokens": int(lengths[i]),
|
|
"output_tokens": args.output_fixed,
|
|
"arrival": args.arrival,
|
|
"token_seed": int(args.workload_seed * 1000003 + i),
|
|
}
|
|
if args.kind == "prefix-pool":
|
|
row.update(
|
|
{
|
|
"prefix_id": int(prefix_ids[i]),
|
|
"num_prefixes": args.num_prefixes,
|
|
"prefix_tokens": args.prefix_len,
|
|
}
|
|
)
|
|
rows.append(row)
|
|
|
|
out = Path(args.out)
|
|
atomic_jsonl(out, rows, mode=0o600)
|
|
summary = manifest_summary(rows)
|
|
summary.update({"sha256": sha256_file(out), "path": str(out)})
|
|
atomic_json(out.with_suffix(out.suffix + ".summary.json"), summary, mode=0o600)
|
|
print(json.dumps(summary, sort_keys=True))
|
|
return summary
|
|
|
|
|
|
def materialize_private(args: argparse.Namespace) -> dict[str, Any]:
|
|
from transformers import AutoTokenizer
|
|
|
|
source = Path(args.source)
|
|
selected: list[dict[str, Any]] = []
|
|
with source.open(encoding="utf-8") as f:
|
|
for source_index, line in enumerate(f):
|
|
row = json.loads(line)
|
|
if (
|
|
float(row["sampling_u"]) <= args.sampling_u_max
|
|
and int(row["input_length"]) <= args.max_input_tokens
|
|
):
|
|
selected.append(
|
|
{
|
|
"schema": SCHEMA,
|
|
"request_id": f"{args.id}-{len(selected):05d}",
|
|
"pattern_id": args.id,
|
|
"kind": "private-trace",
|
|
"input_tokens": int(row["input_length"]),
|
|
"output_tokens": min(
|
|
int(row["output_length"]), args.output_cap
|
|
),
|
|
"arrival": args.arrival,
|
|
"source_index": source_index,
|
|
"prompt": row["prompt"],
|
|
}
|
|
)
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
|
|
diffs = [
|
|
len(tokenizer.encode(row["prompt"], add_special_tokens=False))
|
|
- row["input_tokens"]
|
|
for row in selected
|
|
]
|
|
exact = sum(diff == 0 for diff in diffs)
|
|
exact_fraction = exact / len(diffs) if diffs else 0.0
|
|
max_abs = max((abs(diff) for diff in diffs), default=-1)
|
|
if exact_fraction < 0.99 or max_abs > 1:
|
|
raise RuntimeError(
|
|
"tokenizer parity gate failed: "
|
|
f"exact_fraction={exact_fraction:.6f} max_abs_error={max_abs}"
|
|
)
|
|
|
|
out = Path(args.out)
|
|
atomic_jsonl(out, selected, mode=0o600)
|
|
summary = manifest_summary(selected)
|
|
summary.update(
|
|
{
|
|
"sha256": sha256_file(out),
|
|
"source_sha256": sha256_file(source),
|
|
"tokenizer_exact_n": exact,
|
|
"tokenizer_exact_fraction": exact_fraction,
|
|
"tokenizer_max_abs_error": max_abs,
|
|
"path": str(out),
|
|
}
|
|
)
|
|
atomic_json(out.with_suffix(out.suffix + ".summary.json"), summary, mode=0o600)
|
|
print(json.dumps(summary, sort_keys=True))
|
|
return summary
|
|
|
|
|
|
def load_manifest(path: Path) -> list[dict[str, Any]]:
|
|
rows = [json.loads(line) for line in path.read_text().splitlines() if line]
|
|
required = {
|
|
"request_id",
|
|
"pattern_id",
|
|
"input_tokens",
|
|
"output_tokens",
|
|
"arrival",
|
|
}
|
|
if not rows:
|
|
raise ValueError("empty manifest")
|
|
for row in rows:
|
|
if not required.issubset(row):
|
|
raise ValueError(f"manifest row lacks {sorted(required - set(row))}")
|
|
if len({row["request_id"] for row in rows}) != len(rows):
|
|
raise ValueError("duplicate request_id")
|
|
return rows
|
|
|
|
|
|
def _token_stream(seed: int, count: int) -> list[int]:
|
|
state = seed & 0xFFFFFFFF
|
|
out = []
|
|
for _ in range(count):
|
|
state = (1664525 * state + 1013904223) & 0xFFFFFFFF
|
|
out.append(TOKEN_BASE + state % TOKEN_SPAN)
|
|
return out
|
|
|
|
|
|
def synthetic_prompt(row: dict[str, Any]) -> list[int]:
|
|
length = int(row["input_tokens"])
|
|
seed = int(row.get("token_seed", 0))
|
|
if row.get("kind") == "prefix-pool":
|
|
prefix_n = int(row["prefix_tokens"])
|
|
tokens = _token_stream(0xA5A50000 + int(row["prefix_id"]), prefix_n)
|
|
tokens += _token_stream(seed, length - prefix_n)
|
|
offset = prefix_n
|
|
else:
|
|
tokens = _token_stream(seed, length)
|
|
offset = 0
|
|
if length - offset >= 3:
|
|
index = int(row["request_id"].rsplit("-", 1)[1])
|
|
tokens[offset : offset + 3] = [
|
|
TOKEN_BASE + index % 100,
|
|
TOKEN_BASE + (index // 100) % 100,
|
|
TOKEN_BASE + (index // 10000) % 100,
|
|
]
|
|
return tokens
|
|
|
|
|
|
@dataclass
|
|
class RunContext:
|
|
args: argparse.Namespace
|
|
rows: list[dict[str, Any]]
|
|
t0: float
|
|
clean_end: float
|
|
stop_event: asyncio.Event
|
|
lock: asyncio.Lock
|
|
next_index: int = 0
|
|
in_flight: int = 0
|
|
max_in_flight: int = 0
|
|
exhausted: bool = False
|
|
admission_stop_s: float | None = None
|
|
|
|
async def next_row(self) -> dict[str, Any]:
|
|
async with self.lock:
|
|
if self.next_index >= len(self.rows):
|
|
self.exhausted = True
|
|
raise ManifestExhausted(
|
|
f"manifest exhausted after {self.next_index} admissions"
|
|
)
|
|
row = self.rows[self.next_index]
|
|
self.next_index += 1
|
|
return row
|
|
|
|
|
|
async def request_one(
|
|
ctx: RunContext,
|
|
session: aiohttp.ClientSession,
|
|
row: dict[str, Any],
|
|
scheduled: float,
|
|
) -> dict[str, Any]:
|
|
loop = asyncio.get_running_loop()
|
|
admitted = loop.time()
|
|
ctx.in_flight += 1
|
|
ctx.max_in_flight = max(ctx.max_in_flight, ctx.in_flight)
|
|
status = 0
|
|
actual_output: int | None = None
|
|
first_token: float | None = None
|
|
error_kind: str | None = None
|
|
try:
|
|
prompt: str | list[int] = (
|
|
row["prompt"]
|
|
if row.get("kind") == "private-trace"
|
|
else synthetic_prompt(row)
|
|
)
|
|
if not isinstance(prompt, str) and len(prompt) != int(row["input_tokens"]):
|
|
raise AssertionError("synthetic prompt length drift")
|
|
payload = {
|
|
"model": ctx.args.model,
|
|
"prompt": prompt,
|
|
"max_tokens": int(row["output_tokens"]),
|
|
"temperature": ctx.args.temperature,
|
|
"ignore_eos": ctx.args.ignore_eos,
|
|
"stream": True,
|
|
"stream_options": {"include_usage": True},
|
|
"add_special_tokens": False,
|
|
"seed": ctx.args.server_seed,
|
|
}
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
"x-request-id": str(row["request_id"]),
|
|
}
|
|
async with session.post(
|
|
ctx.args.base_url.rstrip("/") + "/v1/completions",
|
|
json=payload,
|
|
headers=headers,
|
|
) as response:
|
|
status = response.status
|
|
if status != 200:
|
|
error_kind = f"http_{status}"
|
|
else:
|
|
buf = b""
|
|
async for chunk in response.content.iter_any():
|
|
buf += chunk
|
|
while b"\n" in buf:
|
|
line, buf = buf.split(b"\n", 1)
|
|
line = line.strip()
|
|
if not line.startswith(b"data:"):
|
|
continue
|
|
data = line[5:].strip()
|
|
if data == b"[DONE]":
|
|
continue
|
|
event = json.loads(data)
|
|
if event.get("choices") and first_token is None:
|
|
first_token = loop.time()
|
|
if event.get("usage") is not None:
|
|
actual_output = int(event["usage"]["completion_tokens"])
|
|
if actual_output is None:
|
|
error_kind = "missing_usage"
|
|
except (aiohttp.ClientError, asyncio.TimeoutError) as exc:
|
|
error_kind = type(exc).__name__
|
|
except Exception as exc:
|
|
error_kind = type(exc).__name__
|
|
finally:
|
|
completed = loop.time()
|
|
ctx.in_flight -= 1
|
|
success = (
|
|
status == 200
|
|
and error_kind is None
|
|
and actual_output == int(row["output_tokens"])
|
|
)
|
|
if status == 200 and actual_output is not None and not success:
|
|
error_kind = "output_token_mismatch"
|
|
return {
|
|
"schema": SCHEMA,
|
|
"request_id": row["request_id"],
|
|
"scheduled_s": scheduled - ctx.t0,
|
|
"admitted_s": admitted - ctx.t0,
|
|
"first_token_s": None if first_token is None else first_token - ctx.t0,
|
|
"completed_s": completed - ctx.t0,
|
|
"input_tokens": int(row["input_tokens"]),
|
|
"requested_output_tokens": int(row["output_tokens"]),
|
|
"actual_output_tokens": actual_output,
|
|
"http_status": status,
|
|
"success": success,
|
|
"error_kind": error_kind,
|
|
}
|
|
|
|
|
|
async def saturation_load(
|
|
ctx: RunContext, session: aiohttp.ClientSession
|
|
) -> list[dict[str, Any]]:
|
|
results: list[dict[str, Any]] = []
|
|
|
|
async def worker() -> None:
|
|
while not ctx.stop_event.is_set():
|
|
try:
|
|
row = await ctx.next_row()
|
|
except ManifestExhausted:
|
|
ctx.stop_event.set()
|
|
return
|
|
results.append(
|
|
await request_one(ctx, session, row, asyncio.get_running_loop().time())
|
|
)
|
|
|
|
tasks = [
|
|
asyncio.create_task(worker()) for _ in range(ctx.args.max_concurrency)
|
|
]
|
|
await asyncio.gather(*tasks)
|
|
return results
|
|
|
|
|
|
async def finite_load(
|
|
ctx: RunContext, session: aiohttp.ClientSession, rate: float
|
|
) -> list[dict[str, Any]]:
|
|
sem = asyncio.Semaphore(ctx.args.max_concurrency)
|
|
tasks: list[asyncio.Task[dict[str, Any]]] = []
|
|
batch = 8 if str(ctx.rows[0]["arrival"]) == "burst:8" else 1
|
|
period = batch / rate
|
|
event_index = 0
|
|
|
|
async def limited(row: dict[str, Any], scheduled: float) -> dict[str, Any]:
|
|
async with sem:
|
|
return await request_one(ctx, session, row, scheduled)
|
|
|
|
while not ctx.stop_event.is_set():
|
|
scheduled = ctx.t0 + event_index * period
|
|
delay = scheduled - asyncio.get_running_loop().time()
|
|
if delay > 0:
|
|
try:
|
|
await asyncio.wait_for(ctx.stop_event.wait(), timeout=delay)
|
|
break
|
|
except asyncio.TimeoutError:
|
|
pass
|
|
if ctx.stop_event.is_set():
|
|
break
|
|
try:
|
|
for _ in range(batch):
|
|
tasks.append(
|
|
asyncio.create_task(limited(await ctx.next_row(), scheduled))
|
|
)
|
|
except ManifestExhausted:
|
|
ctx.stop_event.set()
|
|
break
|
|
event_index += 1
|
|
return await asyncio.gather(*tasks) if tasks else []
|
|
|
|
|
|
async def post_profile(
|
|
session: aiohttp.ClientSession, base_url: str, endpoint: str
|
|
) -> tuple[float, float, int]:
|
|
loop = asyncio.get_running_loop()
|
|
before = loop.time()
|
|
async with session.post(base_url.rstrip("/") + endpoint) as response:
|
|
status = response.status
|
|
await response.read()
|
|
return before, loop.time(), status
|
|
|
|
|
|
def _trace_loadable(path: Path) -> bool:
|
|
try:
|
|
opener = gzip.open if path.suffix == ".gz" else open
|
|
with opener(path, "rt", encoding="utf-8") as f:
|
|
parsed = json.load(f)
|
|
return isinstance(parsed, dict) and isinstance(parsed.get("traceEvents"), list)
|
|
except (OSError, EOFError, json.JSONDecodeError):
|
|
return False
|
|
|
|
|
|
async def wait_new_trace(
|
|
trace_dir: Path, before: set[Path], timeout: float
|
|
) -> Path:
|
|
deadline = asyncio.get_running_loop().time() + timeout
|
|
while asyncio.get_running_loop().time() < deadline:
|
|
for path in sorted(set(trace_dir.glob("*.pt.trace.json*")) - before):
|
|
if _trace_loadable(path):
|
|
return path
|
|
await asyncio.sleep(0.25)
|
|
raise TimeoutError(f"no new loadable trace within {timeout}s")
|
|
|
|
|
|
async def timeline(
|
|
ctx: RunContext, session: aiohttp.ClientSession
|
|
) -> list[dict[str, Any]]:
|
|
args = ctx.args
|
|
profiles: list[dict[str, Any]] = []
|
|
await asyncio.sleep(max(0, ctx.clean_end - asyncio.get_running_loop().time()))
|
|
if args.profile_after_clean:
|
|
trace_dir = Path(args.profile_trace_dir)
|
|
for window in range(args.num_profile_windows):
|
|
prior = set(trace_dir.glob("*.pt.trace.json*"))
|
|
start_before, start_after, start_status = await post_profile(
|
|
session, args.base_url, "/start_profile"
|
|
)
|
|
trace = await wait_new_trace(
|
|
trace_dir, prior, args.profile_timeout_seconds
|
|
)
|
|
trace_ready = asyncio.get_running_loop().time()
|
|
stop_before, stop_after, stop_status = await post_profile(
|
|
session, args.base_url, "/stop_profile"
|
|
)
|
|
profiles.append(
|
|
{
|
|
"window": window + 1,
|
|
"start_call_s": start_before - ctx.t0,
|
|
"start_return_s": start_after - ctx.t0,
|
|
"trace_ready_s": trace_ready - ctx.t0,
|
|
"stop_call_s": stop_before - ctx.t0,
|
|
"stop_return_s": stop_after - ctx.t0,
|
|
"start_status": start_status,
|
|
"stop_status": stop_status,
|
|
"trace_file": trace.name,
|
|
"trace_sha256": sha256_file(trace),
|
|
}
|
|
)
|
|
if start_status != 200 or stop_status != 200:
|
|
raise RuntimeError("profile endpoint returned non-200")
|
|
await asyncio.sleep(args.recovery_seconds)
|
|
else:
|
|
await asyncio.sleep(args.post_clean_seconds)
|
|
ctx.admission_stop_s = asyncio.get_running_loop().time() - ctx.t0
|
|
ctx.stop_event.set()
|
|
return profiles
|
|
|
|
|
|
def segment_summary(
|
|
records: list[dict[str, Any]], start: float, end: float
|
|
) -> dict[str, Any]:
|
|
admitted = [r for r in records if start <= r["admitted_s"] < end]
|
|
completed = [r for r in records if start <= r["completed_s"] < end]
|
|
successes = [r for r in completed if r["success"]]
|
|
duration = end - start
|
|
return {
|
|
"start_s": start,
|
|
"end_s": end,
|
|
"duration_s": duration,
|
|
"admitted": len(admitted),
|
|
"completed": len(successes),
|
|
"failed": len(completed) - len(successes),
|
|
"offered_rps": len(admitted) / duration,
|
|
"completed_throughput_rps": len(successes) / duration,
|
|
"input_tokens": sum(r["input_tokens"] for r in successes),
|
|
"output_tokens": sum(r["actual_output_tokens"] or 0 for r in successes),
|
|
}
|
|
|
|
|
|
async def run_load(args: argparse.Namespace) -> dict[str, Any]:
|
|
manifest = Path(args.manifest)
|
|
rows = load_manifest(manifest)
|
|
arrivals = {row["arrival"] for row in rows}
|
|
if len(arrivals) != 1:
|
|
raise ValueError("a manifest must have one arrival class")
|
|
if args.load_point == "saturation":
|
|
if args.request_rate != "inf":
|
|
raise ValueError("saturation requires --request-rate inf")
|
|
rate = math.inf
|
|
else:
|
|
if not args.saturation_result:
|
|
raise ValueError("moderate requires --saturation-result")
|
|
sat = json.loads(Path(args.saturation_result).read_text())
|
|
rate = args.rate_fraction * float(sat["clean"]["completed_throughput_rps"])
|
|
if not math.isfinite(rate) or rate <= 0:
|
|
raise ValueError("derived moderate rate must be positive and finite")
|
|
|
|
loop = asyncio.get_running_loop()
|
|
t0 = loop.time()
|
|
clean_seconds = args.clean_segment_seconds * args.num_clean_segments
|
|
ctx = RunContext(
|
|
args=args,
|
|
rows=rows,
|
|
t0=t0,
|
|
clean_end=t0 + args.warmup_seconds + clean_seconds,
|
|
stop_event=asyncio.Event(),
|
|
lock=asyncio.Lock(),
|
|
)
|
|
timeout = aiohttp.ClientTimeout(total=None, connect=30, sock_read=600)
|
|
connector = aiohttp.TCPConnector(limit=args.max_concurrency)
|
|
async with aiohttp.ClientSession(timeout=timeout, connector=connector) as session:
|
|
profile_task = asyncio.create_task(timeline(ctx, session))
|
|
load_task = asyncio.create_task(
|
|
saturation_load(ctx, session)
|
|
if math.isinf(rate)
|
|
else finite_load(ctx, session, rate)
|
|
)
|
|
try:
|
|
profiles = await profile_task
|
|
except Exception:
|
|
ctx.stop_event.set()
|
|
await load_task
|
|
raise
|
|
records = await load_task
|
|
|
|
clean_start = args.warmup_seconds
|
|
clean_end = clean_start + clean_seconds
|
|
clean = segment_summary(records, clean_start, clean_end)
|
|
segments = []
|
|
for i in range(args.num_clean_segments):
|
|
start = clean_start + i * args.clean_segment_seconds
|
|
segments.append(
|
|
{
|
|
"name": chr(ord("A") + i),
|
|
**segment_summary(
|
|
records, start, start + args.clean_segment_seconds
|
|
),
|
|
}
|
|
)
|
|
successful = [r for r in records if r["success"]]
|
|
elapsed_seconds = loop.time() - t0
|
|
if ctx.admission_stop_s is None:
|
|
raise RuntimeError("admission stop timestamp was not recorded")
|
|
drain_seconds = elapsed_seconds - ctx.admission_stop_s
|
|
result = {
|
|
"schema": SCHEMA,
|
|
"manifest_sha256": sha256_file(manifest),
|
|
"manifest_rows": len(rows),
|
|
"manifest_admitted": ctx.next_index,
|
|
"manifest_wrapped": False,
|
|
"manifest_exhausted": ctx.exhausted,
|
|
"load_point": args.load_point,
|
|
"request_rate": "inf" if math.isinf(rate) else rate,
|
|
"rate_fraction": None if math.isinf(rate) else args.rate_fraction,
|
|
"arrival": next(iter(arrivals)),
|
|
"warmup_seconds": args.warmup_seconds,
|
|
"clean_segment_seconds": args.clean_segment_seconds,
|
|
"num_clean_segments": args.num_clean_segments,
|
|
"elapsed_seconds": elapsed_seconds,
|
|
"admission_stop_s": ctx.admission_stop_s,
|
|
"drain_seconds": drain_seconds,
|
|
"max_in_flight": ctx.max_in_flight,
|
|
"records": len(records),
|
|
"successful_records": len(successful),
|
|
"failed_records": len(records) - len(successful),
|
|
"clean": clean,
|
|
"segments": segments,
|
|
"profiles": profiles,
|
|
}
|
|
sanity = {
|
|
"schema": SCHEMA,
|
|
"numeric": {
|
|
"input_tokens": numeric_sanity([r["input_tokens"] for r in records]),
|
|
"requested_output_tokens": numeric_sanity(
|
|
[r["requested_output_tokens"] for r in records]
|
|
),
|
|
"actual_output_tokens": numeric_sanity(
|
|
[
|
|
r["actual_output_tokens"]
|
|
for r in records
|
|
if r["actual_output_tokens"] is not None
|
|
]
|
|
),
|
|
"scheduled_s": numeric_sanity([r["scheduled_s"] for r in records]),
|
|
"admitted_s": numeric_sanity([r["admitted_s"] for r in records]),
|
|
"completed_s": numeric_sanity([r["completed_s"] for r in records]),
|
|
},
|
|
"invariants": {
|
|
"clean_duration_exact": math.isclose(clean["duration_s"], clean_seconds),
|
|
"segment_count_exact": len(segments) == args.num_clean_segments,
|
|
"manifest_no_wrap": ctx.next_index <= len(rows),
|
|
"manifest_not_exhausted": not ctx.exhausted,
|
|
"concurrency_bounded": ctx.max_in_flight <= args.max_concurrency,
|
|
"drain_within_timeout": drain_seconds <= args.drain_timeout_seconds,
|
|
"output_tokens_exact": all(
|
|
r["actual_output_tokens"] == r["requested_output_tokens"]
|
|
for r in successful
|
|
),
|
|
"clean_failures_zero": clean["failed"] == 0,
|
|
"profile_count_exact": len(profiles)
|
|
== (args.num_profile_windows if args.profile_after_clean else 0),
|
|
"profile_status_ok": all(
|
|
p["start_status"] == 200 and p["stop_status"] == 200
|
|
for p in profiles
|
|
),
|
|
},
|
|
}
|
|
if not math.isinf(rate):
|
|
sanity["invariants"]["moderate_offered_within_5pct"] = (
|
|
abs(clean["offered_rps"] / rate - 1) <= 0.05
|
|
)
|
|
out = Path(args.result_dir)
|
|
out.mkdir(parents=True, exist_ok=True)
|
|
atomic_jsonl(out / "requests.jsonl", sorted(records, key=lambda r: r["admitted_s"]))
|
|
atomic_jsonl(out / "segments.jsonl", segments)
|
|
atomic_json(out / "result.json", result)
|
|
atomic_json(out / "sanity.json", sanity)
|
|
if ctx.exhausted:
|
|
raise ManifestExhausted("manifest exhausted; result retained for diagnosis")
|
|
failed = [name for name, ok in sanity["invariants"].items() if not ok]
|
|
if failed:
|
|
raise RuntimeError(f"client sanity failure: {failed}")
|
|
return result
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser()
|
|
sub = parser.add_subparsers(dest="command", required=True)
|
|
mat = sub.add_parser("materialize")
|
|
mat.add_argument("--id", required=True)
|
|
mat.add_argument("--kind", choices=("synthetic", "prefix-pool"), required=True)
|
|
group = mat.add_mutually_exclusive_group()
|
|
group.add_argument("--input-uniform")
|
|
group.add_argument("--input-fixed", type=int)
|
|
group.add_argument("--input-mixture")
|
|
mat.add_argument("--output-fixed", type=int, required=True)
|
|
mat.add_argument("--prefix", default="none")
|
|
mat.add_argument("--arrival", required=True)
|
|
mat.add_argument("--num-requests", type=int, required=True)
|
|
mat.add_argument("--workload-seed", type=int, required=True)
|
|
mat.add_argument("--num-prefixes", type=int, default=0)
|
|
mat.add_argument("--prefix-len", type=int, default=0)
|
|
mat.add_argument("--suffix-fixed", type=int, default=0)
|
|
mat.add_argument("--out", required=True)
|
|
private = sub.add_parser("materialize-private")
|
|
private.add_argument("--id", required=True)
|
|
private.add_argument("--source", required=True)
|
|
private.add_argument("--sampling-u-max", type=float, required=True)
|
|
private.add_argument("--max-input-tokens", type=int, required=True)
|
|
private.add_argument("--output-cap", type=int, required=True)
|
|
private.add_argument("--preserve-prompts", action="store_true", required=True)
|
|
private.add_argument("--disable-shuffle", action="store_true", required=True)
|
|
private.add_argument("--arrival", required=True)
|
|
private.add_argument("--model", required=True)
|
|
private.add_argument("--out", required=True)
|
|
run = sub.add_parser("run")
|
|
run.add_argument("--manifest", required=True)
|
|
run.add_argument("--base-url", required=True)
|
|
run.add_argument("--model", required=True)
|
|
run.add_argument("--load-point", choices=("saturation", "moderate"), required=True)
|
|
run.add_argument("--request-rate")
|
|
run.add_argument("--saturation-result")
|
|
run.add_argument("--rate-fraction", type=float, default=0.60)
|
|
run.add_argument("--max-concurrency", type=int, default=256)
|
|
run.add_argument("--ignore-eos", action="store_true")
|
|
run.add_argument("--temperature", type=float, default=0.0)
|
|
run.add_argument("--warmup-seconds", type=float, default=60)
|
|
run.add_argument("--clean-segment-seconds", type=float, default=80)
|
|
run.add_argument("--num-clean-segments", type=int, default=3)
|
|
run.add_argument("--profile-after-clean", action="store_true")
|
|
run.add_argument("--num-profile-windows", type=int, default=0)
|
|
run.add_argument("--profile-warmup-iterations", type=int, default=2)
|
|
run.add_argument("--profile-active-iterations", type=int, default=8)
|
|
run.add_argument("--profile-trace-dir")
|
|
run.add_argument("--profile-timeout-seconds", type=float, default=120)
|
|
run.add_argument("--recovery-seconds", type=float, default=30)
|
|
run.add_argument("--post-clean-seconds", type=float, default=0)
|
|
run.add_argument("--drain-timeout-seconds", type=float, default=120)
|
|
run.add_argument("--workload-seed", type=int, default=20260712)
|
|
run.add_argument("--server-seed", type=int, default=20260712)
|
|
run.add_argument("--result-dir", required=True)
|
|
return parser
|
|
|
|
|
|
def main() -> None:
|
|
args = build_parser().parse_args()
|
|
if args.command == "materialize":
|
|
materialize(args)
|
|
elif args.command == "materialize-private":
|
|
materialize_private(args)
|
|
else:
|
|
if args.profile_after_clean and not args.profile_trace_dir:
|
|
raise ValueError("--profile-after-clean requires --profile-trace-dir")
|
|
print(json.dumps(asyncio.run(run_load(args)), sort_keys=True))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|