merge(search-with-proof): cluster-unavailable handling [skip ci]

No image rebuild: VM-side changes only (operator asked to stop per-merge rebuilds).

Co-Authored-By: Claude Fable 5.1 <[email protected]>

# Conflicts:
#	benchmarks/ax/batch.py

benchmarks/ax/batch.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 import hashlib
 import ipaddress
 import json
+import logging
 import os
 import queue
 import random
@@ -40,6 +41,10 @@ class NoFreeWorkers(Exception):
     """Capacity back-pressure, never a sandbox or instrument failure."""
 
 
+class ClusterUnavailable(Exception):
+    """No running workers were observed in the worker namespace."""
+
+
 class CapacityDeadline(Exception):
     """No worker became available before the task deadline."""
 
@@ -51,6 +56,7 @@ AX_SLOTS = threading.BoundedSemaphore(2)
 
 def apply_when_available(lane, document, end):
     started = lane.clock()
+    next_health_check = started + 60
     while lane.clock() < end:
         try:
             lane.ax("apply", "-f", "-", manifest=document, timeout=min(30, end - lane.clock()))
@@ -62,6 +68,9 @@ def apply_when_available(lane, document, end):
                 lane.ax("delete", "task", document["metadata"]["name"])
             except Exception as exc:
                 raise RuntimeError("capacity cleanup failed: " + exception_reason(exc)) from None
+            if lane.clock() >= next_health_check:
+                lane.check_cluster_health(timeout=min(10, max(0.01, end - lane.clock())))
+                next_health_check = lane.clock() + 60
             # Substrate's golden-snapshot copy holds the other worker for about a second after each new task
             # (operator, 2026-10-05): retry quickly for the first minute, then back off while real tasks run.
             delay = random.uniform(2, 5) if lane.clock() - started < 60 else random.uniform(15, 30)
@@ -163,7 +172,8 @@ def bounded_read(read, timeout):
 
 
 class AXLane:
-    def __init__(self, args, workload, store=None, pause=time.sleep, clock=time.monotonic):
+    def __init__(self, args, workload, store=None, pause=time.sleep, clock=time.monotonic, runner=None):
+        self.runner = runner or subprocess.run
         self.workload = workload
         self.args = args
         self.pause, self.clock = pause, clock
@@ -204,11 +214,35 @@ class AXLane:
         record = json.loads((output / "result.json").read_text())
         return {**record, "worker_ip": worker, "ax_attempts": attempts}
 
+    def check_cluster_health(self, *, timeout=10):
+        env = dict(os.environ)
+        if self.args.kubeconfig:
+            env["KUBECONFIG"] = str(Path(self.args.kubeconfig).expanduser())
+        done = self.runner(
+            ["kubectl", "-n", "ax-workers", "get", "pods", "-o", "json"],
+            env=env,
+            capture_output=True,
+            text=True,
+            timeout=timeout,
+        )
+        if done.returncode:
+            raise RuntimeError(f"worker pod health check failed (exit {done.returncode})")
+        try:
+            pods = json.loads(done.stdout)["items"]
+            if not isinstance(pods, list):
+                raise ValueError("invalid pod list")
+            phases = [pod["status"]["phase"] for pod in pods]
+        except (ValueError, KeyError, TypeError):
+            raise RuntimeError("worker pod health check returned invalid pod data") from None
+        running = phases.count("Running")
+        if not running:
+            raise ClusterUnavailable(f"worker pods: total={len(pods)}, running=0, pending={phases.count('Pending')}")
+
     def ax(self, *args, manifest=None, timeout=30):
         env = dict(os.environ)
         if self.args.kubeconfig:
             env["KUBECONFIG"] = str(Path(self.args.kubeconfig).expanduser())
-        done = subprocess.run(
+        done = self.runner(
             [str(Path(self.args.ax_bin).expanduser()), *args, *(["--atespace", self.args.atespace] if args[0] != "apply" else [])],
             input=json.dumps(manifest) if manifest is not None else None,
             env=env,
@@ -632,7 +666,7 @@ class SearchExecutorTask:
         self.end = time.monotonic() + self.deadline
         if not isinstance(self.lane, DockerLane):
             if not AX_SLOTS.acquire(timeout=self.deadline):
-                raise CapacityDeadline
+                raise CapacityDeadline("no free workers available before deadline")
             self.slot = True
         expires = timedelta(seconds=self.deadline + 600)
         store = self.lane.store
@@ -677,8 +711,8 @@ class SearchExecutorTask:
                 raise
             self.client = ObjectStoreExecutor(command_put, result_get)
             return self
-        except Exception:
-            self.__exit__(None, None, None)
+        except Exception as exc:
+            self.__exit__(type(exc), exc, exc.__traceback__)
             raise
 
     def prepare(self, spec):
@@ -692,7 +726,7 @@ class SearchExecutorTask:
             raise ValueError("executor archive missing or digest mismatch")
         return stored[0]
 
-    def __exit__(self, *exc):
+    def __exit__(self, exc_type, exc, traceback):
         errors = []
         if self.submitted:
             try:
@@ -708,4 +742,10 @@ class SearchExecutorTask:
             AX_SLOTS.release()
             self.slot = False
         if errors:
-            raise RuntimeError("executor cleanup failed: " + ", ".join(errors))
+            reason = "executor cleanup failed: " + ", ".join(errors)
+            logging.getLogger(__name__).error("%s", reason)
+            if exc is not None:
+                exc.executor_cleanup_reason = reason
+                exc.args = (str(exc) + "; " + reason,)
+                return False
+            raise RuntimeError(reason)

benchmarks/search/pilot_run.py

@@ -17,7 +17,7 @@ from contextlib import contextmanager
 from dataclasses import replace
 from pathlib import Path
 
-from benchmarks.ax.batch import AXLane, SearchExecutorTask, Workload, validate_image
+from benchmarks.ax.batch import AXLane, ClusterUnavailable, SearchExecutorTask, Workload, validate_image
 from benchmarks.ax.search_verify import AXSideDispatcher, validate_result
 from benchmarks.search.executor_smoke import CLONE
 from openultrasast.config import load_dotenv
@@ -65,6 +65,10 @@ def scripted(schema):
     )
 
 
+class RepositoryFailure(ValueError):
+    """The selected checkout or candidate cannot be read."""
+
+
 class Pilot:
     def __init__(self, args, row, index, root, store, run_budget, client, prices, lane=None):
         self.args, self.root, self.lane = args, root, lane
@@ -156,14 +160,14 @@ class Pilot:
             )
             if result.get("exit_code") != 0 or result.get("error"):
                 self.failure("acquisition", "checkout_command_failed")
-                raise ValueError("checkout failed")
+                raise RepositoryFailure("checkout failed")
             remaining = deadline - (time.monotonic() - acquired_at)
             if remaining < 1:
                 raise BudgetExhausted("acquisition deadline", usd=0, calls=0)
             result = task.client.submit("read_file", {"path": self.row["file"]}, timeout_seconds=min(30, remaining))
             if not isinstance(result.get("input_bytes"), int) or result["input_bytes"] <= 0:
                 self.failure("acquisition", "candidate_input_unreadable")
-                raise ValueError("unreadable input")
+                raise RepositoryFailure("unreadable input")
             self.record["input_bytes"] += result["input_bytes"]
             yield task
 
@@ -281,6 +285,8 @@ class Pilot:
                 raise BudgetExhausted("run ceiling", usd=0, calls=0)
             if result.get("status") != "ok":
                 self.failure(task.step, result.get("failure", "worker_execution_failed"))
+                if self.worker.metrics["model_calls"] == 0:
+                    raise RuntimeError(self.record["failures"][-1]["reason"])
                 if not self.verifications:
                     self.record["outcome"] = "inconclusive"
             return result
@@ -304,6 +310,8 @@ class Pilot:
                 else "inconclusive"
             )
             self.record["end_reason"] = "instrument_failure"
+            if isinstance(exc, ClusterUnavailable) or (self.worker.metrics["model_calls"] == 0 and not isinstance(exc, RepositoryFailure)):
+                self.record["aborted_reason"] = self.record["failures"][-1]["reason"]
         metrics = self.worker.metrics
         self.record.update(
             spend={k: metrics[k] for k in ("reserved", "settled")},
@@ -425,15 +433,20 @@ def main(argv=None):
         client = OpenRouterChatClient(
             api_key=os.environ[DEEPSEEK_KEY_ENV], base_url=os.environ.get(DEEPSEEK_BASE_ENV, DEEPSEEK_BASE_URL), max_attempts=1
         )
-        lane = AXLane(args, Workload(args.image, frozenset({"VERIFY_INPUT_URL", "RESULT_URL"}), 900, validate_result))
     root = args.private_root / uuid.uuid4().hex
     root.mkdir(parents=True, mode=0o700)
-    store = open_store(args.board_memory) if args.board_memory else FileStore(root / "boards")
     records = []
     started = time.monotonic()
     args.out.parent.mkdir(parents=True, exist_ok=True)
     with args.out.open("x") as output:
-        for index in indices:
+        aborted_reason = None
+        try:
+            store = open_store(args.board_memory) if args.board_memory else FileStore(root / "boards")
+            if not args.dry_run:
+                lane = AXLane(args, Workload(args.image, frozenset({"VERIFY_INPUT_URL", "RESULT_URL"}), 900, validate_result))
+        except Exception as exc:
+            aborted_reason = exception_reason(exc)
+        for index in indices if aborted_reason is None else []:
             local = root / str(index)
             local.mkdir()
             if args.dry_run:
@@ -447,8 +460,14 @@ def main(argv=None):
             records.append(record)
             output.write(json.dumps(record, allow_nan=False) + "\n")
             output.flush()
+            if record.get("aborted_reason"):
+                aborted_reason = record["aborted_reason"]
+                break
         summary = dict(
             type="summary",
+            end_reason="instrument_failure" if aborted_reason else "completed",
+            aborted_reason=aborted_reason,
+            skipped_pairs=indices[len(records) :],
             searches=len(records),
             dry_run=args.dry_run,
             side=args.side,
@@ -462,7 +481,7 @@ def main(argv=None):
             wall_seconds=time.monotonic() - started,
         )
         output.write(json.dumps(summary, allow_nan=False) + "\n")
-    return int(any(r["end_reason"] == "instrument_failure" for r in records))
+    return int(bool(aborted_reason) or any(r["end_reason"] == "instrument_failure" for r in records))
 
 
 if __name__ == "__main__":

src/openultrasast/search/task_storage.py

@@ -116,4 +116,12 @@ def diagnostic(value: object, *, maximum: int = 300, private: tuple[str, ...] =
 
 def exception_reason(exc: BaseException, *, private: tuple[str, ...] = ()) -> str:
     prefix = type(exc).__name__ + ": "
-    return prefix + diagnostic(exc, maximum=max(0, 300 - len(prefix)), private=private)
+    maximum = max(0, 300 - len(prefix))
+    cleanup = getattr(exc, "executor_cleanup_reason", None)
+    if cleanup:
+        original = str(exc).removesuffix("; " + cleanup)
+        # Reserve space for both causes: a long teardown must not erase the
+        # original failure from the bounded public record.
+        half = max(0, (maximum - 2) // 2)
+        return prefix + diagnostic(original, maximum=half, private=private) + "; " + diagnostic(cleanup, maximum=half, private=private)
+    return prefix + diagnostic(exc, maximum=maximum, private=private)

tests/test_ax_batch.py

@@ -246,6 +246,8 @@ def test_capacity_backpressure_retries_same_manifest_until_deadline(tmp_path, mo
     busy = [True]
 
     def run(argv, **kwargs):
+        if argv[0] == "kubectl":
+            return subprocess.CompletedProcess(argv, 0, json.dumps({"items": [{"status": {"phase": "Running"}}]}), "")
         if argv[1] == "apply":
             applies.append(kwargs["input"])
             return subprocess.CompletedProcess(argv, 1 if busy[0] else 0, "", "no free workers available secret-url")
@@ -326,3 +328,60 @@ def test_shared_capacity_across_workloads(tmp_path, monkeypatch):
         results = list(pool.map(lambda pair: pair[0](item(), pair[1], 0), zip(lanes, paths, strict=True)))
     assert all(r["status"] == "ok" for r in results)
     assert peak[0] == 2
+
+
+@pytest.mark.parametrize("phases", [[], ["Pending", "Pending"], ["Running", "Running"]])
+def test_capacity_health_uses_lane_runner_and_throttles(monkeypatch, phases):
+    import subprocess
+    from types import SimpleNamespace
+
+    from benchmarks.ax import batch
+
+    now, checks, applies = [0.0], [], []
+    monkeypatch.setenv("S3_ENDPOINT", "https://files.example")
+
+    def jitter(low, high):
+        assert (low, high) == ((2, 5) if now[0] < 60 else (15, 30))
+        return 5 if high == 5 else 20
+
+    monkeypatch.setattr(batch.random, "uniform", jitter)
+
+    def runner(argv, **kwargs):
+        if argv[0] == "kubectl":
+            checks.append(now[0])
+            assert argv == ["kubectl", "-n", "ax-workers", "get", "pods", "-o", "json"]
+            assert kwargs["env"]["KUBECONFIG"] == "/tmp/test-kubeconfig"
+            assert kwargs["timeout"] <= 10
+            return subprocess.CompletedProcess(
+                argv, 0, json.dumps({"items": [{"metadata": {"name": "private-name"}, "status": {"phase": phase}} for phase in phases]}), ""
+            )
+        if argv[1] == "apply":
+            applies.append(kwargs["input"])
+            return subprocess.CompletedProcess(argv, 1, "", "no free workers available")
+        return subprocess.CompletedProcess(argv, 0, "", "")
+
+    lane = batch.AXLane(
+        SimpleNamespace(ax_bin="ax", kubeconfig="/tmp/test-kubeconfig", atespace="default"),
+        workload(),
+        store=object(),
+        runner=runner,
+        clock=lambda: now[0],
+        pause=lambda seconds: now.__setitem__(0, now[0] + seconds),
+    )
+    expected = batch.CapacityDeadline if "Running" in phases else batch.ClusterUnavailable
+    with pytest.raises(expected) as caught:
+        batch.apply_when_available(lane, {"metadata": {"name": "same-task"}}, 190)
+    assert checks == ([60, 120, 180] if "Running" in phases else [60])
+    assert len(set(applies)) == 1
+    assert "private-name" not in str(caught.value)
+    assert str(caught.value)
+
+
+def test_executor_slot_deadline_has_reason(monkeypatch):
+    from types import SimpleNamespace
+
+    from benchmarks.ax import batch
+
+    monkeypatch.setattr(batch, "AX_SLOTS", SimpleNamespace(acquire=lambda **kw: False))
+    with pytest.raises(batch.CapacityDeadline, match="no free workers available before deadline"):
+        batch.SearchExecutorTask(object(), b"archive").__enter__()

tests/test_search_executor_smoke.py

@@ -261,3 +261,28 @@ def test_dispatch_rejects_duplicate_items_before_launch(tmp_path):
 
     with pytest.raises(ValueError, match="duplicate dispatch item"):
         batch.dispatch([item, item], batch.Workload("image", frozenset(), 1, None), tmp_path, {"ax": unexpected}, lanes=("ax",))
+
+
+@pytest.mark.parametrize("stage", ["enter", "body", "normal"])
+def test_cleanup_preserves_original_and_logs(tmp_path, monkeypatch, caplog, stage):
+    _, lane, _ = setup(tmp_path, monkeypatch)
+    original = ValueError("original apply failure")
+
+    def ax(*args, **kwargs):
+        if args[0] == "apply" and stage == "enter":
+            raise original
+        if args[0] == "delete":
+            raise ValueError("delete failed https://private.example/?token=secret")
+
+    lane.ax = ax
+    # Retained reservations must not leak into other tests.
+    monkeypatch.setattr(batch, "AX_SLOTS", __import__("threading").BoundedSemaphore(2))
+    with pytest.raises(RuntimeError if stage == "normal" else ValueError) as caught, batch.SearchExecutorTask(lane, b"archive"):
+        if stage == "body":
+            raise original
+    if stage != "normal":
+        assert caught.value is original
+        assert str(caught.value).startswith("original apply failure; executor cleanup failed:")
+    assert "delete failed" in str(caught.value)
+    assert "executor cleanup failed" in caplog.text
+    assert "private.example" not in caplog.text + str(caught.value)

tests/test_search_pilot_run.py

@@ -362,3 +362,99 @@ def test_explore_exception_keeps_sanitized_message(tmp_path, monkeypatch, throug
     assert "app.py" not in failure["reason"]
     assert "https://" not in failure["reason"]
     assert "private-token-1234" not in failure["reason"]
+
+
+@pytest.mark.parametrize(
+    "kind,calls,aborts", [("cluster", 0, True), ("cluster", 1, True), ("infra", 0, True), ("infra", 1, False), ("repo", 0, False)]
+)
+def test_pilot_classifies_run_abort(tmp_path, monkeypatch, kind, calls, aborts):
+    from benchmarks.ax.batch import ClusterUnavailable
+    from benchmarks.search import pilot_run
+
+    pilot = make_pilot(tmp_path)
+    error = {
+        "cluster": ClusterUnavailable("worker pods: total=2, running=0, pending=2"),
+        "infra": RuntimeError("port-forward unavailable"),
+        "repo": pilot_run.RepositoryFailure("checkout failed"),
+    }[kind]
+
+    def dispatch(task):
+        pilot.worker.metrics["model_calls"] = calls
+        raise error
+
+    monkeypatch.setattr(pilot, "dispatch", dispatch)
+    record = pilot.run()
+    assert bool(record.get("aborted_reason")) is aborts
+    assert record["end_reason"] == "instrument_failure"
+    if aborts:
+        assert record["aborted_reason"] == record["failures"][-1]["reason"]
+
+
+@pytest.mark.parametrize("setup_failure", [False, True])
+def test_live_driver_stops_remaining_pairs_on_infrastructure_failure(tmp_path, monkeypatch, setup_failure):
+    from benchmarks.ax.batch import ClusterUnavailable
+    from benchmarks.search import pilot_run
+
+    monkeypatch.chdir(tmp_path)
+    root = Path("benchmarks/search/private")
+    root.mkdir(parents=True)
+    row = dict(
+        repo="https://github.com/private/project",
+        vulnerable="a" * 40,
+        fixed="b" * 40,
+        family="path",
+        language="python",
+        file="app.py",
+        function="main",
+    )
+    manifest = root / "pilot.json"
+    manifest.write_text(json.dumps({"pairs": [row] * 7}))
+    monkeypatch.setattr(pilot_run, "load_dotenv", lambda: None)
+    monkeypatch.setenv(pilot_run.DEEPSEEK_KEY_ENV, "fake-test-key")
+    monkeypatch.setattr(pilot_run, "OpenRouterChatClient", lambda **kw: __import__("types").SimpleNamespace(max_attempts=1))
+
+    def lane(*args):
+        if setup_failure:
+            raise RuntimeError("store unavailable")
+        return object()
+
+    monkeypatch.setattr(pilot_run, "AXLane", lane)
+    attempts = []
+
+    def enter(task):
+        attempts.append(task.name)
+        raise ClusterUnavailable("worker pods: total=2, running=0, pending=2")
+
+    monkeypatch.setattr(pilot_run.SearchExecutorTask, "__enter__", enter)
+    output = tmp_path / "records.jsonl"
+    status = main(
+        [
+            "--side",
+            "vulnerable",
+            "--pairs",
+            "2,5,6",
+            "--ceiling-usd",
+            "1",
+            "--verify-lane",
+            "ax",
+            "--image",
+            "image@sha256:" + "a" * 64,
+            "--out",
+            str(output),
+        ]
+    )
+    records = [json.loads(line) for line in output.read_text().splitlines()]
+    assert status == 1
+    if setup_failure:
+        assert not attempts and len(records) == 1
+        assert records[0]["aborted_reason"] == "RuntimeError: store unavailable"
+        assert records[0]["skipped_pairs"] == [2, 5, 6]
+        assert records[0]["searches"] == 0
+        return
+    assert len(attempts) == 1
+    search, summary = records
+    assert search["pair_index"] == 2 and search["model_calls"] == 0
+    assert summary["searches"] == 1 and summary["skipped_pairs"] == [5, 6]
+    assert summary["end_reason"] == "instrument_failure"
+    assert summary["aborted_reason"] == search["aborted_reason"]
+    assert "running=0" in summary["aborted_reason"]

tests/test_search_task_storage.py

@@ -222,3 +222,21 @@ def test_scratch_failure_preserves_exit_and_stderr(workspace, monkeypatch):
     assert result["exit_code"] is not None
     assert result["stderr"] == "pip build detail"
     assert result["reason"].startswith("ScratchLimit:")
+
+
+def test_capacity_deadline_plain_diagnostic_survives():
+    from benchmarks.ax.batch import CapacityDeadline
+
+    message = "no free workers available before deadline"
+    assert task_storage.diagnostic(message) == message
+    assert task_storage.exception_reason(CapacityDeadline(message)) == "CapacityDeadline: " + message
+
+
+def test_cleanup_diagnostic_retains_both_causes():
+    error = ValueError("original apply failed")
+    error.executor_cleanup_reason = "executor cleanup failed: " + "details " * 100 + "delete failed"
+    error.args = (str(error) + "; " + error.executor_cleanup_reason,)
+    reason = task_storage.exception_reason(error)
+    assert reason.startswith("ValueError: original apply failed;")
+    assert reason.endswith("delete failed")
+    assert len(reason) <= 300