303 lines
12 KiB
Python
303 lines
12 KiB
Python
#!/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()
|