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 <[email protected]>

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

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 {"<none>", "-", ""} 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)

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": []}))