diff --git a/runs/frontier-s3-real-v0/remap_hash_blocks.py b/runs/frontier-s3-real-v0/remap_hash_blocks.py new file mode 100644 index 0000000..c07834d --- /dev/null +++ b/runs/frontier-s3-real-v0/remap_hash_blocks.py @@ -0,0 +1,493 @@ +#!/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() diff --git a/runs/frontier-s3-real-v0/test_remap_hash_blocks.py b/runs/frontier-s3-real-v0/test_remap_hash_blocks.py new file mode 100644 index 0000000..03c48e3 --- /dev/null +++ b/runs/frontier-s3-real-v0/test_remap_hash_blocks.py @@ -0,0 +1,214 @@ +#!/usr/bin/env python3 + +from __future__ import annotations + +import math +import json +import sys +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from trace_utils import ( + common_prefix_length, + expand_hash_ids, + synthetic_tokens, + token_block, + token_block_identities, +) +from remap_hash_blocks import materialize, prompt_token_ids + + +class CharacterTokenizer: + def __call__(self, text: str, **_: object) -> dict[str, list[int]]: + return {"input_ids": [ord(character) for character in text]} + + +class StrictRemapTest(unittest.TestCase): + def test_block_count_is_ceil_input_over_16(self) -> None: + for input_tokens in (1, 15, 16, 17, 63, 64, 65, 127, 128, 129): + source = list(range(math.ceil(input_tokens / 64))) + remapped = expand_hash_ids(source, input_tokens) + self.assertEqual(len(remapped), math.ceil(input_tokens / 16)) + + def test_parent_sequence_remains_prefix(self) -> None: + parent_length = 80 + child_length = 144 + parent_source = [101, 202] + child_source = [101, 202, 303] + parent = expand_hash_ids(parent_source, parent_length) + child = expand_hash_ids(child_source, child_length) + self.assertEqual(child[: len(parent)], parent) + parent_tokens = synthetic_tokens(parent, parent_length, vocab_size=151936) + child_tokens = synthetic_tokens(child, child_length, vocab_size=151936) + self.assertEqual(child_tokens[: len(parent_tokens)], parent_tokens) + + def test_synthetic_tokens_preserve_hash_hit_structure(self) -> None: + request_a_source = [11, 22, 33] + request_b_source = [11, 22, 44] + request_a = expand_hash_ids(request_a_source, 180) + request_b = expand_hash_ids(request_b_source, 190) + source_hits = common_prefix_length(request_a_source, request_b_source) + remapped_hits = common_prefix_length(request_a, request_b) + self.assertEqual(remapped_hits, source_hits * 4) + + token_blocks_a = [tuple(token_block(value, vocab_size=151936)) for value in request_a] + token_blocks_b = [tuple(token_block(value, vocab_size=151936)) for value in request_b] + self.assertEqual(common_prefix_length(token_blocks_a, token_blocks_b), remapped_hits) + self.assertEqual(len(set(token_blocks_a)), len(set(request_a))) + + def test_mapping_rejects_non_strict_source_block_count(self) -> None: + with self.assertRaisesRegex(ValueError, "strict 64-block contract"): + expand_hash_ids([1], 65) + + def test_512_to_16_mapping_and_prefill_only_mode(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + trace = root / "code-window.jsonl" + trace.write_text( + json.dumps( + { + "chat_id": "code-1", + "parent_chat_id": -1, + "session_root": "code-1", + "turn": 1, + "timestamp": 10.0, + "input_length": 513, + "output_length": 128, + "hash_ids": [101, 202], + "sampling_u": 0.1, + } + ) + + "\n" + ) + output = root / "mapped" + args = SimpleNamespace( + input=trace, + prompt=None, + tokenizer=None, + input_is_remapped=False, + output_root=output, + rho=1.0, + start_timestamp=None, + duration_s=None, + vocab_size=151936, + token_offset=1024, + served_model="test-model", + source_block_size=512, + workload_mode="prefill_only", + validate_parents=True, + max_total_tokens=514, + frontier_only=False, + ) + manifest = materialize(args) + request = json.loads((output / "real_requests.jsonl").read_text()) + mapped = json.loads((output / "selected-remapped.jsonl").read_text()) + + self.assertEqual(manifest["source_block_size"], 512) + self.assertEqual(manifest["workload_mode"], "prefill_only") + self.assertEqual(manifest["target_16_blocks"], math.ceil(513 / 16)) + self.assertEqual(request["input_length"], 513) + self.assertEqual(request["output_length"], 1) + self.assertEqual(request["body"]["max_tokens"], 1) + self.assertEqual(mapped["output_length"], 1) + self.assertEqual(len(mapped["hash_ids_16"]), math.ceil(513 / 16)) + + def test_real_prompt_tokens_preserve_parent_and_four_to_one_contract(self) -> None: + tokenizer = CharacterTokenizer() + parent_prompt = "A" * 64 + child_prompt = parent_prompt + "B" * 16 + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + trace = root / "chat-raw-window.jsonl" + prompts = root / "chat-prompt-window.jsonl" + rows = [ + { + "chat_id": "parent", + "parent_chat_id": -1, + "session_root": "parent", + "turn": 1, + "timestamp": 10.0, + "input_length": 64, + "output_length": 8, + "hash_ids": [101], + "sampling_u": 0.1, + }, + { + "chat_id": "child", + "parent_chat_id": "parent", + "session_root": "parent", + "turn": 2, + "timestamp": 11.0, + "input_length": 80, + "output_length": 8, + "hash_ids": [101, 202], + "sampling_u": 0.1, + }, + ] + prompt_rows = [ + {"chat_id": "parent", "turn": 1, "prompt": parent_prompt}, + {"chat_id": "child", "turn": 2, "prompt": child_prompt}, + ] + trace.write_text("".join(json.dumps(row) + "\n" for row in rows)) + prompts.write_text("".join(json.dumps(row) + "\n" for row in prompt_rows)) + output = root / "mapped" + args = SimpleNamespace( + input=trace, + prompt=prompts, + tokenizer=None, + input_is_remapped=False, + output_root=output, + rho=1.0, + start_timestamp=None, + duration_s=None, + vocab_size=151936, + token_offset=1024, + served_model="test-model", + validate_parents=True, + max_total_tokens=None, + frontier_only=False, + ) + manifest = materialize(args, tokenizer=tokenizer) + with (output / "selected-remapped.jsonl").open() as stream: + mapped = [json.loads(line) for line in stream] + filtered_output = root / "filtered" + filtered_args = SimpleNamespace( + input=output / "selected-remapped.jsonl", + prompt=None, + tokenizer=None, + input_is_remapped=True, + output_root=filtered_output, + rho=1.0, + start_timestamp=None, + duration_s=None, + vocab_size=151936, + token_offset=1024, + served_model="test-model", + validate_parents=True, + max_total_tokens=None, + frontier_only=False, + ) + filtered_manifest = materialize(filtered_args) + + expected_parent = token_block_identities(prompt_token_ids(tokenizer, parent_prompt)) + expected_child = token_block_identities(prompt_token_ids(tokenizer, child_prompt)) + self.assertEqual(mapped[0]["hash_ids_16"], expected_parent) + self.assertEqual(mapped[1]["hash_ids_16"], expected_child) + self.assertEqual(mapped[1]["hash_ids_16"][:4], mapped[0]["hash_ids_16"]) + self.assertEqual(len(mapped[0]["hash_ids_16"]), math.ceil(64 / 16)) + self.assertEqual(len(mapped[1]["hash_ids_16"]), math.ceil(80 / 16)) + self.assertEqual(manifest["block_contract"]["source_to_runtime_relation_conflicts"], 0) + self.assertEqual(manifest["block_contract"]["runtime_to_source_relation_conflicts"], 0) + self.assertEqual(manifest["prompt_contract"]["real_prompt_requests"], 2) + self.assertEqual(manifest["prompt_contract"]["synthetic_fallback_requests"], 0) + self.assertEqual(mapped[0]["prompt"], prompt_token_ids(tokenizer, parent_prompt)) + self.assertEqual(filtered_manifest["block_contract"], manifest["block_contract"]) + self.assertIsNotNone(filtered_manifest["upstream_remap_manifest_sha256"]) + self.assertEqual(filtered_manifest["paired_row_vector_sha256"], manifest["paired_row_vector_sha256"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/runs/frontier-s3-real-v0/trace_utils.py b/runs/frontier-s3-real-v0/trace_utils.py new file mode 100644 index 0000000..04f538c --- /dev/null +++ b/runs/frontier-s3-real-v0/trace_utils.py @@ -0,0 +1,269 @@ +#!/usr/bin/env python3 +"""Shared, dependency-free helpers for the S3-real preflight.""" + +from __future__ import annotations + +import hashlib +import json +import math +from pathlib import Path +from typing import Any, Iterable, Iterator, Sequence + + +SOURCE_BLOCK_SIZE = 64 +TARGET_BLOCK_SIZE = 16 +SUBBLOCKS_PER_SOURCE = SOURCE_BLOCK_SIZE // TARGET_BLOCK_SIZE + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1 << 20), b""): + digest.update(chunk) + return digest.hexdigest() + + +def write_json(path: Path, payload: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") + + +def iter_jsonl(path: Path) -> Iterator[dict[str, Any]]: + with path.open() as stream: + for line_number, line in enumerate(stream, 1): + if not line.strip(): + continue + try: + row = json.loads(line) + except json.JSONDecodeError as error: + raise ValueError(f"{path}:{line_number}: invalid JSON") from error + if not isinstance(row, dict): + raise ValueError(f"{path}:{line_number}: row must be an object") + yield row + + +def timestamp(row: dict[str, Any]) -> float: + return float(row["timestamp"]) + + +def input_length(row: dict[str, Any]) -> int: + return int(row["input_length"]) + + +def output_length(row: dict[str, Any]) -> int: + return max(1, int(row["output_length"])) + + +def parse_hash_ids(value: Any) -> list[int]: + if isinstance(value, list): + raw = value + elif isinstance(value, str): + stripped = value.strip() + if not stripped: + return [] + if stripped.startswith("["): + raw = json.loads(stripped) + else: + delimiter = "|" if "|" in stripped else "," + raw = [part.strip() for part in stripped.split(delimiter) if part.strip()] + elif value is None: + return [] + else: + raw = [value] + try: + return [int(item) for item in raw] + except (TypeError, ValueError) as error: + raise ValueError(f"hash_ids must contain integers, got {value!r}") from error + + +def percentile(values: Sequence[float | int], fraction: float) -> float | None: + if not values: + return None + if not 0 <= fraction <= 1: + raise ValueError("percentile fraction must be in [0, 1]") + ordered = sorted(float(value) for value in values) + if len(ordered) == 1: + return ordered[0] + position = (len(ordered) - 1) * fraction + lower = math.floor(position) + upper = math.ceil(position) + if lower == upper: + return ordered[lower] + return ordered[lower] * (upper - position) + ordered[upper] * (position - lower) + + +def distribution(values: Sequence[float | int]) -> dict[str, float | int | None]: + return { + "count": len(values), + "min": min(values) if values else None, + "p50": percentile(values, 0.50), + "p90": percentile(values, 0.90), + "p95": percentile(values, 0.95), + "p99": percentile(values, 0.99), + "max": max(values) if values else None, + "mean": sum(values) / len(values) if values else None, + } + + +def common_prefix_length(left: Sequence[Any], right: Sequence[Any]) -> int: + count = 0 + for a, b in zip(left, right): + if a != b: + break + count += 1 + return count + + +def zigzag(value: int) -> int: + """Map every signed integer injectively to a non-negative integer.""" + return 2 * value if value >= 0 else -2 * value - 1 + + +def expanded_block_id( + source_hash: int, + subblock_index: int, + *, + subblocks_per_source: int = SUBBLOCKS_PER_SOURCE, +) -> int: + if not 0 <= subblock_index < subblocks_per_source: + raise ValueError(f"subblock index out of range: {subblock_index}") + return zigzag(int(source_hash)) * subblocks_per_source + subblock_index + + +def expand_hash_ids( + source_hashes: Sequence[int], + num_tokens: int, + *, + source_block_size: int = SOURCE_BLOCK_SIZE, + target_block_size: int = TARGET_BLOCK_SIZE, +) -> list[int]: + if num_tokens <= 0: + raise ValueError(f"input_length must be positive, got {num_tokens}") + if source_block_size <= 0 or target_block_size <= 0: + raise ValueError("block sizes must be positive") + if source_block_size % target_block_size: + raise ValueError("source block size must be divisible by target block size") + subblocks_per_source = source_block_size // target_block_size + required_source_blocks = math.ceil(num_tokens / source_block_size) + if len(source_hashes) != required_source_blocks: + raise ValueError( + f"strict {source_block_size}-block contract failed: " + f"input_length={num_tokens} requires {required_source_blocks} source blocks, " + f"got {len(source_hashes)}" + ) + target_blocks = math.ceil(num_tokens / target_block_size) + expanded = [ + expanded_block_id( + source_hash, + subblock, + subblocks_per_source=subblocks_per_source, + ) + for source_hash in source_hashes + for subblock in range(subblocks_per_source) + ] + return expanded[:target_blocks] + + +def token_payload(tokens: Sequence[int]) -> bytes: + return len(tokens).to_bytes(2, "little") + b"".join( + int(token).to_bytes(4, "little", signed=False) for token in tokens + ) + + +def token_block_identity_records( + token_ids: Sequence[int], block_size: int = TARGET_BLOCK_SIZE +) -> list[tuple[int, bytes]]: + """Return parent-sensitive block IDs and independent collision witnesses.""" + if block_size <= 0: + raise ValueError("block_size must be positive") + parent = b"FRONTIER_EXACT_TRACE_ROOT" + records: list[tuple[int, bytes]] = [] + 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 token_block_identities( + token_ids: Sequence[int], block_size: int = TARGET_BLOCK_SIZE +) -> list[int]: + return [identity for identity, _ in token_block_identity_records(token_ids, block_size)] + + +def token_block(block_id: int, *, vocab_size: int, token_offset: int = 1024) -> list[int]: + """Encode one block identity injectively as exactly sixteen valid token IDs.""" + base = vocab_size - token_offset + if base < 256: + raise ValueError("vocab has too few non-reserved token IDs") + value = int(block_id) + if value < 0: + raise ValueError("expanded block IDs must be non-negative") + digits = [] + for _ in range(TARGET_BLOCK_SIZE): + digits.append(token_offset + value % base) + value //= base + if value: + raise ValueError("block identity exceeds the injective 16-token code space") + return digits + + +def synthetic_tokens( + block_ids: Sequence[int], num_tokens: int, *, vocab_size: int, token_offset: int = 1024 +) -> list[int]: + tokens = [ + token + for block_id in block_ids + for token in token_block(block_id, vocab_size=vocab_size, token_offset=token_offset) + ][:num_tokens] + if len(tokens) != num_tokens: + raise ValueError(f"synthetic token length mismatch: {len(tokens)} != {num_tokens}") + return tokens + + +def stable_int(value: Any) -> int: + raw = json.dumps(value, sort_keys=True, separators=(",", ":")).encode() + return int.from_bytes(hashlib.sha256(raw).digest()[:8], "big") & ((1 << 63) - 1) + + +def resolve_session_root(row: dict[str, Any], root_of: dict[Any, Any]) -> Any: + chat = row.get("chat_id") + parent = row.get("parent_chat_id") + has_parent = parent not in (None, "", -1, "-1") + root = root_of.get(parent, parent) if has_parent else chat + if chat is not None: + root_of[chat] = root + return root + + +def session_uniform(*, seed: int, window_id: str, session_root: Any) -> float: + payload = json.dumps( + {"seed": seed, "window_id": window_id, "session_root": session_root}, + sort_keys=True, + separators=(",", ":"), + ).encode() + digest = hashlib.blake2b(payload, digest_size=8).digest() + return int.from_bytes(digest, "big") / float(1 << 64) + + +def session_id(row: dict[str, Any]) -> int: + return stable_int( + row.get("session_root", row.get("chat_id", row.get("parent_chat_id", "missing-chat"))) + ) + + +def batched(iterable: Iterable[Any], size: int) -> Iterator[list[Any]]: + batch: list[Any] = [] + for item in iterable: + batch.append(item) + if len(batch) == size: + yield batch + batch = [] + if batch: + yield batch