merge(engine): AX-task executor for the kube-ax cluster

Co-Authored-By: Claude Opus 5.5 <[email protected]>

.github/workflows/engine-image.yml

@@ -6,8 +6,12 @@ on:
     paths:
       - 'src/**'
       - 'benchmarks/learn/engine_trace*.py'
+      - 'benchmarks/learn/engine_task_entry.py'
       - 'benchmarks/learn/engine_worker_service.py'
       - 'Dockerfile'
+      - 'plane/Dockerfile.engine-task'
+      - 'ops/frontend-retention/install.py'
+      - 'pyproject.toml'
       - '.github/workflows/engine-image.yml'
 permissions:
   contents: read
@@ -18,6 +22,14 @@ concurrency:
 jobs:
   image:
     runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - package: openultrasast
+            dockerfile: Dockerfile
+          - package: ousast-engine-task
+            dockerfile: plane/Dockerfile.engine-task
     steps:
       - uses: actions/checkout@v5
       - uses: docker/setup-buildx-action@v3
@@ -29,14 +41,16 @@ jobs:
       - id: names
         env:
           OWNER: ${{ github.repository_owner }}
+          PACKAGE: ${{ matrix.package }}
         run: |
-          echo "image=ghcr.io/${OWNER,,}/openultrasast" >> "$GITHUB_OUTPUT"
+          echo "image=ghcr.io/${OWNER,,}/$PACKAGE" >> "$GITHUB_OUTPUT"
           echo "sha=sha-${GITHUB_SHA:0:7}" >> "$GITHUB_OUTPUT"
       - id: build
         uses: docker/build-push-action@v6
         with:
           context: .
-          file: Dockerfile
+          file: ${{ matrix.dockerfile }}
+          platforms: linux/amd64
           push: true
           tags: |
             ${{ steps.names.outputs.image }}:${{ steps.names.outputs.sha }}
@@ -44,4 +58,6 @@ jobs:
       - env:
           IMAGE: ${{ steps.names.outputs.image }}
           DIGEST: ${{ steps.build.outputs.digest }}
-        run: echo "$IMAGE@$DIGEST" >> "$GITHUB_STEP_SUMMARY"
+        run: |
+          echo "$IMAGE@$DIGEST"
+          echo "$IMAGE@$DIGEST" >> "$GITHUB_STEP_SUMMARY"

README.md

@@ -363,3 +363,81 @@ own terms:
 - Joern (Apache-2.0) is downloaded at image build time.
 - The tree-sitter wheels of the `semantic` extra are MIT.
 - The published benchmark cases were licence-checked one by one. v1's OpenCVE case is BUSL-1.1, see [RELEASE_NOTES.md](RELEASE_NOTES.md).
+
+## Engine traces on kube-ax
+
+`benchmarks/learn/engine_trace.py --executor ax` submits AX v0.3.1 Tasks with
+`ax apply -f -` and removes them with `ax delete task NAME`. The operator supplies
+kubeconfig and AX access; the dispatcher creates no Kubernetes resources or secrets.
+The `ousast-engine-task` GHCR package **must be made Public once** in its package
+settings so the actors can pull it without registry credentials.
+
+The task image contains the current engine and worker code. Rebuild it after code
+changes and pass its digest with `--image`; no analyzer download is needed. The host
+materializes each pin as usual and uploads source and prepared questions under
+`engine-queue/<run>/<pin>/`. Only object-scoped presigned URLs and deadline parameters
+enter Task env. Credentials stay on the coordinator. Configure `OUSAST_MEMORY=s3://…`
+and the store credentials in the coordinator's `.env` or environment, with
+`S3_ENDPOINT=https://files.because-security.com`. Presigned links expire after the
+pin deadline plus ten minutes. IP and HTTP endpoints are rejected: the operator's
+actor EgressPolicy must allow **files.because-security.com:443**.
+
+Completion is the result object, even while AX reports Running. Failed phases and
+host deadlines produce explicit failure/timeout rows. Task and staging-object cleanup
+runs on success and failure; cleanup failure stops further dispatch. Per-pin files,
+`--rerun-status`, resume and `progress.json` retain the existing trace-runner behavior.
+AX results include `executor: "ax"`, `worker_ip`, and the image digest. Placement is
+random; `--parallel` defaults to two AX tasks and cannot exceed two.
+
+AX workers use `--heap-profile cluster` with
+`-Xms1700m -Xmx1700m -XX:MaxMetaspaceSize=256m -XX:ReservedCodeCacheSize=128m -XX:MaxDirectMemorySize=256m -Xss512k -XX:ActiveProcessorCount=1 -XX:+UseG1GC`
+on every frontend/query JVM; instruments record the flags. A `JVM out of memory`
+result is automatically re-queued to one local Docker lane with `--heap-profile vm`
+(2560m), recording `fallback: "vm-oom"` and the AX attempt. `mixed-ax` also uses this
+one Docker lane for overflow, prioritizing pending OOM fallbacks. Thus a mixed run has
+up to two AX tasks plus one Docker container. Use `--docker-image` for the locally
+available VM image; Docker never pulls it implicitly. A STOP file prevents new
+launches; interrupted fallback queues are recovered on resume.
+
+Coordinator commands, from this worktree (the AX binary and kubeconfig remain
+configurable). Build/push the actual uncommitted tree and read its digest:
+
+```bash
+cd /home/mc/Source/OpenUltraSAST/.claude/worktrees/engine-ax
+export PATH=/home/mc/Source/OpenUltraSAST/.venv/bin:$PATH
+export PYTHONPATH="$PWD/src"
+TASK_TAG="ghcr.io/norandom/ousast-engine-task:sha-$(git rev-parse --short=7 HEAD)"
+docker buildx build --platform linux/amd64 -f plane/Dockerfile.engine-task \
+  -t "$TASK_TAG" -t ghcr.io/norandom/ousast-engine-task:main \
+  --push --metadata-file /tmp/ousast-engine-task-build.json .
+export TASK_IMAGE="ghcr.io/norandom/ousast-engine-task@$(python -c 'import json; print(json.load(open("/tmp/ousast-engine-task-build.json"))["containerimage.digest"])')"
+printf '%s\n' "$TASK_IMAGE"
+docker build -f Dockerfile -t openultrasast:engine-ax .
+```
+
+The Engine image workflow also builds/pushes both images and prints each digest;
+CI needs the changes published before dispatch. The local build above includes
+uncommitted changes without publishing a git commit.
+
+With `KUBECONFIG` pointing to the operator's kube-ax config and store settings in
+`.env`, run a one-pin smoke using the default experiment units and local example
+memory (or add `--examples /path/to/frozen-examples.jsonl`):
+
+```bash
+: "${KUBECONFIG:?Set KUBECONFIG to the operator kube-ax configuration}"
+python benchmarks/learn/engine_trace.py --executor ax \
+  --ax-bin "$HOME/.local/ax-v0.3.1/ax" --kubeconfig "$KUBECONFIG" \
+  --atespace default --image "$TASK_IMAGE" --docker-image openultrasast:engine-ax \
+  --parallel 1 --limit 1 --deadline 900 \
+  --out "$HOME/ousast-results/plane/engine-ax-smoke"
+```
+
+Then run the mixed lanes (repeat to resume; defaults re-run failed/timeout pins):
+
+```bash
+python benchmarks/learn/engine_trace.py --executor mixed-ax \
+  --ax-bin "$HOME/.local/ax-v0.3.1/ax" --kubeconfig "$KUBECONFIG" \
+  --atespace default --image "$TASK_IMAGE" --docker-image openultrasast:engine-ax \
+  --parallel 2 --deadline 900 --rerun-status failed,timeout \
+  --out "$HOME/ousast-results/plane/engine-ax-mixed"
+```

benchmarks/learn/engine_task_entry.py

@@ -0,0 +1,154 @@
+"""AX task entry: object-scoped transfers and stdlib only; result PUT is completion."""
+
+from __future__ import annotations
+
+import http.client
+import json
+import os
+import shutil
+import subprocess
+import sys
+import tarfile
+import time
+import urllib.error
+import urllib.request
+from pathlib import Path
+
+
+def transfer(request, destination=None, *, readiness_budget=None):
+    """Wait for actor egress on startup; bound later transfers to three attempts."""
+    deadline = time.monotonic() + readiness_budget if readiness_budget is not None else None
+    delay = 0.5
+    attempt = 0
+    while True:
+        remaining = deadline - time.monotonic() if deadline is not None else 60
+        if remaining <= 0:
+            raise RuntimeError("egress not open after 60 s")
+        attempt += 1
+        try:
+            with urllib.request.urlopen(request, timeout=min(60, remaining)) as response:
+                if destination is not None:
+                    with destination.open("wb") as output:
+                        shutil.copyfileobj(response, output)
+            return
+        except urllib.error.HTTPError as exc:
+            exc.close()
+            if exc.code != 403 and not 500 <= exc.code < 600:
+                raise
+            if deadline is None and attempt >= 3:
+                raise
+        except (urllib.error.URLError, OSError, http.client.HTTPException):
+            if deadline is None and attempt >= 3:
+                raise
+        if deadline is not None:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                raise RuntimeError("egress not open after 60 s") from None
+            time.sleep(min(delay, remaining))
+            delay = min(delay * 2, 8)
+        else:
+            time.sleep(attempt)
+
+
+def download(url, destination, *, readiness_budget=None):
+    transfer(url, destination, readiness_budget=readiness_budget)
+
+
+def unpack(source, destination):
+    """Reject links and traversal; source archives contain only regular files/directories."""
+    with tarfile.open(source) as archive:
+        for member in archive.getmembers():
+            target = (destination / member.name).resolve()
+            if not target.is_relative_to(destination.resolve()) or not (member.isfile() or member.isdir()):
+                raise ValueError("unsafe source archive member")
+        archive.extractall(destination, filter="data")
+
+
+def failure_record(pin, reason):
+    return {
+        **pin,
+        "done": True,
+        "status": "failed",
+        "reason": reason,
+        "units": [
+            {
+                **unit,
+                "status": "failed" if unit.get("supported", True) else "unsupported",
+                "reason": reason,
+                "instrument": {"files": [], "bytes": 0, "jvm": []},
+                "questions": {},
+                "questions_asked": [],
+                "questions_completed": [],
+                "witness_rows": [],
+                "traces": [],
+            }
+            for unit in pin.get("units", [])
+        ],
+    }
+
+
+def main(*, work_root=Path("/tmp"), worker_script=None, readiness_budget=60):
+    pin = {}
+    failed = False
+    output = work_root / "out"
+    result = output / "result.json"
+    try:
+        output.mkdir(parents=True, exist_ok=True)
+        result.unlink(missing_ok=True)
+        questions = work_root / "questions.json"
+        download(os.environ["QUESTIONS_URL"], questions, readiness_budget=readiness_budget)
+        pin = json.loads(questions.read_text())
+        if not isinstance(pin, dict):
+            pin = {}
+            raise ValueError("questions object must contain a prepared pin")
+        source = work_root / "source.tar"
+        download(os.environ["SOURCE_URL"], source)
+        case = work_root / "case"
+        case.mkdir(parents=True, exist_ok=True)
+        unpack(source, case)
+        deadline = int(os.environ.get("DEADLINE", "900"))
+        command = [
+            sys.executable,
+            str(worker_script or Path(__file__).with_name("engine_trace_worker.py")),
+            "--worker",
+            str(questions),
+            "--root",
+            str(case),
+            "--output",
+            str(result),
+            "--heap-profile",
+            "cluster",
+            "--deadline",
+            str(deadline),
+            "--question-deadline",
+            os.environ.get("QUESTION_DEADLINE", "120"),
+        ]
+        completed = subprocess.run(command, timeout=deadline + 5, check=False)
+        if completed.returncode:
+            raise RuntimeError(f"worker exited {completed.returncode}")
+        if not result.is_file() or not json.loads(result.read_text()).get("done"):
+            raise RuntimeError("worker did not write a completed result")
+    except Exception as exc:
+        failed = True
+        # Avoid including exception text containing a presigned URL in durable output.
+        reason = str(exc) if isinstance(exc, (RuntimeError, ValueError)) else f"task entry failed: {type(exc).__name__}"
+        try:
+            output.mkdir(parents=True, exist_ok=True)
+            result.write_text(json.dumps(failure_record(pin, reason)) + "\n")
+        except Exception:
+            print("task entry could not write failure result", file=sys.stderr)
+    try:
+        archive = work_root / "result.tar"
+        with tarfile.open(archive, "w") as tar:
+            for path in sorted(output.iterdir()):
+                tar.add(path, arcname=path.name)
+        data = archive.read_bytes()
+        request = urllib.request.Request(os.environ["RESULT_URL"], data=data, method="PUT", headers={"Content-Type": "application/x-tar"})
+        transfer(request, readiness_budget=readiness_budget if failed else None)
+    except Exception:
+        print("task entry could not upload result; dispatcher deadline remains authoritative", file=sys.stderr)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

benchmarks/learn/engine_trace.py

@@ -21,7 +21,7 @@ import textwrap
 import threading
 import time
 import tomllib
-from collections import Counter, defaultdict
+from collections import Counter, defaultdict, deque
 from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
 from pathlib import Path
 
@@ -46,6 +46,8 @@ SNAPSHOT_FILES = (
     "benchmarks/learn/engine_trace_worker.py",
     "benchmarks/learn/engine_trace_k8s.py",
     "benchmarks/learn/engine_trace_queue.py",
+    "benchmarks/learn/engine_trace_ax.py",
+    "benchmarks/learn/engine_task_entry.py",
     "benchmarks/learn/engine_worker_service.py",
     "benchmarks/learn/engine_trace_parse.py",
     "src/openultrasast/model/taint.py",
@@ -140,6 +142,7 @@ def make_plan(
             recorded = {r["unit"]: r.get("input_digest") for r in old.get("units", [])}
             if (
                 old.get("done")
+                and not old.get("fallback_pending")
                 and old.get("materialization_version") == MATERIALIZATION_VERSION
                 and not any(r.get("status") in rerun_status for r in old.get("units", []) if r["unit"] in {u["unit"] for u in rows})
                 and all(recorded.get(r["unit"]) == r["input_digest"] for r in rows)
@@ -582,6 +585,8 @@ def docker_command(source, checkout, out, name, image, deadline, question_deadli
         "/out/pin.json",
         "--deadline",
         str(deadline),
+        "--heap-profile",
+        "vm",
         "--question-deadline",
         str(question_deadline),
     ]
@@ -636,6 +641,10 @@ def run_pin(pin, source, provenance, inputs, out, args):
                 status, reason = "failed", done.stderr
                 if done.returncode:
                     (output / "container.log").write_text(reason)
+            elif executor == "ax":
+                done, node = args.ax_dispatcher.execute(checkout, output, json.loads((output / "pin.json").read_text()), pin_name(pin)[:-5])
+                status = "timeout" if done is None else "failed"
+                reason = f"AX deadline {args.deadline}s" if done is None else done.stderr
             elif executor == "queue":
                 done, node = args.dispatcher.execute(checkout, output, json.loads((output / "pin.json").read_text()), pin_name(pin)[:-5])
                 status, reason = "timeout", f"queue wait deadline {args.queue_timeout}s"
@@ -683,7 +692,7 @@ def run_pin(pin, source, provenance, inputs, out, args):
             record.update(
                 done=True,
                 seconds=time.monotonic() - started,
-                analyzer=provenance,
+                analyzer={"image": args.image, "code": "baked-in-image"} if executor == "ax" else provenance,
                 image=record.get("image", args.image) if executor == "queue" else args.image,
             )
             if (output / "container.log").exists():
@@ -695,6 +704,11 @@ def run_pin(pin, source, provenance, inputs, out, args):
                 "analyzer": provenance,
                 "units": [unit_record(u, "failed" if u["supported"] else "unsupported", str(exc)) for u in pin["units"]],
             }
+        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 getattr(args, "fallback", None):
+            record.update(executor="docker", fallback=args.fallback, ax_attempt=args.ax_attempt, fallback_pending=False)
         if executor in {"k8s", "k8s-jobs"}:
             record.update(executor="k8s", node=node, image=args.image)
         if executor == "queue":
@@ -760,6 +774,8 @@ def summary(out):
 def run_plan(plan, source, provenance, inputs, args):
     """One task per fixed executor lane; progress is written only by the scheduler."""
     lanes = ["k8s"] * (args.parallel - 1) + ["docker"] if args.executor == "mixed" else [args.executor] * args.parallel
+    if args.executor in {"ax", "mixed-ax"}:
+        lanes = ["ax"] * args.parallel + ["docker"]
     args.input_lock = threading.Lock()
     progress = {
         "pins_done": 0,
@@ -771,7 +787,18 @@ def run_plan(plan, source, provenance, inputs, args):
         "active_lanes": {},
     }
     counts = Counter()
-    pending = iter(plan)
+    pending = deque(plan)
+    fallbacks = deque()
+    if args.executor in {"ax", "mixed-ax"}:
+        pending.clear()
+        for pin in plan:
+            path = args.out / pin_name(pin)
+            old = json.loads(path.read_text()) if path.exists() else {}
+            recorded = {u["unit"]: u.get("input_digest") for u in old.get("units", [])}
+            if old.get("fallback_pending") and all(recorded.get(u["unit"]) == u.get("input_digest") for u in pin["units"]):
+                fallbacks.append((pin, old))
+            else:
+                pending.append(pin)
     active = {}
     failures = False
 
@@ -779,7 +806,7 @@ def run_plan(plan, source, provenance, inputs, args):
         progress.update(pins_in_flight=len(active), active_lanes=dict(Counter(lane for _, lane in active.values())))
         write_json(args.out / "progress.json", progress)
 
-    def execute(pin, lane):
+    def execute(pin, lane, attempt=None):
         if not any(u["supported"] for u in pin["units"]):
             save_pin(
                 args.out / pin_name(pin), {**pin, "done": True, "units": [unit_record(u, "unsupported", u["reason"]) for u in pin["units"]]}
@@ -787,6 +814,9 @@ def run_plan(plan, source, provenance, inputs, args):
             return False
         options = copy.copy(args)
         options.executor = lane
+        if attempt is not None:
+            options.fallback = "vm-oom"
+            options.ax_attempt = attempt
         if lane == "docker":
             options.image = getattr(args, "docker_image", args.image)
         return run_pin(pin, source, provenance, inputs, args.out, options)
@@ -794,10 +824,16 @@ def run_plan(plan, source, provenance, inputs, args):
     with ThreadPoolExecutor(max_workers=len(lanes)) as pool:
 
         def launch(lane):
-            if not (args.out / "STOP").exists():
-                pin = next(pending, None)
-                if pin is not None:
-                    active[pool.submit(execute, pin, lane)] = (pin, lane)
+            if (args.out / "STOP").exists():
+                return
+            attempt = None
+            if lane == "docker" and fallbacks:
+                pin, attempt = fallbacks.popleft()
+            elif pending and not (args.executor == "ax" and lane == "docker"):
+                pin = pending.popleft()
+            else:
+                return
+            active[pool.submit(execute, pin, lane, attempt)] = (pin, lane)
 
         for lane in lanes:
             launch(lane)
@@ -806,14 +842,27 @@ def run_plan(plan, source, provenance, inputs, args):
             completed, _ = wait(active, return_when=FIRST_COMPLETED)
             for future in completed:
                 pin, lane = active.pop(future)
-                failures |= future.result()
+                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)
+                    write_json(args.out / pin_name(pin), result)
+                    fallbacks.append((pin, result))
+                    continue
+                failures |= failed
+                selected = {u["unit"] for u in pin["units"]}
                 counts.update(u["status"] for u in result["units"] if u["unit"] in selected)
                 progress.update(
                     pins_done=progress["pins_done"] + 1, last_pin={"repo": pin["repo"], "pin": pin["pin"]}, statuses=dict(counts)
                 )
-                launch(lane)
+            idle = Counter(lanes) - Counter(lane for _, lane in active.values())
+            # Give a waiting fallback priority over overflow on the single VM lane.
+            for lane, count in idle.items():
+                for _ in range(count):
+                    launch(lane)
+            if args.executor in {"ax", "mixed-ax"}:
+                progress["fallbacks_pending"] = len(fallbacks)
             persist()
     return failures
 
@@ -826,10 +875,14 @@ def main():
     parser.add_argument("--cache", type=Path, default=Path.home() / ".cache/openultrasast")
     parser.add_argument("--out", type=Path, default=DEFAULT_OUT)
     parser.add_argument("--image", help="default: openultrasast:dev for Docker; public 2.0.1 image for k8s/mixed")
-    parser.add_argument("--executor", choices=("docker", "k8s", "k8s-jobs", "mixed", "queue"), default="docker")
+    parser.add_argument("--executor", choices=("docker", "k8s", "k8s-jobs", "mixed", "queue", "ax", "mixed-ax"), default="docker")
     parser.add_argument("--queue-status", action="store_true")
     parser.add_argument("--queue-timeout", type=int, default=3600, help="total queue wait including execution, seconds")
-    parser.add_argument("--parallel", type=int, default=1, help="total pins in flight; mixed reserves one lane for Docker")
+    parser.add_argument("--parallel", type=int, help="AX lanes (default 2, maximum 2); otherwise total lanes (default 1)")
+    parser.add_argument("--ax-bin", default="ax")
+    parser.add_argument("--kubeconfig", type=Path, help="AX kubeconfig; otherwise inherited KUBECONFIG/default CLI config")
+    parser.add_argument("--atespace", default="default")
+    parser.add_argument("--docker-image", default="openultrasast:dev", help="local VM image for AX OOM fallback/mixed overflow")
     parser.add_argument("--kube-context")
     parser.add_argument("--kube-namespace", default="ousast-engine")
     parser.add_argument("--keep-jobs", action="store_true", help="retain Jobs, Secrets and staging objects for debugging")
@@ -847,6 +900,13 @@ def main():
     parser.add_argument("--dry-run", action="store_true")
     parser.add_argument("--summary", action="store_true")
     args = parser.parse_args()
+    if args.parallel is None:
+        args.parallel = 2 if args.executor in {"ax", "mixed-ax"} else 1
+    if args.executor in {"ax", "mixed-ax"}:
+        if args.parallel > 2:
+            parser.error("AX allows at most 2 engine tasks in flight")
+        if not args.image or not re.fullmatch(r"[^\s@]+@sha256:[0-9a-f]{64}", args.image):
+            parser.error("AX requires --image by sha256 digest")
     if args.queue_status:
         print(json.dumps(Dispatcher(args).status(), indent=2))
         return 0
@@ -924,7 +984,8 @@ def main():
     except BlockingIOError:
         parser.error("another trace runner holds this results root")
     failures = False
-    args.docker_image = args.image
+    if args.executor not in {"ax", "mixed-ax"}:
+        args.docker_image = args.image
     if args.executor in {"docker", "mixed"} and any(u["supported"] for pin in plan for u in pin["units"]):
         image = subprocess.run(
             ["docker", "image", "inspect", "--format", "{{.Id}}", args.image], capture_output=True, text=True, check=True
@@ -937,6 +998,10 @@ def main():
         args.kubernetes.namespace()
     if args.executor == "queue":
         args.dispatcher = Dispatcher(args)
+    if args.executor in {"ax", "mixed-ax"}:
+        from engine_trace_ax import AX
+
+        args.ax_dispatcher = AX(args)
     with tempfile.TemporaryDirectory(prefix="trace-source-", dir=args.out) as scratch:
         source = Path(scratch) / "frozen"
         provenance = {"image": args.image, "executor": "queue"} if args.executor == "queue" else freeze_source(source)

benchmarks/learn/engine_trace_ax.py

@@ -0,0 +1,282 @@
+"""Host-only AX Task transport; credentials never enter task manifests."""
+
+from __future__ import annotations
+
+import hashlib
+import io
+import ipaddress
+import json
+import os
+import queue
+import re
+import shutil
+import subprocess
+import tarfile
+import tempfile
+import threading
+import time
+import uuid
+from datetime import timedelta
+from pathlib import Path
+from urllib.parse import urlsplit
+
+import yaml
+from engine_trace_k8s import pack
+
+from openultrasast.config import load_dotenv
+from openultrasast.plane.memory import S3Store, open_store
+
+
+class EgressError(ValueError):
+    """Safe diagnostic without presigned URL contents."""
+
+
+def public_endpoint(url):
+    """Reject endpoints the actor's hostname-based HTTPS egress policy cannot reach."""
+    parts = urlsplit(url)
+    host = parts.hostname or ""
+    try:
+        ipaddress.ip_address(host)
+        is_ip = True
+    except ValueError:
+        is_ip = False
+    if (
+        parts.scheme != "https"
+        or not host
+        or is_ip
+        or "." not in host
+        or host.endswith((".localhost", ".local", ".internal"))
+        or parts.port not in (None, 443)
+        or parts.username
+        or parts.password
+    ):
+        raise EgressError(
+            "AX requires a public HTTPS DNS hostname in S3_ENDPOINT and presigned URLs: "
+            "the egress policy allows the store hostname on port 443 only (no IP or plain HTTP)"
+        )
+    return host.lower()
+
+
+def task_manifest(name, args, urls):
+    if not re.fullmatch(r"[^\s@]+@sha256:[0-9a-f]{64}", args.image):
+        raise ValueError("AX --image must be pinned by sha256 digest")
+    if set(urls) != {"SOURCE_URL", "QUESTIONS_URL", "RESULT_URL"}:
+        raise ValueError("AX Task accepts only source, questions and result presigned URLs")
+    for url in urls.values():
+        public_endpoint(url)
+    env = [{"name": key, "value": value} for key, value in urls.items()]
+    env.extend(
+        {"name": key, "value": str(value)}
+        for key, value in {"DEADLINE": args.deadline, "QUESTION_DEADLINE": args.question_deadline}.items()
+    )
+    if len(json.dumps(env).encode()) >= 32768:
+        raise ValueError("AX Task env exceeds the 32 KB limit")
+    return {
+        "apiVersion": "ax.io/v1alpha1",
+        "kind": "Task",
+        "metadata": {"name": name, "atespace": args.atespace},
+        "spec": {"image": args.image, "env": env, "command": ["python3", "/app/benchmarks/learn/engine_task_entry.py"]},
+    }
+
+
+def task_state(text, name):
+    """Accept AX YAML/JSON output and its human-readable tasks table."""
+    try:
+        doc = yaml.safe_load(text)
+    except yaml.YAMLError:
+        doc = None
+    rows = doc.get("items", doc.get("tasks", [doc])) if isinstance(doc, dict) else doc if isinstance(doc, list) else []
+    for row in rows:
+        if not isinstance(row, dict) or row.get("metadata", {}).get("name", row.get("name")) != name:
+            continue
+        status = row.get("status", {})
+        ip = next(
+            (status.get(key) or row.get(key) for key in ("workerIP", "workerIp", "worker_ip", "ip") if status.get(key) or row.get(key)),
+            None,
+        )
+        worker = status.get("worker")
+        if not ip and isinstance(worker, dict):
+            ip = worker.get("ip") or worker.get("address")
+        return str(status.get("phase", row.get("phase", ""))), ip
+    lines = [line.split() for line in text.splitlines() if line.strip()]
+    if lines:
+        headers = [re.sub(r"[^a-z]", "", cell.lower()) for cell in lines[0]]
+        for cells in lines[1:]:
+            if name not in cells:
+                continue
+            row = dict(zip(headers, cells, strict=False))
+            phase = row.get("phase", row.get("status", ""))
+            ip = row.get("workerip") or row.get("ip")
+            if not ip:
+                for cell in cells:
+                    try:
+                        ipaddress.ip_address(cell)
+                        ip = cell
+                        break
+                    except ValueError:
+                        pass
+            return phase, None if ip in {"<none>", "-", ""} else ip
+    return "", None
+
+
+def bounded_read(read, timeout):
+    """Do not let SDK retries postpone task teardown beyond the pin deadline.
+
+    Only the read runs in a daemon thread; it cannot publish a late result to disk.
+    No executor shutdown waits for an SDK call whose socket has stopped responding.
+    """
+    if timeout <= 0:
+        raise TimeoutError("AX result deadline")
+    replies = queue.Queue(maxsize=1)
+
+    def fetch():
+        try:
+            replies.put((True, read()))
+        except Exception as exc:
+            replies.put((False, exc))
+
+    threading.Thread(target=fetch, daemon=True, name="engine-ax-result-read").start()
+    try:
+        ok, result = replies.get(timeout=timeout)
+    except queue.Empty:
+        raise TimeoutError("AX result deadline") from None
+    if not ok:
+        raise result
+    return result
+
+
+class AX:
+    def __init__(self, args, store=None, pause=time.sleep, clock=time.monotonic):
+        self.args = args
+        self.pause, self.clock = pause, clock
+        self.run = uuid.uuid4().hex
+        load_dotenv()
+        self.host = public_endpoint(os.environ.get("S3_ENDPOINT", ""))
+        # Validate the image before opening a store, whose constructor probes the bucket.
+        task_manifest("validate", args, {key: f"https://{self.host}/" for key in ("SOURCE_URL", "QUESTIONS_URL", "RESULT_URL")})
+        self.store = store if store is not None else open_store()
+        if store is None and not isinstance(self.store, S3Store):
+            raise ValueError("AX requires OUSAST_MEMORY=s3://bucket[/prefix]")
+
+    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(
+            [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,
+            capture_output=True,
+            text=True,
+            timeout=max(0.01, timeout),
+        )
+        if done.returncode:
+            # AX errors may echo manifest URLs; report only verb and exit code.
+            raise ValueError(f"ax {args[0]} failed (exit {done.returncode})")
+        return done
+
+    def receive(self, data, output, pin):
+        # Validate in isolation: malformed results must never become usable checkpoints.
+        with tempfile.TemporaryDirectory(prefix="engine-ax-result-") as temporary:
+            root = Path(temporary)
+            with tarfile.open(fileobj=io.BytesIO(data)) as archive:
+                archive.extractall(root, filter="data")
+            record = json.loads((root / "result.json").read_text())
+            units = record.get("units", [])
+            if record.get("done") and record.get("status") == "failed" and not units:
+                # Questions download can fail before the actor learns any unit identities.
+                # Keep that explicit failure, filling identities only from the host input.
+                record = {
+                    **pin,
+                    **record,
+                    "units": [
+                        {
+                            **unit,
+                            "status": "failed" if unit["supported"] else "unsupported",
+                            "reason": record.get("reason", "AX task entry failed"),
+                            "instrument": {"files": [], "bytes": 0, "cpg_path": None, "cpg_bytes": 0, "jvm": []},
+                            "questions": {},
+                            "questions_asked": [],
+                            "questions_completed": [],
+                            "witness_rows": [],
+                            "traces": [],
+                        }
+                        for unit in pin["units"]
+                    ],
+                }
+                units = record["units"]
+                (root / "result.json").write_text(json.dumps(record))
+            expected = [u["unit"] for u in pin["units"]]
+            actual = [u["unit"] for u in units]
+            if not record.get("done") or sorted(actual) != sorted(expected):
+                raise ValueError("AX Task returned incomplete or mismatched units")
+            output.mkdir(parents=True, exist_ok=True)
+            # Keep the artifacts without interpreting any remote paths as host destinations.
+            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]
+        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}/"
+        keys = [prefix + suffix for suffix in ("source.tar", "questions.json", "result.tar")]
+        expires = timedelta(seconds=self.args.deadline + 600)
+        worker_ip = None
+        submitted = False
+        try:
+            self.store._put(keys[0], pack(checkout))
+            self.store._put(keys[1], json.dumps(pin).encode())
+            urls = {
+                "SOURCE_URL": self.store.presign_get(keys[0], expires),
+                "QUESTIONS_URL": self.store.presign_get(keys[1], expires),
+                "RESULT_URL": self.store.presign_put(keys[2], expires),
+            }
+            if any(public_endpoint(url) != self.host for url in urls.values()):
+                raise EgressError("AX presign hostname must match S3_ENDPOINT for the store egress policy")
+            manifest = task_manifest(name, self.args, urls)
+            end = self.clock() + self.args.deadline
+            # An apply timeout can still have created the task; cleanup must cover it.
+            submitted = True
+            self.ax("apply", "-f", "-", manifest=manifest, timeout=min(30, self.args.deadline))
+            while self.clock() < end:
+                result = bounded_read(lambda: self.store._get(keys[2]), end - self.clock())
+                if self.clock() >= end:
+                    return None, worker_ip
+                phase = ""
+                try:
+                    state = self.ax("get", "tasks", timeout=min(30, end - self.clock())).stdout
+                    phase, ip = task_state(state, name)
+                    worker_ip = ip or worker_ip
+                except (ValueError, OSError, subprocess.SubprocessError):
+                    # The object is authoritative even when AX status is temporarily unavailable.
+                    pass
+                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
+                if phase.lower() == "failed":
+                    return subprocess.CompletedProcess([], 1, "", "AX Task phase Failed"), worker_ip
+                self.pause(min(2, max(0, end - self.clock())))
+            return None, worker_ip
+        except (subprocess.TimeoutExpired, TimeoutError):
+            return None, worker_ip
+        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:
+            errors = []
+            if submitted:
+                try:
+                    self.ax("delete", "task", name)
+                except Exception as exc:
+                    errors.append("task: " + type(exc).__name__)
+            for key in keys:
+                try:
+                    self.store._delete(key)
+                except Exception as exc:
+                    errors.append("object: " + type(exc).__name__)
+            if errors:
+                raise RuntimeError("AX cleanup failed; stop dispatching: " + ", ".join(errors))

benchmarks/learn/engine_trace_worker.py

@@ -41,6 +41,10 @@ def safe_path(root, name):
 FRONTENDS = {"php2cpg", "pysrc2cpg", "jssrc2cpg", "javasrc2cpg", "c2cpg", "joern-parse"}
 MIN_CPG_BYTES = 1024
 HEAP_MB = 2560
+CLUSTER_JVM_FLAGS = (
+    "-Xms1700m -Xmx1700m -XX:MaxMetaspaceSize=256m -XX:ReservedCodeCacheSize=128m "
+    "-XX:MaxDirectMemorySize=256m -Xss512k -XX:ActiveProcessorCount=1 -XX:+UseG1GC"
+)
 ENGINE_BUG = "engine bug: jssrc2cpg ObjectPropertyCallLinker"
 EXCLUDED_DIRS = frozenset({"tests", "test", "docs", "node_modules", "vendor", "dist", "build", "templates", "migrations"})
 
@@ -66,8 +70,16 @@ def instrument_failure(proof):
 class MeasuredBackend(JoernBackend):
     """Use the shipped backend with a process watchdog, no concurrent query threads."""
 
-    def __init__(self, deadline, question_deadline, checkpoint):
-        super().__init__(build_timeout=deadline, query_timeout=deadline, session_transport=False, heap_mb=HEAP_MB)
+    def __init__(self, deadline, question_deadline, checkpoint, heap_profile="vm"):
+        if heap_profile not in {"vm", "cluster"}:
+            raise ValueError(f"unknown heap profile: {heap_profile}")
+        self.heap_profile = heap_profile
+        super().__init__(
+            build_timeout=deadline,
+            query_timeout=deadline,
+            session_transport=False,
+            heap_mb=1700 if heap_profile == "cluster" else HEAP_MB,
+        )
         self.deadline = time.monotonic() + deadline
         self.question_deadline = question_deadline
         self.checkpoint = checkpoint
@@ -81,6 +93,15 @@ class MeasuredBackend(JoernBackend):
         self.retry_pending = False
         self.retries = []
 
+    def _jvm_env(self):
+        env = super()._jvm_env()
+        if self.heap_profile == "cluster":
+            # Bound every JVM, including the child spawned by the Joern launcher.
+            env["JAVA_TOOL_OPTIONS"] = CLUSTER_JVM_FLAGS
+            env.pop("_JAVA_OPTIONS", None)
+            env.pop("JDK_JAVA_OPTIONS", None)
+        return env
+
     def _apply_overlays(self, cpg_path, scratch):
         result = super()._apply_overlays(cpg_path, scratch)
         # An optional overlay failure must not prevent the taint-only attempt.
@@ -250,14 +271,18 @@ def unit_record(unit, status, reason="", **extra):
     }
 
 
-def worker_language(pin, root, out, deadline, question_deadline, persist=write_json):
+def worker_language(pin, root, out, deadline, question_deadline, persist=write_json, heap_profile="vm"):
     proof = {
         "files": [],
         "bytes": 0,
         "cpg_path": None,
         "cpg_bytes": 0,
         "jvm": [],
-        "heap_mb": HEAP_MB,
+        "heap_mb": 1700 if heap_profile == "cluster" else HEAP_MB,
+        "heap_profile": heap_profile,
+        "java_tool_options": CLUSTER_JVM_FLAGS
+        if heap_profile == "cluster"
+        else f"{os.environ.get('JAVA_TOOL_OPTIONS', '').strip()} -Xmx{HEAP_MB}m".strip(),
         "excludes": sorted(EXCLUDED_DIRS),
         "materialization": pin.get("materialization", {}),
     }
@@ -284,7 +309,7 @@ def worker_language(pin, root, out, deadline, question_deadline, persist=write_j
                 row["witness_rows"] = [r for r in backend.answers[rid] if isinstance(r, dict) and "kind" not in r and "sink" in r]
         persist(out, record)
 
-    backend = MeasuredBackend(deadline, question_deadline, checkpoint)
+    backend = MeasuredBackend(deadline, question_deadline, checkpoint, **({"heap_profile": heap_profile} if heap_profile != "vm" else {}))
     proof["jvm"] = backend.stages
     proof["retries"] = getattr(backend, "retries", [])
     backend.labelled_files = sorted({u["file"] for u in pin["units"] if u["supported"]})
@@ -398,11 +423,11 @@ def worker_language(pin, root, out, deadline, question_deadline, persist=write_j
     return record
 
 
-def worker(pin, root, out, deadline, question_deadline):
+def worker(pin, root, out, deadline, question_deadline, heap_profile="vm"):
     """Run language partitions sequentially under one shared container deadline."""
     languages = sorted({"javascript" if u["language"] == "typescript" else u["language"] for u in pin["units"] if u["supported"]})
     if len(languages) <= 1:
-        return worker_language(pin, root, out, deadline, question_deadline)
+        return worker_language(pin, root, out, deadline, question_deadline, heap_profile=heap_profile)
     end = time.monotonic() + deadline
     record = {
         **pin,
@@ -431,7 +456,9 @@ def worker(pin, root, out, deadline, question_deadline):
             u for u in pin["units"] if u["supported"] and ("javascript" if u["language"] == "typescript" else u["language"]) == language
         ]
         questions = {u["unit"]: pin["questions"][u["unit"]] for u in units}
-        worker_language({**pin, "units": units, "questions": questions}, root, out, remaining, question_deadline, persist)
+        worker_language(
+            {**pin, "units": units, "questions": questions}, root, out, remaining, question_deadline, persist, heap_profile=heap_profile
+        )
     record["done"] = True
     write_json(out, record)
     return record
@@ -442,10 +469,13 @@ def main():
     parser.add_argument("--worker", type=Path, required=True)
     parser.add_argument("--deadline", type=int, default=900)
     parser.add_argument("--question-deadline", type=int, default=120)
+    parser.add_argument("--heap-profile", choices=("cluster", "vm"), default="vm")
+    parser.add_argument("--root", type=Path, default=Path("/case"))
+    parser.add_argument("--output", type=Path, default=Path("/out/result.json"))
     args = parser.parse_args()
     if min(args.deadline, args.question_deadline) <= 0:
         parser.error("deadlines must be positive")
-    worker(json.loads(args.worker.read_text()), Path("/case"), Path("/out/result.json"), args.deadline, args.question_deadline)
+    worker(json.loads(args.worker.read_text()), args.root, args.output, args.deadline, args.question_deadline, args.heap_profile)
     return 0
 
 

plane/Dockerfile.engine-task

@@ -0,0 +1,52 @@
+# Build from the repository root. AX v0.3.1 starts spec.command immediately;
+# preserve its runner as PID 1 and put the engine command in the Task manifest.
+FROM eclipse-temurin:21-jre-jammy AS jre
+
+FROM ghcr.io/norandom/ax-task-runner:v0.3.1
+LABEL org.opencontainers.image.licenses="Apache-2.0" \
+      org.opencontainers.image.source="https://github.com/norandom/OpenUltraSAST"
+
+USER root
+ARG JOERN_VERSION=v4.0.625
+ARG ARCHIVE=joern-cli-linux-x86_64.zip
+
+# The AX runner is Python 3.12 / Debian slim. Copy the Java 21 runtime from
+# Jammy (older glibc than Debian bookworm) instead of relying on distro JRE versions.
+COPY --from=jre /opt/java/openjdk /opt/java/openjdk
+ENV JAVA_HOME=/opt/java/openjdk
+ENV PATH="/venv/bin:/opt/java/openjdk/bin:/opt/joern-cli:${PATH}"
+RUN apt-get update \
+ && apt-get install -y --no-install-recommends \
+      php-cli unzip curl ca-certificates libstdc++6 zlib1g \
+ && rm -rf /var/lib/apt/lists/*
+
+RUN set -eux; \
+    cd /tmp; \
+    curl -fsSL -O "https://github.com/joernio/joern/releases/download/${JOERN_VERSION}/${ARCHIVE}"; \
+    curl -fsSL -O "https://github.com/joernio/joern/releases/download/${JOERN_VERSION}/${ARCHIVE}.sha512"; \
+    sed 's@  target/@  @' "${ARCHIVE}.sha512" > checksum.sha512; \
+    sha512sum -c checksum.sha512; \
+    unzip -q "${ARCHIVE}" -d /opt; \
+    rm -f "${ARCHIVE}" "${ARCHIVE}.sha512" checksum.sha512; \
+    /opt/joern-cli/joern-parse --help > /dev/null
+
+RUN if [ ! -e /opt/joern-cli/jssrc2cpg ] && [ -x /opt/joern-cli/jssrc2cpg.sh ]; then \
+      ln -s jssrc2cpg.sh /opt/joern-cli/jssrc2cpg; \
+    fi
+COPY ops/frontend-retention/install.py /tmp/frontend-retention-install.py
+RUN python3 /tmp/frontend-retention-install.py /opt/joern-cli \
+ && rm /tmp/frontend-retention-install.py
+
+WORKDIR /app
+COPY pyproject.toml README.md LICENSE ./
+COPY src ./src
+COPY benchmarks/learn/engine_task_entry.py \
+     benchmarks/learn/engine_trace_worker.py \
+     benchmarks/learn/engine_trace_parse.py ./benchmarks/learn/
+# Worker imports are standard-library plus engine modules. Isolate them from
+# the AX runner's Python environment and omit PyYAML, boto3 and all extras.
+RUN python3 -m venv /venv \
+ && /venv/bin/pip install --no-cache-dir --no-deps . \
+ && /venv/bin/python -I -c 'import pathlib, runpy, sys; sys.modules.update(yaml=None, boto3=None); sys.path.insert(0, "/app/benchmarks/learn"); p = pathlib.Path("/app/benchmarks/learn/engine_trace_worker.py"); print("Worker import smoke bytes:", len(p.read_bytes())); runpy.run_path(str(p), run_name="engine_worker_import_smoke")'
+
+# Do not set ENTRYPOINT or CMD: both belong to the upstream AX runner.

tests/test_ax_executor.py

@@ -0,0 +1,277 @@
+"""AX transport controls using an executable shim and an in-memory object store."""
+
+import io
+import json
+import sys
+import tarfile
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+with pytest.MonkeyPatch.context() as patch:
+    patch.syspath_prepend(str(Path(__file__).resolve().parents[1] / "benchmarks/learn"))
+    import engine_trace_ax as ax
+
+
+IMAGE = "ghcr.io/norandom/ousast-engine-task@sha256:" + "a" * 64
+PIN = {"repo": "owner/repo", "pin": "123", "units": [{"unit": "a", "supported": True}], "questions": {"a": {}}}
+
+
+def archive(record):
+    buffer = io.BytesIO()
+    with tarfile.open(fileobj=buffer, mode="w") as tar:
+        data = json.dumps(record).encode()
+        info = tarfile.TarInfo("result.json")
+        info.size = len(data)
+        tar.addfile(info, io.BytesIO(data))
+    return buffer.getvalue()
+
+
+class Store:
+    def __init__(self, result=None):
+        self.result = result
+        self.objects = {}
+        self.deleted = []
+        self.expiries = []
+
+    def _put(self, key, data):
+        self.objects[key] = data
+
+    def _get(self, key):
+        return (self.result, None) if self.result else None
+
+    def _delete(self, key):
+        self.deleted.append(key)
+        self.objects.pop(key, None)
+
+    def presign_get(self, key, expiry):
+        self.expiries.append(expiry.total_seconds())
+        return "https://files.because-security.com/bucket/" + key + "?signature=secret"
+
+    presign_put = presign_get
+
+
+@pytest.fixture
+def setup(tmp_path, monkeypatch):
+    log = tmp_path / "calls.jsonl"
+    shim = tmp_path / "fake-ax"
+    shim.write_text("#!/bin/sh\nexec " + sys.executable + " " + str(tmp_path / "shim.py") + ' "$@"\n')
+    shim.chmod(0o755)
+    (tmp_path / "shim.py").write_text("""import json, os, sys
+from pathlib import Path
+args = sys.argv[1:]
+with Path(os.environ["AX_LOG"]).open("a") as log:
+    log.write(json.dumps({"args": args, "input": sys.stdin.read() if args[0] == "apply" else "",
+        "kubeconfig": os.environ.get("KUBECONFIG")}) + "\\n")
+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"])
+    print(json.dumps({"items": [{"metadata": manifest["metadata"], "status": {
+        "phase": os.environ.get("AX_PHASE", "Running"), "workerIP": "10.0.0.2"}}]}))
+if args[0] == "delete" and os.environ.get("AX_DELETE_FAIL"):
+    sys.exit(1)
+""")
+    monkeypatch.setenv("AX_LOG", str(log))
+    monkeypatch.setenv("S3_ENDPOINT", "https://files.because-security.com")
+    args = SimpleNamespace(
+        ax_bin=str(shim), kubeconfig="/operator/kubeconfig", atespace="default", image=IMAGE, deadline=4, question_deadline=2
+    )
+    source = tmp_path / "case"
+    source.mkdir()
+    (source / "app.py").write_text("print(42)\n")
+    output = tmp_path / "out"
+    output.mkdir()
+    now = [0]
+
+    def pause(seconds):
+        now[0] += seconds
+
+    return args, source, output, log, pause, lambda: now[0]
+
+
+def test_ax_executor_success_while_running_and_cleanup(setup):
+    args, source, output, log, pause, clock = setup
+    store = Store(archive({**PIN, "done": True}))
+    executor = ax.AX(args, store, pause=pause, clock=clock)
+    done, ip = executor.execute(source, output, PIN, "owner/repo/123")
+    assert done.returncode == 0
+    assert ip == "10.0.0.2"
+    assert json.loads((output / "result.json").read_text())["done"]
+    calls = [json.loads(line) for line in log.read_text().splitlines()]
+    assert calls[0]["args"] == ["apply", "-f", "-"]
+    assert calls[-1]["args"][:2] == ["delete", "task"]
+    assert calls[-1]["args"][-2:] == ["--atespace", args.atespace]
+    assert all(c["kubeconfig"] == args.kubeconfig for c in calls)
+    manifest = json.loads(calls[0]["input"])
+    assert manifest["apiVersion"] == "ax.io/v1alpha1"
+    assert manifest["kind"] == "Task"
+    assert len(manifest["metadata"]["name"]) <= 49
+    assert manifest["spec"]["image"] == IMAGE
+    env = {e["name"]: e["value"] for e in manifest["spec"]["env"]}
+    assert set(env) == {"SOURCE_URL", "QUESTIONS_URL", "RESULT_URL", "DEADLINE", "QUESTION_DEADLINE"}
+    assert len(json.dumps(manifest["spec"]["env"]).encode()) < 32768
+    assert "AWS" not in calls[0]["input"]
+    assert store.expiries == [604] * 3
+    assert len(store.deleted) == 3 and not store.objects
+
+
+def test_ax_executor_timeout_deletes_task(setup):
+    args, source, output, log, pause, clock = setup
+    store = Store()
+    done, _ = ax.AX(args, store, pause=pause, clock=clock).execute(source, output, PIN, "pin")
+    assert done is None
+    assert json.loads(log.read_text().splitlines()[-1])["args"][:2] == ["delete", "task"]
+    assert len(store.deleted) == 3
+
+
+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 json.loads(log.read_text().splitlines()[-1])["args"][:2] == ["delete", "task"]
+
+
+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")
+    store = Store(archive({**PIN, "done": True}))
+    with pytest.raises(RuntimeError, match="cleanup"):
+        ax.AX(args, store, pause=pause, clock=clock).execute(source, output, PIN, "pin")
+    assert len(store.deleted) == 3
+
+
+@pytest.mark.parametrize("endpoint", ["http://files.because-security.com", "https://127.0.0.1", "https://[::1]", "https://localhost"])
+def test_ax_executor_endpoint_rejected_before_store_open(setup, monkeypatch, endpoint):
+    args = setup[0]
+    monkeypatch.setenv("S3_ENDPOINT", endpoint)
+    monkeypatch.setattr(ax, "open_store", lambda: pytest.fail("store opened before validation"))
+    with pytest.raises(ValueError, match="egress policy"):
+        ax.AX(args)
+
+
+def test_ax_executor_rejects_tag_and_large_env(setup):
+    args = setup[0]
+    urls = {key: "https://files.because-security.com/" for key in ("SOURCE_URL", "QUESTIONS_URL", "RESULT_URL")}
+    args.image = "image:latest"
+    with pytest.raises(ValueError, match="digest"):
+        ax.task_manifest("name", args, urls)
+    args.image = IMAGE
+    urls["SOURCE_URL"] += "x" * 32768
+    with pytest.raises(ValueError, match="32 KB"):
+        ax.task_manifest("name", args, urls)
+
+
+def test_ax_executor_rejects_mismatched_result(setup):
+    args, source, output, _, pause, clock = setup
+    store = Store(archive({"done": True, "units": []}))
+    done, _ = ax.AX(args, store, pause=pause, clock=clock).execute(source, output, PIN, "pin")
+    assert done.returncode == 1
+    assert not (output / "result.json").exists()
+
+
+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_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": []}))
+    done, _ = ax.AX(args, store, pause=pause, clock=clock).execute(source, output, PIN, "pin")
+    assert done.returncode == 0
+    unit = json.loads((output / "result.json").read_text())["units"][0]
+    assert unit["unit"] == "a" and unit["status"] == "failed"
+    assert unit["reason"] == "questions download failed"
+
+
+def test_ax_executor_presign_host_must_match_egress_policy(setup):
+    args, source, output, log, pause, clock = setup
+    store = Store()
+    store.presign_put = lambda key, expiry: "https://another-store.example/result"
+    done, _ = ax.AX(args, store, pause=pause, clock=clock).execute(source, output, PIN, "pin")
+    assert done.returncode == 1 and "egress policy" in done.stderr
+    assert not log.exists()
+    assert len(store.deleted) == 3
+
+
+def test_ax_executor_get_failure_does_not_lose_result(setup, monkeypatch):
+    args, source, output, _, pause, clock = setup
+    monkeypatch.setenv("AX_GET_FAIL", "1")
+    done, _ = ax.AX(args, Store(archive({**PIN, "done": True})), pause=pause, clock=clock).execute(source, output, PIN, "pin")
+    assert done.returncode == 0
+
+
+def test_ax_executor_absent_worker_ip():
+    assert ax.task_state("NAME ATESPACE PHASE ACTOR WORKER-IP AGE\nours other Pending actor <none> 1s\n", "ours") == ("Pending", None)
+
+
+def test_ax_executor_late_store_result_times_out(setup):
+    args, source, output, log, pause, clock = setup
+    store = Store(archive({**PIN, "done": True}))
+    original_get = store._get
+
+    def late_get(key):
+        pause(args.deadline)
+        return original_get(key)
+
+    store._get = late_get
+    done, _ = ax.AX(args, store, pause=pause, clock=clock).execute(source, output, PIN, "pin")
+    assert done is None
+    assert not (output / "result.json").exists()
+    assert json.loads(log.read_text().splitlines()[-1])["args"][:2] == ["delete", "task"]
+
+
+def test_ax_executor_late_status_times_out(setup):
+    args, source, output, _, pause, clock = setup
+    executor = ax.AX(args, Store(archive({**PIN, "done": True})), pause=pause, clock=clock)
+    original_ax = executor.ax
+
+    def late_ax(*command, **kwargs):
+        result = original_ax(*command, **kwargs)
+        if command[0] == "get":
+            pause(args.deadline)
+        return result
+
+    executor.ax = late_ax
+    done, _ = executor.execute(source, output, PIN, "pin")
+    assert done is None
+    assert not (output / "result.json").exists()
+
+
+def test_ax_executor_blocking_store_read_cannot_delay_task_deletion(setup):
+    import threading
+    import time
+
+    args, source, output, log, _, _ = setup
+    args.deadline = 0.15
+    release = threading.Event()
+    store = Store()
+    store._get = lambda key: release.wait(5)
+    started = time.monotonic()
+    try:
+        done, _ = ax.AX(args, store).execute(source, output, PIN, "pin")
+        assert done is None
+        assert time.monotonic() - started < 2
+        assert json.loads(log.read_text().splitlines()[-1])["args"][:2] == ["delete", "task"]
+    finally:
+        release.set()
+
+
+def test_ax_task_names_reserve_watcher_prefix_and_fit_actor_limit(setup):
+    args, source, output, log, pause, clock = setup
+    executor = ax.AX(args, Store(archive({**PIN, "done": True})), pause=pause, clock=clock)
+    # Sanitization and truncation must not collapse distinct pins.
+    for pin_id in ["A/B", "a-b", "x" * 100 + "1", "x" * 100 + "2", "☃"]:
+        done, _ = executor.execute(source, output, PIN, pin_id)
+        assert done.returncode == 0
+    manifests = [json.loads(call["input"]) for line in log.read_text().splitlines() if (call := json.loads(line))["args"][0] == "apply"]
+    names = [manifest["metadata"]["name"] for manifest in manifests]
+    assert len(set(names)) == 5
+    for manifest, name in zip(manifests, names, strict=True):
+        assert manifest["kind"] == "Task"
+        assert manifest["metadata"]["atespace"] == "default"
+        assert name.startswith("ousast-engine-")
+        assert len(name.encode()) <= 49
+        assert len((name + "-tmpl-12345678").encode()) <= 63

tests/test_engine_task_entry.py

@@ -0,0 +1,224 @@
+"""Offline AX worker entry controls; HTTP is loopback only."""
+
+import importlib.util
+import io
+import json
+import tarfile
+import threading
+import time
+import urllib.error
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from pathlib import Path
+
+import pytest
+
+SCRIPTS = Path(__file__).resolve().parents[1] / "benchmarks/learn"
+
+
+def load(name):
+    with pytest.MonkeyPatch.context() as patch:
+        patch.syspath_prepend(str(SCRIPTS))
+        spec = importlib.util.spec_from_file_location(name, SCRIPTS / f"{name}.py")
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+        return module
+
+
+def test_engine_cluster_heap_flags(monkeypatch):
+    worker = load("engine_trace_worker")
+    monkeypatch.setenv("JAVA_TOOL_OPTIONS", "-Xmx9999m")
+    backend = worker.MeasuredBackend(30, 10, lambda: None, heap_profile="cluster")
+    flags = (
+        "-Xms1700m -Xmx1700m -XX:MaxMetaspaceSize=256m -XX:ReservedCodeCacheSize=128m "
+        "-XX:MaxDirectMemorySize=256m -Xss512k -XX:ActiveProcessorCount=1 -XX:+UseG1GC"
+    )
+    assert backend._jvm_env()["JAVA_TOOL_OPTIONS"] == flags
+    assert backend._heap_flag() == "-J-Xmx1700m"
+    assert worker.MeasuredBackend(30, 10, lambda: None)._heap_flag() == "-J-Xmx2560m"
+
+
+@pytest.mark.parametrize("readiness", ["ready", "delayed", "closed"])
+@pytest.mark.parametrize("failure", [False, True])
+@pytest.mark.parametrize("transport", ["memory", "loopback"])
+def test_task_entry_http_roundtrip(tmp_path, monkeypatch, failure, transport, readiness):
+    entry = load("engine_task_entry")
+    archive = io.BytesIO()
+    with tarfile.open(fileobj=archive, mode="w") as tar:
+        data = b"print('input proof')\n"
+        info = tarfile.TarInfo("app.py")
+        info.size = len(data)
+        tar.addfile(info, io.BytesIO(data))
+    pin = {"pin": "abc", "units": [{"unit": "u", "supported": True}], "questions": {"u": {}}}
+    uploads = []
+    gets = []
+    puts = []
+    started = time.monotonic()
+
+    class Handler(BaseHTTPRequestHandler):
+        def do_GET(self):
+            gets.append((self.path, time.monotonic() - started))
+            if readiness == "closed" or (readiness == "delayed" and gets[-1][1] < 3):
+                self.send_response(403 if len(gets) % 2 else 503)
+                self.end_headers()
+                return
+            data = archive.getvalue() if self.path == "/source" else json.dumps(pin).encode()
+            self.send_response(200)
+            self.end_headers()
+            self.wfile.write(data)
+
+        def do_PUT(self):
+            body = self.rfile.read(int(self.headers["Content-Length"]))
+            puts.append(time.monotonic())
+            if readiness == "closed" and len(puts) == 1:
+                self.send_response(403)
+                self.end_headers()
+                return
+            uploads.append(body)
+            self.send_response(200)
+            self.end_headers()
+
+        def log_message(self, *args):
+            pass
+
+    fake_worker = tmp_path / "worker.py"
+    fake_worker.write_text(
+        "import argparse,json,pathlib,sys\np=argparse.ArgumentParser()\n"
+        "for key in ['worker','root','output','deadline','question-deadline','heap-profile']: p.add_argument('--'+key)\n"
+        "a=p.parse_args()\nassert a.heap_profile == 'cluster'\n"
+        "assert (pathlib.Path(a.root)/'app.py').stat().st_size > 0\n"
+        + ("sys.exit(7)\n" if failure else "pathlib.Path(a.output).write_text(json.dumps({'done':True,'units':[]}))\n")
+    )
+    if transport == "loopback":
+        try:
+            server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
+        except PermissionError:
+            pytest.skip("sandbox prohibits loopback sockets; in-memory HTTP roundtrip still runs")
+        thread = threading.Thread(target=server.serve_forever, daemon=True)
+        thread.start()
+        base = f"http://127.0.0.1:{server.server_port}"
+    else:
+        base = "https://fake-store.invalid"
+
+        def urlopen(request, timeout):
+            method = request.get_method() if hasattr(request, "get_method") else "GET"
+            url = request.full_url if hasattr(request, "full_url") else request
+            body = request.data or b"" if method == "PUT" else b""
+            raw = f"{method} {url.removeprefix(base)} HTTP/1.0\r\nContent-Length: {len(body)}\r\n\r\n".encode() + body
+
+            class Socket:
+                def __init__(self):
+                    self.sent = bytearray()
+
+                def makefile(self, *args):
+                    return io.BytesIO(raw)
+
+                def sendall(self, data):
+                    self.sent.extend(data)
+
+            connection = Socket()
+            Handler(connection, ("127.0.0.1", 0), None)
+            headers, payload = bytes(connection.sent).split(b"\r\n\r\n", 1)
+            status = int(headers.split()[1])
+            if status != 200:
+                raise urllib.error.HTTPError(url, status, "denied", {}, io.BytesIO(payload))
+            return io.BytesIO(payload)
+
+        monkeypatch.setattr(entry.urllib.request, "urlopen", urlopen)
+    try:
+        for name, suffix in [("SOURCE_URL", "source"), ("QUESTIONS_URL", "questions"), ("RESULT_URL", "result")]:
+            monkeypatch.setenv(name, base + "/" + suffix)
+        budget = 0.7 if readiness == "closed" else 60
+        assert entry.main(work_root=tmp_path / "work", worker_script=fake_worker, readiness_budget=budget) == 0
+        if readiness == "delayed":
+            assert len(gets) >= 5
+            assert gets[-1][1] >= 3
+        if readiness == "closed":
+            assert 0.7 <= puts[0] - started < 1.5
+            assert len(puts) == 2
+            assert puts[1] - puts[0] >= 0.5
+            assert all(path == "/questions" for path, _ in gets)
+    finally:
+        if transport == "loopback":
+            server.shutdown()
+            thread.join()
+            server.server_close()
+    assert len(uploads) == 1
+    with tarfile.open(fileobj=io.BytesIO(uploads[0])) as tar:
+        result = json.load(tar.extractfile("result.json"))
+    assert result["done"]
+    if readiness == "closed":
+        assert result["status"] == "failed"
+        assert result["reason"] == "egress not open after 60 s"
+        assert json.loads((tmp_path / "work/out/result.json").read_text()) == result
+    elif failure:
+        assert result["units"][0]["status"] == "failed"
+        assert "7" in result["units"][0]["reason"]
+
+
+def test_engine_cluster_profile_recorded_on_input_failure(tmp_path):
+    worker = load("engine_trace_worker")
+    pin = {
+        "units": [{"unit": "u", "file": "missing.py", "language": "python", "supported": True, "reason": ""}],
+        "questions": {"u": {}},
+    }
+    result = worker.worker(pin, tmp_path, tmp_path / "result.json", 30, 10, heap_profile="cluster")
+    assert result["units"][0]["status"] == "failed"
+    assert result["instrument"]["heap_profile"] == "cluster"
+    assert result["instrument"]["heap_mb"] == 1700
+    assert result["instrument"]["java_tool_options"] == worker.CLUSTER_JVM_FLAGS
+
+
+def test_task_entry_rejects_tar_traversal(tmp_path):
+    entry = load("engine_task_entry")
+    source = tmp_path / "source.tar"
+    with tarfile.open(source, "w") as archive:
+        member = tarfile.TarInfo("../escape")
+        member.size = 1
+        archive.addfile(member, io.BytesIO(b"x"))
+    with pytest.raises(ValueError, match="unsafe"):
+        entry.unpack(source, tmp_path / "case")
+    assert not (tmp_path / "escape").exists()
+
+
+@pytest.mark.parametrize("error", [urllib.error.URLError("refused"), ConnectionResetError("reset")])
+def test_readiness_connection_errors_back_off_to_cap(tmp_path, monkeypatch, error):
+    entry = load("engine_task_entry")
+    now = [0.0]
+    pauses = []
+    timeouts = []
+
+    def pause(seconds):
+        pauses.append(seconds)
+        now[0] += seconds
+
+    def refused(request, timeout):
+        timeouts.append(timeout)
+        raise error
+
+    monkeypatch.setattr(entry.time, "monotonic", lambda: now[0])
+    monkeypatch.setattr(entry.time, "sleep", pause)
+    monkeypatch.setattr(entry.urllib.request, "urlopen", refused)
+    with pytest.raises(RuntimeError, match="^egress not open after 60 s$"):
+        entry.download("https://fake.invalid", tmp_path / "input", readiness_budget=60)
+    assert now[0] == 60
+    assert pauses[:5] == [0.5, 1, 2, 4, 8]
+    assert max(pauses) == 8
+    assert timeouts[-1] == pauses[-1]
+
+
+@pytest.mark.parametrize("status, attempts", [(403, 3), (503, 3), (404, 1)])
+def test_later_transfers_keep_short_retry(monkeypatch, status, attempts):
+    entry = load("engine_task_entry")
+    calls = []
+    pauses = []
+
+    def refused(request, timeout):
+        calls.append(timeout)
+        raise urllib.error.HTTPError(request, status, "denied", {}, io.BytesIO())
+
+    monkeypatch.setattr(entry.time, "sleep", pauses.append)
+    monkeypatch.setattr(entry.urllib.request, "urlopen", refused)
+    with pytest.raises(urllib.error.HTTPError):
+        entry.transfer("https://fake.invalid")
+    assert len(calls) == attempts
+    assert pauses == list(range(1, attempts))

tests/test_engine_trace.py

@@ -6,6 +6,7 @@ import importlib.util
 import json
 import subprocess
 import sys
+from collections import Counter
 from dataclasses import replace
 from pathlib import Path
 
@@ -1452,3 +1453,146 @@ def test_engine_queue_pin_persistence_resume_and_progress(tmp_path, monkeypatch,
         assert trace.make_plan([unit()], tmp_path, rerun_status=("failed", "timeout"))
     progress = json.loads((tmp_path / "progress.json").read_text())
     assert progress["pins_done"] == 1 and progress["pins_in_flight"] == 0
+
+
+@pytest.mark.parametrize("executor", ["ax", "mixed-ax"])
+def test_ax_scheduler_requeues_oom_to_single_vm_lane(tmp_path, monkeypatch, executor):
+    import threading
+    import time
+    from types import SimpleNamespace
+
+    plan = trace.make_plan([unit(str(i), str(i)) for i in range(7)], tmp_path)
+    args = SimpleNamespace(executor=executor, parallel=2, out=tmp_path, image="cluster", docker_image="vm")
+    lock = threading.Lock()
+    active = Counter()
+    peak = Counter()
+    attempts = []
+
+    def run(pin, source, provenance, inputs, out, options):
+        with lock:
+            active[options.executor] += 1
+            peak[options.executor] = max(peak[options.executor], active[options.executor])
+            attempts.append((pin["pin"], options.executor, options.image))
+        time.sleep(0.01)
+        oom = options.executor == "ax"
+        record = {
+            **pin,
+            "done": True,
+            "units": [worker.unit_record(u, "failed" if oom else "path", "JVM out of memory" if oom else "") for u in pin["units"]],
+        }
+        if getattr(options, "fallback", None):
+            record["fallback"] = options.fallback
+        trace.write_json(out / trace.pin_name(pin), record)
+        with lock:
+            active[options.executor] -= 1
+        return oom
+
+    monkeypatch.setattr(trace, "run_pin", run)
+    assert not trace.run_plan(plan, None, {}, None, args)
+    assert peak["ax"] <= 2 and peak["docker"] == 1
+    for pin in plan:
+        record = json.loads((tmp_path / trace.pin_name(pin)).read_text())
+        if (pin["pin"], "ax", "cluster") in attempts:
+            assert record["fallback"] == "vm-oom"
+            assert (pin["pin"], "docker", "vm") in attempts
+    progress = json.loads((tmp_path / "progress.json").read_text())
+    assert progress["pins_done"] == 7
+    assert progress["statuses"] == {"path": 7}
+    assert progress["pins_in_flight"] == 0
+
+
+@pytest.mark.parametrize("timeout", [False, True])
+def test_ax_run_pin_persists_worker_ip_timeout_and_resume(tmp_path, monkeypatch, timeout):
+    pin = trace.make_plan([unit()], tmp_path)[0]
+    calls = []
+
+    class Inputs:
+        def materialize(self, prepared, target):
+            target.mkdir()
+            (target / "app.py").write_text("def run(): pass")
+            return {"files": 1, "bytes": 15}
+
+    class Dispatcher:
+        def execute(self, checkout, output, prepared, name):
+            assert (checkout / "app.py").read_bytes() == b"def run(): pass"
+            assert prepared["questions"] == {"a": {}}
+            assert prepared["materialization"]["bytes"] == 15
+            calls.append(name)
+            if timeout:
+                return None, "10.0.0.12"
+            trace.write_json(
+                output / "result.json",
+                {**prepared, "done": True, "units": [worker.unit_record(u, "asked-nothing") for u in prepared["units"]]},
+            )
+            return subprocess.CompletedProcess([], 0, "", ""), "10.0.0.12"
+
+    args = k8s_args(tmp_path, executor="ax", ax_dispatcher=Dispatcher(), image="registry/engine@sha256:" + "a" * 64)
+    monkeypatch.setattr(trace, "prepare_questions", lambda *_: {"a": {}})
+    monkeypatch.setattr(subprocess, "run", lambda *_a, **_k: pytest.fail("AX run_pin must use dispatcher only"))
+    assert trace.run_pin(pin, None, {"host": "unused"}, Inputs(), tmp_path, args) is timeout
+    result = json.loads((tmp_path / trace.pin_name(pin)).read_text())
+    assert calls == [trace.pin_name(pin)[:-5]]
+    assert result["done"] and result["executor"] == "ax"
+    assert result["worker_ip"] == result["worker"] == "10.0.0.12"
+    assert result["analyzer"] == {"image": args.image, "code": "baked-in-image"}
+    assert result["units"][0]["status"] == ("timeout" if timeout else "asked-nothing")
+    assert result["container_exit"] == (None if timeout else 0)
+    assert trace.make_plan([unit()], tmp_path, rerun_status=()) == []
+    assert bool(trace.make_plan([unit()], tmp_path, rerun_status=("timeout",))) is timeout
+
+
+@pytest.mark.parametrize("executor", ["ax", "mixed-ax"])
+def test_ax_pending_fallback_resumes_in_docker_and_retains_attempt(tmp_path, monkeypatch, executor):
+    pin = trace.make_plan([unit()], tmp_path)[0]
+    attempt = {
+        **pin,
+        "done": True,
+        "executor": "ax",
+        "worker_ip": "10.0.0.12",
+        "fallback": "vm-oom",
+        "fallback_pending": True,
+        "units": [worker.unit_record(u, "failed", "JVM out of memory") for u in pin["units"]],
+    }
+    trace.write_json(tmp_path / trace.pin_name(pin), attempt)
+    resumed = trace.make_plan([unit()], tmp_path, rerun_status=())
+    assert len(resumed) == 1
+    calls = []
+
+    class Inputs:
+        def materialize(self, prepared, target):
+            target.mkdir()
+            (target / "app.py").write_text("def run(): pass")
+            return {"files": 1, "bytes": 15}
+
+    class Dispatcher:
+        def execute(self, *_):
+            pytest.fail("pending OOM fallback must not dispatch AX again")
+
+    def docker(command, **kwargs):
+        calls.append(command)
+        assert command[:2] in (["docker", "run"], ["docker", "rm"])
+        if command[1] == "run":
+            assert command[command.index("--heap-profile") + 1] == "vm"
+            assert "vm-image" in command
+            output = Path(next(arg.removesuffix(":/out") for arg in command if arg.endswith(":/out")))
+            prepared = json.loads((output / "pin.json").read_text())
+            trace.write_json(
+                output / "result.json",
+                {**prepared, "done": True, "units": [worker.unit_record(u, "asked-nothing") for u in prepared["units"]]},
+            )
+        return subprocess.CompletedProcess(command, 0, "", "")
+
+    args = k8s_args(tmp_path, executor=executor, ax_dispatcher=Dispatcher(), docker_image="vm-image")
+    monkeypatch.setattr(trace, "prepare_questions", lambda *_: {"a": {}})
+    monkeypatch.setattr(subprocess, "run", docker)
+    assert not trace.run_plan(resumed, tmp_path / "frozen", {"code": "frozen"}, Inputs(), args)
+    result = json.loads((tmp_path / trace.pin_name(pin)).read_text())
+    assert [command[1] for command in calls] == ["run", "rm"]
+    assert result["executor"] == "docker" and result["image"] == "vm-image"
+    assert result["fallback"] == "vm-oom" and not result["fallback_pending"]
+    assert result["ax_attempt"] == attempt
+    assert result["units"][0]["status"] == "asked-nothing"
+    assert trace.make_plan([unit()], tmp_path, rerun_status=()) == []
+    progress = json.loads((tmp_path / "progress.json").read_text())
+    assert progress["pins_done"] == 1 and progress["pins_in_flight"] == 0
+    assert progress["fallbacks_pending"] == 0 and progress["statuses"] == {"asked-nothing": 1}