From 02b2cbf99c8693ce3f195e9f896715ad82c911da Mon Sep 17 00:00:00 2001 From: Marius Ciepluch <11855163+norandom@users.noreply.github.com> Date: Sun, 4 Oct 2026 17:14:20 +0200 Subject: [PATCH] engine: retry a pin once when its AX sandbox dies, then hand it to the VM lane Two engine actors died at gVisor level during a 12-pin run (AX phase Failed, no result). Implemented with Codex. Co-Authored-By: Claude Opus 5.5 --- benchmarks/learn/engine_trace.py | 14 ++- benchmarks/learn/engine_trace_ax.py | 49 +++++++++-- tests/test_ax_executor.py | 131 +++++++++++++++++++++++++++- 3 files changed, 181 insertions(+), 13 deletions(-) diff --git a/benchmarks/learn/engine_trace.py b/benchmarks/learn/engine_trace.py index 0af7c05..15d56a0 100644 --- a/benchmarks/learn/engine_trace.py +++ b/benchmarks/learn/engine_trace.py @@ -668,7 +668,7 @@ def run_pin(pin, source, provenance, inputs, out, args): if done is None and removed.returncode != 0: raise RuntimeError(f"cannot remove timed-out container {name}; refusing another launch") result_path = output / "result.json" - if not result_path.exists(): + if not result_path.exists() and reason != "AX sandbox failed twice": reason = f"container produced no result: {reason}" record = json.loads(result_path.read_text()) if result_path.exists() else {**pin, "units": []} if not record.get("done"): @@ -707,8 +707,12 @@ def run_pin(pin, source, provenance, inputs, out, args): if executor == "ax": record.update(executor="ax", worker=node, worker_ip=node, image=args.image) record["analyzer"] = {"image": args.image, "code": "baked-in-image"} + if (output / "ax_attempts.json").exists(): + record["ax_attempts"] = json.loads((output / "ax_attempts.json").read_text()) if getattr(args, "fallback", None): record.update(executor="docker", fallback=args.fallback, ax_attempt=args.ax_attempt, fallback_pending=False) + if "ax_attempts" in args.ax_attempt: + record["ax_attempts"] = args.ax_attempt["ax_attempts"] if executor in {"k8s", "k8s-jobs"}: record.update(executor="k8s", node=node, image=args.image) if executor == "queue": @@ -815,7 +819,7 @@ def run_plan(plan, source, provenance, inputs, args): options = copy.copy(args) options.executor = lane if attempt is not None: - options.fallback = "vm-oom" + options.fallback = attempt.get("fallback", "vm-oom") options.ax_attempt = attempt if lane == "docker": options.image = getattr(args, "docker_image", args.image) @@ -845,8 +849,10 @@ def run_plan(plan, source, provenance, inputs, args): failed = future.result() result = json.loads((args.out / pin_name(pin)).read_text()) selected = {u["unit"] for u in pin["units"]} - if lane == "ax" and any("JVM out of memory" in u.get("reason", "") for u in result["units"] if u["unit"] in selected): - result.update(fallback="vm-oom", fallback_pending=True) + reasons = [u.get("reason", "") for u in result["units"] if u["unit"] in selected] + sandbox_failed = args.executor == "mixed-ax" and result.get("container_exit") == 1 and "AX sandbox failed twice" in reasons + if lane == "ax" and (sandbox_failed or any("JVM out of memory" in reason for reason in reasons)): + result.update(fallback="vm-sandbox" if sandbox_failed else "vm-oom", fallback_pending=True) write_json(args.out / pin_name(pin), result) fallbacks.append((pin, result)) continue diff --git a/benchmarks/learn/engine_trace_ax.py b/benchmarks/learn/engine_trace_ax.py index 1c16cf2..d6317ce 100644 --- a/benchmarks/learn/engine_trace_ax.py +++ b/benchmarks/learn/engine_trace_ax.py @@ -31,6 +31,10 @@ class EgressError(ValueError): """Safe diagnostic without presigned URL contents.""" +class SandboxFailure(Exception): + """A terminal AX task without an authoritative result object.""" + + def public_endpoint(url): """Reject endpoints the actor's hostname-based HTTPS egress policy cannot reach.""" parts = urlsplit(url) @@ -116,7 +120,8 @@ def task_state(text, name): except ValueError: pass return phase, None if ip in {"", "-", ""} else ip - return "", None + listed = isinstance(doc, (dict, list)) or bool(lines and "name" in headers and {"phase", "status"}.intersection(headers)) + return "Missing" if listed else "", None def bounded_read(read, timeout): @@ -215,15 +220,35 @@ class AX: shutil.copytree(root, output, dirs_exist_ok=True) def execute(self, checkout, output, pin, pin_id): - digest = hashlib.sha256(f"{self.run}/{pin_id}".encode()).hexdigest()[:12] + attempts = [] + try: + for attempt in range(2): + try: + return self._execute_attempt(checkout, output, pin, pin_id, attempt, attempts) + except SandboxFailure: + if attempt == 0: + self.pause(10) + return subprocess.CompletedProcess([], 1, "", "AX sandbox failed twice"), None + finally: + (output / "ax_attempts.json").write_text(json.dumps(attempts)) + result_path = output / "result.json" + if result_path.exists(): + record = json.loads(result_path.read_text()) + record["ax_attempts"] = attempts + result_path.write_text(json.dumps(record)) + + def _execute_attempt(self, checkout, output, pin, pin_id, attempt, attempts): + digest = hashlib.sha256(f"{self.run}/{pin_id}/attempt-{attempt}".encode()).hexdigest()[:12] slug = re.sub(r"[^a-z0-9]+", "-", str(pin_id).lower()).strip("-")[:22] name = f"ousast-engine-{slug or 'pin'}-{digest}" key_id = hashlib.sha256(str(pin_id).encode()).hexdigest() - prefix = f"engine-queue/{self.run}/{key_id}/" + prefix = f"engine-queue/{self.run}/{key_id}/attempt-{attempt}/" keys = [prefix + suffix for suffix in (f"source{PRESIGNED_ARCHIVE_SUFFIX}", "questions.json", f"result{PRESIGNED_ARCHIVE_SUFFIX}")] expires = timedelta(seconds=self.args.deadline + 600) worker_ip = None submitted = False + started = self.clock() + phase = "" try: self.store._put(keys[0], pack(checkout)) self.store._put(keys[1], json.dumps(pin).encode()) @@ -256,23 +281,35 @@ class AX: if result is not None: self.receive(result[0], output, pin) return subprocess.CompletedProcess([], 0, "", ""), worker_ip - if phase.lower() == "failed": - return subprocess.CompletedProcess([], 1, "", "AX Task phase Failed"), worker_ip + if phase.lower() in {"failed", "missing"}: + # The worker may have uploaded while we were reading AX status. + result = bounded_read(lambda: self.store._get(keys[2]), end - self.clock()) + if self.clock() >= end: + return None, worker_ip + if result is not None: + self.receive(result[0], output, pin) + return subprocess.CompletedProcess([], 0, "", ""), worker_ip + raise SandboxFailure self.pause(min(2, max(0, end - self.clock()))) return None, worker_ip except (subprocess.TimeoutExpired, TimeoutError): return None, worker_ip + except SandboxFailure: + raise except Exception as exc: # Never expose exception strings which may contain object-scoped bearer URLs. detail = str(exc) if isinstance(exc, EgressError) else type(exc).__name__ return subprocess.CompletedProcess([], 1, "", "AX transport failed: " + detail), worker_ip finally: + attempts.append({"name": name, "phase": phase, "seconds": self.clock() - started}) errors = [] if submitted: try: self.ax("delete", "task", name) except Exception as exc: - errors.append("task: " + type(exc).__name__) + # A successful list already confirmed a disappeared task is gone. + if phase != "Missing": + errors.append("task: " + type(exc).__name__) for key in keys: try: self.store._delete(key) diff --git a/tests/test_ax_executor.py b/tests/test_ax_executor.py index 978493a..abff6bb 100644 --- a/tests/test_ax_executor.py +++ b/tests/test_ax_executor.py @@ -2,6 +2,7 @@ import io import json +import subprocess import sys import tarfile from pathlib import Path @@ -67,9 +68,13 @@ with Path(os.environ["AX_LOG"]).open("a") as log: if args[0] == "get" and os.environ.get("AX_GET_FAIL"): sys.exit(1) if args[0] == "get": - manifest = json.loads(json.loads(Path(os.environ["AX_LOG"]).read_text().splitlines()[0])["input"]) + applies = [json.loads(line) for line in Path(os.environ["AX_LOG"]).read_text().splitlines() + if json.loads(line)["args"][0] == "apply"] + manifest = json.loads(applies[-1]["input"]) + phases = os.environ.get("AX_PHASE", "Running").split(",") + phase = phases[min(len(applies) - 1, len(phases) - 1)] print(json.dumps({"items": [{"metadata": manifest["metadata"], "status": { - "phase": os.environ.get("AX_PHASE", "Running"), "workerIP": "10.0.0.2"}}]})) + "phase": phase, "workerIP": "10.0.0.2"}}] if phase != "Missing" else []})) if args[0] == "delete" and os.environ.get("AX_DELETE_FAIL"): sys.exit(1) """) @@ -134,10 +139,115 @@ def test_ax_executor_failed_phase_deletes_task(setup, monkeypatch): args, source, output, log, pause, clock = setup monkeypatch.setenv("AX_PHASE", "Failed") done, _ = ax.AX(args, Store(), pause=pause, clock=clock).execute(source, output, PIN, "pin") - assert done.returncode == 1 and "Failed" in done.stderr + assert done.returncode == 1 and done.stderr == "AX sandbox failed twice" + attempts = json.loads((output / "ax_attempts.json").read_text()) + assert [a["phase"] for a in attempts] == ["Failed", "Failed"] assert json.loads(log.read_text().splitlines()[-1])["args"][:2] == ["delete", "task"] +@pytest.mark.parametrize("phase", ["Failed", "Missing"]) +def test_ax_executor_retries_with_new_name_links_and_backoff(setup, monkeypatch, phase): + args, source, output, log, pause, clock = setup + monkeypatch.setenv("AX_PHASE", phase + ",Succeeded") + store = Store() + store._get = lambda key: (archive({**PIN, "done": True}), None) if "/attempt-1/" in key else None + pauses = [] + + def backoff(seconds): + calls = [json.loads(line) for line in log.read_text().splitlines()] + assert calls[-1]["args"][:2] == ["delete", "task"] + pauses.append(seconds) + pause(seconds) + + done, _ = ax.AX(args, store, pause=backoff, clock=clock).execute(source, output, PIN, "x" * 100) + assert done.returncode == 0 + record = json.loads((output / "result.json").read_text()) + attempts = record["ax_attempts"] + assert record["done"] and [a["phase"] for a in attempts] == [phase, "Succeeded"] + assert len({a["name"] for a in attempts}) == 2 + assert all(a["name"].startswith("ousast-engine-") and len(a["name"]) <= 49 and a["seconds"] >= 0 for a in attempts) + assert pauses == [10] + calls = [json.loads(line) for line in log.read_text().splitlines()] + manifests = [json.loads(c["input"]) for c in calls if c["args"][0] == "apply"] + links = [{e["name"]: e["value"] for e in m["spec"]["env"]} for m in manifests] + assert all(links[0][key] != links[1][key] for key in ("SOURCE_URL", "QUESTIONS_URL", "RESULT_URL")) + assert len(store.expiries) == 6 and len(store.deleted) == 6 + + +@pytest.mark.parametrize("late", [False, True]) +def test_ax_executor_failed_result_is_never_retried(setup, monkeypatch, late): + args, source, output, log, pause, clock = setup + monkeypatch.setenv("AX_PHASE", "Failed") + store = Store(archive({**PIN, "done": True, "status": "failed"})) + original_get = store._get + reads = [] + + def get(key): + reads.append(key) + return None if late and len(reads) == 1 else original_get(key) + + store._get = get + done, _ = ax.AX(args, store, pause=pause, clock=clock).execute(source, output, PIN, "pin") + record = json.loads((output / "result.json").read_text()) + assert done.returncode == 0 and record["status"] == "failed" + assert len(record["ax_attempts"]) == 1 + assert sum(json.loads(line)["args"][0] == "apply" for line in log.read_text().splitlines()) == 1 + + +@pytest.mark.parametrize("mode", ["ax", "mixed-ax"]) +def test_ax_sandbox_failed_twice_scheduler(setup, monkeypatch, mode): + with monkeypatch.context() as patch: + patch.syspath_prepend(str(Path(__file__).resolve().parents[1] / "benchmarks/learn")) + import engine_trace as trace + + args, source, output, log, pause, clock = setup + monkeypatch.setenv("AX_PHASE", "Failed") + args.executor, args.parallel, args.out, args.docker_image = mode, 1, output, "vm-image" + args.ax_dispatcher = ax.AX(args, Store(), pause=pause, clock=clock) + monkeypatch.setattr(trace, "prepare_questions", lambda *_: PIN["questions"]) + + class Inputs: + def materialize(self, pin, target): + target.mkdir() + data = (source / "app.py").read_bytes() + (target / "app.py").write_bytes(data) + return {"files": 1, "bytes": len(data)} + + original_run = subprocess.run + docker_calls = [] + + def run(command, **kwargs): + if command[0] == args.ax_bin: + return original_run(command, **kwargs) + assert command[:2] in (["docker", "run"], ["docker", "rm"]) + docker_calls.append(command) + if command[1] == "run": + assert command[command.index("--heap-profile") + 1] == "vm" + assert "vm-image" in command + target = Path(next(arg.removesuffix(":/out") for arg in command if arg.endswith(":/out"))) + prepared = json.loads((target / "pin.json").read_text()) + trace.write_json( + target / "result.json", + {**prepared, "done": True, "units": [trace.unit_record(u, "asked-nothing") for u in prepared["units"]]}, + ) + return subprocess.CompletedProcess(command, 0, "", "") + + monkeypatch.setattr(subprocess, "run", run) + assert trace.run_plan([PIN], source, {}, Inputs(), args) is (mode == "ax") + record = json.loads((output / trace.pin_name(PIN)).read_text()) + assert [a["phase"] for a in record["ax_attempts"]] == ["Failed", "Failed"] + assert sum(json.loads(line)["args"][0] == "apply" for line in log.read_text().splitlines()) == 2 + if mode == "mixed-ax": + assert [command[1] for command in docker_calls] == ["run", "rm"] + assert record["executor"] == "docker" and record["fallback"] == "vm-sandbox" + assert record["units"][0]["status"] == "asked-nothing" + assert record["ax_attempt"]["ax_attempts"] == record["ax_attempts"] + else: + assert not docker_calls + assert record["units"][0]["status"] == "failed" + assert record["units"][0]["reason"] == "AX sandbox failed twice" + + def test_ax_executor_delete_failure_is_fatal_and_still_cleans_objects(setup, monkeypatch): args, source, output, _, pause, clock = setup monkeypatch.setenv("AX_DELETE_FAIL", "1") @@ -180,6 +290,21 @@ def test_ax_executor_table_worker_ip(): assert ax.task_state("NAME ATESPACE PHASE WORKER-IP\nours default Running 10.2.3.4\n", "ours") == ("Running", "10.2.3.4") +def test_ax_executor_missing_requires_readable_task_list(): + assert ax.task_state('{"items": []}', "ours") == ("Missing", None) + assert ax.task_state("NAME ATESPACE PHASE WORKER-IP\n", "ours") == ("Missing", None) + assert ax.task_state("", "ours") == ("", None) + assert ax.task_state("temporarily unavailable", "ours") == ("", None) + + +def test_ax_executor_status_query_failure_is_not_disappearance(setup, monkeypatch): + args, source, output, log, pause, clock = setup + monkeypatch.setenv("AX_GET_FAIL", "1") + done, _ = ax.AX(args, Store(), pause=pause, clock=clock).execute(source, output, PIN, "pin") + assert done is None + assert sum(json.loads(line)["args"][0] == "apply" for line in log.read_text().splitlines()) == 1 + + def test_ax_executor_entry_download_failure_restores_host_units(setup): args, source, output, _, pause, clock = setup store = Store(archive({"done": True, "status": "failed", "reason": "questions download failed", "units": []}))