engine: AX-task executor for kube-ax (engine-task image FROM ax-task-runner v0.3.1, presigned HTTPS I/O, ousast-engine- names, 60 s egress-ready retry, 1.7 GB cluster heap, VM OOM fallback)

Implemented with Codex. Live smoke pending the operator egress watcher.

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

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}