learn: Kubernetes executor for the engine trace runner (Jobs, presigned source/result via the S3 store, mixed lanes, --k8s-check)

Implemented with Codex; live cluster check pending.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
This commit is contained in:
Marius Ciepluch
2026-10-04 10:35:59 +02:00
co-authored by Claude Opus 5.5
parent 431ab1e738
commit 2e6ba93aee
3 changed files with 616 additions and 40 deletions
+117 -39
View File
@@ -1,4 +1,4 @@
"""Resumable, serial labelled-function traces; host fetches pinned sources, Docker stays offline. """Resumable labelled-function traces with Docker and Kubernetes execution lanes.
Pair blob identities resolve through catalogs to real source commits. The analyzer Pair blob identities resolve through catalogs to real source commits. The analyzer
is git archive HEAD src plus a hashed snapshot of this uncommitted tool. is git archive HEAD src plus a hashed snapshot of this uncommitted tool.
@@ -8,6 +8,7 @@ from __future__ import annotations
import argparse import argparse
import ast import ast
import copy
import fcntl import fcntl
import hashlib import hashlib
import json import json
@@ -17,11 +18,14 @@ import shutil
import subprocess import subprocess
import tempfile import tempfile
import textwrap import textwrap
import threading
import time import time
import tomllib import tomllib
from collections import Counter, defaultdict from collections import Counter, defaultdict
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
from pathlib import Path from pathlib import Path
from engine_trace_k8s import IMAGE, Kubernetes
from engine_trace_worker import EXCLUDED_DIRS, safe_path, unit_record, write_json from engine_trace_worker import EXCLUDED_DIRS, safe_path, unit_record, write_json
from openultrasast.learn.examples import read_jsonl from openultrasast.learn.examples import read_jsonl
@@ -39,6 +43,7 @@ DEFAULT_OUT = Path.home() / "ousast-results/plane/engine-trace"
SNAPSHOT_FILES = ( SNAPSHOT_FILES = (
"benchmarks/learn/engine_trace.py", "benchmarks/learn/engine_trace.py",
"benchmarks/learn/engine_trace_worker.py", "benchmarks/learn/engine_trace_worker.py",
"benchmarks/learn/engine_trace_k8s.py",
"src/openultrasast/model/trace.py", "src/openultrasast/model/trace.py",
"src/openultrasast/model/taint.py", "src/openultrasast/model/taint.py",
"src/openultrasast/cpg/queries/taint.sc", "src/openultrasast/cpg/queries/taint.sc",
@@ -611,32 +616,42 @@ def prepare_questions(pin, root, question_deadline):
def run_pin(pin, source, provenance, inputs, out, args): def run_pin(pin, source, provenance, inputs, out, args):
executor = getattr(args, "executor", "docker")
node = None
with tempfile.TemporaryDirectory(prefix="trace-pin-", dir=out) as scratch: with tempfile.TemporaryDirectory(prefix="trace-pin-", dir=out) as scratch:
scratch = Path(scratch) scratch = Path(scratch)
checkout, output = scratch / "case", scratch / "output" checkout, output = scratch / "case", scratch / "output"
output.mkdir() output.mkdir()
try: try:
pin = {**pin, "materialization": inputs.materialize(pin, checkout)} with getattr(args, "input_lock", threading.Lock()):
pin = {**pin, "materialization": inputs.materialize(pin, checkout)}
write_json(output / "pin.json", {**pin, "questions": prepare_questions(pin, checkout, args.question_deadline)}) write_json(output / "pin.json", {**pin, "questions": prepare_questions(pin, checkout, args.question_deadline)})
name = "ousast-trace-" + pin_name(pin)[:16] name = "ousast-trace-" + pin_name(pin)[:16]
command = docker_command(source, checkout, output, name, args.image, args.deadline, args.question_deadline)
started = time.monotonic() started = time.monotonic()
done = None if executor == "k8s":
try: done, node = args.kubernetes.execute(source, checkout, output, json.loads((output / "pin.json").read_text()))
done = subprocess.run(command, capture_output=True, text=True, timeout=args.deadline) status, reason = "failed", done.stderr
status, reason = "failed", f"container exit {done.returncode}: {done.stderr[-1000:]}" if done.returncode:
(output / "container.log").write_text(done.stdout + "\n" + done.stderr) (output / "container.log").write_text(reason)
except subprocess.TimeoutExpired: else:
command = docker_command(source, checkout, output, name, args.image, args.deadline, args.question_deadline)
started = time.monotonic()
done = None done = None
status, reason = "timeout", f"container deadline {args.deadline}s"
finally:
# Never start the next pin unless this container is known to be gone.
try: try:
removed = subprocess.run(["docker", "rm", "-f", name], capture_output=True, timeout=20) done = subprocess.run(command, capture_output=True, text=True, timeout=args.deadline)
except subprocess.SubprocessError as exc: status, reason = "failed", f"container exit {done.returncode}: {done.stderr[-1000:]}"
raise RuntimeError(f"cannot confirm container {name} stopped; refusing another launch") from exc (output / "container.log").write_text(done.stdout + "\n" + done.stderr)
if done is None and removed.returncode != 0: except subprocess.TimeoutExpired:
raise RuntimeError(f"cannot remove timed-out container {name}; refusing another launch") done = None
status, reason = "timeout", f"container deadline {args.deadline}s"
finally:
# Never start the next pin unless this container is known to be gone.
try:
removed = subprocess.run(["docker", "rm", "-f", name], capture_output=True, timeout=20)
except subprocess.SubprocessError as exc:
raise RuntimeError(f"cannot confirm container {name} stopped; refusing another launch") from exc
if done is None and removed.returncode != 0:
raise RuntimeError(f"cannot remove timed-out container {name}; refusing another launch")
result_path = output / "result.json" result_path = output / "result.json"
if not result_path.exists(): if not result_path.exists():
reason = f"container produced no result: {reason}" reason = f"container produced no result: {reason}"
@@ -666,6 +681,8 @@ def run_pin(pin, source, provenance, inputs, out, args):
"analyzer": provenance, "analyzer": provenance,
"units": [unit_record(u, "failed" if u["supported"] else "unsupported", str(exc)) for u in pin["units"]], "units": [unit_record(u, "failed" if u["supported"] else "unsupported", str(exc)) for u in pin["units"]],
} }
if executor == "k8s":
record.update(executor="k8s", node=node, image=args.image)
save_pin(out / pin_name(pin), record) save_pin(out / pin_name(pin), record)
print( print(
json.dumps({"repo": pin["repo"], "pin": pin["pin"], "statuses": dict(Counter(u["status"] for u in record["units"]))}), json.dumps({"repo": pin["repo"], "pin": pin["pin"], "statuses": dict(Counter(u["status"] for u in record["units"]))}),
@@ -724,6 +741,67 @@ 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
args.input_lock = threading.Lock()
progress = {
"pins_done": 0,
"pins_total": len(plan),
"last_pin": None,
"statuses": {},
"lanes": dict(Counter(lanes)),
"pins_in_flight": 0,
"active_lanes": {},
}
counts = Counter()
pending = iter(plan)
active = {}
failures = False
def persist():
progress.update(pins_in_flight=len(active), active_lanes=dict(Counter(lane for _, lane in active.values())))
write_json(args.out / "progress.json", progress)
def execute(pin, lane):
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"]]}
)
return False
options = copy.copy(args)
options.executor = lane
if lane == "docker":
options.image = getattr(args, "docker_image", args.image)
return run_pin(pin, source, provenance, inputs, args.out, options)
with ThreadPoolExecutor(max_workers=len(lanes)) as pool:
def 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)
for lane in lanes:
launch(lane)
persist()
while active:
completed, _ = wait(active, return_when=FIRST_COMPLETED)
for future in completed:
pin, lane = active.pop(future)
failures |= future.result()
result = json.loads((args.out / pin_name(pin)).read_text())
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)
persist()
return failures
def main(): def main():
parser = argparse.ArgumentParser(description=__doc__) parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--units", type=Path, default=ROOT / "plane/experiments/exp-005-graph-slice.units.jsonl") parser.add_argument("--units", type=Path, default=ROOT / "plane/experiments/exp-005-graph-slice.units.jsonl")
@@ -731,7 +809,13 @@ def main():
parser.add_argument("--memory", type=Path, default=Path.home() / "ousast-results/plane/memory") parser.add_argument("--memory", type=Path, default=Path.home() / "ousast-results/plane/memory")
parser.add_argument("--cache", type=Path, default=Path.home() / ".cache/openultrasast") parser.add_argument("--cache", type=Path, default=Path.home() / ".cache/openultrasast")
parser.add_argument("--out", type=Path, default=DEFAULT_OUT) parser.add_argument("--out", type=Path, default=DEFAULT_OUT)
parser.add_argument("--image", default="openultrasast:dev") 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", "mixed"), default="docker")
parser.add_argument("--parallel", type=int, default=1, help="total pins in flight; mixed reserves one lane for Docker")
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")
parser.add_argument("--k8s-check", action="store_true")
parser.add_argument("--deadline", type=int, default=900) parser.add_argument("--deadline", type=int, default=900)
parser.add_argument("--question-deadline", type=int, default=120) parser.add_argument("--question-deadline", type=int, default=120)
parser.add_argument("--limit", type=int, default=0) parser.add_argument("--limit", type=int, default=0)
@@ -747,6 +831,13 @@ def main():
args = parser.parse_args() args = parser.parse_args()
if min(args.deadline, args.question_deadline) <= 0 or args.limit < 0: if min(args.deadline, args.question_deadline) <= 0 or args.limit < 0:
parser.error("deadlines must be positive; limit must be nonnegative") parser.error("deadlines must be positive; limit must be nonnegative")
if args.parallel < 1 or (args.executor == "mixed" and args.parallel < 2):
parser.error("--parallel must be positive; mixed needs at least 2 total lanes")
if (args.executor != "docker" or args.k8s_check) and not args.kube_context:
parser.error("--kube-context is required for Kubernetes")
args.image = args.image or (IMAGE if args.executor != "docker" or args.k8s_check else "openultrasast:dev")
if args.k8s_check:
return Kubernetes(args).check()
rerun_status = set(filter(None, args.rerun_status.split(","))) rerun_status = set(filter(None, args.rerun_status.split(",")))
if rerun_status - set(STATUSES): if rerun_status - set(STATUSES):
parser.error("unknown --rerun-status value") parser.error("unknown --rerun-status value")
@@ -806,34 +897,21 @@ def main():
except BlockingIOError: except BlockingIOError:
parser.error("another trace runner holds this results root") parser.error("another trace runner holds this results root")
failures = False failures = False
if any(u["supported"] for pin in plan for u in pin["units"]): 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( image = subprocess.run(
["docker", "image", "inspect", "--format", "{{.Id}}", args.image], capture_output=True, text=True, check=True ["docker", "image", "inspect", "--format", "{{.Id}}", args.image], capture_output=True, text=True, check=True
) )
args.image = image.stdout.strip() args.docker_image = image.stdout.strip()
if not args.image: if not args.docker_image:
raise ValueError("Docker returned no local image identity") raise ValueError("Docker returned no local image identity")
if args.executor != "docker" and any(u["supported"] for pin in plan for u in pin["units"]):
args.kubernetes = Kubernetes(args)
args.kubernetes.namespace()
with tempfile.TemporaryDirectory(prefix="trace-source-", dir=args.out) as scratch: with tempfile.TemporaryDirectory(prefix="trace-source-", dir=args.out) as scratch:
source = Path(scratch) / "frozen" source = Path(scratch) / "frozen"
provenance = freeze_source(source) provenance = freeze_source(source)
progress = {"pins_done": 0, "pins_total": len(plan), "last_pin": None, "statuses": {}} failures = run_plan(plan, source, provenance, inputs, args)
counts = Counter()
write_json(args.out / "progress.json", progress)
for pin in plan:
if (args.out / "STOP").exists():
break
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"]]},
)
else:
failures |= run_pin(pin, source, provenance, inputs, args.out, args)
result = json.loads((args.out / pin_name(pin)).read_text())
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))
write_json(args.out / "progress.json", progress)
return 2 if failures else 0 return 2 if failures else 0
+259
View File
@@ -0,0 +1,259 @@
"""Host Kubernetes transport; the embedded bootstrap uses only the standard library."""
from __future__ import annotations
import hashlib
import io
import json
import re
import subprocess
import tarfile
import time
import uuid
from datetime import timedelta
from openultrasast.plane.memory import S3Store, open_store
IMAGE = "ghcr.io/norandom/openultrasast:2.0.1"
BOOTSTRAP = r"""
import io, json, os, pathlib, subprocess, sys, tarfile, urllib.request
def get(name):
with urllib.request.urlopen(os.environ[name], timeout=120) as response:
return response.read()
def put(data):
request = urllib.request.Request(os.environ['RESULT_URL'], data=data, method='PUT')
with urllib.request.urlopen(request, timeout=120) as response:
response.read()
try:
if os.environ.get('CHECK') == '1':
data = get('SOURCE_URL')
assert len(data) == 1024
subprocess.run([sys.executable, '-c', 'print(42)'], check=True)
put(data)
else:
for name, target in [('SOURCE_URL', '/case'), ('ANALYZER_URL', '/frozen')]:
pathlib.Path(target).mkdir(parents=True, exist_ok=True)
with tarfile.open(fileobj=io.BytesIO(get(name))) as archive:
archive.extractall(target, filter='data')
out = pathlib.Path('/out')
out.mkdir(exist_ok=True)
(out / 'pin.json').write_text(os.environ['PIN'])
env = dict(os.environ, PYTHONPATH='/frozen/src')
for key in ['SOURCE_URL', 'ANALYZER_URL', 'RESULT_URL']:
env.pop(key, None)
with (out / 'container.log').open('w') as log:
done = subprocess.run([sys.executable, '/frozen/benchmarks/learn/engine_trace_worker.py',
'--worker', '/out/pin.json', '--deadline', os.environ['DEADLINE'],
'--question-deadline', os.environ['QUESTION_DEADLINE']], env=env, stdout=log, stderr=log)
(out / 'exit.json').write_text(json.dumps(done.returncode))
buffer = io.BytesIO()
with tarfile.open(fileobj=buffer, mode='w') as archive:
for path in out.iterdir():
archive.add(path, arcname=path.name)
put(buffer.getvalue())
sys.exit(done.returncode)
except Exception as exc:
# urllib exceptions can contain bearer URLs. Never print exception text.
print('engine bootstrap failed: ' + type(exc).__name__, file=sys.stderr)
sys.exit(1)
"""
def pack(path):
buffer = io.BytesIO()
with tarfile.open(fileobj=buffer, mode="w") as archive:
for child in sorted(path.iterdir()):
archive.add(child, arcname=child.name)
return buffer.getvalue()
def manifests(name, run, args, urls, pin, check=False):
labels = {"engine-trace-run": run, "engine-trace-job": name, "engine-trace": "true"}
secret = {
"apiVersion": "v1",
"kind": "Secret",
"metadata": {"name": name, "namespace": args.kube_namespace, "labels": labels},
"type": "Opaque",
"stringData": urls,
}
env = [{"name": key, "valueFrom": {"secretKeyRef": {"name": name, "key": key}}} for key in urls]
env += [
{"name": key, "value": value}
for key, value in {
"PIN": json.dumps(pin),
"DEADLINE": str(args.deadline),
"QUESTION_DEADLINE": str(args.question_deadline),
"CHECK": "1" if check else "0",
}.items()
]
job = {
"apiVersion": "batch/v1",
"kind": "Job",
"metadata": {"name": name, "namespace": args.kube_namespace, "labels": labels},
"spec": {
"activeDeadlineSeconds": args.deadline,
"backoffLimit": 0,
"ttlSecondsAfterFinished": 86400,
"template": {
"metadata": {"labels": labels},
"spec": {
"restartPolicy": "Never",
"automountServiceAccountToken": False,
"securityContext": {"runAsUser": 1000, "runAsGroup": 1000, "fsGroup": 1000},
"volumes": [{"name": name, "emptyDir": {}} for name in ("case", "frozen", "out")],
"affinity": {
"podAntiAffinity": {
"requiredDuringSchedulingIgnoredDuringExecution": [
{
"labelSelector": {"matchLabels": {"engine-trace": "true"}},
"topologyKey": "kubernetes.io/hostname",
}
]
}
},
"containers": [
{
"name": "engine",
"image": args.image,
"command": ["python", "-c", BOOTSTRAP],
"env": env,
"volumeMounts": [{"name": name, "mountPath": "/" + name} for name in ("case", "frozen", "out")],
"resources": {"requests": {"memory": "3Gi", "cpu": "1"}, "limits": {"memory": "3Gi"}},
}
],
},
},
},
}
return job, secret
def redact(text):
return re.sub(r"https?://[^\s\"'<>]+", "[redacted URL]", text)
class Kubernetes:
def __init__(self, args, store=None):
self.args = args
self.run = uuid.uuid4().hex
self.store = store if store is not None else open_store()
if not isinstance(self.store, S3Store):
raise ValueError("Kubernetes requires OUSAST_MEMORY=s3://bucket[/prefix]")
def kubectl(self, *args, manifest=None, check=True):
result = subprocess.run(
["kubectl", "--context", self.args.kube_context, "--namespace", self.args.kube_namespace, *args],
input=json.dumps(manifest) if manifest is not None else None,
capture_output=True,
text=True,
timeout=30,
)
if check and result.returncode:
# kubectl can echo submitted Secret data on errors; never expose stderr.
raise ValueError(f"kubectl {args[0]} failed (exit {result.returncode})")
return result
def namespace(self):
self.kubectl("get", "nodes", "-o", "json")
self.kubectl("apply", "-f", "-", manifest={"apiVersion": "v1", "kind": "Namespace", "metadata": {"name": self.args.kube_namespace}})
def execute(self, source, checkout, output, pin, check=False):
name = "engine-" + uuid.uuid4().hex[:24]
pin_key = hashlib.sha256(json.dumps([pin.get("repo"), pin.get("pin")]).encode()).hexdigest()
prefix = f"engine-jobs/{self.run}/{pin_key}/"
keys = [prefix + file for file in ("source.tar", "analyzer.tar", "result.tar")]
expiry = timedelta(seconds=self.args.deadline + 900)
node = None
try:
self.store._put(keys[0], b"x" * 1024 if check else pack(checkout))
self.store._put(keys[1], b"" if check else pack(source))
urls = {
"SOURCE_URL": self.store.presign_get(keys[0], expiry),
"ANALYZER_URL": self.store.presign_get(keys[1], expiry),
"RESULT_URL": self.store.presign_put(keys[2], expiry),
}
job, secret = manifests(name, self.run, self.args, urls, pin, check)
# create avoids kubectl's last-applied annotation duplicating Secret values.
created = json.loads(self.kubectl("create", "-f", "-", "-o", "json", manifest=job).stdout)
secret["metadata"]["ownerReferences"] = [
{
"apiVersion": "batch/v1",
"kind": "Job",
"name": name,
"uid": created["metadata"]["uid"],
}
]
self.kubectl("create", "-f", "-", manifest=secret)
end = time.monotonic() + self.args.deadline + 60
reason = "host deadline waiting for Job"
while time.monotonic() < end:
state = json.loads(self.kubectl("get", "job", name, "-o", "json").stdout).get("status", {})
pods = json.loads(self.kubectl("get", "pods", "-l", f"engine-trace-job={name}", "-o", "json").stdout)
node = next((p.get("spec", {}).get("nodeName") for p in pods.get("items", []) if p.get("spec", {}).get("nodeName")), node)
result = self.store._get(keys[2])
failed = next((c for c in state.get("conditions", []) if c.get("type") == "Failed" and c.get("status") == "True"), None)
if failed or state.get("failed"):
reason = (failed or {}).get("reason", "Job failed") + ": " + (failed or {}).get("message", "")
break
if state.get("succeeded") and result is not None:
if check:
if result[0] != b"x" * 1024:
raise ValueError("round-trip bytes mismatch")
else:
with tarfile.open(fileobj=io.BytesIO(result[0])) as archive:
archive.extractall(output, filter="data")
if not (output / "result.json").is_file():
raise ValueError("Job produced no result.json")
record = json.loads((output / "result.json").read_text())
if not record.get("done") or {u["unit"] for u in record["units"]} != {u["unit"] for u in pin["units"]}:
raise ValueError("Job returned incomplete or mismatched units")
return subprocess.CompletedProcess([], 0, "", ""), node
time.sleep(2)
logs = self.kubectl("logs", f"job/{name}", "--tail=40", check=False)
reason = redact(reason + "\n" + logs.stdout + logs.stderr)
return subprocess.CompletedProcess([], 1, "", reason), node
except Exception as exc:
reason = "Kubernetes transport failed: " + type(exc).__name__ + ": " + redact(str(exc))
try:
logs = self.kubectl("logs", f"job/{name}", "--tail=40", check=False)
reason += "\n" + redact(logs.stdout + logs.stderr)
except (OSError, ValueError, subprocess.SubprocessError):
pass
return subprocess.CompletedProcess([], 1, "", reason), node
finally:
if not self.args.keep_jobs:
# Attempt every cleanup even if an earlier deletion fails.
errors = []
for kind in ("job", "secret"):
try:
self.kubectl("delete", kind, name, "--ignore-not-found", "--wait=true", "--timeout=30s")
except Exception as exc:
errors.append(type(exc).__name__)
for key in keys:
try:
self.store._delete(key)
except Exception as exc:
errors.append(type(exc).__name__)
if errors:
raise RuntimeError("engine Job cleanup failed: " + ", ".join(errors))
def check(self):
self.namespace()
nodes = json.loads(self.kubectl("get", "nodes", "-o", "json").stdout)
for node in nodes["items"]:
print(json.dumps({"node": node["metadata"]["name"], "allocatable_memory": node["status"]["allocatable"]["memory"]}))
done, node = self.execute(None, None, None, {}, check=True)
print(
json.dumps(
{
"k8s_check": "passed" if done.returncode == 0 else "failed",
"node": node,
"reason": done.stderr,
"round_trip_bytes": 1024 if done.returncode == 0 else 0,
}
)
)
return done.returncode
+240 -1
View File
@@ -531,7 +531,15 @@ def test_trace_real_unsupported_run_then_typescript_opt_in_replans(tmp_path, mon
first = json.loads(result_file.read_text()) first = json.loads(result_file.read_text())
assert first["done"] and first["units"][0]["status"] == "unsupported" assert first["done"] and first["units"][0]["status"] == "unsupported"
progress = json.loads((out / "progress.json").read_text()) progress = json.loads((out / "progress.json").read_text())
assert progress == {"pins_done": 1, "pins_total": 1, "last_pin": {"repo": u["repo"], "pin": u["pin"]}, "statuses": {"unsupported": 1}} assert progress == {
"pins_done": 1,
"pins_total": 1,
"last_pin": {"repo": u["repo"], "pin": u["pin"]},
"statuses": {"unsupported": 1},
"lanes": {"docker": 1},
"pins_in_flight": 0,
"active_lanes": {},
}
assert trace.make_plan(trace.join_units([u], [example]), out) == [] assert trace.make_plan(trace.join_units([u], [example]), out) == []
called = [] called = []
@@ -1158,3 +1166,234 @@ def test_overlay_engine_bug_allows_one_taint_only_attempt(tmp_path, monkeypatch)
monkeypatch.setattr(JoernBackend, "_batch_once", batch) monkeypatch.setattr(JoernBackend, "_batch_once", batch)
assert backend._batch_once(tmp_path / "cpg.bin", "taint", {"a": {}}) == {"a": []} assert backend._batch_once(tmp_path / "cpg.bin", "taint", {"a": {}}) == {"a": []}
assert len(backend.retries) == 1 assert len(backend.retries) == 1
def k8s_args(tmp_path, **overrides):
from types import SimpleNamespace
return SimpleNamespace(
**{
"executor": "k8s",
"parallel": 2,
"kube_context": "test-context",
"kube_namespace": "ousast-engine",
"image": "ghcr.io/norandom/openultrasast:2.0.1",
"deadline": 900,
"question_deadline": 120,
"keep_jobs": False,
"out": tmp_path,
**overrides,
}
)
def test_kubernetes_manifests_keep_urls_in_secret(tmp_path):
from engine_trace_k8s import BOOTSTRAP, manifests
args = k8s_args(tmp_path)
urls = {name: "https://store/" + name + "?signature=secret" for name in ("SOURCE_URL", "ANALYZER_URL", "RESULT_URL")}
job, secret = manifests("job", "run", args, urls, {})
spec = job["spec"]
pod = spec["template"]["spec"]
container = pod["containers"][0]
assert spec["activeDeadlineSeconds"] == 900
assert spec["backoffLimit"] == 0 and spec["ttlSecondsAfterFinished"] > 0
assert pod["restartPolicy"] == "Never" and not pod["automountServiceAccountToken"]
assert pod["affinity"]["podAntiAffinity"]["requiredDuringSchedulingIgnoredDuringExecution"]
assert container["image"] == args.image
assert pod["securityContext"] == {"runAsUser": 1000, "runAsGroup": 1000, "fsGroup": 1000}
assert {m["mountPath"] for m in container["volumeMounts"]} == {"/case", "/frozen", "/out"}
assert all(v["emptyDir"] == {} for v in pod["volumes"])
assert container["resources"] == {"requests": {"memory": "3Gi", "cpu": "1"}, "limits": {"memory": "3Gi"}}
assert secret["stringData"] == urls
assert all(url not in json.dumps(job) for url in urls.values())
assert "boto3" not in BOOTSTRAP and "yaml" not in BOOTSTRAP
compile(BOOTSTRAP, "<bootstrap>", "exec")
@pytest.mark.parametrize("failure", [False, True])
def test_kubernetes_host_pin_success_failure_and_cleanup(tmp_path, monkeypatch, failure):
import io
import tarfile
from engine_trace_k8s import Kubernetes
from openultrasast.plane.memory import S3Store
objects, calls, expiries = {}, [], []
class Client:
def put(self, key, data, labels):
objects[key] = data
def get(self, key):
return (objects[key], None) if key in objects else None
def delete(self, key):
objects.pop(key, None)
def presign(self, method, key, expires):
expiries.append(expires.total_seconds())
return "https://store/" + key + "?secret"
store = object.__new__(S3Store)
store.client, store.prefix = Client(), "prefix"
args = k8s_args(tmp_path)
transport = Kubernetes(args, store)
args.kubernetes = transport
pin = trace.make_plan([unit()], tmp_path)[0]
source = tmp_path / "frozen"
source.mkdir()
(source / "worker.py").write_text("print(42)")
class Inputs:
def materialize(self, pin, target):
target.mkdir()
(target / "app.py").write_text("def run(): pass")
return {"bytes": 15, "files": 1}
def kubectl(command, **kwargs):
assert command[:5] == ["kubectl", "--context", "test-context", "--namespace", "ousast-engine"]
calls.append(command[5:])
action = command[5:]
stdout = "{}"
if action[0] == "create":
manifest = json.loads(kwargs["input"])
if manifest["kind"] == "Secret":
assert manifest["metadata"]["ownerReferences"][0]["uid"] == "job-uid"
if manifest["kind"] == "Job":
stdout = json.dumps({"metadata": {"uid": "job-uid"}})
env = manifest["spec"]["template"]["spec"]["containers"][0]["env"]
prepared = json.loads(next(e["value"] for e in env if e["name"] == "PIN"))
result = {**prepared, "done": True, "units": [worker.unit_record(u, "asked-nothing", "") for u in prepared["units"]]}
result_bytes = json.dumps(result).encode()
stream = io.BytesIO()
with tarfile.open(fileobj=stream, mode="w") as archive:
entry = tarfile.TarInfo("result.json")
entry.size = len(result_bytes)
archive.addfile(entry, io.BytesIO(result_bytes))
source_key = next(key for key in objects if key.endswith("source.tar"))
with tarfile.open(fileobj=io.BytesIO(objects[source_key])) as archive:
assert archive.extractfile("app.py").read() == b"def run(): pass"
objects[source_key.replace("source.tar", "result.tar")] = stream.getvalue()
if action[:2] == ["get", "job"]:
stdout = json.dumps(
{"status": {"conditions": [{"type": "Failed", "status": "True", "reason": "DeadlineExceeded"}]}}
if failure
else {"status": {"succeeded": 1}}
)
if action[:2] == ["get", "pods"]:
stdout = json.dumps({"items": [{"spec": {"nodeName": "worker-1"}}]})
if action[0] == "logs":
assert "--tail=40" in action
stdout = "last worker line https://store/private?secret"
return subprocess.CompletedProcess(command, 0, stdout, "")
monkeypatch.setattr(subprocess, "run", kubectl)
monkeypatch.setattr(trace, "prepare_questions", lambda *a: {"a": {}})
assert trace.run_pin(pin, source, {"head": "frozen"}, Inputs(), tmp_path, args) is failure
result = json.loads((tmp_path / trace.pin_name(pin)).read_text())
assert result["done"] and result["executor"] == "k8s" and result["node"] == "worker-1"
assert result["analyzer"] == {"head": "frozen"}
assert result["container_exit"] == int(failure)
assert result["units"][0]["status"] == ("failed" if failure else "asked-nothing")
if failure:
assert "DeadlineExceeded" in result["units"][0]["reason"]
assert "last worker line" in result["units"][0]["reason"]
assert "?secret" not in json.dumps(result)
else:
assert trace.make_plan([unit()], tmp_path) == []
assert not objects
assert any(c[:2] == ["delete", "secret"] for c in calls)
assert any(c[:2] == ["delete", "job"] for c in calls)
assert min(expiries) >= args.deadline + 900
@pytest.mark.parametrize("executor,parallel", [("docker", 1), ("k8s", 2), ("mixed", 3)])
def test_scheduler_bounds_lanes_and_progress(tmp_path, monkeypatch, executor, parallel):
import threading
import time
from collections import Counter
active, peak = Counter(), Counter()
lock = threading.Lock()
args = k8s_args(tmp_path, executor=executor, parallel=parallel)
plan = trace.make_plan([unit(str(i), str(i)) for i in range(9)], tmp_path)
def run(pin, source, provenance, inputs, out, options):
with lock:
active[options.executor] += 1
active["total"] += 1
for key, value in active.items():
peak[key] = max(peak[key], value)
time.sleep(0.02)
trace.save_pin(
out / trace.pin_name(pin), {**pin, "done": True, "units": [worker.unit_record(u, "asked-nothing", "") for u in pin["units"]]}
)
with lock:
active[options.executor] -= 1
active["total"] -= 1
return False
monkeypatch.setattr(trace, "run_pin", run)
assert not trace.run_plan(plan, None, {}, None, args)
assert peak["total"] == parallel
if executor == "mixed":
assert peak["docker"] == 1 and peak["k8s"] == parallel - 1
progress = json.loads((tmp_path / "progress.json").read_text())
assert progress["pins_done"] == 9 and progress["pins_in_flight"] == 0
assert sum(progress["lanes"].values()) == parallel
assert trace.make_plan([unit(str(i), str(i)) for i in range(9)], tmp_path) == []
@pytest.mark.parametrize("keep", [False, True])
def test_k8s_check_roundtrip_and_retention(tmp_path, monkeypatch, capsys, keep):
from engine_trace_k8s import Kubernetes
from openultrasast.plane.memory import S3Store
data, manifests_seen, commands = {}, [], []
class Client:
def put(self, key, body, labels):
data[key] = body
def get(self, key):
return (b"x" * 1024, None)
def delete(self, key):
data.pop(key, None)
def presign(self, method, key, expires):
return "https://private/" + key
store = object.__new__(S3Store)
store.client, store.prefix = Client(), ""
runner = Kubernetes(k8s_args(tmp_path, keep_jobs=keep), store)
def fake(command, **kwargs):
commands.append(command)
action = command[5:]
result = {}
if kwargs.get("input"):
manifest = json.loads(kwargs["input"])
manifests_seen.append(manifest)
result = {"metadata": {"uid": "check-uid"}}
if action[:2] == ["get", "nodes"]:
result = {"items": [{"metadata": {"name": "node"}, "status": {"allocatable": {"memory": "6000000Ki"}}}]}
if action[:2] == ["get", "pods"]:
result = {"items": [{"spec": {"nodeName": "node"}}]}
if action[:2] == ["get", "job"]:
result = {"status": {"succeeded": 1}}
return subprocess.CompletedProcess(command, 0, json.dumps(result), "")
monkeypatch.setattr(subprocess, "run", fake)
assert runner.check() == 0
text = capsys.readouterr().out
assert "6000000Ki" in text and "1024" in text and "passed" in text
assert "https://" not in text
assert any(m["kind"] == "Namespace" for m in manifests_seen)
job = next(m for m in manifests_seen if m["kind"] == "Job")
assert {"name": "CHECK", "value": "1"} in job["spec"]["template"]["spec"]["containers"][0]["env"]
assert bool(data) is keep
assert any("delete" in c for c in commands) is not keep