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]>

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
 
 

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