learn: engine_trace container worker imports only the engine; dry-run writes nothing; resume keyed on selection

The 2026-09-14 engine image has no PyYAML or boto3; the worker no longer imports
learn.examples. Implemented with Codex.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

benchmarks/learn/engine_trace_worker.py

@@ -0,0 +1,276 @@
+"""Container-only Joern trace worker; host prepares all corpus and question metadata.
+
+Keep this module independent of learn/plane and their optional dependencies. The
+frozen host invokes this file directly in the shipped, offline analyzer image.
+"""
+
+from __future__ import annotations
+
+import argparse
+import contextlib
+import hashlib
+import json
+import os
+import selectors
+import signal
+import subprocess
+import time
+from collections import Counter
+from pathlib import Path
+
+from openultrasast.cpg.backend import JoernBackend
+from openultrasast.model.trace import parse_trace
+
+
+def write_json(path, value):
+    path.parent.mkdir(parents=True, exist_ok=True)
+    tmp = path.with_suffix(".tmp")
+    tmp.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n")
+    tmp.replace(path)
+
+
+def safe_path(root, name):
+    path = (root / name).resolve()
+    if not path.is_relative_to(root.resolve()):
+        raise ValueError(f"path escapes input root: {name}")
+    return path
+
+
+def instrument_failure(proof):
+    if not proof.get("files") or not proof.get("bytes"):
+        return "instrument read zero files or zero bytes"
+    if not proof.get("cpg_bytes"):
+        return "instrument produced no CPG bytes"
+    for stage in proof.get("jvm", []):
+        if stage["seconds"] < 5:
+            return f"implausibly fast JVM step: {stage['command']} in {stage['seconds']:.3f}s (<5s)"
+    if not proof.get("jvm"):
+        return "no JVM step recorded"
+    return ""
+
+
+class MeasuredBackend(JoernBackend):
+    """Use the shipped backend with a process watchdog, no concurrent query threads."""
+
+    def __init__(self, deadline, question_deadline, checkpoint):
+        super().__init__(build_timeout=deadline, query_timeout=deadline, session_transport=False)
+        self.deadline = time.monotonic() + deadline
+        self.question_deadline = question_deadline
+        self.checkpoint = checkpoint
+        self.stages = []
+        self.asked = []
+        self.answers = {}
+        self.completion_counts = Counter()
+        self.timed_out = False
+        self.fatal = ""
+
+    def _run(self, command, *, timeout, cwd=None):
+        if self.fatal or self.timed_out:
+            return None
+        if time.monotonic() >= self.deadline:
+            self.timed_out = True
+            return None
+        started = time.monotonic()
+        end = min(self.deadline, started + timeout)
+        question_end = end
+        jvm = Path(command[0]).name != "php"
+        output = bytearray()
+        pending = b""
+        process = subprocess.Popen(
+            command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, cwd=cwd, env=self._jvm_env(), start_new_session=True
+        )
+        try:
+            with selectors.DefaultSelector() as selector:
+                selector.register(process.stdout, selectors.EVENT_READ)
+                while selector.get_map():
+                    if time.monotonic() >= min(end, question_end):
+                        self.timed_out = True
+                        break
+                    for key, _ in selector.select(min(0.2, max(0, min(end, question_end) - time.monotonic()))):
+                        chunk = os.read(key.fileobj.fileno(), 65536)
+                        if not chunk:
+                            selector.unregister(key.fileobj)
+                            continue
+                        output.extend(chunk)
+                        pending += chunk
+                        while b"\n" in pending:
+                            line, pending = pending.split(b"\n", 1)
+                            try:
+                                event = json.loads(line)
+                            except ValueError:
+                                continue
+                            if not isinstance(event, dict):
+                                continue
+                            if "__question__" in event:
+                                self.asked.append(event["__question__"])
+                                question_end = min(end, time.monotonic() + self.question_deadline)
+                            if "id" in event and "rows" in event:
+                                self.answers[event["id"]] = event["rows"]
+                                self.completion_counts[event["id"]] += 1
+                                question_end = end
+                            self.checkpoint()
+            if not self.timed_out:
+                process.wait(timeout=max(0.01, end - time.monotonic()))
+        finally:
+            with contextlib.suppress(ProcessLookupError):
+                os.killpg(process.pid, signal.SIGKILL)
+            process.wait(timeout=5)
+            process.stdout.close()
+        elapsed = time.monotonic() - started
+        if jvm:
+            self.stages.append({"command": Path(command[0]).name, "seconds": elapsed, "exit": process.returncode})
+            if elapsed < 5:
+                self.fatal = f"implausibly fast JVM step: {command[0]} in {elapsed:.3f}s (<5s)"
+        self.checkpoint()
+        return subprocess.CompletedProcess(command, process.returncode if not self.fatal else 1, output.decode(errors="replace"), "")
+
+
+def unit_record(unit, status, reason="", **extra):
+    return {
+        **unit,
+        "status": status,
+        "reason": reason,
+        "witness_rows": [],
+        "traces": [],
+        "instrument": {"files": [], "bytes": 0, "cpg_bytes": 0, "jvm": []},
+        "questions": {},
+        "questions_asked": [],
+        "questions_completed": [],
+        **extra,
+    }
+
+
+def worker(pin, root, out, deadline, question_deadline):
+    proof = {"files": [], "bytes": 0, "cpg_bytes": 0, "jvm": []}
+    record = {
+        **pin,
+        "done": False,
+        "instrument": proof,
+        "units": [unit_record(u, "unsupported" if not u["supported"] else "timeout", u["reason"]) for u in pin["units"]],
+    }
+
+    def checkpoint():
+        # Checkpoints retain raw completed rows even if the outer container deadline fires.
+        # They remain timeout until the finished JVM's instrument proof is validated.
+        for row in record["units"]:
+            rid = row["unit"]
+            row["instrument"] = proof
+            row["questions_asked"] = [rid] if rid in backend.asked else []
+            if rid in backend.answers:
+                row["questions_completed"] = [rid]
+                row["witness_rows"] = [r for r in backend.answers[rid] if isinstance(r, dict) and "kind" not in r and "sink" in r]
+        write_json(out, record)
+
+    backend = MeasuredBackend(deadline, question_deadline, checkpoint)
+    proof["jvm"] = backend.stages
+    cpg = None
+    try:
+        active = [u for u in pin["units"] if u["supported"]]
+        for file in sorted({u["file"] for u in active}):
+            data = safe_path(root, file).read_bytes()
+            proof["files"].append({"file": file, "bytes": len(data), "sha256": hashlib.sha256(data).hexdigest()})
+            proof["bytes"] += len(data)
+            print(f"read {file}: {len(data)} bytes", flush=True)
+        if not proof["bytes"]:
+            raise ValueError("instrument read zero files or zero bytes")
+        languages = {"javascript" if u["language"] == "typescript" else u["language"] for u in active}
+        if len(languages) != 1:
+            raise ValueError(f"one pin has multiple frontend languages: {sorted(languages)}")
+        language = next(iter(languages))
+        requests = pin["questions"]
+        if set(requests) != {u["unit"] for u in active}:
+            raise ValueError("prepared questions do not match supported units")
+        record["questions"] = requests
+        for unit_row in record["units"]:
+            unit_row["questions"] = {unit_row["unit"]: requests[unit_row["unit"]]} if unit_row["unit"] in requests else {}
+        write_json(out, record)
+        cpg = backend.build(root, language=language)
+        if cpg is None:
+            raise ValueError(backend.fatal or backend.last_failure or ("pin deadline" if backend.timed_out else "CPG build failed"))
+        proof["cpg_bytes"] = sum(p.stat().st_size for p in cpg.cpg_path.parent.rglob("*.bin"))
+        failure = instrument_failure(proof)
+        if failure:
+            raise ValueError(failure)
+        if cpg.run_batch_once is None:
+            raise ValueError("backend offers no batch operation")
+        answers = cpg.run_batch_once("taint", requests) or {}
+        record["census"] = answers.get("__census__", [])
+        failure = backend.fatal or instrument_failure(proof)
+        census = record["census"]
+        if census and int(census[0].get("files", 0)) <= 0:
+            failure = failure or "query census read zero files"
+        elif not census and not backend.timed_out:
+            failure = failure or "query census is missing"
+        for index, unit in enumerate(pin["units"]):
+            if not unit["supported"]:
+                continue
+            rid = unit["unit"]
+            extra = {
+                "questions": {rid: requests[rid]},
+                "questions_asked": [rid] if rid in backend.asked else [],
+                "questions_completed": [rid] if rid in answers else [],
+                "instrument": proof,
+            }
+            rows = answers.get(rid)
+            # A batch across backend shards is complete only when every shard answered.
+            # The legacy merger can return a union even when one shard was killed.
+            shards = int(census[0].get("shards", 1)) if census else 1
+            shard_incomplete = shards > 1 and getattr(backend, "completion_counts", {}).get(rid, 0) < shards
+            if rows is not None:
+                raw_flows = [r for r in rows if isinstance(r, dict) and "kind" not in r and "sink" in r]
+                extra.update(witness_rows=raw_flows, traces=[parse_trace(r) for r in raw_flows])
+            if shard_incomplete:
+                extra["questions_completed"] = []
+            dropped = any(p.endswith(unit["file"]) for p in cpg.unparsed)
+            names = census[0].get("file_names", []) if census else []
+            file_missing = not any(str(p).removeprefix("./") == unit["file"] or str(p).endswith("/" + unit["file"]) for p in names)
+            if file_missing and rows is not None:
+                dropped = True
+            if failure or dropped:
+                status, reason = "failed", failure or "labelled file was dropped or absent from query census"
+            elif rows is None or shard_incomplete:
+                status, reason = ("timeout" if backend.timed_out else "failed"), "question did not complete"
+            else:
+                flows = [r for r in rows if isinstance(r, dict) and "kind" not in r and "sink" in r]
+                if any(not row.get("trace") for row in flows):
+                    raise ValueError("taint witness returned without requested trace")
+                extra.update(witness_rows=flows, traces=[parse_trace(r) for r in flows])
+                status, reason = ("path" if flows else "asked-nothing"), ""
+            record["units"][index] = unit_record(unit, status, reason, **extra)
+    except (OSError, ValueError, RuntimeError) as exc:
+        for index, unit in enumerate(pin["units"]):
+            if unit["supported"]:
+                record["units"][index] = unit_record(
+                    unit,
+                    "timeout" if backend.timed_out and not backend.fatal else "failed",
+                    str(exc),
+                    instrument=proof,
+                    questions={unit["unit"]: record["questions"][unit["unit"]]} if unit["unit"] in record.get("questions", {}) else {},
+                    witness_rows=[r for r in backend.answers.get(unit["unit"], []) if isinstance(r, dict) and "sink" in r],
+                    questions_completed=[unit["unit"]] if unit["unit"] in backend.answers else [],
+                    questions_asked=[unit["unit"]] if unit["unit"] in backend.asked else [],
+                )
+    finally:
+        if cpg and cpg.cleanup:
+            cpg.cleanup()
+        record["questions_asked"] = list(dict.fromkeys(backend.asked))
+        record["questions_completed"] = [rid for u in record["units"] for rid in u["questions_completed"]]
+        record["done"] = True
+        write_json(out, record)
+    return record
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--worker", type=Path, required=True)
+    parser.add_argument("--deadline", type=int, default=900)
+    parser.add_argument("--question-deadline", type=int, default=120)
+    args = parser.parse_args()
+    if min(args.deadline, args.question_deadline) <= 0:
+        parser.error("deadlines must be positive")
+    worker(json.loads(args.worker.read_text()), Path("/case"), Path("/out/result.json"), args.deadline, args.question_deadline)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())