#!/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()