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]>

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