Add fidelity-aware verification pilot

runs/fidelity-headroom/test_pilot_tools.py

@@ -0,0 +1,85 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import json
+import math
+import sys
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+
+
+HERE = Path(__file__).resolve().parent
+sys.path.insert(0, str(HERE))
+
+import pilot_controller as controller  # noqa: E402
+import prepare_pilot as prepare  # noqa: E402
+
+
+@dataclass
+class Request:
+    row_id: str
+    sampling_u: float
+    arrival_s: float = 0.0
+    prompt_tokens_hint: int = 1
+
+
+def main() -> None:
+    requests = [
+        Request("a", 0.1),
+        Request("b", 0.2),
+        Request("c", 0.2),
+        Request("d", 0.9),
+    ]
+    anchor, selected = prepare.attainable_anchor(requests, target_count=2)
+    assert anchor == 0.2
+    assert [request.row_id for request in selected] == ["a", "b", "c"]
+
+    with tempfile.TemporaryDirectory() as temporary:
+        root = Path(temporary)
+        source = root / "source.jsonl"
+        rows = []
+        for index, role in enumerate(prepare.ROLES):
+            rows.append(
+                {
+                    "request_id": role,
+                    "timestamp": float(index),
+                    "sampling_u": (index + 0.5) / len(prepare.ROLES),
+                    "input_length": 16 + index,
+                    "messages": [{"role": "user", "content": role}],
+                }
+            )
+        source.write_text(
+            "".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8"
+        )
+        windows, stats = prepare.materialize_bands(
+            source,
+            {
+                "window_id": "source",
+                "trace_type": "chat",
+                "window_start": 0.0,
+                "window_end": 600.0,
+            },
+            root / "private",
+        )
+        assert windows.is_file()
+        assert all(stats[role]["rows"] == 1 for role in prepare.ROLES)
+        for role in prepare.ROLES:
+            row = json.loads((root / "private" / "traces" / f"{role}.jsonl").read_text())
+            assert row["fidelity_pilot_band"] == role
+            assert abs(float(row["sampling_u"]) - 0.5) < 1e-12
+
+    assert len(controller.ORDER) == 6
+    assert set(controller.ORDER) == set(prepare.CELLS)
+    assert math.isclose(
+        sum(
+            controller.CELL_ESTIMATE_H20_HOURS[int(config["tp"])]
+            for config in prepare.CELLS.values()
+        ) + controller.SAFETY_H20_HOURS,
+        3.0,
+    )
+    print("fidelity pilot tools: PASS")
+
+
+if __name__ == "__main__":
+    main()