Revise action-aware pilot after token overload

runs/action-aware-v0/analyze_pilot.py

@@ -454,7 +454,10 @@ def evaluate_decisions(
 
 def analyze(run_root: Path, manifest_path: Path) -> dict[str, Any]:
     manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
-    if manifest.get("schema") != "action-aware-constraint-pilot-manifest-v0":
+    if manifest.get("schema") not in {
+        "action-aware-constraint-pilot-manifest-v0",
+        "action-aware-constraint-pilot-manifest-v1",
+    }:
         raise ValueError("unexpected manifest schema")
     duration_s = float(manifest["engine"]["duration_s"])
     phase_fractions = [float(value) for value in manifest["gates"]["phase_fractions"]]

runs/action-aware-v0/pilot-manifest-v1.json

@@ -0,0 +1,227 @@
+{
+  "budget": {
+    "expected_h20_hours": [
+      6.0,
+      7.2
+    ],
+    "expected_wall_minutes": [
+      90,
+      110
+    ],
+    "global_hard_cap_h20_hours": 8.0,
+    "hard_cap_h20_hours": 7.614013100465138,
+    "prior_attempt_artifact": "/home/admin/cpfs/wjh/action-aware-constraint-v0-20260714/operational-stop-v0.json",
+    "prior_attempt_h20_hours": 0.38598689953486126,
+    "safety_h20_hours": 0.25,
+    "session_estimate_h20_hours": 1.35
+  },
+  "burnin": {
+    "anchor": 0.18919793755240089,
+    "arrival_order_sha256": "6c0ac4cb9a30ef501eeeacc8e6cc631c345e976db5ccf530ea5a1ec706d62a24",
+    "input_length_order_sha256": "7939cc20e1a00d1031d27d71508789f38decbbbb6ea59a1df18b2ec342fd2ef8",
+    "offered_req_s": 8.5,
+    "offered_req_s_per_gpu": 2.125,
+    "request_id_order_sha256": "84f4809acbc8acd3b1d14dfa357134a1dc0b9287341624b33f598dafeef54dc7",
+    "selected_count": 510,
+    "study": "/home/admin/cpfs/wjh/fidelity-prefix-pilot-20260714/private/studies/burnin-tp4.json",
+    "study_sha256": "5d6c2098042909a863efd3112818fbee9bafe96f22898ac98b66846dbe1fef0f"
+  },
+  "configs": [
+    {
+      "id": "b_base",
+      "mbbt": 2048,
+      "mns": 64,
+      "repetition_order": [
+        1,
+        2,
+        3
+      ]
+    },
+    {
+      "id": "a_base",
+      "mbbt": 8192,
+      "mns": 16,
+      "repetition_order": [
+        2,
+        3,
+        1
+      ]
+    },
+    {
+      "id": "shared",
+      "mbbt": 8192,
+      "mns": 64,
+      "repetition_order": [
+        3,
+        1,
+        2
+      ]
+    },
+    {
+      "id": "b_mns",
+      "mbbt": 2048,
+      "mns": 128,
+      "repetition_order": [
+        1,
+        3,
+        2
+      ]
+    },
+    {
+      "id": "a_mbbt",
+      "mbbt": 16384,
+      "mns": 16,
+      "repetition_order": [
+        2,
+        1,
+        3
+      ]
+    }
+  ],
+  "engine": {
+    "burnin_max_elapsed_s": 90.0,
+    "client_timeout_s": 450.0,
+    "disable_slo_early_stop": true,
+    "duration_s": 300.0,
+    "tp": 4
+  },
+  "gates": {
+    "material_kv_usage": 0.9,
+    "minimum_exclusive_fraction": 0.1,
+    "minimum_exclusive_ratio": 5.0,
+    "minimum_relative_winner_margin": 0.1,
+    "phase_fractions": [
+      0.25,
+      0.5,
+      0.75,
+      1.0
+    ]
+  },
+  "regimes": {
+    "A": {
+      "actions": {
+        "mbbt": "a_mbbt",
+        "mns": "shared"
+      },
+      "source": "a_base"
+    },
+    "B": {
+      "actions": {
+        "mbbt": "shared",
+        "mns": "b_mns"
+      },
+      "source": "b_base"
+    }
+  },
+  "repetitions": {
+    "1": {
+      "merged_trace": {
+        "bytes": 337429767,
+        "path": "/home/admin/cpfs/wjh/intervention-response-v3-20260714/private/traces/rep1.jsonl",
+        "request_id_scheme": "sha256(source_sha256:line_number:original_id)",
+        "rows": 9420,
+        "sha256": "68983266aa0e66aa589562f7c08edbd966f9ba4405e20c105adb43777d2dfbf5",
+        "source_sha256": [
+          "b242d1d9086df3accab57b4c92445d5edd581e12f47e12cea227aa63964c6930",
+          "d23b549f7b69af3647308677bbf76f818a3c226a1c98f9a9f93f09ceee46be87"
+        ],
+        "sources": [
+          "/home/admin/cpfs/wjh/fidelity-prefix-pilot-20260714/private/traces/low1.jsonl",
+          "/home/admin/cpfs/wjh/fidelity-prefix-pilot-20260714/private/traces/high1.jsonl"
+        ]
+      },
+      "selection": {
+        "anchor": 0.48686986110831465,
+        "arrival_order_sha256": "c2ad99986ce558da5901a9c5ec0a00bd69f198c981d8779235f2773a5c87f1c0",
+        "input_length_order_sha256": "9442bfebdc3fab5062dc1f4d688dc28c02afe3fd806c56dd8159f0ac7e6d0b94",
+        "offered_req_s": 8.5,
+        "offered_req_s_per_gpu": 2.125,
+        "request_id_order_sha256": "0bb61dbc9c26875e991d0d4f984134910d37463e5063f86ee960cf4f8aafb771",
+        "selected_count": 2550,
+        "target_count": 2550,
+        "target_req_s_per_gpu": 2.125
+      },
+      "study": "/home/admin/cpfs/wjh/intervention-response-v3-20260714/private/studies/rep1-tp4.json",
+      "study_sha256": "ecfff96e33d458eb1e3b9a6d24386f00cc6f1b19ff926e2ec6320b3f671a7ae3"
+    },
+    "2": {
+      "merged_trace": {
+        "bytes": 337509330,
+        "path": "/home/admin/cpfs/wjh/intervention-response-v3-20260714/private/traces/rep2.jsonl",
+        "request_id_scheme": "sha256(source_sha256:line_number:original_id)",
+        "rows": 9457,
+        "sha256": "f38e8938f6a481fc6725b71b21aa04ff7eaf79783cdfd6e41aa2f074156f00c2",
+        "source_sha256": [
+          "4cbb0baac082bd54af562ce2f39104c5c23b4671672da365a67b1e8c146adf9f",
+          "bb0bcd2564a88000f435f12feb21c7c902eafc9ea5fe916adfe9d1eae47f3f9a"
+        ],
+        "sources": [
+          "/home/admin/cpfs/wjh/fidelity-prefix-pilot-20260714/private/traces/low2.jsonl",
+          "/home/admin/cpfs/wjh/fidelity-prefix-pilot-20260714/private/traces/high2.jsonl"
+        ]
+      },
+      "selection": {
+        "anchor": 0.4825698948735577,
+        "arrival_order_sha256": "b9fc12cf3f86bc8a79bee65296e65aa2b8bf2aeca46b2887094c669adcbb9a00",
+        "input_length_order_sha256": "d8d4bd6fc8ba852a45605b673b6b3e4f33b58f459e69f2a032d226ee175b074e",
+        "offered_req_s": 8.5,
+        "offered_req_s_per_gpu": 2.125,
+        "request_id_order_sha256": "56a0616b6b54abafd37875c7cb25f8639afef2706ccc55dfbe568f45859ea382",
+        "selected_count": 2550,
+        "target_count": 2550,
+        "target_req_s_per_gpu": 2.125
+      },
+      "study": "/home/admin/cpfs/wjh/intervention-response-v3-20260714/private/studies/rep2-tp4.json",
+      "study_sha256": "d92a576db031db24bb58f354ea725d7f7567cb76699d387117ac5a6c9317bbb9"
+    },
+    "3": {
+      "merged_trace": {
+        "bytes": 337450256,
+        "path": "/home/admin/cpfs/wjh/intervention-response-v3-20260714/private/traces/rep3.jsonl",
+        "request_id_scheme": "sha256(source_sha256:line_number:original_id)",
+        "rows": 9431,
+        "sha256": "3094084b0bb20cc02eecf465091a5c919b4e5b112f704cdc36a563d1efdcee46",
+        "source_sha256": [
+          "1f7ececb142f9a363d2d1ca25eb7b8488b2cc319a51b55faa384f2a3d51f2142",
+          "6f326234791e1cff4ff866bface0d097d0d6e3844eebb1c97653d8e9c35e9397"
+        ],
+        "sources": [
+          "/home/admin/cpfs/wjh/fidelity-prefix-pilot-20260714/private/traces/low3.jsonl",
+          "/home/admin/cpfs/wjh/fidelity-prefix-pilot-20260714/private/traces/high3.jsonl"
+        ]
+      },
+      "selection": {
+        "anchor": 0.48664343020532463,
+        "arrival_order_sha256": "efce7339e22d3618cb4d55e6b55bfddb2c563c18faba2a992d5829c13e3f55e9",
+        "input_length_order_sha256": "0792b05fff6729fbd92ab2bb4cb6d31bea7799e232ad42772936bc06efbafb54",
+        "offered_req_s": 8.5,
+        "offered_req_s_per_gpu": 2.125,
+        "request_id_order_sha256": "2a2fabe2c4cf176aeb7e0d32fb8e7dbb1f27429a2e7a0cd18d7d186f23096f19",
+        "selected_count": 2550,
+        "target_count": 2550,
+        "target_req_s_per_gpu": 2.125
+      },
+      "study": "/home/admin/cpfs/wjh/intervention-response-v3-20260714/private/studies/rep3-tp4.json",
+      "study_sha256": "fb8ffe256dace32f4ca8a8d49b662d98c3b69b94ecc8fa826e43068b238884ab"
+    }
+  },
+  "sanity": {
+    "invariants": {
+      "all_repetition_orders_are_permutations": true,
+      "five_unique_configs": true,
+      "same_load_all_repetitions": true,
+      "shared_endpoint_reused_by_both_regimes": true,
+      "three_disjoint_repetitions": true
+    },
+    "red_flags": []
+  },
+  "schema": "action-aware-constraint-pilot-manifest-v1",
+  "source": {
+    "base_manifest": "/home/gahow/phd/aituner/runs/intervention-response-v2/pilot-manifest-v3.json",
+    "base_manifest_sha256": "273db1181dcc9d6b64439650d0642ebe553b12e6aa9adebfbe3758a7977e5611",
+    "source_trace": "/home/admin/cpfs/wjh/aituner/aituner/trace_windows/traces/chat_w20260312_1000.jsonl",
+    "source_trace_sha256": "875ba869775deb78086477919f03b322da14e2673c7d070e26528c4190912757",
+    "window_id": "chat_w20260312_1000"
+  },
+  "status": "PASS"
+}

runs/action-aware-v0/pilot_controller.py

@@ -56,7 +56,10 @@ def configure(args: argparse.Namespace, manifest: Mapping[str, Any]) -> None:
 
 
 def validate_inputs(args: argparse.Namespace, manifest: Mapping[str, Any]) -> None:
-    if manifest.get("schema") != "action-aware-constraint-pilot-manifest-v0":
+    if manifest.get("schema") not in {
+        "action-aware-constraint-pilot-manifest-v0",
+        "action-aware-constraint-pilot-manifest-v1",
+    }:
         raise RuntimeError("unexpected action-aware manifest schema")
     if manifest.get("status") != "PASS":
         raise RuntimeError("action-aware manifest did not pass preflight")
@@ -455,7 +458,7 @@ def execute_session(
         first = manifest["repetitions"][str(config["repetition_order"][0])]
         session_state["status"] = "warmup"
         atomic_json(state_path, state)
-        run_client(
+        burnin_result = run_client(
             entry=entry,
             config=config,
             role="warmup",
@@ -479,6 +482,19 @@ def execute_session(
             state=state,
             timeout_s=float(manifest["engine"]["client_timeout_s"]),
         )
+        burnin_elapsed_s = float(burnin_result["interval"]["elapsed_s"])
+        session_state["burnin"] = {
+            "elapsed_s": burnin_elapsed_s,
+            "pass_rate": float(burnin_result["pass_rate"]),
+            "feasible": bool(burnin_result["feasible"]),
+        }
+        atomic_json(state_path, state)
+        burnin_limit_s = float(manifest["engine"]["burnin_max_elapsed_s"])
+        if burnin_elapsed_s > burnin_limit_s:
+            raise RuntimeError(
+                f"burnin throughput gate failed: {burnin_elapsed_s:.3f}s > "
+                f"{burnin_limit_s:.3f}s"
+            )
         session_state["status"] = "measured"
         atomic_json(state_path, state)
         for repetition in config["repetition_order"]:

runs/action-aware-v0/prepare_pilot.py

@@ -11,14 +11,28 @@ from pathlib import Path
 from typing import Any
 
 
-SCHEMA = "action-aware-constraint-pilot-manifest-v0"
-CONFIGS = (
-    {"id": "b_base", "mns": 64, "mbbt": 256, "repetition_order": [1, 2, 3]},
-    {"id": "a_base", "mns": 16, "mbbt": 8192, "repetition_order": [2, 3, 1]},
-    {"id": "shared", "mns": 64, "mbbt": 8192, "repetition_order": [3, 1, 2]},
-    {"id": "b_mns", "mns": 128, "mbbt": 256, "repetition_order": [1, 3, 2]},
-    {"id": "a_mbbt", "mns": 16, "mbbt": 16384, "repetition_order": [2, 1, 3]},
-)
+SCHEMA_V0 = "action-aware-constraint-pilot-manifest-v0"
+SCHEMA_V1 = "action-aware-constraint-pilot-manifest-v1"
+
+
+def configs(token_source_mbbt: int) -> tuple[dict[str, Any], ...]:
+    return (
+        {
+            "id": "b_base",
+            "mns": 64,
+            "mbbt": token_source_mbbt,
+            "repetition_order": [1, 2, 3],
+        },
+        {"id": "a_base", "mns": 16, "mbbt": 8192, "repetition_order": [2, 3, 1]},
+        {"id": "shared", "mns": 64, "mbbt": 8192, "repetition_order": [3, 1, 2]},
+        {
+            "id": "b_mns",
+            "mns": 128,
+            "mbbt": token_source_mbbt,
+            "repetition_order": [1, 3, 2],
+        },
+        {"id": "a_mbbt", "mns": 16, "mbbt": 16384, "repetition_order": [2, 1, 3]},
+    )
 
 
 def sha256_file(path: Path) -> str:
@@ -38,7 +52,17 @@ def atomic_json(path: Path, payload: Any) -> None:
     os.replace(temporary, path)
 
 
-def build(base_path: Path) -> dict[str, Any]:
+def build(
+    base_path: Path,
+    *,
+    token_source_mbbt: int = 256,
+    prior_attempt_h20_hours: float = 0.0,
+    prior_attempt_artifact: str | None = None,
+) -> dict[str, Any]:
+    if token_source_mbbt <= 0:
+        raise ValueError("token source MBBT must be positive")
+    if prior_attempt_h20_hours < 0.0 or prior_attempt_h20_hours >= 8.0:
+        raise ValueError("prior attempt cost must be in [0, 8)")
     base = json.loads(base_path.read_text(encoding="utf-8"))
     if base.get("schema") != "intervention-response-phase-aware-pilot-manifest-v3":
         raise ValueError("unexpected base manifest schema")
@@ -60,9 +84,15 @@ def build(base_path: Path) -> dict[str, Any]:
             "merged_trace": source["merged_trace"],
         }
 
-    config_ids = [str(config["id"]) for config in CONFIGS]
+    frozen_configs = configs(token_source_mbbt)
+    config_ids = [str(config["id"]) for config in frozen_configs]
+    schema = (
+        SCHEMA_V0
+        if token_source_mbbt == 256 and prior_attempt_h20_hours == 0.0
+        else SCHEMA_V1
+    )
     payload = {
-        "schema": SCHEMA,
+        "schema": schema,
         "status": "PASS",
         "source": {
             "base_manifest": str(base_path.resolve()),
@@ -76,10 +106,11 @@ def build(base_path: Path) -> dict[str, Any]:
             "duration_s": 300.0,
             "disable_slo_early_stop": True,
             "client_timeout_s": 450.0,
+            "burnin_max_elapsed_s": 90.0,
         },
         "burnin": base["burnin"],
         "repetitions": repetitions,
-        "configs": [dict(config) for config in CONFIGS],
+        "configs": [dict(config) for config in frozen_configs],
         "regimes": {
             "A": {
                 "source": "a_base",
@@ -91,7 +122,10 @@ def build(base_path: Path) -> dict[str, Any]:
             },
         },
         "budget": {
-            "hard_cap_h20_hours": 8.0,
+            "global_hard_cap_h20_hours": 8.0,
+            "hard_cap_h20_hours": 8.0 - prior_attempt_h20_hours,
+            "prior_attempt_h20_hours": prior_attempt_h20_hours,
+            "prior_attempt_artifact": prior_attempt_artifact,
             "session_estimate_h20_hours": 1.35,
             "safety_h20_hours": 0.25,
             "expected_h20_hours": [6.0, 7.2],
@@ -117,7 +151,7 @@ def build(base_path: Path) -> dict[str, Any]:
                 == 1,
                 "all_repetition_orders_are_permutations": all(
                     sorted(config["repetition_order"]) == [1, 2, 3]
-                    for config in CONFIGS
+                    for config in frozen_configs
                 ),
             }
         },
@@ -141,8 +175,16 @@ def main() -> None:
     parser = argparse.ArgumentParser()
     parser.add_argument("--base-manifest", type=Path, required=True)
     parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument("--token-source-mbbt", type=int, default=256)
+    parser.add_argument("--prior-attempt-h20-hours", type=float, default=0.0)
+    parser.add_argument("--prior-attempt-artifact")
     args = parser.parse_args()
-    payload = build(args.base_manifest)
+    payload = build(
+        args.base_manifest,
+        token_source_mbbt=args.token_source_mbbt,
+        prior_attempt_h20_hours=args.prior_attempt_h20_hours,
+        prior_attempt_artifact=args.prior_attempt_artifact,
+    )
     atomic_json(args.output, payload)
     print(json.dumps(payload["sanity"], sort_keys=True))
     if payload["status"] != "PASS":

runs/action-aware-v0/test_pilot.py

@@ -184,6 +184,23 @@ def main() -> None:
     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")