Add crossed-constraint action-aware pilot

runs/action-aware-v0/test_pilot.py

@@ -0,0 +1,191 @@
+#!/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"]
+    print("action-aware constraint pilot: PASS")
+
+
+if __name__ == "__main__":
+    main()