diff --git a/benchmarks/learn/engine_trace.py b/benchmarks/learn/engine_trace.py index e4d0a7d..f121ef2 100644 --- a/benchmarks/learn/engine_trace.py +++ b/benchmarks/learn/engine_trace.py @@ -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 diff --git a/benchmarks/learn/engine_trace_k8s.py b/benchmarks/learn/engine_trace_k8s.py new file mode 100644 index 0000000..528e31b --- /dev/null +++ b/benchmarks/learn/engine_trace_k8s.py @@ -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 diff --git a/tests/test_engine_trace.py b/tests/test_engine_trace.py index a4f99df..27075a2 100644 --- a/tests/test_engine_trace.py +++ b/tests/test_engine_trace.py @@ -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, "", "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