diff --git a/benchmarks/ax/batch.py b/benchmarks/ax/batch.py index 24708e0..47eeb6d 100644 --- a/benchmarks/ax/batch.py +++ b/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) diff --git a/benchmarks/search/pilot_run.py b/benchmarks/search/pilot_run.py index a36ea0f..477dc35 100644 --- a/benchmarks/search/pilot_run.py +++ b/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__": diff --git a/src/openultrasast/search/task_storage.py b/src/openultrasast/search/task_storage.py index 071ac78..456d5f4 100644 --- a/src/openultrasast/search/task_storage.py +++ b/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) diff --git a/tests/test_ax_batch.py b/tests/test_ax_batch.py index 94e902f..33505b9 100644 --- a/tests/test_ax_batch.py +++ b/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__() diff --git a/tests/test_search_executor_smoke.py b/tests/test_search_executor_smoke.py index 9d81e31..ca93bc8 100644 --- a/tests/test_search_executor_smoke.py +++ b/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) diff --git a/tests/test_search_pilot_run.py b/tests/test_search_pilot_run.py index 1a64915..63f62a3 100644 --- a/tests/test_search_pilot_run.py +++ b/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"] diff --git a/tests/test_search_task_storage.py b/tests/test_search_task_storage.py index 3a45e4f..abf3e26 100644 --- a/tests/test_search_task_storage.py +++ b/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