aituner/runs/intervention-response-v2/prepare_pilot.py

#!/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_ROLE_PAIRS = (("low1", "high1"), ("low2", "high2"), ("low3", "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 resolve_role_trace(base: dict[str, Any], role: str) -> Path:
    windows_path = Path(base["private"]["windows"])
    windows = json.loads(windows_path.read_text(encoding="utf-8"))["windows"]
    window_id = f"fidelity_pilot_{role}"
    record = next(item for item in windows if item["window_id"] == window_id)
    trace = Path(record["trace_file"])
    if not trace.is_absolute():
        trace = (windows_path.parent / trace).resolve()
    return trace


def merge_role_traces(sources: tuple[Path, Path], target: Path) -> dict[str, Any]:
    target.parent.mkdir(parents=True, exist_ok=True)
    temporary = target.with_suffix(target.suffix + ".tmp")
    rows = 0
    with temporary.open("w", encoding="utf-8") as output:
        for source in sources:
            with source.open(encoding="utf-8") as input_file:
                for line in input_file:
                    if not line.strip():
                        continue
                    json.loads(line)
                    output.write(line if line.endswith("\n") else line + "\n")
                    rows += 1
    os.replace(temporary, target)
    return {
        "path": str(target),
        "sha256": sha256_file(target),
        "bytes": target.stat().st_size,
        "rows": rows,
        "sources": [str(source) for source in sources],
        "source_sha256": [sha256_file(source) for source in sources],
    }


def materialize_study(
    source: Path, target: Path, *, replicate: int, trace_override: Path
) -> 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["trace_file_override"] = str(trace_override)
    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]] = {}
    merged_traces = {}
    for replicate, roles in enumerate(REPLICATE_ROLE_PAIRS, start=1):
        source_study = Path(base["private"]["studies"][roles[1]]["tp4"])
        source_traces = tuple(resolve_role_trace(base, role) for role in roles)
        merged_trace = private_root / "traces" / f"rep{replicate}.jsonl"
        merged_traces[str(replicate)] = merge_role_traces(
            source_traces, merged_trace
        )
        target_study = private_root / "studies" / f"rep{replicate}-tp4.json"
        materialize_study(
            source_study,
            target_study,
            replicate=replicate,
            trace_override=merged_trace,
        )
        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_roles": list(roles),
            "merged_trace": merged_traces[str(replicate)],
            "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,
        "private": {"merged_traces": merged_traces},
        "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()