135 lines
4.9 KiB
Python
135 lines
4.9 KiB
Python
#!/usr/bin/env python3
|
|
"""Project base workload cases to a constant per-GPU request-rate contract."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--source-manifest", type=Path, required=True)
|
|
parser.add_argument("--output-root", type=Path, required=True)
|
|
parser.add_argument("--tp", type=int, action="append", default=None)
|
|
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 rho_label(rho: float) -> str:
|
|
return format(rho, ".12g").replace(".", "p")
|
|
|
|
|
|
def project_case(
|
|
case: dict[str, Any], output_root: Path, tp: int
|
|
) -> dict[str, Any]:
|
|
source_public = Path(case["public_csv"])
|
|
source_private = Path(case["private_jsonl"])
|
|
if sha256(source_public) != case["public_csv_sha256"]:
|
|
raise ValueError(f"source public digest mismatch: {source_public}")
|
|
if sha256(source_private) != case["private_jsonl_sha256"]:
|
|
raise ValueError(f"source private digest mismatch: {source_private}")
|
|
|
|
case_root = output_root / f"tp{tp}" / case["family"] / f"rho{rho_label(case['rho'])}"
|
|
public_root = case_root / "public"
|
|
private_root = case_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 source_public.open(newline="") as source:
|
|
reader = csv.DictReader(source)
|
|
rows = list(reader)
|
|
fieldnames = reader.fieldnames
|
|
if not rows or fieldnames is None:
|
|
raise ValueError(f"empty public source: {source_public}")
|
|
with public_path.open("w", newline="") as output:
|
|
writer = csv.DictWriter(output, fieldnames=fieldnames, lineterminator="\n")
|
|
writer.writeheader()
|
|
for row in rows:
|
|
row["arrived_at"] = f"{float(row['arrived_at']) / tp:.12f}"
|
|
writer.writerow(row)
|
|
|
|
private_rows = [
|
|
json.loads(line) for line in source_private.open() if line.strip()
|
|
]
|
|
if len(private_rows) != len(rows):
|
|
raise ValueError(f"public/private row mismatch: {source_public}")
|
|
with private_path.open("w") as output:
|
|
for row in private_rows:
|
|
row["arrived_at"] = float(row["arrived_at"]) / tp
|
|
output.write(json.dumps(row, separators=(",", ":")) + "\n")
|
|
|
|
projected = dict(case)
|
|
projected.update(
|
|
{
|
|
"load_contract": "constant_per_gpu_request_rate",
|
|
"tp": tp,
|
|
"rho_per_gpu": case["rho"],
|
|
"per_gpu_offered_request_rate": case["global_offered_request_rate"],
|
|
"global_offered_request_rate": case["global_offered_request_rate"] * tp,
|
|
"empirical_interarrival_rate": case["empirical_interarrival_rate"] * tp,
|
|
"decode_offered_tokens_per_second_per_gpu": case[
|
|
"decode_offered_tokens_per_second"
|
|
],
|
|
"decode_offered_tokens_per_second": case[
|
|
"decode_offered_tokens_per_second"
|
|
]
|
|
* tp,
|
|
"last_arrival_s": case["last_arrival_s"] / tp,
|
|
"source_case_public_csv": str(source_public),
|
|
"source_case_public_csv_sha256": case["public_csv_sha256"],
|
|
"source_case_private_jsonl": str(source_private),
|
|
"source_case_private_jsonl_sha256": case["private_jsonl_sha256"],
|
|
"public_csv": str(public_path.resolve()),
|
|
"public_csv_sha256": sha256(public_path),
|
|
"private_jsonl": str(private_path.resolve()),
|
|
"private_jsonl_sha256": sha256(private_path),
|
|
}
|
|
)
|
|
(public_root / "manifest.json").write_text(
|
|
json.dumps(projected, indent=2, sort_keys=True) + "\n"
|
|
)
|
|
return projected
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
tps = args.tp or [1, 2, 4]
|
|
if any(tp not in (1, 2, 4) for tp in tps):
|
|
raise ValueError("TP must be 1, 2, or 4")
|
|
suite = json.loads(args.source_manifest.read_text())
|
|
projected = [
|
|
project_case(case, args.output_root, tp)
|
|
for tp in tps
|
|
for case in suite["cases"]
|
|
]
|
|
output = {
|
|
"schema": "frontier-workload-regime-per-gpu-v1",
|
|
"load_contract": "constant_per_gpu_request_rate",
|
|
"source_manifest": str(args.source_manifest.resolve()),
|
|
"source_manifest_sha256": sha256(args.source_manifest),
|
|
"tensor_parallel_sizes": tps,
|
|
"cases": projected,
|
|
}
|
|
args.output_root.mkdir(parents=True, exist_ok=True)
|
|
path = args.output_root / "manifest.json"
|
|
path.write_text(json.dumps(output, indent=2, sort_keys=True) + "\n")
|
|
print(path)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|