unseen-eval groups 4-6: generic AX batch dispatcher (engine_trace on it), replay task and image stage, scorer with adoption gate

Implemented with Codex.

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

benchmarks/learn/engine_trace.py

@@ -16,13 +16,13 @@ import os
 import re
 import shutil
 import subprocess
+import sys
 import tempfile
 import textwrap
 import threading
 import time
 import tomllib
-from collections import Counter, defaultdict, deque
-from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
+from collections import Counter, defaultdict
 from pathlib import Path
 
 from engine_trace_k8s import IMAGE, Kubernetes
@@ -39,6 +39,9 @@ from openultrasast.plane.memory import FileStore
 from openultrasast.plane.tasks.alerts import engine_languages
 from openultrasast.preprocess import detect_language
 
+sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
+from benchmarks.ax.batch import schedule
+
 ROOT = Path(__file__).resolve().parents[2]
 DEFAULT_OUT = Path.home() / "ousast-results/plane/engine-trace"
 SNAPSHOT_FILES = (
@@ -47,6 +50,7 @@ SNAPSHOT_FILES = (
     "benchmarks/learn/engine_trace_k8s.py",
     "benchmarks/learn/engine_trace_queue.py",
     "benchmarks/learn/engine_trace_ax.py",
+    "benchmarks/ax/batch.py",
     "benchmarks/learn/engine_task_entry.py",
     "benchmarks/learn/engine_worker_service.py",
     "benchmarks/learn/engine_trace_parse.py",
@@ -791,8 +795,8 @@ def run_plan(plan, source, provenance, inputs, args):
         "active_lanes": {},
     }
     counts = Counter()
-    pending = deque(plan)
-    fallbacks = deque()
+    pending = list(plan)
+    fallbacks = []
     if args.executor in {"ax", "mixed-ax"}:
         pending.clear()
         for pin in plan:
@@ -803,12 +807,6 @@ def run_plan(plan, source, provenance, inputs, args):
                 fallbacks.append((pin, old))
             else:
                 pending.append(pin)
-    active = {}
-    failures = False
-
-    def persist():
-        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, attempt=None):
         if not any(u["supported"] for u in pin["units"]):
@@ -825,52 +823,20 @@ def run_plan(plan, source, provenance, inputs, args):
             options.image = getattr(args, "docker_image", args.image)
         return run_pin(pin, source, provenance, inputs, args.out, options)
 
-    with ThreadPoolExecutor(max_workers=len(lanes)) as pool:
+    def complete(pin, lane, failed, progress):
+        result = json.loads((args.out / pin_name(pin)).read_text())
+        selected = {u["unit"] for u in pin["units"]}
+        reasons = [u.get("reason", "") for u in result["units"] if u["unit"] in selected]
+        sandbox_failed = args.executor == "mixed-ax" and result.get("container_exit") == 1 and "AX sandbox failed twice" in reasons
+        if lane == "ax" and (sandbox_failed or any("JVM out of memory" in reason for reason in reasons)):
+            result.update(fallback="vm-sandbox" if sandbox_failed else "vm-oom", fallback_pending=True)
+            write_json(args.out / pin_name(pin), result)
+            return False, result
+        counts.update(u["status"] for u in result["units"] if u["unit"] in selected)
+        progress.update(last_pin={"repo": pin["repo"], "pin": pin["pin"]}, statuses=dict(counts))
+        return failed, None
 
-        def launch(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)
-        persist()
-        while active:
-            completed, _ = wait(active, return_when=FIRST_COMPLETED)
-            for future in completed:
-                pin, lane = active.pop(future)
-                failed = future.result()
-                result = json.loads((args.out / pin_name(pin)).read_text())
-                selected = {u["unit"] for u in pin["units"]}
-                reasons = [u.get("reason", "") for u in result["units"] if u["unit"] in selected]
-                sandbox_failed = args.executor == "mixed-ax" and result.get("container_exit") == 1 and "AX sandbox failed twice" in reasons
-                if lane == "ax" and (sandbox_failed or any("JVM out of memory" in reason for reason in reasons)):
-                    result.update(fallback="vm-sandbox" if sandbox_failed else "vm-oom", fallback_pending=True)
-                    write_json(args.out / pin_name(pin), result)
-                    fallbacks.append((pin, result))
-                    continue
-                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)
-                )
-            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
+    return schedule(pending, lanes, args.out, execute, complete, fallbacks=fallbacks, overflow=args.executor != "ax", progress=progress)
 
 
 def main():

benchmarks/learn/engine_trace_ax.py

@@ -1,184 +1,50 @@
-"""Host-only AX Task transport; credentials never enter task manifests."""
+"""Engine workload adapter for the shared AX transport."""
 
 from __future__ import annotations
 
-import hashlib
 import io
-import ipaddress
 import json
-import os
-import queue
-import re
 import shutil
 import subprocess
+import sys
 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
+sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
 from engine_trace_k8s import pack
 
-from openultrasast.config import load_dotenv
-from openultrasast.plane.memory import PRESIGNED_ARCHIVE_SUFFIX, S3Store, open_store
-
-
-class EgressError(ValueError):
-    """Safe diagnostic without presigned URL contents."""
-
-
-class SandboxFailure(Exception):
-    """A terminal AX task without an authoritative result object."""
-
-
-def public_endpoint(url):
-    """Reject endpoints the actor's hostname-based HTTPS egress policy cannot reach."""
-    parts = urlsplit(url)
-    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()
+from benchmarks.ax.batch import AXLane, Item, SandboxFailure, Workload, manifest
+from benchmarks.ax.batch import task_state as task_state
+from openultrasast.plane.memory import open_store as open_store
 
 
 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()
+    return manifest(
+        name, engine_workload(args, lambda *_: None), engine_item({}, "pin", b"", args), urls, args.atespace, validate_name=False
     )
-    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
-    listed = isinstance(doc, (dict, list)) or bool(lines and "name" in headers and {"phase", "status"}.intersection(headers))
-    return "Missing" if listed else "", None
+def engine_workload(args, validate):
+    return Workload(args.image, frozenset({"SOURCE_URL", "QUESTIONS_URL", "RESULT_URL"}), args.deadline, validate)
 
 
-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
+def engine_item(pin, pin_id, source, args):
+    return Item(
+        pin_id,
+        {"SOURCE_URL": source, "QUESTIONS_URL": json.dumps(pin).encode()},
+        "result.dat",
+        ("python3", "/app/benchmarks/learn/engine_task_entry.py"),
+        {"DEADLINE": str(args.deadline), "QUESTION_DEADLINE": str(args.question_deadline)},
+        "",
+        {"SOURCE_URL": "source.dat", "QUESTIONS_URL": "questions.json"},
+    )
 
 
-class AX:
+class AX(AXLane):
     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
+        super().__init__(args, engine_workload(args, lambda *_: None), store, pause, clock)
 
     def receive(self, data, output, pin):
         # Validate in isolation: malformed results must never become usable checkpoints.
@@ -220,11 +86,15 @@ class AX:
             shutil.copytree(root, output, dirs_exist_ok=True)
 
     def execute(self, checkout, output, pin, pin_id):
+        import copy
+
+        self = copy.copy(self)
+        self.workload = engine_workload(self.args, lambda data, target: self.receive(data, target, pin))
         attempts = []
         try:
             for attempt in range(2):
                 try:
-                    return self._execute_attempt(checkout, output, pin, pin_id, attempt, attempts)
+                    return self._execute_attempt(engine_item(pin, pin_id, pack(checkout), self.args), output, attempt, attempts)
                 except SandboxFailure:
                     if attempt == 0:
                         self.pause(10)
@@ -236,84 +106,3 @@ class AX:
                 record = json.loads(result_path.read_text())
                 record["ax_attempts"] = attempts
                 result_path.write_text(json.dumps(record))
-
-    def _execute_attempt(self, checkout, output, pin, pin_id, attempt, attempts):
-        digest = hashlib.sha256(f"{self.run}/{pin_id}/attempt-{attempt}".encode()).hexdigest()[:12]
-        slug = re.sub(r"[^a-z0-9]+", "-", str(pin_id).lower()).strip("-")[:22]
-        name = f"ousast-engine-{slug or 'pin'}-{digest}"
-        key_id = hashlib.sha256(str(pin_id).encode()).hexdigest()
-        prefix = f"engine-queue/{self.run}/{key_id}/attempt-{attempt}/"
-        keys = [prefix + suffix for suffix in (f"source{PRESIGNED_ARCHIVE_SUFFIX}", "questions.json", f"result{PRESIGNED_ARCHIVE_SUFFIX}")]
-        expires = timedelta(seconds=self.args.deadline + 600)
-        worker_ip = None
-        submitted = False
-        started = self.clock()
-        phase = ""
-        try:
-            self.store._put(keys[0], pack(checkout))
-            self.store._put(keys[1], json.dumps(pin).encode())
-            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() in {"failed", "missing"}:
-                    # The worker may have uploaded while we were reading AX status.
-                    result = bounded_read(lambda: self.store._get(keys[2]), end - self.clock())
-                    if self.clock() >= end:
-                        return None, worker_ip
-                    if result is not None:
-                        self.receive(result[0], output, pin)
-                        return subprocess.CompletedProcess([], 0, "", ""), worker_ip
-                    raise SandboxFailure
-                self.pause(min(2, max(0, end - self.clock())))
-            return None, worker_ip
-        except (subprocess.TimeoutExpired, TimeoutError):
-            return None, worker_ip
-        except SandboxFailure:
-            raise
-        except Exception as exc:
-            # Never expose exception strings which may contain object-scoped bearer URLs.
-            detail = str(exc) if isinstance(exc, EgressError) else type(exc).__name__
-            return subprocess.CompletedProcess([], 1, "", "AX transport failed: " + detail), worker_ip
-        finally:
-            attempts.append({"name": name, "phase": phase, "seconds": self.clock() - started})
-            errors = []
-            if submitted:
-                try:
-                    self.ax("delete", "task", name)
-                except Exception as exc:
-                    # A successful list already confirmed a disappeared task is gone.
-                    if phase != "Missing":
-                        errors.append("task: " + type(exc).__name__)
-            for key in keys:
-                try:
-                    self.store._delete(key)
-                except Exception as exc:
-                    errors.append("object: " + type(exc).__name__)
-            if errors:
-                raise RuntimeError("AX cleanup failed; stop dispatching: " + ", ".join(errors))