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)