From abca40b072eed6fef72cdfcae6beb780f95b3565 Mon Sep 17 00:00:00 2001 From: Marius Ciepluch <11855163+norandom@users.noreply.github.com> Date: Sun, 4 Oct 2026 15:22:44 +0200 Subject: [PATCH] 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 --- .github/workflows/engine-image.yml | 22 +- README.md | 78 +++++++ benchmarks/learn/engine_task_entry.py | 154 +++++++++++++ benchmarks/learn/engine_trace.py | 91 ++++++-- benchmarks/learn/engine_trace_ax.py | 282 ++++++++++++++++++++++++ benchmarks/learn/engine_trace_worker.py | 48 +++- plane/Dockerfile.engine-task | 52 +++++ tests/test_ax_executor.py | 277 +++++++++++++++++++++++ tests/test_engine_task_entry.py | 224 +++++++++++++++++++ tests/test_engine_trace.py | 144 ++++++++++++ 10 files changed, 1347 insertions(+), 25 deletions(-) create mode 100644 benchmarks/learn/engine_task_entry.py create mode 100644 benchmarks/learn/engine_trace_ax.py create mode 100644 plane/Dockerfile.engine-task create mode 100644 tests/test_ax_executor.py create mode 100644 tests/test_engine_task_entry.py diff --git a/.github/workflows/engine-image.yml b/.github/workflows/engine-image.yml index 0b9ad37..dca961c 100644 --- a/.github/workflows/engine-image.yml +++ b/.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" diff --git a/README.md b/README.md index dceb884..1ab3832 100644 --- a/README.md +++ b/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///`. 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" +``` diff --git a/benchmarks/learn/engine_task_entry.py b/benchmarks/learn/engine_task_entry.py new file mode 100644 index 0000000..965de87 --- /dev/null +++ b/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()) diff --git a/benchmarks/learn/engine_trace.py b/benchmarks/learn/engine_trace.py index 5d30cfd..0af7c05 100644 --- a/benchmarks/learn/engine_trace.py +++ b/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) diff --git a/benchmarks/learn/engine_trace_ax.py b/benchmarks/learn/engine_trace_ax.py new file mode 100644 index 0000000..3a3e539 --- /dev/null +++ b/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 {"", "-", ""} 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)) diff --git a/benchmarks/learn/engine_trace_worker.py b/benchmarks/learn/engine_trace_worker.py index ca8e9ed..6304a86 100644 --- a/benchmarks/learn/engine_trace_worker.py +++ b/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 diff --git a/plane/Dockerfile.engine-task b/plane/Dockerfile.engine-task new file mode 100644 index 0000000..5c059ef --- /dev/null +++ b/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. diff --git a/tests/test_ax_executor.py b/tests/test_ax_executor.py new file mode 100644 index 0000000..9ae4030 --- /dev/null +++ b/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 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 diff --git a/tests/test_engine_task_entry.py b/tests/test_engine_task_entry.py new file mode 100644 index 0000000..58c00fc --- /dev/null +++ b/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)) diff --git a/tests/test_engine_trace.py b/tests/test_engine_trace.py index 914f117..80afc24 100644 --- a/tests/test_engine_trace.py +++ b/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}