#!/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 # A per-step stream may have a submit gap above one second when the # preceding model execution itself spans that interval. Such a gap is # covered telemetry, not a dropped-record interval. asynchronous = [ {"submit_mono_ns": 0, "complete_mono_ns": 1_200_000_000}, {"submit_mono_ns": 1_100_000_000, "complete_mono_ns": 1_300_000_000}, ] coverage, covered = analysis.telemetry_coverage( asynchronous, start_ns=0, end_ns=1_100_000_000 ) assert coverage["max_internal_submit_gap_s"] == 1.1 assert coverage["max_uncovered_gap_s"] == 0.0 assert covered missing = copy.deepcopy(asynchronous) missing[0]["complete_mono_ns"] = 0 assert not analysis.telemetry_coverage( missing, start_ns=0, end_ns=1_100_000_000 )[1] mechanism = analysis.mechanism_summary( [ { "model_executed": True, "submit_mono_ns": 0, "complete_mono_ns": 2_000_000, "prefill_tokens": 8, "prefill_requests": 2, "chunked_prefill": { "first": 1, "middle": 0, "final": 0, "unsplit": 1, "tokens": 8, }, "prefix": {"local": {"queries": 10, "hits": 2}}, }, { "model_executed": True, "submit_mono_ns": 2_000_000, "complete_mono_ns": 3_000_000, "prefill_tokens": 0, "prefill_requests": 0, "chunked_prefill": { "first": 0, "middle": 0, "final": 0, "unsplit": 0, "tokens": 0, }, "prefix": {"local": {"queries": 0, "hits": 0}}, }, ] ) assert mechanism["prefill"]["requests_per_step"] == 2.0 assert mechanism["prefill"]["chunks"]["first"] == 1 assert mechanism["prefix"]["hit_rate"] == 0.2 assert all(mechanism["sanity"]["invariants"].values()) 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"] ) accepted_burnin = { "kind": "anchor", "selection": {"count": 510}, "interval": {"elapsed_s": 61.25}, "pass_rate": 0.5, "feasible": False, } assert controller.burnin_gate( accepted_burnin, expected_count=510, maximum_elapsed_s=90.0 )["elapsed_s"] == 61.25 warmup = copy.deepcopy(accepted_burnin) warmup["kind"] = "warmup" try: controller.burnin_gate(warmup, expected_count=510, maximum_elapsed_s=90.0) except RuntimeError as error: assert "non-anchor" in str(error) else: raise AssertionError("warmup incorrectly passed the burnin gate") slow = copy.deepcopy(accepted_burnin) slow["interval"]["elapsed_s"] = 91.0 try: controller.burnin_gate(slow, expected_count=510, maximum_elapsed_s=90.0) except RuntimeError as error: assert "throughput gate failed" in str(error) else: raise AssertionError("slow burnin incorrectly passed the throughput gate") print("action-aware constraint pilot: PASS") if __name__ == "__main__": main()