241 lines
9.4 KiB
Python
241 lines
9.4 KiB
Python
#!/usr/bin/env python3
|
|
"""Derive a per-GPU-normalized real/Frontier Qwen30 trace pair.
|
|
|
|
Arrival times change as ``t' = t / TP``. The private JSONL preserves the
|
|
complete prompt identity vector for real replay; the Frontier CSV exports only
|
|
the completed 16-token prompt blocks that vLLM can legally reuse from prefix
|
|
cache. Prompt bodies remain exclusively in the private output.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import hashlib
|
|
import json
|
|
import math
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
CSV_FIELDS = (
|
|
"arrived_at",
|
|
"num_prefill_tokens",
|
|
"num_decode_tokens",
|
|
"session_id",
|
|
"block_hash_ids",
|
|
)
|
|
|
|
|
|
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 update_digest(digest: Any, values: list[Any]) -> None:
|
|
digest.update(json.dumps(values, separators=(",", ":")).encode())
|
|
digest.update(b"\n")
|
|
|
|
|
|
def full_prompt_block_ids(
|
|
runtime_block_ids: list[int], input_tokens: int, block_size: int
|
|
) -> list[int]:
|
|
"""Project exact prompt identities to vLLM-cacheable full blocks only."""
|
|
if input_tokens <= 0 or block_size <= 0:
|
|
raise ValueError("input_tokens and block_size must be positive")
|
|
expected_runtime_blocks = math.ceil(input_tokens / block_size)
|
|
if len(runtime_block_ids) != expected_runtime_blocks:
|
|
raise ValueError(
|
|
"runtime block identity count does not match the exact prompt: "
|
|
f"got {len(runtime_block_ids)}, expected {expected_runtime_blocks}"
|
|
)
|
|
return runtime_block_ids[: input_tokens // block_size]
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--base-public-csv", type=Path, required=True)
|
|
parser.add_argument("--base-private-jsonl", type=Path, required=True)
|
|
parser.add_argument("--tp", type=int, required=True)
|
|
parser.add_argument("--output-root", type=Path, required=True)
|
|
parser.add_argument("--base-window-seconds", type=float, default=600.0)
|
|
parser.add_argument("--runtime-block-size", type=int, default=16)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
if (
|
|
args.tp <= 0
|
|
or args.runtime_block_size <= 0
|
|
or not math.isfinite(args.base_window_seconds)
|
|
or args.base_window_seconds <= 0
|
|
):
|
|
raise ValueError("TP, runtime block size, and base window must be positive")
|
|
|
|
public_rows = list(csv.DictReader(args.base_public_csv.open(newline="")))
|
|
private_rows = [
|
|
json.loads(line) for line in args.base_private_jsonl.open() if line.strip()
|
|
]
|
|
if not public_rows or len(public_rows) != len(private_rows):
|
|
raise ValueError("public/private request counts must match and be non-zero")
|
|
if set(public_rows[0]) != set(CSV_FIELDS):
|
|
raise ValueError("unexpected public trace schema")
|
|
|
|
normalized_digest = hashlib.sha256()
|
|
semantic_digest = hashlib.sha256()
|
|
normalized_private: list[dict[str, Any]] = []
|
|
normalized_public: list[dict[str, str]] = []
|
|
original_arrivals: list[float] = []
|
|
frontier_prefix_digest = hashlib.sha256()
|
|
partial_prompt_count = 0
|
|
raw_runtime_block_count = 0
|
|
frontier_full_block_count = 0
|
|
|
|
for index, (public, private) in enumerate(
|
|
zip(public_rows, private_rows, strict=True)
|
|
):
|
|
original_arrival = float(public["arrived_at"])
|
|
if not math.isfinite(original_arrival) or original_arrival < 0:
|
|
raise ValueError(f"invalid arrival at row {index}")
|
|
if not math.isclose(original_arrival, float(private["arrived_at"]), abs_tol=1e-9):
|
|
raise ValueError(f"public/private arrival mismatch at row {index}")
|
|
runtime_block_ids = [int(value) for value in private["runtime_block_ids"]]
|
|
expected = (
|
|
int(public["num_prefill_tokens"]),
|
|
int(public["num_decode_tokens"]),
|
|
int(public["session_id"]),
|
|
public["block_hash_ids"],
|
|
)
|
|
actual = (
|
|
int(private["input_length"]),
|
|
int(private["output_length"]),
|
|
int(private["session_id"]),
|
|
"|".join(str(value) for value in runtime_block_ids),
|
|
)
|
|
if expected != actual:
|
|
raise ValueError(f"public/private payload mismatch at row {index}")
|
|
|
|
input_tokens = int(private["input_length"])
|
|
frontier_block_ids = full_prompt_block_ids(
|
|
runtime_block_ids, input_tokens, args.runtime_block_size
|
|
)
|
|
if not frontier_block_ids:
|
|
raise ValueError(
|
|
"Frontier prefix-cache replay requires at least one complete "
|
|
f"runtime block; request {index} has ISL={input_tokens}"
|
|
)
|
|
partial_prompt_count += int(input_tokens % args.runtime_block_size != 0)
|
|
raw_runtime_block_count += len(runtime_block_ids)
|
|
frontier_full_block_count += len(frontier_block_ids)
|
|
|
|
normalized_arrival = original_arrival / args.tp
|
|
updated = dict(private)
|
|
updated["arrived_at"] = normalized_arrival
|
|
normalized_private.append(updated)
|
|
normalized_public.append(
|
|
{
|
|
"arrived_at": f"{normalized_arrival:.12f}",
|
|
"num_prefill_tokens": public["num_prefill_tokens"],
|
|
"num_decode_tokens": public["num_decode_tokens"],
|
|
"session_id": public["session_id"],
|
|
"block_hash_ids": "|".join(str(value) for value in frontier_block_ids),
|
|
}
|
|
)
|
|
update_digest(
|
|
semantic_digest,
|
|
[
|
|
int(private["source_index"]),
|
|
int(private["input_length"]),
|
|
int(private["output_length"]),
|
|
int(private["session_id"]),
|
|
runtime_block_ids,
|
|
],
|
|
)
|
|
update_digest(
|
|
normalized_digest,
|
|
[
|
|
int(private["source_index"]),
|
|
normalized_arrival,
|
|
int(private["input_length"]),
|
|
int(private["output_length"]),
|
|
int(private["session_id"]),
|
|
runtime_block_ids,
|
|
],
|
|
)
|
|
update_digest(
|
|
frontier_prefix_digest,
|
|
[
|
|
int(private["source_index"]),
|
|
normalized_arrival,
|
|
input_tokens,
|
|
frontier_block_ids,
|
|
],
|
|
)
|
|
original_arrivals.append(original_arrival)
|
|
|
|
if any(
|
|
right < left for left, right in zip(original_arrivals, original_arrivals[1:])
|
|
):
|
|
raise ValueError("base arrival order is not monotonic")
|
|
|
|
public_dir = args.output_root / "public"
|
|
private_dir = args.output_root / "private"
|
|
public_dir.mkdir(parents=True, exist_ok=True)
|
|
private_dir.mkdir(parents=True, exist_ok=True)
|
|
public_path = public_dir / "frontier.csv"
|
|
private_path = private_dir / "real_requests.jsonl"
|
|
with public_path.open("w", newline="") as output:
|
|
writer = csv.DictWriter(output, fieldnames=CSV_FIELDS, lineterminator="\n")
|
|
writer.writeheader()
|
|
writer.writerows(normalized_public)
|
|
with private_path.open("w") as output:
|
|
for row in normalized_private:
|
|
output.write(json.dumps(row, separators=(",", ":")) + "\n")
|
|
|
|
manifest = {
|
|
"schema": "qwen30-tp-normalized-trace-v2",
|
|
"privacy": "prompt text exists only under private/ and must not be harvested",
|
|
"transform": "arrival_seconds_prime = arrival_seconds / tensor_parallel_size",
|
|
"tensor_parallel_size": args.tp,
|
|
"requests": len(normalized_private),
|
|
"base_window_seconds": args.base_window_seconds,
|
|
"global_offered_request_rate": len(normalized_private)
|
|
/ args.base_window_seconds
|
|
* args.tp,
|
|
"per_gpu_offered_request_rate": len(normalized_private)
|
|
/ args.base_window_seconds,
|
|
"frontier_prefix_block_projection": {
|
|
"runtime_block_size": args.runtime_block_size,
|
|
"rule": "full_prompt_blocks_only=floor(input_tokens/runtime_block_size)",
|
|
"partial_prompt_count": partial_prompt_count,
|
|
"raw_runtime_block_count": raw_runtime_block_count,
|
|
"frontier_full_block_count": frontier_full_block_count,
|
|
},
|
|
"original_first_arrival_s": original_arrivals[0],
|
|
"original_last_arrival_s": original_arrivals[-1],
|
|
"normalized_first_arrival_s": original_arrivals[0] / args.tp,
|
|
"normalized_last_arrival_s": original_arrivals[-1] / args.tp,
|
|
"base_public_csv": str(args.base_public_csv.resolve()),
|
|
"base_public_csv_sha256": sha256(args.base_public_csv),
|
|
"base_private_jsonl": str(args.base_private_jsonl.resolve()),
|
|
"base_private_jsonl_sha256": sha256(args.base_private_jsonl),
|
|
"public_csv": str(public_path.resolve()),
|
|
"public_csv_sha256": sha256(public_path),
|
|
"private_jsonl": str(private_path.resolve()),
|
|
"private_jsonl_sha256": sha256(private_path),
|
|
"semantic_vector_sha256": semantic_digest.hexdigest(),
|
|
"normalized_row_vector_sha256": normalized_digest.hexdigest(),
|
|
"frontier_prefix_vector_sha256": frontier_prefix_digest.hexdigest(),
|
|
}
|
|
manifest_path = public_dir / "manifest.json"
|
|
manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
|
|
print(json.dumps(manifest, sort_keys=True))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|