Add simulator-aware fidelity pilot audit
This commit is contained in:
331
runs/fidelity-headroom/prepare_pilot_simulator.py
Normal file
331
runs/fidelity-headroom/prepare_pilot_simulator.py
Normal file
@@ -0,0 +1,331 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user