From 93daf291f6fd14d4ab1f732dcc8625faa51e9d71 Mon Sep 17 00:00:00 2001
From: Gahow Wang <gahow.wang@gmail.com>
Date: Tue, 14 Jul 2026 12:58:01 +0800
Subject: [PATCH] Fix prefix pilot warmup validation

---
 runs/fidelity-headroom/pilot_controller.py | 37 +++++++++++++---
 runs/fidelity-headroom/test_pilot_tools.py | 51 ++++++++++++++++++++++
 2 files changed, 82 insertions(+), 6 deletions(-)

diff --git a/runs/fidelity-headroom/pilot_controller.py b/runs/fidelity-headroom/pilot_controller.py
index e63c9dd..6ffb4f1 100644
--- a/runs/fidelity-headroom/pilot_controller.py
+++ b/runs/fidelity-headroom/pilot_controller.py
@@ -233,8 +233,37 @@ def run_client(
             f"client failed: cell={entry['cell']} role={role} rc={process.returncode}"
         )
     result = json.loads((output / "result.json").read_text(encoding="utf-8"))
+    validate_result_selection(
+        result=result,
+        selection=selection,
+        cell=entry["cell"],
+        role=role,
+        warmup=warmup,
+    )
+    entry["results"].append(
+        {"anchor": float(selection["anchor"]), "dir": str(output), "kind": result["kind"]}
+    )
+    return result
+
+
+def validate_result_selection(
+    *,
+    result: dict[str, Any],
+    selection: dict[str, Any],
+    cell: str,
+    role: str,
+    warmup: bool,
+) -> None:
+    if warmup:
+        if result["kind"] != "warmup" or int(result["selection"]["count"]) != 16:
+            raise RuntimeError(f"invalid warmup selection: {cell} {role}")
+        for key in ("warmup_16", "warmup_exact_16", "warmup_long"):
+            if not result["invariants"].get(key, False):
+                raise RuntimeError(f"warmup invariant {key} failed: {cell} {role}")
+        return
+
     if int(result["selection"]["count"]) != int(selection["selected_count"]):
-        raise RuntimeError(f"selection count mismatch: {entry['cell']} {role}")
+        raise RuntimeError(f"selection count mismatch: {cell} {role}")
     for key in (
         "request_id_order_sha256",
         "arrival_order_sha256",
@@ -244,11 +273,7 @@ def run_client(
             "input_length_order_sha256" if key == "raw_length_order_sha256" else key
         )
         if result["selection"][key] != selection[manifest_key]:
-            raise RuntimeError(f"selection hash mismatch {key}: {entry['cell']} {role}")
-    entry["results"].append(
-        {"anchor": float(selection["anchor"]), "dir": str(output), "kind": result["kind"]}
-    )
-    return result
+            raise RuntimeError(f"selection hash mismatch {key}: {cell} {role}")
 
 
 def execute_cell(
diff --git a/runs/fidelity-headroom/test_pilot_tools.py b/runs/fidelity-headroom/test_pilot_tools.py
index a2e1a32..f5a674f 100644
--- a/runs/fidelity-headroom/test_pilot_tools.py
+++ b/runs/fidelity-headroom/test_pilot_tools.py
@@ -78,6 +78,57 @@ def main() -> None:
         ) + controller.SAFETY_H20_HOURS,
         3.0,
     )
+    selection = {
+        "selected_count": 122,
+        "request_id_order_sha256": "request-hash",
+        "arrival_order_sha256": "arrival-hash",
+        "input_length_order_sha256": "length-hash",
+    }
+    warmup = {
+        "kind": "warmup",
+        "selection": {"count": 16},
+        "invariants": {
+            "warmup_16": True,
+            "warmup_exact_16": True,
+            "warmup_long": True,
+        },
+    }
+    controller.validate_result_selection(
+        result=warmup,
+        selection=selection,
+        cell="tp1_mns8",
+        role="burnin",
+        warmup=True,
+    )
+    measured = {
+        "kind": "anchor",
+        "selection": {
+            "count": 122,
+            "request_id_order_sha256": "request-hash",
+            "arrival_order_sha256": "arrival-hash",
+            "raw_length_order_sha256": "length-hash",
+        },
+        "invariants": {},
+    }
+    controller.validate_result_selection(
+        result=measured,
+        selection=selection,
+        cell="tp1_mns8",
+        role="low1",
+        warmup=False,
+    )
+    try:
+        controller.validate_result_selection(
+            result=warmup,
+            selection=selection,
+            cell="tp1_mns8",
+            role="low1",
+            warmup=False,
+        )
+    except RuntimeError as error:
+        assert "selection count mismatch" in str(error)
+    else:
+        raise AssertionError("measured selection accepted a warmup subset")
     print("fidelity pilot tools: PASS")