123 lines
3.8 KiB
Python
123 lines
3.8 KiB
Python
#!/usr/bin/env python3
|
|
"""Materialize deterministic, prefix-disjoint fixed-shape Frontier traces."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import hashlib
|
|
import json
|
|
import math
|
|
from pathlib import Path
|
|
|
|
|
|
FIELDS = (
|
|
"arrived_at",
|
|
"num_prefill_tokens",
|
|
"num_decode_tokens",
|
|
"session_id",
|
|
"block_hash_ids",
|
|
)
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--output-root", type=Path, required=True)
|
|
parser.add_argument("--input-tokens", type=int, required=True)
|
|
parser.add_argument("--output-tokens", type=int, required=True)
|
|
parser.add_argument("--requests", type=int, default=64)
|
|
parser.add_argument("--rate", type=float, action="append", required=True)
|
|
parser.add_argument("--block-size", type=int, default=16)
|
|
return parser.parse_args()
|
|
|
|
|
|
def sha256(path: Path) -> str:
|
|
return hashlib.sha256(path.read_bytes()).hexdigest()
|
|
|
|
|
|
def rate_key(rate: float) -> str:
|
|
return f"r{rate:g}".replace(".", "p")
|
|
|
|
|
|
def materialize(
|
|
root: Path,
|
|
*,
|
|
input_tokens: int,
|
|
output_tokens: int,
|
|
requests: int,
|
|
rates: list[float],
|
|
block_size: int,
|
|
) -> dict[str, object]:
|
|
values = [input_tokens, output_tokens, requests, block_size]
|
|
if any(value <= 0 for value in values):
|
|
raise ValueError("token counts, requests, and block size must be positive")
|
|
if input_tokens + output_tokens > 40960:
|
|
raise ValueError("shape exceeds the frozen max model length")
|
|
if not rates or any(rate <= 0 or not math.isfinite(rate) for rate in rates):
|
|
raise ValueError("rates must be positive and finite")
|
|
if len(set(rates)) != len(rates):
|
|
raise ValueError("rates must be unique")
|
|
root.mkdir(parents=True, exist_ok=True)
|
|
anchors = []
|
|
blocks = math.ceil(input_tokens / block_size)
|
|
for rate in rates:
|
|
path = root / f"{rate_key(rate)}.csv"
|
|
with path.open("w", newline="") as output:
|
|
writer = csv.DictWriter(output, fieldnames=FIELDS, lineterminator="\n")
|
|
writer.writeheader()
|
|
for request_id in range(requests):
|
|
first_block = request_id * blocks + 1
|
|
writer.writerow(
|
|
{
|
|
"arrived_at": f"{request_id / rate:.12f}",
|
|
"num_prefill_tokens": input_tokens,
|
|
"num_decode_tokens": output_tokens,
|
|
"session_id": request_id,
|
|
"block_hash_ids": "|".join(
|
|
str(first_block + offset) for offset in range(blocks)
|
|
),
|
|
}
|
|
)
|
|
anchors.append(
|
|
{
|
|
"label": rate_key(rate),
|
|
"rate": rate,
|
|
"path": path.name,
|
|
"sha256": sha256(path),
|
|
}
|
|
)
|
|
manifest = {
|
|
"schema": "qwen30-fixed-frontier-traces-v1",
|
|
"contract": {
|
|
"input_tokens": input_tokens,
|
|
"output_tokens": output_tokens,
|
|
"requests": requests,
|
|
"block_size": block_size,
|
|
"arrival": "open_loop_uniform",
|
|
"prefix_caching": False,
|
|
"prefix_relation": "all request block identities are disjoint",
|
|
},
|
|
"anchors": anchors,
|
|
}
|
|
(root / "manifest.json").write_text(
|
|
json.dumps(manifest, indent=2, sort_keys=True) + "\n"
|
|
)
|
|
return manifest
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
manifest = materialize(
|
|
args.output_root.resolve(),
|
|
input_tokens=args.input_tokens,
|
|
output_tokens=args.output_tokens,
|
|
requests=args.requests,
|
|
rates=args.rate,
|
|
block_size=args.block_size,
|
|
)
|
|
print(json.dumps(manifest, sort_keys=True))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|