mirror of
https://github.com/norandom/OpenUltraSAST.git
synced 2026-10-16 03:49:26 +00:00
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]>
This commit is contained in:
Marius Ciepluch
co-authored by
Claude Opus 5.5
Claude Opus 5.5
parent
afff080874
commit
abca40b072
@@ -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"
|
||||
|
||||
@@ -363,3 +363,81 @@ own terms:
|
||||
- Joern (Apache-2.0) is downloaded at image build time.
|
||||
- The tree-sitter wheels of the `semantic` extra are MIT.
|
||||
- The published benchmark cases were licence-checked one by one. v1's OpenCVE case is BUSL-1.1, see [RELEASE_NOTES.md](RELEASE_NOTES.md).
|
||||
|
||||
## Engine traces on kube-ax
|
||||
|
||||
`benchmarks/learn/engine_trace.py --executor ax` submits AX v0.3.1 Tasks with
|
||||
`ax apply -f -` and removes them with `ax delete task NAME`. The operator supplies
|
||||
kubeconfig and AX access; the dispatcher creates no Kubernetes resources or secrets.
|
||||
The `ousast-engine-task` GHCR package **must be made Public once** in its package
|
||||
settings so the actors can pull it without registry credentials.
|
||||
|
||||
The task image contains the current engine and worker code. Rebuild it after code
|
||||
changes and pass its digest with `--image`; no analyzer download is needed. The host
|
||||
materializes each pin as usual and uploads source and prepared questions under
|
||||
`engine-queue/<run>/<pin>/`. Only object-scoped presigned URLs and deadline parameters
|
||||
enter Task env. Credentials stay on the coordinator. Configure `OUSAST_MEMORY=s3://…`
|
||||
and the store credentials in the coordinator's `.env` or environment, with
|
||||
`S3_ENDPOINT=https://files.because-security.com`. Presigned links expire after the
|
||||
pin deadline plus ten minutes. IP and HTTP endpoints are rejected: the operator's
|
||||
actor EgressPolicy must allow **files.because-security.com:443**.
|
||||
|
||||
Completion is the result object, even while AX reports Running. Failed phases and
|
||||
host deadlines produce explicit failure/timeout rows. Task and staging-object cleanup
|
||||
runs on success and failure; cleanup failure stops further dispatch. Per-pin files,
|
||||
`--rerun-status`, resume and `progress.json` retain the existing trace-runner behavior.
|
||||
AX results include `executor: "ax"`, `worker_ip`, and the image digest. Placement is
|
||||
random; `--parallel` defaults to two AX tasks and cannot exceed two.
|
||||
|
||||
AX workers use `--heap-profile cluster` with
|
||||
`-Xms1700m -Xmx1700m -XX:MaxMetaspaceSize=256m -XX:ReservedCodeCacheSize=128m -XX:MaxDirectMemorySize=256m -Xss512k -XX:ActiveProcessorCount=1 -XX:+UseG1GC`
|
||||
on every frontend/query JVM; instruments record the flags. A `JVM out of memory`
|
||||
result is automatically re-queued to one local Docker lane with `--heap-profile vm`
|
||||
(2560m), recording `fallback: "vm-oom"` and the AX attempt. `mixed-ax` also uses this
|
||||
one Docker lane for overflow, prioritizing pending OOM fallbacks. Thus a mixed run has
|
||||
up to two AX tasks plus one Docker container. Use `--docker-image` for the locally
|
||||
available VM image; Docker never pulls it implicitly. A STOP file prevents new
|
||||
launches; interrupted fallback queues are recovered on resume.
|
||||
|
||||
Coordinator commands, from this worktree (the AX binary and kubeconfig remain
|
||||
configurable). Build/push the actual uncommitted tree and read its digest:
|
||||
|
||||
```bash
|
||||
cd /home/mc/Source/OpenUltraSAST/.claude/worktrees/engine-ax
|
||||
export PATH=/home/mc/Source/OpenUltraSAST/.venv/bin:$PATH
|
||||
export PYTHONPATH="$PWD/src"
|
||||
TASK_TAG="ghcr.io/norandom/ousast-engine-task:sha-$(git rev-parse --short=7 HEAD)"
|
||||
docker buildx build --platform linux/amd64 -f plane/Dockerfile.engine-task \
|
||||
-t "$TASK_TAG" -t ghcr.io/norandom/ousast-engine-task:main \
|
||||
--push --metadata-file /tmp/ousast-engine-task-build.json .
|
||||
export TASK_IMAGE="ghcr.io/norandom/ousast-engine-task@$(python -c 'import json; print(json.load(open("/tmp/ousast-engine-task-build.json"))["containerimage.digest"])')"
|
||||
printf '%s\n' "$TASK_IMAGE"
|
||||
docker build -f Dockerfile -t openultrasast:engine-ax .
|
||||
```
|
||||
|
||||
The Engine image workflow also builds/pushes both images and prints each digest;
|
||||
CI needs the changes published before dispatch. The local build above includes
|
||||
uncommitted changes without publishing a git commit.
|
||||
|
||||
With `KUBECONFIG` pointing to the operator's kube-ax config and store settings in
|
||||
`.env`, run a one-pin smoke using the default experiment units and local example
|
||||
memory (or add `--examples /path/to/frozen-examples.jsonl`):
|
||||
|
||||
```bash
|
||||
: "${KUBECONFIG:?Set KUBECONFIG to the operator kube-ax configuration}"
|
||||
python benchmarks/learn/engine_trace.py --executor ax \
|
||||
--ax-bin "$HOME/.local/ax-v0.3.1/ax" --kubeconfig "$KUBECONFIG" \
|
||||
--atespace default --image "$TASK_IMAGE" --docker-image openultrasast:engine-ax \
|
||||
--parallel 1 --limit 1 --deadline 900 \
|
||||
--out "$HOME/ousast-results/plane/engine-ax-smoke"
|
||||
```
|
||||
|
||||
Then run the mixed lanes (repeat to resume; defaults re-run failed/timeout pins):
|
||||
|
||||
```bash
|
||||
python benchmarks/learn/engine_trace.py --executor mixed-ax \
|
||||
--ax-bin "$HOME/.local/ax-v0.3.1/ax" --kubeconfig "$KUBECONFIG" \
|
||||
--atespace default --image "$TASK_IMAGE" --docker-image openultrasast:engine-ax \
|
||||
--parallel 2 --deadline 900 --rerun-status failed,timeout \
|
||||
--out "$HOME/ousast-results/plane/engine-ax-mixed"
|
||||
```
|
||||
|
||||
@@ -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())
|
||||
@@ -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)
|
||||
|
||||
@@ -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))
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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.
|
||||
@@ -0,0 +1,277 @@
|
||||
"""AX transport controls using an executable shim and an in-memory object store."""
|
||||
|
||||
import io
|
||||
import json
|
||||
import sys
|
||||
import tarfile
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
with pytest.MonkeyPatch.context() as patch:
|
||||
patch.syspath_prepend(str(Path(__file__).resolve().parents[1] / "benchmarks/learn"))
|
||||
import engine_trace_ax as ax
|
||||
|
||||
|
||||
IMAGE = "ghcr.io/norandom/ousast-engine-task@sha256:" + "a" * 64
|
||||
PIN = {"repo": "owner/repo", "pin": "123", "units": [{"unit": "a", "supported": True}], "questions": {"a": {}}}
|
||||
|
||||
|
||||
def archive(record):
|
||||
buffer = io.BytesIO()
|
||||
with tarfile.open(fileobj=buffer, mode="w") as tar:
|
||||
data = json.dumps(record).encode()
|
||||
info = tarfile.TarInfo("result.json")
|
||||
info.size = len(data)
|
||||
tar.addfile(info, io.BytesIO(data))
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
class Store:
|
||||
def __init__(self, result=None):
|
||||
self.result = result
|
||||
self.objects = {}
|
||||
self.deleted = []
|
||||
self.expiries = []
|
||||
|
||||
def _put(self, key, data):
|
||||
self.objects[key] = data
|
||||
|
||||
def _get(self, key):
|
||||
return (self.result, None) if self.result else None
|
||||
|
||||
def _delete(self, key):
|
||||
self.deleted.append(key)
|
||||
self.objects.pop(key, None)
|
||||
|
||||
def presign_get(self, key, expiry):
|
||||
self.expiries.append(expiry.total_seconds())
|
||||
return "https://files.because-security.com/bucket/" + key + "?signature=secret"
|
||||
|
||||
presign_put = presign_get
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def setup(tmp_path, monkeypatch):
|
||||
log = tmp_path / "calls.jsonl"
|
||||
shim = tmp_path / "fake-ax"
|
||||
shim.write_text("#!/bin/sh\nexec " + sys.executable + " " + str(tmp_path / "shim.py") + ' "$@"\n')
|
||||
shim.chmod(0o755)
|
||||
(tmp_path / "shim.py").write_text("""import json, os, sys
|
||||
from pathlib import Path
|
||||
args = sys.argv[1:]
|
||||
with Path(os.environ["AX_LOG"]).open("a") as log:
|
||||
log.write(json.dumps({"args": args, "input": sys.stdin.read() if args[0] == "apply" else "",
|
||||
"kubeconfig": os.environ.get("KUBECONFIG")}) + "\\n")
|
||||
if args[0] == "get" and os.environ.get("AX_GET_FAIL"):
|
||||
sys.exit(1)
|
||||
if args[0] == "get":
|
||||
manifest = json.loads(json.loads(Path(os.environ["AX_LOG"]).read_text().splitlines()[0])["input"])
|
||||
print(json.dumps({"items": [{"metadata": manifest["metadata"], "status": {
|
||||
"phase": os.environ.get("AX_PHASE", "Running"), "workerIP": "10.0.0.2"}}]}))
|
||||
if args[0] == "delete" and os.environ.get("AX_DELETE_FAIL"):
|
||||
sys.exit(1)
|
||||
""")
|
||||
monkeypatch.setenv("AX_LOG", str(log))
|
||||
monkeypatch.setenv("S3_ENDPOINT", "https://files.because-security.com")
|
||||
args = SimpleNamespace(
|
||||
ax_bin=str(shim), kubeconfig="/operator/kubeconfig", atespace="default", image=IMAGE, deadline=4, question_deadline=2
|
||||
)
|
||||
source = tmp_path / "case"
|
||||
source.mkdir()
|
||||
(source / "app.py").write_text("print(42)\n")
|
||||
output = tmp_path / "out"
|
||||
output.mkdir()
|
||||
now = [0]
|
||||
|
||||
def pause(seconds):
|
||||
now[0] += seconds
|
||||
|
||||
return args, source, output, log, pause, lambda: now[0]
|
||||
|
||||
|
||||
def test_ax_executor_success_while_running_and_cleanup(setup):
|
||||
args, source, output, log, pause, clock = setup
|
||||
store = Store(archive({**PIN, "done": True}))
|
||||
executor = ax.AX(args, store, pause=pause, clock=clock)
|
||||
done, ip = executor.execute(source, output, PIN, "owner/repo/123")
|
||||
assert done.returncode == 0
|
||||
assert ip == "10.0.0.2"
|
||||
assert json.loads((output / "result.json").read_text())["done"]
|
||||
calls = [json.loads(line) for line in log.read_text().splitlines()]
|
||||
assert calls[0]["args"] == ["apply", "-f", "-"]
|
||||
assert calls[-1]["args"][:2] == ["delete", "task"]
|
||||
assert calls[-1]["args"][-2:] == ["--atespace", args.atespace]
|
||||
assert all(c["kubeconfig"] == args.kubeconfig for c in calls)
|
||||
manifest = json.loads(calls[0]["input"])
|
||||
assert manifest["apiVersion"] == "ax.io/v1alpha1"
|
||||
assert manifest["kind"] == "Task"
|
||||
assert len(manifest["metadata"]["name"]) <= 49
|
||||
assert manifest["spec"]["image"] == IMAGE
|
||||
env = {e["name"]: e["value"] for e in manifest["spec"]["env"]}
|
||||
assert set(env) == {"SOURCE_URL", "QUESTIONS_URL", "RESULT_URL", "DEADLINE", "QUESTION_DEADLINE"}
|
||||
assert len(json.dumps(manifest["spec"]["env"]).encode()) < 32768
|
||||
assert "AWS" not in calls[0]["input"]
|
||||
assert store.expiries == [604] * 3
|
||||
assert len(store.deleted) == 3 and not store.objects
|
||||
|
||||
|
||||
def test_ax_executor_timeout_deletes_task(setup):
|
||||
args, source, output, log, pause, clock = setup
|
||||
store = Store()
|
||||
done, _ = ax.AX(args, store, pause=pause, clock=clock).execute(source, output, PIN, "pin")
|
||||
assert done is None
|
||||
assert json.loads(log.read_text().splitlines()[-1])["args"][:2] == ["delete", "task"]
|
||||
assert len(store.deleted) == 3
|
||||
|
||||
|
||||
def test_ax_executor_failed_phase_deletes_task(setup, monkeypatch):
|
||||
args, source, output, log, pause, clock = setup
|
||||
monkeypatch.setenv("AX_PHASE", "Failed")
|
||||
done, _ = ax.AX(args, Store(), pause=pause, clock=clock).execute(source, output, PIN, "pin")
|
||||
assert done.returncode == 1 and "Failed" in done.stderr
|
||||
assert json.loads(log.read_text().splitlines()[-1])["args"][:2] == ["delete", "task"]
|
||||
|
||||
|
||||
def test_ax_executor_delete_failure_is_fatal_and_still_cleans_objects(setup, monkeypatch):
|
||||
args, source, output, _, pause, clock = setup
|
||||
monkeypatch.setenv("AX_DELETE_FAIL", "1")
|
||||
store = Store(archive({**PIN, "done": True}))
|
||||
with pytest.raises(RuntimeError, match="cleanup"):
|
||||
ax.AX(args, store, pause=pause, clock=clock).execute(source, output, PIN, "pin")
|
||||
assert len(store.deleted) == 3
|
||||
|
||||
|
||||
@pytest.mark.parametrize("endpoint", ["http://files.because-security.com", "https://127.0.0.1", "https://[::1]", "https://localhost"])
|
||||
def test_ax_executor_endpoint_rejected_before_store_open(setup, monkeypatch, endpoint):
|
||||
args = setup[0]
|
||||
monkeypatch.setenv("S3_ENDPOINT", endpoint)
|
||||
monkeypatch.setattr(ax, "open_store", lambda: pytest.fail("store opened before validation"))
|
||||
with pytest.raises(ValueError, match="egress policy"):
|
||||
ax.AX(args)
|
||||
|
||||
|
||||
def test_ax_executor_rejects_tag_and_large_env(setup):
|
||||
args = setup[0]
|
||||
urls = {key: "https://files.because-security.com/" for key in ("SOURCE_URL", "QUESTIONS_URL", "RESULT_URL")}
|
||||
args.image = "image:latest"
|
||||
with pytest.raises(ValueError, match="digest"):
|
||||
ax.task_manifest("name", args, urls)
|
||||
args.image = IMAGE
|
||||
urls["SOURCE_URL"] += "x" * 32768
|
||||
with pytest.raises(ValueError, match="32 KB"):
|
||||
ax.task_manifest("name", args, urls)
|
||||
|
||||
|
||||
def test_ax_executor_rejects_mismatched_result(setup):
|
||||
args, source, output, _, pause, clock = setup
|
||||
store = Store(archive({"done": True, "units": []}))
|
||||
done, _ = ax.AX(args, store, pause=pause, clock=clock).execute(source, output, PIN, "pin")
|
||||
assert done.returncode == 1
|
||||
assert not (output / "result.json").exists()
|
||||
|
||||
|
||||
def test_ax_executor_table_worker_ip():
|
||||
assert ax.task_state("NAME ATESPACE PHASE WORKER-IP\nours default Running 10.2.3.4\n", "ours") == ("Running", "10.2.3.4")
|
||||
|
||||
|
||||
def test_ax_executor_entry_download_failure_restores_host_units(setup):
|
||||
args, source, output, _, pause, clock = setup
|
||||
store = Store(archive({"done": True, "status": "failed", "reason": "questions download failed", "units": []}))
|
||||
done, _ = ax.AX(args, store, pause=pause, clock=clock).execute(source, output, PIN, "pin")
|
||||
assert done.returncode == 0
|
||||
unit = json.loads((output / "result.json").read_text())["units"][0]
|
||||
assert unit["unit"] == "a" and unit["status"] == "failed"
|
||||
assert unit["reason"] == "questions download failed"
|
||||
|
||||
|
||||
def test_ax_executor_presign_host_must_match_egress_policy(setup):
|
||||
args, source, output, log, pause, clock = setup
|
||||
store = Store()
|
||||
store.presign_put = lambda key, expiry: "https://another-store.example/result"
|
||||
done, _ = ax.AX(args, store, pause=pause, clock=clock).execute(source, output, PIN, "pin")
|
||||
assert done.returncode == 1 and "egress policy" in done.stderr
|
||||
assert not log.exists()
|
||||
assert len(store.deleted) == 3
|
||||
|
||||
|
||||
def test_ax_executor_get_failure_does_not_lose_result(setup, monkeypatch):
|
||||
args, source, output, _, pause, clock = setup
|
||||
monkeypatch.setenv("AX_GET_FAIL", "1")
|
||||
done, _ = ax.AX(args, Store(archive({**PIN, "done": True})), pause=pause, clock=clock).execute(source, output, PIN, "pin")
|
||||
assert done.returncode == 0
|
||||
|
||||
|
||||
def test_ax_executor_absent_worker_ip():
|
||||
assert ax.task_state("NAME ATESPACE PHASE ACTOR WORKER-IP AGE\nours other Pending actor <none> 1s\n", "ours") == ("Pending", None)
|
||||
|
||||
|
||||
def test_ax_executor_late_store_result_times_out(setup):
|
||||
args, source, output, log, pause, clock = setup
|
||||
store = Store(archive({**PIN, "done": True}))
|
||||
original_get = store._get
|
||||
|
||||
def late_get(key):
|
||||
pause(args.deadline)
|
||||
return original_get(key)
|
||||
|
||||
store._get = late_get
|
||||
done, _ = ax.AX(args, store, pause=pause, clock=clock).execute(source, output, PIN, "pin")
|
||||
assert done is None
|
||||
assert not (output / "result.json").exists()
|
||||
assert json.loads(log.read_text().splitlines()[-1])["args"][:2] == ["delete", "task"]
|
||||
|
||||
|
||||
def test_ax_executor_late_status_times_out(setup):
|
||||
args, source, output, _, pause, clock = setup
|
||||
executor = ax.AX(args, Store(archive({**PIN, "done": True})), pause=pause, clock=clock)
|
||||
original_ax = executor.ax
|
||||
|
||||
def late_ax(*command, **kwargs):
|
||||
result = original_ax(*command, **kwargs)
|
||||
if command[0] == "get":
|
||||
pause(args.deadline)
|
||||
return result
|
||||
|
||||
executor.ax = late_ax
|
||||
done, _ = executor.execute(source, output, PIN, "pin")
|
||||
assert done is None
|
||||
assert not (output / "result.json").exists()
|
||||
|
||||
|
||||
def test_ax_executor_blocking_store_read_cannot_delay_task_deletion(setup):
|
||||
import threading
|
||||
import time
|
||||
|
||||
args, source, output, log, _, _ = setup
|
||||
args.deadline = 0.15
|
||||
release = threading.Event()
|
||||
store = Store()
|
||||
store._get = lambda key: release.wait(5)
|
||||
started = time.monotonic()
|
||||
try:
|
||||
done, _ = ax.AX(args, store).execute(source, output, PIN, "pin")
|
||||
assert done is None
|
||||
assert time.monotonic() - started < 2
|
||||
assert json.loads(log.read_text().splitlines()[-1])["args"][:2] == ["delete", "task"]
|
||||
finally:
|
||||
release.set()
|
||||
|
||||
|
||||
def test_ax_task_names_reserve_watcher_prefix_and_fit_actor_limit(setup):
|
||||
args, source, output, log, pause, clock = setup
|
||||
executor = ax.AX(args, Store(archive({**PIN, "done": True})), pause=pause, clock=clock)
|
||||
# Sanitization and truncation must not collapse distinct pins.
|
||||
for pin_id in ["A/B", "a-b", "x" * 100 + "1", "x" * 100 + "2", "☃"]:
|
||||
done, _ = executor.execute(source, output, PIN, pin_id)
|
||||
assert done.returncode == 0
|
||||
manifests = [json.loads(call["input"]) for line in log.read_text().splitlines() if (call := json.loads(line))["args"][0] == "apply"]
|
||||
names = [manifest["metadata"]["name"] for manifest in manifests]
|
||||
assert len(set(names)) == 5
|
||||
for manifest, name in zip(manifests, names, strict=True):
|
||||
assert manifest["kind"] == "Task"
|
||||
assert manifest["metadata"]["atespace"] == "default"
|
||||
assert name.startswith("ousast-engine-")
|
||||
assert len(name.encode()) <= 49
|
||||
assert len((name + "-tmpl-12345678").encode()) <= 63
|
||||
@@ -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))
|
||||
@@ -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}
|
||||
|
||||
Reference in New Issue
Block a user