Add batch-aware profile and exact trace preparation
This commit is contained in:
302
runs/frontier-fidelity-envelope-v1/prepare_exact_trace.py
Normal file
302
runs/frontier-fidelity-envelope-v1/prepare_exact_trace.py
Normal file
@@ -0,0 +1,302 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Prepare exact Qwen trace replays with block-16 prefix identities.
|
||||
|
||||
Prompt text is written only below ``private/``. Public manifests and Frontier
|
||||
fixtures contain lengths, arrivals, session IDs, and deterministic block IDs.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import platform
|
||||
import socket
|
||||
import time
|
||||
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 threshold_name(value: float) -> str:
|
||||
return f"u{value:.12g}".replace(".", "p")
|
||||
|
||||
|
||||
def token_payload(tokens: list[int]) -> bytes:
|
||||
return len(tokens).to_bytes(2, "little") + b"".join(
|
||||
int(token).to_bytes(4, "little", signed=False) for token in tokens
|
||||
)
|
||||
|
||||
|
||||
def block_identity_records(
|
||||
token_ids: list[int], block_size: int
|
||||
) -> list[tuple[int, bytes]]:
|
||||
"""Return parent-sensitive identities and independent collision witnesses."""
|
||||
parent = b"FRONTIER_EXACT_TRACE_ROOT"
|
||||
records = []
|
||||
for start in range(0, len(token_ids), block_size):
|
||||
payload = token_payload(token_ids[start : start + block_size])
|
||||
identity_input = parent + b"\0" + payload
|
||||
parent = hashlib.blake2b(identity_input, digest_size=16).digest()
|
||||
records.append(
|
||||
(
|
||||
int.from_bytes(parent, "big", signed=False),
|
||||
hashlib.sha256(identity_input).digest(),
|
||||
)
|
||||
)
|
||||
return records
|
||||
|
||||
|
||||
def block_identities(token_ids: list[int], block_size: int) -> list[int]:
|
||||
"""Return parent-sensitive identities with the same prefix equivalence as vLLM."""
|
||||
return [identity for identity, _ in block_identity_records(token_ids, block_size)]
|
||||
|
||||
|
||||
def root_sessions(rows: list[dict[str, Any]]) -> dict[int, int]:
|
||||
roots: dict[int, int] = {}
|
||||
for row in rows:
|
||||
chat_id = int(row["chat_id"])
|
||||
parent = int(row["parent_chat_id"])
|
||||
roots[chat_id] = chat_id if parent == -1 else roots.get(parent, parent)
|
||||
return roots
|
||||
|
||||
|
||||
def update_digest(digest: Any, values: list[Any]) -> None:
|
||||
digest.update(json.dumps(values, separators=(",", ":")).encode())
|
||||
digest.update(b"\n")
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--trace", type=Path, required=True)
|
||||
parser.add_argument("--model", type=Path, required=True)
|
||||
parser.add_argument("--output-root", type=Path, required=True)
|
||||
parser.add_argument("--served-model-name", default="qwen3-30b-a3b")
|
||||
parser.add_argument("--sampling-u-max", type=float, action="append", required=True)
|
||||
parser.add_argument("--max-model-len", type=int, default=40960)
|
||||
parser.add_argument("--source-block-size", type=int, default=64)
|
||||
parser.add_argument("--runtime-block-size", type=int, default=16)
|
||||
parser.add_argument("--batch-size", type=int, default=16)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
thresholds = sorted(set(args.sampling_u_max))
|
||||
if not thresholds or thresholds[0] < 0 or thresholds[-1] > 1:
|
||||
raise ValueError("sampling thresholds must be in [0, 1]")
|
||||
if args.source_block_size % args.runtime_block_size:
|
||||
raise ValueError("source block size must be divisible by runtime block size")
|
||||
if min(args.max_model_len, args.source_block_size, args.runtime_block_size, args.batch_size) <= 0:
|
||||
raise ValueError("length and batch arguments must be positive")
|
||||
|
||||
import transformers
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
rows = [json.loads(line) for line in args.trace.open() if line.strip()]
|
||||
roots = root_sessions(rows)
|
||||
eligible = [
|
||||
(index, row)
|
||||
for index, row in enumerate(rows)
|
||||
if int(row["input_length"]) + int(row["output_length"]) <= args.max_model_len
|
||||
and int(row["output_length"]) > 0
|
||||
]
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
|
||||
prepared: list[dict[str, Any]] = []
|
||||
identity_to_digest: dict[int, bytes] = {}
|
||||
source_to_runtime: dict[int, tuple[int, ...]] = {}
|
||||
runtime_to_source: dict[tuple[int, ...], int] = {}
|
||||
identity_collision_count = 0
|
||||
source_to_runtime_conflicts = 0
|
||||
runtime_to_source_conflicts = 0
|
||||
length_mismatches = 0
|
||||
source_hash_count_mismatches = 0
|
||||
started = time.time()
|
||||
|
||||
for start in range(0, len(eligible), args.batch_size):
|
||||
batch = eligible[start : start + args.batch_size]
|
||||
encoded = tokenizer(
|
||||
[row["prompt"] for _, row in batch],
|
||||
add_special_tokens=False,
|
||||
padding=False,
|
||||
truncation=False,
|
||||
)["input_ids"]
|
||||
for (source_index, row), token_ids in zip(batch, encoded, strict=True):
|
||||
token_ids = [int(token) for token in token_ids]
|
||||
if len(token_ids) != int(row["input_length"]):
|
||||
length_mismatches += 1
|
||||
source_hashes = [int(value) for value in row["hash_ids"]]
|
||||
if len(source_hashes) != math.ceil(len(token_ids) / args.source_block_size):
|
||||
source_hash_count_mismatches += 1
|
||||
identity_records = block_identity_records(token_ids, args.runtime_block_size)
|
||||
runtime_ids = [identity for identity, _ in identity_records]
|
||||
for runtime_id, witness in identity_records:
|
||||
previous = identity_to_digest.setdefault(runtime_id, witness)
|
||||
identity_collision_count += int(previous != witness)
|
||||
|
||||
blocks_per_source = args.source_block_size // args.runtime_block_size
|
||||
for block_index, source_id in enumerate(source_hashes):
|
||||
begin = block_index * blocks_per_source
|
||||
relation = tuple(runtime_ids[begin : begin + blocks_per_source])
|
||||
previous_relation = source_to_runtime.setdefault(source_id, relation)
|
||||
source_to_runtime_conflicts += int(previous_relation != relation)
|
||||
previous_source = runtime_to_source.setdefault(relation, source_id)
|
||||
runtime_to_source_conflicts += int(previous_source != source_id)
|
||||
|
||||
prepared.append(
|
||||
{
|
||||
"source_index": source_index,
|
||||
"chat_id": int(row["chat_id"]),
|
||||
"session_id": roots[int(row["chat_id"])],
|
||||
"arrival": float(row["timestamp"]),
|
||||
"input_length": int(row["input_length"]),
|
||||
"output_length": int(row["output_length"]),
|
||||
"sampling_u": float(row["sampling_u"]),
|
||||
"prompt": row["prompt"],
|
||||
"runtime_block_ids": runtime_ids,
|
||||
}
|
||||
)
|
||||
|
||||
args.output_root.mkdir(parents=True, exist_ok=True)
|
||||
threshold_manifests = []
|
||||
for threshold in thresholds:
|
||||
name = threshold_name(threshold)
|
||||
public = args.output_root / "public" / name
|
||||
private = args.output_root / "private" / name
|
||||
public.mkdir(parents=True, exist_ok=True)
|
||||
private.mkdir(parents=True, exist_ok=True)
|
||||
selected = [row for row in prepared if row["sampling_u"] <= threshold]
|
||||
frontier_path = public / "frontier.csv"
|
||||
real_path = private / "real_requests.jsonl"
|
||||
vector_digest = hashlib.sha256()
|
||||
with frontier_path.open("w", newline="") as output:
|
||||
writer = csv.DictWriter(output, fieldnames=CSV_FIELDS, lineterminator="\n")
|
||||
writer.writeheader()
|
||||
for row in selected:
|
||||
writer.writerow(
|
||||
{
|
||||
"arrived_at": f"{row['arrival']:.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"]
|
||||
),
|
||||
}
|
||||
)
|
||||
update_digest(
|
||||
vector_digest,
|
||||
[
|
||||
row["source_index"],
|
||||
row["arrival"],
|
||||
row["input_length"],
|
||||
row["output_length"],
|
||||
row["session_id"],
|
||||
row["runtime_block_ids"],
|
||||
],
|
||||
)
|
||||
with real_path.open("w") as output:
|
||||
for row in selected:
|
||||
output.write(
|
||||
json.dumps(
|
||||
{
|
||||
"source_index": row["source_index"],
|
||||
"arrived_at": row["arrival"],
|
||||
"input_length": row["input_length"],
|
||||
"output_length": row["output_length"],
|
||||
"session_id": row["session_id"],
|
||||
"runtime_block_ids": row["runtime_block_ids"],
|
||||
"body": {
|
||||
"model": args.served_model_name,
|
||||
"prompt": row["prompt"],
|
||||
"min_tokens": row["output_length"],
|
||||
"max_tokens": row["output_length"],
|
||||
"ignore_eos": True,
|
||||
"stream": True,
|
||||
},
|
||||
},
|
||||
separators=(",", ":"),
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
threshold_manifests.append(
|
||||
{
|
||||
"sampling_u_max": threshold,
|
||||
"selected_requests": len(selected),
|
||||
"request_rate": len(selected) / 600.0,
|
||||
"frontier_csv": str(frontier_path.resolve()),
|
||||
"frontier_csv_sha256": sha256(frontier_path),
|
||||
"real_requests_private": str(real_path.resolve()),
|
||||
"real_requests_private_sha256": sha256(real_path),
|
||||
"row_vector_sha256": vector_digest.hexdigest(),
|
||||
}
|
||||
)
|
||||
|
||||
failures = {
|
||||
"input_length_mismatches": length_mismatches,
|
||||
"source_hash_count_mismatches": source_hash_count_mismatches,
|
||||
"runtime_identity_collisions": identity_collision_count,
|
||||
"source_to_runtime_relation_conflicts": source_to_runtime_conflicts,
|
||||
"runtime_to_source_relation_conflicts": runtime_to_source_conflicts,
|
||||
}
|
||||
manifest = {
|
||||
"schema": "qwen30-exact-trace-block16-v1",
|
||||
"status": "pass" if not any(failures.values()) else "fail",
|
||||
"execution": {
|
||||
"host": socket.gethostname(),
|
||||
"python": platform.python_version(),
|
||||
"transformers": transformers.__version__,
|
||||
"tokenizer": type(tokenizer).__name__,
|
||||
"elapsed_seconds": round(time.time() - started, 3),
|
||||
},
|
||||
"source": {
|
||||
"trace": str(args.trace.resolve()),
|
||||
"trace_sha256": sha256(args.trace),
|
||||
"requests": len(rows),
|
||||
"eligible_requests": len(eligible),
|
||||
"model": str(args.model.resolve()),
|
||||
"max_model_len": args.max_model_len,
|
||||
},
|
||||
"block_contract": {
|
||||
"source_block_size": args.source_block_size,
|
||||
"runtime_block_size": args.runtime_block_size,
|
||||
"identity": "BLAKE2b-128(parent runtime identity, exact token-id block)",
|
||||
"unique_runtime_identities": len(identity_to_digest),
|
||||
"unique_source_relations": len(source_to_runtime),
|
||||
"failures": failures,
|
||||
},
|
||||
"selection_contract": (
|
||||
"eligible universe followed by source sampling_u threshold; no length "
|
||||
"selection or output override; original arrivals/order preserved"
|
||||
),
|
||||
"thresholds": threshold_manifests,
|
||||
"privacy": "prompt text exists only under private/ and must not be harvested",
|
||||
}
|
||||
manifest_path = args.output_root / "public" / "manifest.json"
|
||||
manifest_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
|
||||
print(json.dumps({"status": manifest["status"], "failures": failures, "thresholds": threshold_manifests}, sort_keys=True))
|
||||
if manifest["status"] != "pass":
|
||||
raise SystemExit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user