Files
aituner/runs/frontier-qwen30-vllm020-profile-v1/prepare_profile_comparison.py

105 lines
3.5 KiB
Python

#!/usr/bin/env python3
"""Prepare old/new profile-only Frontier manifests from the frozen P1 probes."""
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("--source", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--mode", choices=("old-profile-only", "new-profile-only"), required=True)
parser.add_argument("--profile-root", type=Path)
return parser.parse_args()
def main() -> None:
args = parse_args()
source = json.loads(args.source.read_text())
if source.get("status") != "PASS" or len(source.get("entries", [])) != 12:
raise SystemExit("source manifest must contain 12 passing P1 probes")
if args.mode == "new-profile-only" and args.profile_root is None:
raise SystemExit("--profile-root is required for new-profile-only")
output = args.output.resolve()
config_root = output / "configs"
cache_root = output / "prediction-cache"
entries: list[dict[str, Any]] = []
profile_hashes: dict[str, str] = {}
if args.profile_root is not None:
profile_root = args.profile_root.resolve()
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)
for entry in source["entries"]:
config_path = Path(entry["config"])
config = json.loads(config_path.read_text())
config["mode"] = args.mode
config["config_id"] = f"{config['cell_id']}__{args.mode}"
config["calibration"]["a_tp"] = 1.0
knobs = config["frontier"]["knobs"]
knobs["cache_dir"] = str(cache_root)
knobs["no_cache"] = False
if args.mode == "new-profile-only":
for key, filename in PROFILE_KEYS.items():
knobs[key] = str((args.profile_root.resolve() / filename))
target_config = config_root / f"{entry['fixture_id']}.json"
write_json(target_config, config)
updated_entry = dict(entry)
updated_entry["config"] = str(target_config)
updated_entry["calibration_scale"] = 1.0
entries.append(updated_entry)
prepared = {
"schema": "frontier-qwen30-profile-comparison-prepared.v1",
"status": "PASS",
"mode": args.mode,
"source": {
"manifest": str(args.source.resolve()),
"sha256": sha256(args.source),
},
"profile_hashes": profile_hashes,
"isolation": {
"calibration_a_tp": 1.0,
"prediction_cache": str(cache_root),
"all_non_profile_knobs_inherited": True,
},
"entries": entries,
}
write_json(output / "prepared-manifest.json", prepared)
print(output / "prepared-manifest.json")
if __name__ == "__main__":
main()