332 lines
13 KiB
Python
332 lines
13 KiB
Python
#!/usr/bin/env python3
|
|
"""Prepare exact Frontier fixtures for the P1 primary low/high probes.
|
|
|
|
Prompt-bearing band traces remain under ``--private-root``. The emitted
|
|
fixtures and public manifest contain token IDs, block IDs, hashes, and
|
|
aggregate metadata, but no prompt text.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import importlib.util
|
|
import json
|
|
import math
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from transformers import AutoTokenizer
|
|
|
|
|
|
HERE = Path(__file__).resolve().parent
|
|
AITUNER_ROOT = HERE.parents[1]
|
|
sys.path.insert(0, str(HERE))
|
|
|
|
import prepare_pilot as pilot # noqa: E402
|
|
|
|
|
|
PRIMARY_ROLES = ("low1", "high1")
|
|
|
|
|
|
def load_module(path: Path):
|
|
spec = importlib.util.spec_from_file_location("simfid_s2rb_prepare", path)
|
|
if spec is None or spec.loader is None:
|
|
raise ImportError(path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[spec.name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def order_hash(values: list[str]) -> str:
|
|
return hashlib.sha256("\n".join(values).encode()).hexdigest()
|
|
|
|
|
|
def sha256_file(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as source:
|
|
for chunk in iter(lambda: source.read(1 << 20), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def git_capture(root: Path, *arguments: str) -> str:
|
|
return subprocess.run(
|
|
["git", "-C", str(root), *arguments],
|
|
check=True,
|
|
text=True,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
).stdout
|
|
|
|
|
|
def raw_rows(path: Path) -> dict[int, dict[str, Any]]:
|
|
result = {}
|
|
with path.open(encoding="utf-8") as source:
|
|
for index, line in enumerate(source):
|
|
if line.strip():
|
|
result[index] = json.loads(line)
|
|
return result
|
|
|
|
|
|
def selected_hashes(
|
|
selected: list[Any], rows: dict[int, dict[str, Any]]
|
|
) -> dict[str, str]:
|
|
identifiers = []
|
|
arrivals = []
|
|
lengths = []
|
|
for item in selected:
|
|
row = rows[item.row_index]
|
|
identifiers.append(str(row.get("request_id") or row.get("id") or item.row_index))
|
|
arrivals.append(f"{float(item.timestamp) * 0.1:.12f}")
|
|
lengths.append(str(int(item.input_length)))
|
|
return {
|
|
"request_id_order_sha256": order_hash(identifiers),
|
|
"arrival_order_sha256": order_hash(arrivals),
|
|
"input_length_order_sha256": order_hash(lengths),
|
|
}
|
|
|
|
|
|
def kv_blocks(raw_root: Path, cell: str) -> int:
|
|
stream = next((raw_root / cell / "opprof").glob("*.jsonl"))
|
|
with stream.open(encoding="utf-8") as source:
|
|
for line in source:
|
|
record = json.loads(line)
|
|
if "step_index" in record:
|
|
return int(record["kv"]["total_blocks"])
|
|
raise ValueError(f"no Layer-1 record for {cell}")
|
|
|
|
|
|
def source_window(windows_path: Path, window_id: str) -> tuple[dict[str, Any], Path]:
|
|
return pilot.resolve_source_trace(windows_path, window_id)
|
|
|
|
|
|
def prepare(args: argparse.Namespace) -> dict[str, Any]:
|
|
simulator = load_module(args.replayserve_root / "tools/simfid_s2rb_prepare.py")
|
|
manifest = json.loads(args.pilot_manifest.read_text(encoding="utf-8"))
|
|
window, trace = source_window(args.source_windows, args.source_window_id)
|
|
if args.band_root is not None:
|
|
role_paths = {
|
|
role: (args.band_root / f"{role}.jsonl").resolve()
|
|
for role in PRIMARY_ROLES
|
|
}
|
|
band_stats = {
|
|
role: manifest["private"]["band_stats"][role]
|
|
for role in PRIMARY_ROLES
|
|
}
|
|
for role, path in role_paths.items():
|
|
if sha256_file(path) != band_stats[role]["sha256"]:
|
|
raise ValueError(f"pre-materialized band hash mismatch: {role}")
|
|
private_windows = None
|
|
else:
|
|
private_windows, all_band_stats = pilot.materialize_bands(
|
|
trace, window, args.private_root
|
|
)
|
|
private_payload = json.loads(private_windows.read_text(encoding="utf-8"))
|
|
role_paths = {
|
|
item["fidelity_pilot_role"]: (
|
|
private_windows.parent / item["trace_file"]
|
|
).resolve()
|
|
for item in private_payload["windows"]
|
|
}
|
|
band_stats = {
|
|
role: all_band_stats[role]
|
|
for role in PRIMARY_ROLES
|
|
}
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
args.tokenizer, local_files_only=True, use_fast=True
|
|
)
|
|
fixture_root = args.output / "fixtures"
|
|
config_root = args.output / "configs"
|
|
fixture_root.mkdir(parents=True, exist_ok=True)
|
|
config_root.mkdir(parents=True, exist_ok=True)
|
|
entries = []
|
|
red_flags = []
|
|
for role in PRIMARY_ROLES:
|
|
trace_path = role_paths[role]
|
|
retained, trace_stats = simulator.scan_trace(trace_path)
|
|
rows = raw_rows(trace_path)
|
|
primary_pool = [retained[(index * len(retained)) // 512] for index in range(512)]
|
|
selections: dict[str, list[Any]] = {}
|
|
selected_union: set[int] = set()
|
|
for cell, cell_manifest in sorted(manifest["cells"].items()):
|
|
level = "low" if role.startswith("low") else "high"
|
|
expected = cell_manifest["targets"][level]["selections"][role]
|
|
pool = retained if int(cell_manifest["tp"]) == 4 else primary_pool
|
|
selected = [item for item in pool if item.sampling_u <= float(expected["anchor"])]
|
|
selections[cell] = selected
|
|
selected_union.update(item.row_index for item in selected)
|
|
hashes = selected_hashes(selected, rows)
|
|
if len(selected) != int(expected["selected_count"]):
|
|
red_flags.append(f"selection_count_{cell}_{role}")
|
|
for key, value in hashes.items():
|
|
if value != expected[key]:
|
|
red_flags.append(f"selection_hash_{cell}_{role}_{key}")
|
|
|
|
token_gates, selected_records, block_stats = simulator.tokenize_and_hash(
|
|
trace=trace_path,
|
|
tokenizer=tokenizer,
|
|
retained=retained,
|
|
selected_union=selected_union,
|
|
)
|
|
if any(gate["status"] != "pass" for gate in token_gates.values()):
|
|
red_flags.append(f"token_gate_{role}")
|
|
for cell, selected in selections.items():
|
|
cell_manifest = manifest["cells"][cell]
|
|
level = "low" if role.startswith("low") else "high"
|
|
expected = cell_manifest["targets"][level]["selections"][role]
|
|
fixture_id = f"fidelity_p1_{cell}_{role}"
|
|
cell_record = {
|
|
"cell_id": cell,
|
|
"tensor_parallel_size": int(cell_manifest["tp"]),
|
|
"max_num_seqs": int(cell_manifest["mns"]),
|
|
"store_role": "companion" if int(cell_manifest["tp"]) == 4 else "primary",
|
|
"kv_capacity": {
|
|
"block_size_tokens": 16,
|
|
"num_blocks": kv_blocks(args.phase6_raw_root, cell),
|
|
},
|
|
}
|
|
probe = {
|
|
"probe_index": 0 if role == "low1" else 1,
|
|
"sampling_u": float(expected["anchor"]),
|
|
}
|
|
fixture = simulator.create_fixture(
|
|
fixture_root=fixture_root,
|
|
fixture_id=fixture_id,
|
|
cell=cell_record,
|
|
probe=probe,
|
|
row_indexes=[item.row_index for item in selected],
|
|
meta_by_index={item.row_index: item for item in retained},
|
|
selected_records=selected_records,
|
|
)
|
|
config_path = config_root / f"{fixture_id}.json"
|
|
config = simulator.build_config(
|
|
path=config_path,
|
|
cell=cell_record,
|
|
mode="frozen-calibrated",
|
|
fixture_ids=[fixture_id],
|
|
frontier_root=args.frontier_root,
|
|
cache_dir=args.cache_dir,
|
|
)
|
|
entries.append(
|
|
{
|
|
"cell": cell,
|
|
"role": role,
|
|
"level": level,
|
|
"anchor": expected["anchor"],
|
|
"selected_count": len(selected),
|
|
"fixture_id": fixture_id,
|
|
"fixture_manifest": str(
|
|
(fixture_root / fixture_id / "fixture_manifest.json").resolve()
|
|
),
|
|
"frontier_csv": fixture["frontier_csv"]["path"],
|
|
"sidecar": fixture["sidecar_jsonl"]["path"],
|
|
"config": str(config_path.resolve()),
|
|
"calibration_scale": config["calibration"]["a_tp"],
|
|
}
|
|
)
|
|
if block_stats["selected_union_records"] != len(selected_union):
|
|
red_flags.append(f"selected_union_{role}")
|
|
if trace_stats["retained_inclusive_0_8192"] < 512:
|
|
red_flags.append(f"retained_too_small_{role}")
|
|
|
|
selected_counts = [int(entry["selected_count"]) for entry in entries]
|
|
calibration = [float(entry["calibration_scale"]) for entry in entries]
|
|
result = {
|
|
"schema": "fidelity-p1-frontier-prepared-v1",
|
|
"status": "PASS" if not red_flags else "STOP",
|
|
"source": {
|
|
"pilot_manifest": str(args.pilot_manifest.resolve()),
|
|
"source_windows": str(args.source_windows.resolve()),
|
|
"source_window_id": args.source_window_id,
|
|
"source_trace": str(trace.resolve()),
|
|
"private_windows": (
|
|
str(private_windows.resolve()) if private_windows is not None else None
|
|
),
|
|
"pre_materialized_band_root": (
|
|
str(args.band_root.resolve()) if args.band_root is not None else None
|
|
),
|
|
"band_stats": band_stats,
|
|
},
|
|
"simulator": {
|
|
"replayserve_root": str(args.replayserve_root.resolve()),
|
|
"frontier_root": str(args.frontier_root.resolve()),
|
|
"tokenizer": str(args.tokenizer.resolve()),
|
|
"mode": "frozen-calibrated",
|
|
},
|
|
"generator": {
|
|
"script": str(Path(__file__).resolve()),
|
|
"script_sha256": sha256_file(Path(__file__).resolve()),
|
|
"aituner_git_head": git_capture(AITUNER_ROOT, "rev-parse", "HEAD").strip(),
|
|
"aituner_git_status_short": git_capture(AITUNER_ROOT, "status", "--short"),
|
|
},
|
|
"entries": entries,
|
|
"sanity": {
|
|
"red_flags": red_flags,
|
|
"n": len(entries),
|
|
"selected_count": {
|
|
"n": len(selected_counts),
|
|
"min": min(selected_counts),
|
|
"max": max(selected_counts),
|
|
"distinct_n": len(set(selected_counts)),
|
|
},
|
|
"calibration_scale": {
|
|
"n": len(calibration),
|
|
"min": min(calibration),
|
|
"max": max(calibration),
|
|
"distinct_n": len(set(calibration)),
|
|
},
|
|
"invariants": {
|
|
"entries_12": len(entries) == 12,
|
|
"roles_2": {entry["role"] for entry in entries} == set(PRIMARY_ROLES),
|
|
"cells_6": len({entry["cell"] for entry in entries}) == 6,
|
|
"selected_nonnegative": all(value > 0 for value in selected_counts),
|
|
"per_config_not_identical": len(set(selected_counts)) > 1,
|
|
},
|
|
},
|
|
}
|
|
args.public_manifest.parent.mkdir(parents=True, exist_ok=True)
|
|
args.public_manifest.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
|
|
return result
|
|
|
|
|
|
def parser() -> argparse.ArgumentParser:
|
|
result = argparse.ArgumentParser()
|
|
result.add_argument("--pilot-manifest", type=Path, required=True)
|
|
result.add_argument("--source-windows", type=Path, required=True)
|
|
result.add_argument("--source-window-id", required=True)
|
|
result.add_argument("--private-root", type=Path, required=True)
|
|
result.add_argument("--band-root", type=Path)
|
|
result.add_argument("--output", type=Path, required=True)
|
|
result.add_argument("--public-manifest", type=Path, required=True)
|
|
result.add_argument("--phase6-raw-root", type=Path, required=True)
|
|
result.add_argument("--replayserve-root", type=Path, required=True)
|
|
result.add_argument("--frontier-root", type=Path, required=True)
|
|
result.add_argument("--cache-dir", type=Path, required=True)
|
|
result.add_argument("--tokenizer", type=Path, required=True)
|
|
return result
|
|
|
|
|
|
def main() -> None:
|
|
result = prepare(parser().parse_args())
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"status": result["status"],
|
|
"entries": len(result["entries"]),
|
|
"red_flags": result["sanity"]["red_flags"],
|
|
},
|
|
sort_keys=True,
|
|
)
|
|
)
|
|
if result["status"] != "PASS":
|
|
raise RuntimeError(result["sanity"]["red_flags"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|