293 lines
9.5 KiB
Python
293 lines
9.5 KiB
Python
#!/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()
|