experiment: add workload regime taxonomy
This commit is contained in:
@@ -0,0 +1,376 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Materialize controlled Qwen30 workload families from one exact trace cohort."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import csv
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Iterable
|
||||
|
||||
|
||||
FIELDS = (
|
||||
"arrived_at",
|
||||
"num_prefill_tokens",
|
||||
"num_decode_tokens",
|
||||
"session_id",
|
||||
"block_hash_ids",
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Family:
|
||||
name: str
|
||||
shape: str
|
||||
arrival: str
|
||||
prefix: bool
|
||||
|
||||
|
||||
FAMILIES = (
|
||||
Family("w0-short-fixed-uniform-none", "short-fixed", "uniform", False),
|
||||
Family("w1-mean-fixed-uniform-none", "mean-fixed", "uniform", False),
|
||||
Family("w2-mean-fixed-trace-none", "mean-fixed", "trace", False),
|
||||
Family("w3-heterogeneous-uniform-none", "heterogeneous", "uniform", False),
|
||||
Family("w4-heterogeneous-trace-none", "heterogeneous", "trace", False),
|
||||
Family("w5-heterogeneous-uniform-prefix", "heterogeneous", "uniform", True),
|
||||
Family("w6-heterogeneous-trace-prefix", "heterogeneous", "trace", True),
|
||||
)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--source-public", type=Path, required=True)
|
||||
parser.add_argument("--source-private", type=Path, required=True)
|
||||
parser.add_argument("--model", type=Path, required=True)
|
||||
parser.add_argument("--output-root", type=Path, required=True)
|
||||
parser.add_argument(
|
||||
"--reference-decode-tokens-per-second", type=float, default=3064.0
|
||||
)
|
||||
parser.add_argument(
|
||||
"--rho", type=float, action="append", default=None,
|
||||
help="Normalized decode offered load; repeat for multiple levels.",
|
||||
)
|
||||
parser.add_argument("--requests", type=int)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def sha256(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as source:
|
||||
for chunk in iter(lambda: source.read(1 << 20), b""):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def vector_sha256(rows: Iterable[dict[str, Any]]) -> str:
|
||||
digest = hashlib.sha256()
|
||||
for row in rows:
|
||||
digest.update(
|
||||
json.dumps(
|
||||
[
|
||||
row["source_index"],
|
||||
row["arrived_at"],
|
||||
row["input_length"],
|
||||
row["output_length"],
|
||||
row["session_id"],
|
||||
row["runtime_block_ids"],
|
||||
],
|
||||
separators=(",", ":"),
|
||||
).encode()
|
||||
)
|
||||
digest.update(b"\n")
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def load_source(public_path: Path, private_path: Path) -> list[dict[str, Any]]:
|
||||
with public_path.open(newline="") as source:
|
||||
public_rows = list(csv.DictReader(source))
|
||||
private_rows = [json.loads(line) for line in private_path.open() if line.strip()]
|
||||
if not public_rows or len(public_rows) != len(private_rows):
|
||||
raise ValueError("source public/private request count mismatch")
|
||||
|
||||
rows = []
|
||||
for index, (public, private) in enumerate(
|
||||
zip(public_rows, private_rows, strict=True)
|
||||
):
|
||||
input_length = int(public["num_prefill_tokens"])
|
||||
output_length = int(public["num_decode_tokens"])
|
||||
if input_length != int(private["input_length"]):
|
||||
raise ValueError(f"source input-length drift at row {index}")
|
||||
if output_length != int(private["output_length"]):
|
||||
raise ValueError(f"source output-length drift at row {index}")
|
||||
runtime_ids = [
|
||||
int(value) for value in public["block_hash_ids"].split("|") if value
|
||||
]
|
||||
if len(runtime_ids) != input_length // 16:
|
||||
raise ValueError(f"incomplete prefix-block projection at row {index}")
|
||||
rows.append(
|
||||
{
|
||||
"source_index": int(private["source_index"]),
|
||||
"source_arrived_at": float(public["arrived_at"]),
|
||||
"input_length": input_length,
|
||||
"output_length": output_length,
|
||||
"session_id": int(private["session_id"]),
|
||||
"runtime_block_ids": runtime_ids,
|
||||
"body": copy.deepcopy(private["body"]),
|
||||
}
|
||||
)
|
||||
if any(
|
||||
right["source_arrived_at"] < left["source_arrived_at"]
|
||||
for left, right in zip(rows, rows[1:])
|
||||
):
|
||||
raise ValueError("source arrival order drift")
|
||||
return rows
|
||||
|
||||
|
||||
def arrivals_for(
|
||||
source_rows: list[dict[str, Any]], arrival: str, target_rate: float
|
||||
) -> list[float]:
|
||||
if target_rate <= 0 or not math.isfinite(target_rate):
|
||||
raise ValueError("target request rate must be finite and positive")
|
||||
if arrival == "uniform":
|
||||
return [index / target_rate for index in range(len(source_rows))]
|
||||
if arrival != "trace":
|
||||
raise ValueError(f"unknown arrival family: {arrival}")
|
||||
if len(source_rows) < 2:
|
||||
raise ValueError("trace arrivals require at least two requests")
|
||||
first = source_rows[0]["source_arrived_at"]
|
||||
relative = [row["source_arrived_at"] - first for row in source_rows]
|
||||
if relative[-1] <= 0:
|
||||
raise ValueError("trace arrival window must be positive")
|
||||
source_rate = (len(relative) - 1) / relative[-1]
|
||||
scale = source_rate / target_rate
|
||||
return [value * scale for value in relative]
|
||||
|
||||
|
||||
def non_special_token_ids(model: Path, requests: int) -> list[int]:
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True)
|
||||
special = set(tokenizer.all_special_ids)
|
||||
candidates = [token for token in range(tokenizer.vocab_size) if token not in special]
|
||||
if len(candidates) < requests + 1:
|
||||
raise ValueError("tokenizer lacks enough non-special token IDs")
|
||||
return candidates[: requests + 1]
|
||||
|
||||
|
||||
def build_family_rows(
|
||||
source_rows: list[dict[str, Any]],
|
||||
family: Family,
|
||||
*,
|
||||
target_rate: float,
|
||||
fixed_token_ids: list[int],
|
||||
) -> list[dict[str, Any]]:
|
||||
arrivals = arrivals_for(source_rows, family.arrival, target_rate)
|
||||
mean_input = round(sum(row["input_length"] for row in source_rows) / len(source_rows))
|
||||
mean_output = round(sum(row["output_length"] for row in source_rows) / len(source_rows))
|
||||
if family.shape == "short-fixed":
|
||||
fixed_shape = (2048, 128)
|
||||
elif family.shape == "mean-fixed":
|
||||
fixed_shape = (mean_input, mean_output)
|
||||
elif family.shape == "heterogeneous":
|
||||
fixed_shape = None
|
||||
else:
|
||||
raise ValueError(f"unknown shape family: {family.shape}")
|
||||
|
||||
rows = []
|
||||
for index, (source, arrived_at) in enumerate(
|
||||
zip(source_rows, arrivals, strict=True)
|
||||
):
|
||||
if fixed_shape is None:
|
||||
input_length = source["input_length"]
|
||||
output_length = source["output_length"]
|
||||
body = copy.deepcopy(source["body"])
|
||||
body.update(
|
||||
{
|
||||
"min_tokens": output_length,
|
||||
"max_tokens": output_length,
|
||||
"ignore_eos": True,
|
||||
}
|
||||
)
|
||||
else:
|
||||
input_length, output_length = fixed_shape
|
||||
body = copy.deepcopy(source["body"])
|
||||
body.update(
|
||||
{
|
||||
"prompt": [
|
||||
fixed_token_ids[index + 1],
|
||||
*([fixed_token_ids[0]] * (input_length - 1)),
|
||||
],
|
||||
"min_tokens": output_length,
|
||||
"max_tokens": output_length,
|
||||
"ignore_eos": True,
|
||||
}
|
||||
)
|
||||
if input_length + output_length > 40960:
|
||||
raise ValueError(f"shape exceeds model limit at row {index}")
|
||||
rows.append(
|
||||
{
|
||||
"source_index": source["source_index"],
|
||||
"arrived_at": arrived_at,
|
||||
"input_length": input_length,
|
||||
"output_length": output_length,
|
||||
"session_id": source["session_id"] if family.prefix else index,
|
||||
"runtime_block_ids": (
|
||||
list(source["runtime_block_ids"]) if family.prefix else []
|
||||
),
|
||||
"body": body,
|
||||
}
|
||||
)
|
||||
return rows
|
||||
|
||||
|
||||
def write_case(
|
||||
root: Path,
|
||||
rows: list[dict[str, Any]],
|
||||
*,
|
||||
family: Family,
|
||||
rho: float,
|
||||
target_rate: float,
|
||||
reference_capacity: float,
|
||||
source_public: Path,
|
||||
source_private: Path,
|
||||
) -> dict[str, Any]:
|
||||
public_root = root / "public"
|
||||
private_root = root / "private"
|
||||
public_root.mkdir(parents=True, exist_ok=True)
|
||||
private_root.mkdir(parents=True, exist_ok=True)
|
||||
public_path = public_root / "frontier.csv"
|
||||
private_path = private_root / "real_requests.jsonl"
|
||||
|
||||
with public_path.open("w", newline="") as output:
|
||||
writer = csv.DictWriter(output, fieldnames=FIELDS, lineterminator="\n")
|
||||
writer.writeheader()
|
||||
for row in rows:
|
||||
writer.writerow(
|
||||
{
|
||||
"arrived_at": f"{row['arrived_at']:.12f}",
|
||||
"num_prefill_tokens": row["input_length"],
|
||||
"num_decode_tokens": row["output_length"],
|
||||
"session_id": row["session_id"],
|
||||
"block_hash_ids": "|".join(
|
||||
str(value) for value in row["runtime_block_ids"]
|
||||
),
|
||||
}
|
||||
)
|
||||
with private_path.open("w") as output:
|
||||
for row in rows:
|
||||
output.write(json.dumps(row, separators=(",", ":")) + "\n")
|
||||
|
||||
input_lengths = [row["input_length"] for row in rows]
|
||||
output_lengths = [row["output_length"] for row in rows]
|
||||
arrivals = [row["arrived_at"] for row in rows]
|
||||
empirical_rate = (len(rows) - 1) / (arrivals[-1] - arrivals[0])
|
||||
manifest = {
|
||||
"schema": "frontier-workload-regime-v1",
|
||||
"family": family.name,
|
||||
"shape_contract": family.shape,
|
||||
"arrival_contract": family.arrival,
|
||||
"prefix_caching": family.prefix,
|
||||
"requests": len(rows),
|
||||
"rho": rho,
|
||||
"reference_decode_tokens_per_second": reference_capacity,
|
||||
"global_offered_request_rate": target_rate,
|
||||
"empirical_interarrival_rate": empirical_rate,
|
||||
"decode_offered_tokens_per_second": target_rate
|
||||
* sum(output_lengths)
|
||||
/ len(output_lengths),
|
||||
"input_tokens": {
|
||||
"mean": sum(input_lengths) / len(input_lengths),
|
||||
"min": min(input_lengths),
|
||||
"max": max(input_lengths),
|
||||
},
|
||||
"output_tokens": {
|
||||
"mean": sum(output_lengths) / len(output_lengths),
|
||||
"min": min(output_lengths),
|
||||
"max": max(output_lengths),
|
||||
},
|
||||
"first_arrival_s": arrivals[0],
|
||||
"last_arrival_s": arrivals[-1],
|
||||
"source_public": str(source_public.resolve()),
|
||||
"source_public_sha256": sha256(source_public),
|
||||
"source_private": str(source_private.resolve()),
|
||||
"source_private_sha256": sha256(source_private),
|
||||
"public_csv": str(public_path.resolve()),
|
||||
"public_csv_sha256": sha256(public_path),
|
||||
"private_jsonl": str(private_path.resolve()),
|
||||
"private_jsonl_sha256": sha256(private_path),
|
||||
"row_vector_sha256": vector_sha256(rows),
|
||||
}
|
||||
manifest_path = public_root / "manifest.json"
|
||||
manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
|
||||
return manifest
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
rhos = args.rho or [0.05, 0.25, 0.50, 0.90, 1.20]
|
||||
if any(not math.isfinite(rho) or rho <= 0 for rho in rhos):
|
||||
raise ValueError("rho values must be finite and positive")
|
||||
if args.reference_decode_tokens_per_second <= 0:
|
||||
raise ValueError("reference decode capacity must be positive")
|
||||
|
||||
source_rows = load_source(args.source_public, args.source_private)
|
||||
if args.requests is not None:
|
||||
if args.requests < 2 or args.requests > len(source_rows):
|
||||
raise ValueError("requests must be between 2 and the source cohort size")
|
||||
source_rows = source_rows[: args.requests]
|
||||
fixed_token_ids = non_special_token_ids(args.model, len(source_rows))
|
||||
|
||||
cases = []
|
||||
for family in FAMILIES:
|
||||
if family.shape == "short-fixed":
|
||||
mean_output = 128.0
|
||||
elif family.shape == "mean-fixed":
|
||||
mean_output = round(
|
||||
sum(row["output_length"] for row in source_rows) / len(source_rows)
|
||||
)
|
||||
else:
|
||||
mean_output = sum(row["output_length"] for row in source_rows) / len(
|
||||
source_rows
|
||||
)
|
||||
for rho in rhos:
|
||||
target_rate = (
|
||||
rho * args.reference_decode_tokens_per_second / mean_output
|
||||
)
|
||||
rows = build_family_rows(
|
||||
source_rows,
|
||||
family,
|
||||
target_rate=target_rate,
|
||||
fixed_token_ids=fixed_token_ids,
|
||||
)
|
||||
rho_label = f"rho{rho:.2f}".replace(".", "p")
|
||||
root = args.output_root / family.name / rho_label
|
||||
cases.append(
|
||||
write_case(
|
||||
root,
|
||||
rows,
|
||||
family=family,
|
||||
rho=rho,
|
||||
target_rate=target_rate,
|
||||
reference_capacity=args.reference_decode_tokens_per_second,
|
||||
source_public=args.source_public,
|
||||
source_private=args.source_private,
|
||||
)
|
||||
)
|
||||
|
||||
experiment_manifest = {
|
||||
"schema": "frontier-workload-regime-suite-v1",
|
||||
"reference_decode_tokens_per_second": args.reference_decode_tokens_per_second,
|
||||
"rhos": rhos,
|
||||
"families": [family.__dict__ for family in FAMILIES],
|
||||
"cases": cases,
|
||||
}
|
||||
args.output_root.mkdir(parents=True, exist_ok=True)
|
||||
output = args.output_root / "manifest.json"
|
||||
output.write_text(json.dumps(experiment_manifest, indent=2, sort_keys=True) + "\n")
|
||||
print(output)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user