gahow/aituner
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Generalize trace remapping for code workloads

Browse Source
This commit is contained in:
Gahow Wang
2026-07-24 00:09:50 +08:00
parent a9e88f14de
commit 029c8991b6
3 changed files with 976 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

493
runs/frontier-s3-real-v0/remap_hash_blocks.py Normal file
View File

@@ -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()

214
runs/frontier-s3-real-v0/test_remap_hash_blocks.py Normal file
View File

@@ -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()

269
runs/frontier-s3-real-v0/trace_utils.py Normal file
View File

@@ -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