aituner/runs/action-aware-v0/test_pilot.py

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

import copy
import importlib.util
from pathlib import Path
from types import SimpleNamespace


HERE = Path(__file__).resolve().parent
ROOT = HERE.parents[1]


def load(name: str, filename: str):
    spec = importlib.util.spec_from_file_location(name, HERE / filename)
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(module)
    return module


def record(*, waiting: int, running: int, tokens: int) -> dict:
    return {
        "queues": {"waiting": waiting, "deferred": 0, "running": running},
        "prefill_tokens": tokens,
        "decode_tokens": 0,
        "kv": {"usage": 0.5},
        "preemptions": 0,
    }


def fake_run(
    config: str,
    repetition: int,
    *,
    goodput: float,
    mns_score: float = 0.0,
    mbbt_score: float = 0.0,
    ambiguous: float = 0.0,
) -> dict:
    binding = {
        "mns_exclusive_fraction": mns_score,
        "mbbt_exclusive_fraction": mbbt_score,
        "both_fraction": ambiguous,
        "waiting_unresolved_fraction": 0.0,
        "kv_usage_max": 0.5,
        "preemptions": 0,
    }
    phases = {
        phase: {
            "mns_exclusive_fraction": mns_score,
            "mbbt_exclusive_fraction": mbbt_score,
        }
        for phase in ("0.25", "0.50", "0.75", "1.00")
    }
    return {
        "config_id": config,
        "repetition": repetition,
        "outcome": {"slo_goodput_req_s": goodput},
        "binding": binding,
        "phases": phases,
    }


def main() -> None:
    analysis = load("action_aware_analysis", "analyze_pilot.py")
    summary = analysis.binding_summary(
        [
            record(waiting=1, running=16, tokens=8),
            record(waiting=1, running=8, tokens=32),
            record(waiting=1, running=16, tokens=32),
            record(waiting=1, running=8, tokens=8),
            record(waiting=0, running=8, tokens=8),
        ],
        mns=16,
        mbbt=32,
    )
    assert summary["mns_exclusive_count"] == 1
    assert summary["mbbt_exclusive_count"] == 1
    assert summary["both_count"] == 1
    assert summary["waiting_unresolved_count"] == 1
    assert summary["waiting_count"] == 4

    manifest = {
        "repetitions": {str(index): {} for index in (1, 2, 3)},
        "regimes": {
            "A": {
                "source": "a_base",
                "actions": {"mns": "shared", "mbbt": "a_mbbt"},
            },
            "B": {
                "source": "b_base",
                "actions": {"mns": "b_mns", "mbbt": "shared"},
            },
        },
        "gates": {
            "minimum_relative_winner_margin": 0.10,
            "minimum_exclusive_fraction": 0.10,
            "minimum_exclusive_ratio": 5.0,
            "material_kv_usage": 0.90,
        },
    }
    runs = []
    for repetition in (1, 2, 3):
        runs.extend(
            [
                fake_run(
                    "a_base",
                    repetition,
                    goodput=1.0,
                    mns_score=0.8,
                    mbbt_score=0.01,
                ),
                fake_run(
                    "b_base",
                    repetition,
                    goodput=1.0,
                    mns_score=0.01,
                    mbbt_score=0.7,
                ),
                fake_run("shared", repetition, goodput=3.0),
                fake_run("a_mbbt", repetition, goodput=1.5),
                fake_run("b_mns", repetition, goodput=1.2),
            ]
        )
    result = analysis.evaluate_decisions(runs, manifest)
    assert result["decision"] == "STOP_NO_NEW_INSTRUMENTATION_NEEDED"
    assert result["baselines"] == {
        "always_mns_correct": 3,
        "always_mbbt_correct": 3,
        "binding_correct": 6,
        "decision_count": 6,
    }

    ambiguous = copy.deepcopy(runs)
    for run in ambiguous:
        if run["config_id"] == "b_base":
            run["binding"]["both_fraction"] = 0.8
    assert (
        analysis.evaluate_decisions(ambiguous, manifest)["decision"]
        == "OPEN_EXACT_ATTRIBUTION_ABLATION"
    )

    wrong = copy.deepcopy(runs)
    for run in wrong:
        if run["config_id"] == "b_base":
            run["binding"]["mns_exclusive_fraction"] = 0.8
            run["binding"]["mbbt_exclusive_fraction"] = 0.01
            for phase in run["phases"].values():
                phase["mns_exclusive_fraction"] = 0.8
                phase["mbbt_exclusive_fraction"] = 0.01
    assert (
        analysis.evaluate_decisions(wrong, manifest)["decision"]
        == "STOP_BINDING_NOT_PREDICTIVE"
    )

    prepare = load("action_aware_prepare", "prepare_pilot.py")
    frozen = prepare.build(
        ROOT / "runs/intervention-response-v2/pilot-manifest-v3.json"
    )
    assert frozen["status"] == "PASS"
    assert frozen["sanity"]["red_flags"] == []
    assert [config["id"] for config in frozen["configs"]] == [
        "b_base",
        "a_base",
        "shared",
        "b_mns",
        "a_mbbt",
    ]

    controller = load("action_aware_controller", "pilot_controller.py")
    args = SimpleNamespace(
        manifest=Path("/tmp/manifest.json"),
        run_root=Path("/tmp/action-aware"),
        aituner_root=Path("/tmp/aituner"),
        vllm_source=Path("/tmp/vllm"),
        venv=Path("/tmp/venv"),
        model=Path("/tmp/model"),
        client=Path("/tmp/client.py"),
    )
    controller.configure(args, frozen)
    plan = controller.dry_run_plan(args, frozen)
    assert plan["status"] == "PASS"
    assert len(plan["sessions"]) == 5
    assert plan["projected_h20_hours"] == 7.0
    assert "--max-num-batched-tokens 256" in plan["sessions"][0]["commands"]["server"]
    revised = prepare.build(
        ROOT / "runs/intervention-response-v2/pilot-manifest-v3.json",
        token_source_mbbt=2048,
        prior_attempt_h20_hours=0.38598689953486126,
        prior_attempt_artifact="/tmp/operational-stop-v0.json",
    )
    assert revised["schema"] == "action-aware-constraint-pilot-manifest-v1"
    assert revised["configs"][0]["mbbt"] == 2048
    assert revised["configs"][3]["mbbt"] == 2048
    assert revised["budget"]["hard_cap_h20_hours"] < 8.0
    controller.configure(args, revised)
    revised_plan = controller.dry_run_plan(args, revised)
    assert revised_plan["projected_h20_hours"] < revised_plan["hard_cap_h20_hours"]
    assert (
        "--max-num-batched-tokens 2048"
        in revised_plan["sessions"][0]["commands"]["server"]
    )
    print("action-aware constraint pilot: PASS")


if __name__ == "__main__":
    main()