aituner/runs/active-intervention-v0/test_prospective.py

#!/usr/bin/env python3
from __future__ import annotations

import importlib.util
import json
import sys
import tempfile
from pathlib import Path


HERE = Path(__file__).resolve().parent


def load(name: str, path: Path):
    spec = importlib.util.spec_from_file_location(name, path)
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def write_json(path: Path, payload) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload) + "\n", encoding="utf-8")


def main() -> None:
    prepare = load("active_intervention_prepare_test", HERE / "prepare_prospective.py")
    decision_module = load(
        "active_intervention_decision_test", HERE / "prospective_decision.py"
    )
    analyzer = load("active_intervention_audit_test", HERE / "analyze_prospective.py")
    with tempfile.TemporaryDirectory() as temporary:
        root = Path(temporary)
        source = root / "source.jsonl"
        source.write_text(
            "".join(
                json.dumps(
                    {
                        "request_id": f"request-{index}",
                        "timestamp": float(index),
                        "sampling_u": index / 100.0,
                    }
                )
                + "\n"
                for index in range(60)
            ),
            encoding="utf-8",
        )
        partition = prepare.partition_trace(source, root / "partitions")
        assert sum(item["rows"] for item in partition["partitions"].values()) == 60
        ids = []
        for item in partition["partitions"].values():
            assert item["rows"] > 0
            ids.extend(
                json.loads(line)["request_id"]
                for line in Path(item["path"]).read_text(encoding="utf-8").splitlines()
            )
        assert len(ids) == len(set(ids)) == 60

        checkpoints = [
            {
                "phase": "0.25",
                "cutoff_s": 75.0,
                "selected_action": "joint",
                "confident": True,
                "candidates": [
                    {"action_id": "joint", "upper": 0.5, "prediction": {"mean": 0.4}},
                    {"action_id": "mns", "upper": 0.2, "prediction": {"mean": 0.1}},
                    {"action_id": "mbbt", "upper": 0.1, "prediction": {"mean": 0.05}},
                    {"action_id": "noop", "upper": 0.0, "prediction": {"mean": 0.0}},
                ],
            },
            {
                "phase": "0.50",
                "cutoff_s": 150.0,
                "selected_action": "joint",
                "confident": True,
                "candidates": [
                    {"action_id": "joint", "upper": 0.45, "prediction": {"mean": 0.4}},
                    {"action_id": "mns", "upper": 0.2, "prediction": {"mean": 0.1}},
                    {"action_id": "mbbt", "upper": 0.1, "prediction": {"mean": 0.05}},
                    {"action_id": "noop", "upper": 0.0, "prediction": {"mean": 0.0}},
                ],
            },
        ]
        selected = decision_module.apply_measurement_and_acquisition(checkpoints)
        assert selected["selected_cutoff_s"] == 150.0
        assert selected["selected_action"] == "joint"

        configs = prepare.configs()
        repetitions = {
            str(rep): {
                "selection": {
                    "offered_req_s_per_gpu": 0.25,
                    "request_id_order_sha256": f"hash-{rep}",
                }
            }
            for rep in (1, 2, 3)
        }
        manifest = {
            "schema": "active-intervention-prospective-manifest-v0",
            "engine": {"duration_s": 300.0, "tp": 4},
            "repetitions": repetitions,
            "configs": configs,
            "source_config_id": "source_mns32_mbbt4096",
            "actions": {
                "noop": "source_mns32_mbbt4096",
                "mns": "mns64_mbbt4096",
                "mbbt": "mns32_mbbt8192",
                "joint": "joint_mns64_mbbt8192",
            },
            "gates": {
                "acceptable_regret": 0.02,
                "confirmation_trigger_gpu_cost_reduction": 0.10,
                "contribution_gpu_cost_reduction": 0.20,
            },
        }
        manifest_path = root / "manifest.json"
        write_json(manifest_path, manifest)
        run_root = root / "run"
        scores = {
            "source_mns32_mbbt4096": 0.5,
            "mns64_mbbt4096": 0.8,
            "mns32_mbbt8192": 0.7,
            "joint_mns64_mbbt8192": 1.0,
        }
        sessions = {}
        for config in configs:
            config_id = config["id"]
            sessions[config_id] = {"status": "complete", "gpu_hours": 1.2}
            for repetition in (1, 2, 3):
                result = {
                    "selection": {
                        "request_id_order_sha256": f"hash-{repetition}"
                    },
                    "slo_pass_count": round(scores[config_id] * 300),
                    "pass_rate": scores[config_id],
                    "interval": {"elapsed_s": 300.0},
                }
                write_json(
                    run_root
                    / "sessions"
                    / config_id
                    / f"rep{repetition}"
                    / "result.json",
                    result,
                )
        state = {
            "status": "complete",
            "gpu_hours_total": 4.8,
            "sessions": sessions,
        }
        write_json(run_root / "controller-state.json", state)
        mode_base = {
            "selected_cutoff_s": 300.0,
            "selected_action": "mns",
            "decision_kind": "exploit",
            "intervention_order": ["mns", "mbbt", "joint", "noop"],
        }
        mode_telemetry = {
            "selected_cutoff_s": 150.0,
            "selected_action": "joint",
            "decision_kind": "exploit",
            "intervention_order": ["joint", "mns", "mbbt", "noop"],
        }
        decision = {
            "schema": "active-intervention-prospective-decision-v0",
            "manifest_sha256": analyzer.sha256_file(manifest_path),
            "decisions": {
                "outcome_only": mode_base,
                "telemetry": mode_telemetry,
            },
        }
        decision_path = root / "decision.json"
        write_json(decision_path, decision)
        audit = analyzer.build_audit(
            manifest_path=manifest_path,
            decision_path=decision_path,
            run_root=run_root,
        )
        assert audit["status"] == "TRIGGER_ACTUAL_EARLY_STOP_CONFIRMATION"
        assert audit["comparison"]["telemetry_gpu_cost_reduction_fraction"] > 0.10
        assert not audit["sanity"]["red_flags"]
    print("active intervention prospective pipeline: PASS")


if __name__ == "__main__":
    main()