mirror of
https://github.com/norandom/OpenUltraSAST.git
synced 2026-10-16 03:49:26 +00:00
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:
co-authored by
Claude Opus 5.5
parent
431ab1e738
commit
2e6ba93aee
@@ -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
|
||||
is git archive HEAD src plus a hashed snapshot of this uncommitted tool.
|
||||
@@ -8,6 +8,7 @@ from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import ast
|
||||
import copy
|
||||
import fcntl
|
||||
import hashlib
|
||||
import json
|
||||
@@ -17,11 +18,14 @@ import shutil
|
||||
import subprocess
|
||||
import tempfile
|
||||
import textwrap
|
||||
import threading
|
||||
import time
|
||||
import tomllib
|
||||
from collections import Counter, defaultdict
|
||||
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
|
||||
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 openultrasast.learn.examples import read_jsonl
|
||||
@@ -39,6 +43,7 @@ DEFAULT_OUT = Path.home() / "ousast-results/plane/engine-trace"
|
||||
SNAPSHOT_FILES = (
|
||||
"benchmarks/learn/engine_trace.py",
|
||||
"benchmarks/learn/engine_trace_worker.py",
|
||||
"benchmarks/learn/engine_trace_k8s.py",
|
||||
"src/openultrasast/model/trace.py",
|
||||
"src/openultrasast/model/taint.py",
|
||||
"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):
|
||||
executor = getattr(args, "executor", "docker")
|
||||
node = None
|
||||
with tempfile.TemporaryDirectory(prefix="trace-pin-", dir=out) as scratch:
|
||||
scratch = Path(scratch)
|
||||
checkout, output = scratch / "case", scratch / "output"
|
||||
output.mkdir()
|
||||
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)})
|
||||
name = "ousast-trace-" + pin_name(pin)[:16]
|
||||
command = docker_command(source, checkout, output, name, args.image, args.deadline, args.question_deadline)
|
||||
started = time.monotonic()
|
||||
done = None
|
||||
try:
|
||||
done = subprocess.run(command, capture_output=True, text=True, timeout=args.deadline)
|
||||
status, reason = "failed", f"container exit {done.returncode}: {done.stderr[-1000:]}"
|
||||
(output / "container.log").write_text(done.stdout + "\n" + done.stderr)
|
||||
except subprocess.TimeoutExpired:
|
||||
if executor == "k8s":
|
||||
done, node = args.kubernetes.execute(source, checkout, output, json.loads((output / "pin.json").read_text()))
|
||||
status, reason = "failed", done.stderr
|
||||
if done.returncode:
|
||||
(output / "container.log").write_text(reason)
|
||||
else:
|
||||
command = docker_command(source, checkout, output, name, args.image, args.deadline, args.question_deadline)
|
||||
started = time.monotonic()
|
||||
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")
|
||||
done = subprocess.run(command, capture_output=True, text=True, timeout=args.deadline)
|
||||
status, reason = "failed", f"container exit {done.returncode}: {done.stderr[-1000:]}"
|
||||
(output / "container.log").write_text(done.stdout + "\n" + done.stderr)
|
||||
except subprocess.TimeoutExpired:
|
||||
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"
|
||||
if not result_path.exists():
|
||||
reason = f"container produced no result: {reason}"
|
||||
@@ -666,6 +681,8 @@ 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 == "k8s":
|
||||
record.update(executor="k8s", node=node, image=args.image)
|
||||
save_pin(out / pin_name(pin), record)
|
||||
print(
|
||||
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():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
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("--cache", type=Path, default=Path.home() / ".cache/openultrasast")
|
||||
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("--question-deadline", type=int, default=120)
|
||||
parser.add_argument("--limit", type=int, default=0)
|
||||
@@ -747,6 +831,13 @@ def main():
|
||||
args = parser.parse_args()
|
||||
if min(args.deadline, args.question_deadline) <= 0 or args.limit < 0:
|
||||
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(",")))
|
||||
if rerun_status - set(STATUSES):
|
||||
parser.error("unknown --rerun-status value")
|
||||
@@ -806,34 +897,21 @@ def main():
|
||||
except BlockingIOError:
|
||||
parser.error("another trace runner holds this results root")
|
||||
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(
|
||||
["docker", "image", "inspect", "--format", "{{.Id}}", args.image], capture_output=True, text=True, check=True
|
||||
)
|
||||
args.image = image.stdout.strip()
|
||||
if not args.image:
|
||||
args.docker_image = image.stdout.strip()
|
||||
if not args.docker_image:
|
||||
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:
|
||||
source = Path(scratch) / "frozen"
|
||||
provenance = freeze_source(source)
|
||||
progress = {"pins_done": 0, "pins_total": len(plan), "last_pin": None, "statuses": {}}
|
||||
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)
|
||||
failures = run_plan(plan, source, provenance, inputs, args)
|
||||
return 2 if failures else 0
|
||||
|
||||
|
||||
|
||||
@@ -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
@@ -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())
|
||||
assert first["done"] and first["units"][0]["status"] == "unsupported"
|
||||
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) == []
|
||||
|
||||
called = []
|
||||
@@ -1158,3 +1166,234 @@ def test_overlay_engine_bug_allows_one_taint_only_attempt(tmp_path, monkeypatch)
|
||||
monkeypatch.setattr(JoernBackend, "_batch_once", batch)
|
||||
assert backend._batch_once(tmp_path / "cpg.bin", "taint", {"a": {}}) == {"a": []}
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user