134 lines
4.8 KiB
Python
134 lines
4.8 KiB
Python
#!/usr/bin/env python3
|
|
"""Prepare the 92 frozen S2-R-b probes with a replacement profile bundle."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
PROFILE_KEYS = {
|
|
"linear_op_input_file": "linear_op.csv",
|
|
"atten_input_file": "attention.csv",
|
|
"moe_input_file": "moe.csv",
|
|
}
|
|
|
|
|
|
def sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as handle:
|
|
for chunk in iter(lambda: handle.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 parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--resolved-plan", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--profile-root", type=Path, required=True)
|
|
parser.add_argument("--cache-root", type=Path)
|
|
parser.add_argument("--shards", type=int, default=1)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
if args.shards < 1:
|
|
raise SystemExit("--shards must be positive")
|
|
source = json.loads(args.resolved_plan.read_text())
|
|
source_rows = [row for row in source["runs"] if row["mode"] == "uncalibrated"]
|
|
if len(source_rows) != 92:
|
|
raise SystemExit(f"expected 92 uncalibrated probes, found {len(source_rows)}")
|
|
|
|
output = args.output.resolve()
|
|
profile_root = args.profile_root.resolve()
|
|
cache_root = (args.cache_root or (output / "prediction-cache")).resolve()
|
|
profile_hashes: dict[str, str] = {}
|
|
for filename in PROFILE_KEYS.values():
|
|
path = profile_root / filename
|
|
if not path.is_file():
|
|
raise SystemExit(f"missing frozen profile: {path}")
|
|
profile_hashes[str(path)] = sha256(path)
|
|
|
|
entries: list[dict[str, Any]] = []
|
|
for row in source_rows:
|
|
fixture_dir = Path(row["fixture_dir"]).resolve()
|
|
fixture_manifest_path = fixture_dir / "fixture_manifest.json"
|
|
fixture = json.loads(fixture_manifest_path.read_text())
|
|
config = json.loads(Path(row["config_path"]).read_text())
|
|
config["mode"] = "new-profile-only"
|
|
config["config_id"] = f"{row['cell_id']}__new-profile-only"
|
|
config["calibration"]["a_tp"] = 1.0
|
|
knobs = config["frontier"]["knobs"]
|
|
knobs["cache_dir"] = str(cache_root)
|
|
knobs["no_cache"] = False
|
|
for key, filename in PROFILE_KEYS.items():
|
|
knobs[key] = str(profile_root / filename)
|
|
|
|
target_config = output / "configs" / f"{row['fixture_id']}.json"
|
|
write_json(target_config, config)
|
|
entries.append(
|
|
{
|
|
"fixture_id": row["fixture_id"],
|
|
"cell": row["cell_id"],
|
|
"role": f"probe-{int(row['probe_index']):02d}",
|
|
"anchor": float(row["sampling_u"]),
|
|
"selected_count": int(row["request_count"]),
|
|
"probe_index": int(row["probe_index"]),
|
|
"sampling_u": float(row["sampling_u"]),
|
|
"tensor_parallel_size": int(row["tensor_parallel_size"]),
|
|
"config": str(target_config),
|
|
"fixture_manifest": str(fixture_manifest_path),
|
|
"frontier_csv": str(fixture_dir / "frontier.csv"),
|
|
"sidecar": str(fixture_dir / "sidecar.jsonl"),
|
|
"calibration_scale": 1.0,
|
|
"request_count": int(fixture["request_count"]),
|
|
}
|
|
)
|
|
|
|
base = {
|
|
"schema": "frontier-qwen30-s2-profile-comparison-prepared.v1",
|
|
"status": "PASS",
|
|
"mode": "new-profile-only",
|
|
"source": {
|
|
"resolved_plan": str(args.resolved_plan.resolve()),
|
|
"sha256": sha256(args.resolved_plan),
|
|
},
|
|
"profile_hashes": profile_hashes,
|
|
"isolation": {
|
|
"calibration_a_tp": 1.0,
|
|
"prediction_cache": str(cache_root),
|
|
"all_non_profile_knobs_inherited": True,
|
|
},
|
|
"suite_total_runs": len(entries),
|
|
}
|
|
write_json(
|
|
output / "prepared-manifest.json",
|
|
{**base, "expected_runs": len(entries), "entries": entries},
|
|
)
|
|
for shard in range(args.shards):
|
|
shard_entries = entries[shard :: args.shards]
|
|
write_json(
|
|
output / f"prepared-manifest-shard-{shard:02d}-of-{args.shards:02d}.json",
|
|
{
|
|
**base,
|
|
"shard": {"index": shard, "count": args.shards},
|
|
"expected_runs": len(shard_entries),
|
|
"entries": shard_entries,
|
|
},
|
|
)
|
|
print(output / "prepared-manifest.json")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|