#!/usr/bin/env python3 """Remap source hashes to real-token 16-token blocks and replay prompts.""" from __future__ import annotations import argparse import csv import hashlib import json import math from contextlib import ExitStack from itertools import zip_longest from pathlib import Path from typing import Any, Iterator from trace_utils import ( TARGET_BLOCK_SIZE, expand_hash_ids, input_length, iter_jsonl, output_length, parse_hash_ids, session_id, sha256, synthetic_tokens, token_block_identity_records, timestamp, write_json, ) CSV_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("--input", type=Path, required=True) parser.add_argument("--prompt", type=Path) parser.add_argument("--tokenizer", type=Path) parser.add_argument("--input-is-remapped", action="store_true") parser.add_argument("--output-root", type=Path, required=True) parser.add_argument("--rho", type=float, default=1.0) parser.add_argument("--start-timestamp", type=float) parser.add_argument("--duration-s", type=float) parser.add_argument("--vocab-size", type=int, default=151936) parser.add_argument("--token-offset", type=int, default=1024) parser.add_argument("--served-model", default="qwen3-30b-s3-real") parser.add_argument("--source-block-size", type=int, default=64) parser.add_argument( "--workload-mode", choices=("prefill_decode", "prefill_only"), default="prefill_decode", ) parser.add_argument("--validate-parents", action="store_true") parser.add_argument("--max-total-tokens", type=int) parser.add_argument("--frontier-only", action="store_true") return parser.parse_args() def keep_row(row: dict[str, Any], rho: float) -> bool: if not 0 < rho <= 1: raise ValueError("rho must be in (0, 1]") if rho == 1: return True if "sampling_u" not in row: raise ValueError("sampling_u is required for rho < 1") return float(row["sampling_u"]) <= rho def row_identity(row: dict[str, Any], source_index: int) -> Any: return row.get("chat_id", source_index) def iter_source_prompt_rows( source: Path, prompt: Path | None, input_is_remapped: bool ) -> Iterator[tuple[int, dict[str, Any], dict[str, Any] | None]]: if input_is_remapped: for source_index, row in enumerate(iter_jsonl(source)): yield source_index, row, None return if prompt is None: for source_index, row in enumerate(iter_jsonl(source)): yield source_index, row, None return with source.open() as source_stream, prompt.open() as prompt_stream: for source_index, pair in enumerate(zip_longest(source_stream, prompt_stream)): source_line, prompt_line = pair if source_line is None or prompt_line is None: raise ValueError("trace/prompt window row-count mismatch") if not source_line.strip() and not prompt_line.strip(): continue if not source_line.strip() or not prompt_line.strip(): raise ValueError(f"trace/prompt blank-line mismatch at row {source_index + 1}") row = json.loads(source_line) prompt_row = json.loads(prompt_line) if ( row.get("chat_id") != prompt_row.get("chat_id") or row.get("turn") != prompt_row.get("turn") ): raise ValueError(f"trace/prompt alignment failed at row {source_index + 1}") yield source_index, row, prompt_row def load_tokenizer(path: Path | None) -> Any: if path is None: raise ValueError("--tokenizer is required for real prompt tokenization") from transformers import AutoTokenizer return AutoTokenizer.from_pretrained(path, trust_remote_code=True, local_files_only=True) def prompt_token_ids(tokenizer: Any, prompt: str) -> list[int]: encoded = tokenizer( prompt, add_special_tokens=False, padding=False, truncation=False, )["input_ids"] return [int(token) for token in encoded] def materialize(args: argparse.Namespace, *, tokenizer: Any | None = None) -> dict[str, Any]: if args.output_root.exists(): raise ValueError(f"refusing to overwrite {args.output_root}") args.output_root.mkdir(parents=True) frontier_path = args.output_root / "frontier.csv" real_path = args.output_root / "real_requests.jsonl" selected_path = args.output_root / "selected-remapped.jsonl" start = args.start_timestamp end = start + args.duration_s if start is not None and args.duration_s is not None else None first_selected_timestamp: float | None = None request_count = 0 source_rows = 0 prefill_tokens = 0 decode_tokens = 0 target_blocks = 0 source_blocks = 0 parent_sequences: dict[Any, list[int]] = {} parent_links_checked = 0 parent_links_outside = 0 parent_prefix_common_blocks = 0 parent_tail_nonreused_blocks = 0 excluded_over_max_tokens = 0 real_prompt_requests = 0 synthetic_fallback_requests = 0 input_length_mismatches = 0 source_hash_count_mismatches = 0 identity_collisions = 0 source_to_runtime_conflicts = 0 runtime_to_source_conflicts = 0 identity_to_witness: dict[int, bytes] = {} source_to_runtime: dict[int, tuple[int, ...]] = {} runtime_to_source: dict[tuple[int, ...], int] = {} row_vector_digest = hashlib.sha256() upstream_manifest_path: Path | None = None upstream_block_contract: dict[str, Any] | None = None source_block_size = int(getattr(args, "source_block_size", 64)) workload_mode = str(getattr(args, "workload_mode", "prefill_decode")) if source_block_size <= 0 or source_block_size % TARGET_BLOCK_SIZE: raise ValueError("source block size must be a positive multiple of 16") subblocks_per_source = source_block_size // TARGET_BLOCK_SIZE if not args.input_is_remapped and tokenizer is None and args.prompt is not None: tokenizer = load_tokenizer(args.tokenizer) if args.input_is_remapped: upstream_manifest_path = args.input.parent / "manifest.json" if not upstream_manifest_path.is_file(): raise FileNotFoundError(f"pre-remapped input requires sibling manifest: {upstream_manifest_path}") upstream = json.loads(upstream_manifest_path.read_text()) source_block_size = int(upstream["source_block_size"]) subblocks_per_source = source_block_size // TARGET_BLOCK_SIZE if int(upstream["target_block_size"]) != TARGET_BLOCK_SIZE: raise ValueError("upstream remap target block size is not 16") upstream_block_contract = dict(upstream["block_contract"]) with ExitStack() as stack: frontier_stream = stack.enter_context(frontier_path.open("w", newline="")) selected_stream = ( None if args.input_is_remapped and args.frontier_only else stack.enter_context(selected_path.open("w")) ) real_stream = None if args.frontier_only else stack.enter_context(real_path.open("w")) writer = csv.DictWriter(frontier_stream, fieldnames=CSV_FIELDS, lineterminator="\n") writer.writeheader() for source_index, row, prompt_row in iter_source_prompt_rows( args.input, args.prompt, args.input_is_remapped ): source_rows += 1 ts = timestamp(row) if start is not None and ts < start: continue if end is not None and ts >= end: continue if not keep_row(row, args.rho): continue isl = input_length(row) osl = 1 if workload_mode == "prefill_only" else output_length(row) if args.max_total_tokens is not None and isl + osl > args.max_total_tokens: excluded_over_max_tokens += 1 continue original_hashes = parse_hash_ids(row.get("hash_ids")) if args.input_is_remapped: block_ids = [int(value) for value in row["hash_ids_16"]] if len(block_ids) != math.ceil(isl / TARGET_BLOCK_SIZE): raise ValueError(f"pre-remapped 16-block count mismatch for chat_id={row.get('chat_id')!r}") if len(original_hashes) != math.ceil(isl / source_block_size): raise ValueError( f"pre-remapped {source_block_size}-block count mismatch " f"for chat_id={row.get('chat_id')!r}" ) request_prompt = row["prompt"] prompt_kind = str(row["prompt_source"]) is_real_prompt = prompt_kind in ("real_text", "real_tokens") real_prompt_requests += int(is_real_prompt) synthetic_fallback_requests += int(not is_real_prompt) else: prompt_text = prompt_row.get("prompt") if prompt_row is not None else None if isinstance(prompt_text, str) and prompt_text: assert tokenizer is not None tokens = prompt_token_ids(tokenizer, prompt_text) input_length_mismatches += int(len(tokens) != isl) if len(tokens) != isl: raise ValueError( f"prompt token length mismatch for chat_id={row.get('chat_id')!r}: " f"trace={isl}, tokenizer={len(tokens)}" ) expected_source_blocks = math.ceil(len(tokens) / source_block_size) source_hash_count_mismatches += int(len(original_hashes) != expected_source_blocks) if len(original_hashes) != expected_source_blocks: raise ValueError( f"strict {source_block_size}-block contract failed " f"for chat_id={row.get('chat_id')!r}: " f"expected={expected_source_blocks}, got={len(original_hashes)}" ) records = token_block_identity_records(tokens) block_ids = [identity for identity, _ in records] for runtime_id, witness in records: previous = identity_to_witness.setdefault(runtime_id, witness) identity_collisions += int(previous != witness) for block_index, source_hash in enumerate(original_hashes): begin = block_index * subblocks_per_source relation = tuple( block_ids[begin : begin + subblocks_per_source] ) previous_relation = source_to_runtime.setdefault(source_hash, relation) source_to_runtime_conflicts += int(previous_relation != relation) previous_source = runtime_to_source.setdefault(relation, source_hash) runtime_to_source_conflicts += int(previous_source != source_hash) request_prompt = tokens prompt_kind = "real_tokens" real_prompt_requests += 1 else: block_ids = expand_hash_ids( original_hashes, isl, source_block_size=source_block_size, ) request_prompt = synthetic_tokens( block_ids, isl, vocab_size=args.vocab_size, token_offset=args.token_offset, ) prompt_kind = "synthetic_missing_prompt_fallback" synthetic_fallback_requests += 1 parent = row.get("parent_chat_id") if args.validate_parents and parent not in (None, "", -1, "-1"): parent_blocks = parent_sequences.get(parent) if parent_blocks is None: parent_links_outside += 1 else: parent_links_checked += 1 # Real multi-turn chats re-tokenize the full turn-1 history into # the turn-2 prompt. Because source blocks are not necessarily # aligned to turn boundaries, the parent's trailing partial block # does not survive as a child prefix: vLLM/Frontier reuse the # parent's complete leading blocks and recompute from the first # unaligned block. So we require a contiguous common prefix and # only tolerate a bounded mismatch confined to the parent tail. common = 0 for a, b in zip(block_ids, parent_blocks): if a != b: break common += 1 # Tail slack is one possibly-partial source block plus one # runtime block for tokenizer boundary effects. max_parent_tail_slack = subblocks_per_source + 1 if len(parent_blocks) - common > max_parent_tail_slack: raise ValueError( f"parent prefix violation beyond tail slack: " f"parent={parent!r}, child={row.get('chat_id')!r}, " f"common={common}, parent_blocks={len(parent_blocks)}" ) parent_prefix_common_blocks += common parent_tail_nonreused_blocks += len(parent_blocks) - common parent_sequences[row_identity(row, source_index)] = block_ids if first_selected_timestamp is None: first_selected_timestamp = ts arrival_origin = start if start is not None else first_selected_timestamp arrived_at = ts - arrival_origin sid = session_id(row) writer.writerow( { "arrived_at": f"{arrived_at:.9f}", "num_prefill_tokens": isl, "num_decode_tokens": osl, "session_id": sid, "block_hash_ids": "|".join(str(value) for value in block_ids), } ) row_vector_digest.update( json.dumps( [ row.get("source_index", source_index), arrived_at, isl, osl, sid, block_ids, ], separators=(",", ":"), ).encode() ) row_vector_digest.update(b"\n") real_request = { "source_index": row.get("source_index", source_index), "chat_id": row.get("chat_id"), "parent_chat_id": parent, "arrived_at": arrived_at, "input_length": isl, "output_length": osl, "session_id": sid, "runtime_block_ids": block_ids, "body": { "model": args.served_model, "prompt": request_prompt, "min_tokens": osl, "max_tokens": osl, "ignore_eos": True, "stream": True, }, } if real_stream is not None: real_stream.write(json.dumps(real_request, separators=(",", ":")) + "\n") if selected_stream is not None: selected_stream.write( json.dumps( { **row, "output_length": osl, "source_index": row.get("source_index", source_index), "arrived_at": arrived_at, "hash_ids_16": block_ids, "prompt": request_prompt, "prompt_source": prompt_kind, }, separators=(",", ":"), ) + "\n" ) request_count += 1 prefill_tokens += isl decode_tokens += osl source_blocks += len(original_hashes) target_blocks += len(block_ids) if request_count == 0: raise ValueError("rho/window selected no requests") contract_failures = { "input_length_mismatches": input_length_mismatches, "source_hash_count_mismatches": source_hash_count_mismatches, "runtime_identity_collisions": identity_collisions, "source_to_runtime_relation_conflicts": source_to_runtime_conflicts, "runtime_to_source_relation_conflicts": runtime_to_source_conflicts, } if upstream_block_contract is not None: inherited_failures = { key: value for key, value in upstream_block_contract.items() if key.endswith("mismatches") or key.endswith("collisions") or key.endswith("conflicts") } if any(inherited_failures.values()): raise ValueError(f"upstream block contract failed: {inherited_failures}") contract_failures = inherited_failures if any(contract_failures.values()): raise ValueError(f"source-to-16 real-token block contract failed: {contract_failures}") manifest = { "schema": "frontier-s3-real-remap-v2", "source": str(args.input.resolve()), "source_sha256": sha256(args.input), "prompt_window": str(args.prompt.resolve()) if args.prompt else None, "prompt_window_sha256": sha256(args.prompt) if args.prompt else None, "upstream_remap_manifest": str(upstream_manifest_path.resolve()) if upstream_manifest_path else None, "upstream_remap_manifest_sha256": sha256(upstream_manifest_path) if upstream_manifest_path else None, "rho": args.rho, "sampling_rule": "keep iff sampling_u <= rho (SimFid convention; rho=1 keeps all rows)", "window": {"start_timestamp": start, "duration_s": args.duration_s}, "requests": request_count, "source_rows_scanned": source_rows, "excluded_over_max_total_tokens": excluded_over_max_tokens, "max_total_tokens_filter": args.max_total_tokens, "total_prefill_tokens": prefill_tokens, "total_decode_tokens": decode_tokens, "source_blocks": source_blocks, "target_16_blocks": target_blocks, "paired_row_vector_sha256": row_vector_digest.hexdigest(), "source_block_size": source_block_size, "target_block_size": TARGET_BLOCK_SIZE, "workload_mode": workload_mode, "mapping": ( "real prompt tokens: BLAKE2b-128(parent runtime identity, exact 16-token block); " "missing prompt fallback: expanded_id = zigzag(source_hash) * " "subblocks_per_source + subblock_index" ), "prompt_contract": { "real_prompt_requests": real_prompt_requests, "synthetic_fallback_requests": synthetic_fallback_requests, "input_is_pre_remapped": args.input_is_remapped, "tokenizer": str(args.tokenizer.resolve()) if args.tokenizer else None, }, "block_contract": { **contract_failures, "unique_runtime_identities": ( upstream_block_contract.get("unique_runtime_identities", 0) if upstream_block_contract is not None else len(identity_to_witness) ), "unique_source_hash_relations": ( upstream_block_contract.get("unique_source_hash_relations", 0) if upstream_block_contract is not None else len(source_to_runtime) ), "unique_runtime_relations": ( upstream_block_contract.get("unique_runtime_relations", 0) if upstream_block_contract is not None else len(runtime_to_source) ), "source_runtime_granularity": ( f"one {source_block_size}-token source hash to " f"{subblocks_per_source} consecutive 16-token identities; " "final partial source block may have fewer" ), }, "synthetic_fallback_contract": { "vocab_size": args.vocab_size, "token_offset": args.token_offset, "encoding": "injective base-(vocab_size-token_offset), 16 little-endian digits", "same_16_block_hash_same_tokens": True, "different_16_block_hash_different_tokens": True, }, "parent_validation": { "enabled": args.validate_parents, "links_checked": parent_links_checked, "links_outside_selected_stream": parent_links_outside, "rule": ( "contiguous common prefix required; mismatch tolerated only " f"within parent trailing partial {source_block_size}-token block " f"(<= {subblocks_per_source + 1} runtime blocks)" ), "parent_prefix_common_blocks": parent_prefix_common_blocks, "parent_tail_nonreused_blocks": parent_tail_nonreused_blocks, "mean_tail_nonreused_per_link": ( parent_tail_nonreused_blocks / parent_links_checked if parent_links_checked else 0.0 ), }, "frontier_csv": str(frontier_path.resolve()), "frontier_csv_sha256": sha256(frontier_path), "real_requests": str(real_path.resolve()) if real_path.is_file() else None, "real_requests_sha256": sha256(real_path) if real_path.is_file() else None, "selected_remapped": str(selected_path.resolve()) if selected_path.is_file() else None, "selected_remapped_sha256": sha256(selected_path) if selected_path.is_file() else None, } write_json(args.output_root / "manifest.json", manifest) return manifest def main() -> None: args = parse_args() print(json.dumps(materialize(args), sort_keys=True)) if __name__ == "__main__": main()