501 lines
23 KiB
Python
501 lines
23 KiB
Python
#!/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:
|
|
content_block_ids = expand_hash_ids(
|
|
original_hashes,
|
|
isl,
|
|
source_block_size=source_block_size,
|
|
)
|
|
request_prompt = synthetic_tokens(
|
|
content_block_ids,
|
|
isl,
|
|
vocab_size=args.vocab_size,
|
|
token_offset=args.token_offset,
|
|
)
|
|
records = token_block_identity_records(request_prompt)
|
|
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)
|
|
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: source hashes first define deterministic 16-token "
|
|
"content blocks, then the same parent-sensitive BLAKE2b-128 runtime identity "
|
|
"contract is applied"
|
|
),
|
|
"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,
|
|
"runtime_identity_is_parent_sensitive": 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()
|