196 lines
6.8 KiB
Python
196 lines
6.8 KiB
Python
#!/usr/bin/env python3
|
|
"""Freeze the crossed-constraint action-aware development pilot."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
SCHEMA_V0 = "action-aware-constraint-pilot-manifest-v0"
|
|
SCHEMA_V1 = "action-aware-constraint-pilot-manifest-v1"
|
|
|
|
|
|
def configs(token_source_mbbt: int) -> tuple[dict[str, Any], ...]:
|
|
return (
|
|
{
|
|
"id": "b_base",
|
|
"mns": 64,
|
|
"mbbt": token_source_mbbt,
|
|
"repetition_order": [1, 2, 3],
|
|
},
|
|
{"id": "a_base", "mns": 16, "mbbt": 8192, "repetition_order": [2, 3, 1]},
|
|
{"id": "shared", "mns": 64, "mbbt": 8192, "repetition_order": [3, 1, 2]},
|
|
{
|
|
"id": "b_mns",
|
|
"mns": 128,
|
|
"mbbt": token_source_mbbt,
|
|
"repetition_order": [1, 3, 2],
|
|
},
|
|
{"id": "a_mbbt", "mns": 16, "mbbt": 16384, "repetition_order": [2, 1, 3]},
|
|
)
|
|
|
|
|
|
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 atomic_json(path: Path, payload: Any) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary = path.with_suffix(path.suffix + ".tmp")
|
|
temporary.write_text(
|
|
json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8"
|
|
)
|
|
os.replace(temporary, path)
|
|
|
|
|
|
def build(
|
|
base_path: Path,
|
|
*,
|
|
token_source_mbbt: int = 256,
|
|
prior_attempt_h20_hours: float = 0.0,
|
|
prior_attempt_artifact: str | None = None,
|
|
) -> dict[str, Any]:
|
|
if token_source_mbbt <= 0:
|
|
raise ValueError("token source MBBT must be positive")
|
|
if prior_attempt_h20_hours < 0.0 or prior_attempt_h20_hours >= 8.0:
|
|
raise ValueError("prior attempt cost must be in [0, 8)")
|
|
base = json.loads(base_path.read_text(encoding="utf-8"))
|
|
if base.get("schema") != "intervention-response-phase-aware-pilot-manifest-v3":
|
|
raise ValueError("unexpected base manifest schema")
|
|
if base.get("status") != "PASS":
|
|
raise ValueError("base manifest did not pass its preflight")
|
|
if sorted(int(key) for key in base["repetitions"]) != [1, 2, 3]:
|
|
raise ValueError("base manifest must contain exactly three repetitions")
|
|
|
|
repetitions = {}
|
|
selection_hashes = []
|
|
for repetition in (1, 2, 3):
|
|
source = base["repetitions"][str(repetition)]
|
|
selection = dict(source["selections"]["mid"])
|
|
selection_hashes.append(selection["request_id_order_sha256"])
|
|
repetitions[str(repetition)] = {
|
|
"study": source["study"],
|
|
"study_sha256": source["study_sha256"],
|
|
"selection": selection,
|
|
"merged_trace": source["merged_trace"],
|
|
}
|
|
|
|
frozen_configs = configs(token_source_mbbt)
|
|
config_ids = [str(config["id"]) for config in frozen_configs]
|
|
schema = (
|
|
SCHEMA_V0
|
|
if token_source_mbbt == 256 and prior_attempt_h20_hours == 0.0
|
|
else SCHEMA_V1
|
|
)
|
|
payload = {
|
|
"schema": schema,
|
|
"status": "PASS",
|
|
"source": {
|
|
"base_manifest": str(base_path.resolve()),
|
|
"base_manifest_sha256": sha256_file(base_path),
|
|
"window_id": base["source"]["window_id"],
|
|
"source_trace": base["source"]["source_trace"],
|
|
"source_trace_sha256": base["source"]["source_trace_sha256"],
|
|
},
|
|
"engine": {
|
|
"tp": 4,
|
|
"duration_s": 300.0,
|
|
"disable_slo_early_stop": True,
|
|
"client_timeout_s": 450.0,
|
|
"burnin_max_elapsed_s": 90.0,
|
|
},
|
|
"burnin": base["burnin"],
|
|
"repetitions": repetitions,
|
|
"configs": [dict(config) for config in frozen_configs],
|
|
"regimes": {
|
|
"A": {
|
|
"source": "a_base",
|
|
"actions": {"mns": "shared", "mbbt": "a_mbbt"},
|
|
},
|
|
"B": {
|
|
"source": "b_base",
|
|
"actions": {"mns": "b_mns", "mbbt": "shared"},
|
|
},
|
|
},
|
|
"budget": {
|
|
"global_hard_cap_h20_hours": 8.0,
|
|
"hard_cap_h20_hours": 8.0 - prior_attempt_h20_hours,
|
|
"prior_attempt_h20_hours": prior_attempt_h20_hours,
|
|
"prior_attempt_artifact": prior_attempt_artifact,
|
|
"session_estimate_h20_hours": 1.35,
|
|
"safety_h20_hours": 0.25,
|
|
"expected_h20_hours": [6.0, 7.2],
|
|
"expected_wall_minutes": [90, 110],
|
|
},
|
|
"gates": {
|
|
"minimum_relative_winner_margin": 0.10,
|
|
"minimum_exclusive_fraction": 0.10,
|
|
"minimum_exclusive_ratio": 5.0,
|
|
"phase_fractions": [0.25, 0.50, 0.75, 1.0],
|
|
"material_kv_usage": 0.90,
|
|
},
|
|
"sanity": {
|
|
"invariants": {
|
|
"five_unique_configs": len(config_ids) == len(set(config_ids)) == 5,
|
|
"three_disjoint_repetitions": len(set(selection_hashes)) == 3,
|
|
"same_load_all_repetitions": len(
|
|
{
|
|
float(item["selection"]["offered_req_s_per_gpu"])
|
|
for item in repetitions.values()
|
|
}
|
|
)
|
|
== 1,
|
|
"all_repetition_orders_are_permutations": all(
|
|
sorted(config["repetition_order"]) == [1, 2, 3]
|
|
for config in frozen_configs
|
|
),
|
|
}
|
|
},
|
|
}
|
|
payload["sanity"]["invariants"]["shared_endpoint_reused_by_both_regimes"] = (
|
|
payload["regimes"]["A"]["actions"]["mns"]
|
|
== payload["regimes"]["B"]["actions"]["mbbt"]
|
|
== "shared"
|
|
)
|
|
payload["sanity"]["red_flags"] = [
|
|
name
|
|
for name, passed in payload["sanity"]["invariants"].items()
|
|
if not passed
|
|
]
|
|
if payload["sanity"]["red_flags"]:
|
|
payload["status"] = "FAIL"
|
|
return payload
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--base-manifest", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--token-source-mbbt", type=int, default=256)
|
|
parser.add_argument("--prior-attempt-h20-hours", type=float, default=0.0)
|
|
parser.add_argument("--prior-attempt-artifact")
|
|
args = parser.parse_args()
|
|
payload = build(
|
|
args.base_manifest,
|
|
token_source_mbbt=args.token_source_mbbt,
|
|
prior_attempt_h20_hours=args.prior_attempt_h20_hours,
|
|
prior_attempt_artifact=args.prior_attempt_artifact,
|
|
)
|
|
atomic_json(args.output, payload)
|
|
print(json.dumps(payload["sanity"], sort_keys=True))
|
|
if payload["status"] != "PASS":
|
|
raise SystemExit("manifest preflight failed")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|