413 lines
15 KiB
Python
413 lines
15 KiB
Python
#!/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,
|
|
help="Optional Frontier CSV; otherwise project complete blocks from private rows.",
|
|
)
|
|
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 rho_label(rho: float) -> str:
|
|
return f"rho{format(rho, '.12g').replace('.', 'p')}"
|
|
|
|
|
|
def load_source(
|
|
public_path: Path | None, private_path: Path
|
|
) -> list[dict[str, Any]]:
|
|
private_rows = [json.loads(line) for line in private_path.open() if line.strip()]
|
|
if not private_rows:
|
|
raise ValueError("empty private source")
|
|
if public_path is None:
|
|
public_rows = [
|
|
{
|
|
"arrived_at": row["arrived_at"],
|
|
"num_prefill_tokens": row["input_length"],
|
|
"num_decode_tokens": row["output_length"],
|
|
"block_hash_ids": "|".join(
|
|
str(value)
|
|
for value in row["runtime_block_ids"][
|
|
: int(row["input_length"]) // 16
|
|
]
|
|
),
|
|
}
|
|
for row in private_rows
|
|
]
|
|
else:
|
|
with public_path.open(newline="") as source:
|
|
public_rows = list(csv.DictReader(source))
|
|
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 | None,
|
|
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()) if source_public is not None else None
|
|
),
|
|
"source_public_sha256": (
|
|
sha256(source_public) if source_public is not None else None
|
|
),
|
|
"frontier_block_projection": (
|
|
"source_public_complete_blocks"
|
|
if source_public is not None
|
|
else "private_runtime_ids_first_floor_isl_div_16_blocks"
|
|
),
|
|
"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,
|
|
)
|
|
root = args.output_root / family.name / rho_label(rho)
|
|
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()
|