268 lines
9.5 KiB
Python
268 lines
9.5 KiB
Python
#!/usr/bin/env python3
|
|
"""Prepare the uncensored 300-second TP4 matched pilot on dash0."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
AITUNER_ROOT = Path(os.environ.get("AITUNER_ROOT", Path(__file__).resolve().parents[2]))
|
|
sys.path.insert(0, str(AITUNER_ROOT / "src"))
|
|
|
|
from aituner.spec import load_study_spec # noqa: E402
|
|
from aituner.trace import load_trace_requests, select_requests_for_threshold # noqa: E402
|
|
|
|
|
|
SCHEMA = "intervention-response-phase-aware-pilot-manifest-v2"
|
|
TP = 4
|
|
MNS_ENDPOINTS = (16, 64)
|
|
REPLAY_TIME_SCALE = 0.5
|
|
EXPECTED_DURATION_S = 300.0
|
|
SAFETY_DEADLINE_S = 360.0
|
|
LOADS_REQ_S_GPU = {"low": 1.5, "mid": 2.125, "high": 3.125}
|
|
REPLICATE_ROLES = ("high1", "high2", "high3")
|
|
LOAD_ORDERS = {
|
|
1: ("low", "mid", "high"),
|
|
2: ("high", "low", "mid"),
|
|
3: ("mid", "high", "low"),
|
|
}
|
|
SESSION_ORDER = (
|
|
(1, 16),
|
|
(1, 64),
|
|
(2, 64),
|
|
(2, 16),
|
|
(3, 16),
|
|
(3, 64),
|
|
)
|
|
|
|
|
|
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")
|
|
os.replace(temporary, path)
|
|
|
|
|
|
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 order_hash(values: list[str]) -> str:
|
|
return hashlib.sha256("\n".join(values).encode()).hexdigest()
|
|
|
|
|
|
def attainable_anchor(requests: list[Any], target_count: int) -> tuple[float, list[Any]]:
|
|
ordered = sorted(float(request.sampling_u) for request in requests)
|
|
if not ordered:
|
|
raise ValueError("no requests remain after study filtering")
|
|
if target_count <= 0 or target_count > len(ordered):
|
|
raise ValueError(
|
|
f"target count {target_count} is outside available range 1..{len(ordered)}"
|
|
)
|
|
indices = sorted(
|
|
{
|
|
max(0, min(len(ordered) - 1, target_count - 1)),
|
|
max(0, min(len(ordered) - 1, target_count)),
|
|
}
|
|
)
|
|
candidates = []
|
|
for index in indices:
|
|
anchor = ordered[index]
|
|
selected = select_requests_for_threshold(requests, threshold=anchor)
|
|
candidates.append((abs(len(selected) - target_count), len(selected), anchor, selected))
|
|
_error, _count, anchor, selected = min(
|
|
candidates, key=lambda item: (item[0], item[1], item[2])
|
|
)
|
|
return anchor, selected
|
|
|
|
|
|
def selection_record(selected: list[Any], *, duration_s: float) -> dict[str, Any]:
|
|
return {
|
|
"anchor": max(float(request.sampling_u) for request in selected),
|
|
"selected_count": len(selected),
|
|
"offered_req_s": len(selected) / duration_s,
|
|
"offered_req_s_per_gpu": len(selected) / duration_s / TP,
|
|
"request_id_order_sha256": order_hash([request.row_id for request in selected]),
|
|
"arrival_order_sha256": order_hash(
|
|
[f"{request.arrival_s:.12f}" for request in selected]
|
|
),
|
|
"input_length_order_sha256": order_hash(
|
|
[str(request.prompt_tokens_hint) for request in selected]
|
|
),
|
|
}
|
|
|
|
|
|
def materialize_study(source: Path, target: Path, *, replicate: int) -> Path:
|
|
payload = json.loads(source.read_text(encoding="utf-8"))
|
|
payload["study_id"] = f"phase-aware-telemetry-v2-rep{replicate}"
|
|
payload["hardware"]["host_candidates"] = ["dash0"]
|
|
payload["engine"]["engine_version"] = "0.24.1.dev3+opprof"
|
|
trace = payload["trace"]
|
|
trace["replay_time_scale"] = REPLAY_TIME_SCALE
|
|
trace["early_stop_max_lag_s"] = None
|
|
trace["early_stop_max_elapsed_s"] = SAFETY_DEADLINE_S
|
|
trace["restart_engine_after_early_stop"] = False
|
|
trace["adaptive_stop"] = {"enabled": False}
|
|
atomic_json(target, payload)
|
|
return target
|
|
|
|
|
|
def build_manifest(
|
|
*, base_manifest_path: Path, private_root: Path
|
|
) -> dict[str, Any]:
|
|
base = json.loads(base_manifest_path.read_text(encoding="utf-8"))
|
|
if base.get("schema") != "fidelity-prefix-pilot-manifest-v1":
|
|
raise ValueError("unexpected base P1 manifest schema")
|
|
|
|
repetitions = {}
|
|
selection_hashes = []
|
|
selected_ids_by_replicate: dict[int, set[str]] = {}
|
|
for replicate, role in enumerate(REPLICATE_ROLES, start=1):
|
|
source_study = Path(base["private"]["studies"][role]["tp4"])
|
|
target_study = private_root / "studies" / f"rep{replicate}-tp4.json"
|
|
materialize_study(source_study, target_study, replicate=replicate)
|
|
study = load_study_spec(target_study)
|
|
window, requests = load_trace_requests(study, study_spec_path=target_study)
|
|
duration_s = float(window.window_end - window.window_start)
|
|
if not math.isclose(
|
|
duration_s, EXPECTED_DURATION_S, rel_tol=0.0, abs_tol=1e-9
|
|
):
|
|
raise ValueError(
|
|
f"rep{replicate}: replay duration {duration_s} != {EXPECTED_DURATION_S}"
|
|
)
|
|
selections = {}
|
|
selected_ids_by_level: dict[str, set[str]] = {}
|
|
for level, target_rate in LOADS_REQ_S_GPU.items():
|
|
target_count = round(target_rate * duration_s * TP)
|
|
anchor, selected = attainable_anchor(requests, target_count)
|
|
record = selection_record(selected, duration_s=duration_s)
|
|
record.update(
|
|
{
|
|
"anchor": anchor,
|
|
"target_count": target_count,
|
|
"target_req_s_per_gpu": target_rate,
|
|
}
|
|
)
|
|
selections[level] = record
|
|
selected_ids_by_level[level] = {request.row_id for request in selected}
|
|
selection_hashes.append(record["request_id_order_sha256"])
|
|
selected_ids_by_replicate[replicate] = selected_ids_by_level["high"]
|
|
repetitions[str(replicate)] = {
|
|
"source_role": role,
|
|
"study": str(target_study),
|
|
"study_sha256": sha256_file(target_study),
|
|
"duration_s": duration_s,
|
|
"load_order": list(LOAD_ORDERS[replicate]),
|
|
"selections": selections,
|
|
}
|
|
|
|
burnin = base["cells"]["tp4_mns16"]["targets"]["low"]["selections"][
|
|
"burnin"
|
|
]
|
|
burnin = dict(burnin)
|
|
burnin["study_sha256"] = sha256_file(Path(burnin["study"]))
|
|
sessions = [
|
|
{
|
|
"session": f"rep{replicate}-mns{mns}",
|
|
"replicate": replicate,
|
|
"mns": mns,
|
|
}
|
|
for replicate, mns in SESSION_ORDER
|
|
]
|
|
invariants = {
|
|
"three_repetitions": len(repetitions) == 3,
|
|
"six_sessions": len(sessions) == 6,
|
|
"load_levels_three": all(
|
|
len(item["selections"]) == 3 for item in repetitions.values()
|
|
),
|
|
"selection_hashes_unique": len(selection_hashes) == len(set(selection_hashes)),
|
|
"selection_sets_disjoint_across_repetitions": all(
|
|
not selected_ids_by_replicate[left] & selected_ids_by_replicate[right]
|
|
for left in selected_ids_by_replicate
|
|
for right in selected_ids_by_replicate
|
|
if left < right
|
|
),
|
|
"all_counts_positive": all(
|
|
selection["selected_count"] > 0
|
|
for item in repetitions.values()
|
|
for selection in item["selections"].values()
|
|
),
|
|
}
|
|
red_flags = [name for name, passed in invariants.items() if not passed]
|
|
return {
|
|
"schema": SCHEMA,
|
|
"status": "PASS" if not red_flags else "STOP",
|
|
"source": {
|
|
"base_manifest": str(base_manifest_path),
|
|
"base_manifest_sha256": sha256_file(base_manifest_path),
|
|
"window_id": base["source"]["window_id"],
|
|
"source_trace": base["source"]["trace"],
|
|
"source_trace_sha256": base["source"]["trace_sha256"],
|
|
},
|
|
"engine": {
|
|
"tp": TP,
|
|
"mns_endpoints": list(MNS_ENDPOINTS),
|
|
"replay_time_scale": REPLAY_TIME_SCALE,
|
|
"duration_s": EXPECTED_DURATION_S,
|
|
"safety_deadline_s": SAFETY_DEADLINE_S,
|
|
"disable_slo_early_stop": True,
|
|
},
|
|
"burnin": burnin,
|
|
"repetitions": repetitions,
|
|
"sessions": sessions,
|
|
"checkpoints": {
|
|
"fractions": [0.1, 0.25, 0.5, 0.75, 1.0],
|
|
"seconds": [30.0, 75.0, 150.0, 225.0, 300.0],
|
|
},
|
|
"budget": {
|
|
"hard_cap_h20_hours": 8.0,
|
|
"expected_wall_minutes": [95, 120],
|
|
"expected_h20_hours": [6.3, 8.0],
|
|
},
|
|
"sanity": {
|
|
"red_flags": red_flags,
|
|
"invariants": invariants,
|
|
"selected_sets": len(selection_hashes),
|
|
"distinct_selected_sets": len(set(selection_hashes)),
|
|
},
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--base-manifest", type=Path, required=True)
|
|
parser.add_argument("--private-root", type=Path, required=True)
|
|
parser.add_argument("--public-manifest", type=Path, required=True)
|
|
args = parser.parse_args()
|
|
manifest = build_manifest(
|
|
base_manifest_path=args.base_manifest, private_root=args.private_root
|
|
)
|
|
atomic_json(args.public_manifest, manifest)
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"status": manifest["status"],
|
|
"manifest": str(args.public_manifest),
|
|
"sanity": manifest["sanity"],
|
|
},
|
|
sort_keys=True,
|
|
)
|
|
)
|
|
if manifest["status"] != "PASS":
|
|
raise RuntimeError(f"phase-aware pilot preflight failed: {manifest['sanity']}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|