Add OpProf campaign: protocols, results, patches, run evidence (P0-P6)
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>
This commit is contained in:
784
runs/opprof-phase3-ea/provenance/opprof_phase3_client.py
Normal file
784
runs/opprof-phase3-ea/provenance/opprof_phase3_client.py
Normal file
@@ -0,0 +1,784 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user