mirror of
https://github.com/norandom/OpenUltraSAST.git
synced 2026-10-16 03:49:26 +00:00
merge(engine): Kubernetes engine workers (Deployment from the image, S3 queue); leak-audit records carry digests, not repository names
The Joern leak-audit record named spent population v1/v2 repositories, which the reserved-repository guard forbids outside benchmarks/independent/ (no v3 repository was involved: 0 matches). Records now carry counts and SHA-256 digests. With Codex. Co-Authored-By: Claude Opus 5.5 <[email protected]>
This commit is contained in:
Marius Ciepluch
@@ -0,0 +1,290 @@
|
||||
"""Offline queue protocol, image imports and deployment contract controls."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
with pytest.MonkeyPatch.context() as patch:
|
||||
patch.syspath_prepend(str(ROOT / "benchmarks/learn"))
|
||||
service = importlib.import_module("engine_worker_service")
|
||||
transport = importlib.import_module("engine_trace_queue")
|
||||
|
||||
|
||||
class FakeS3:
|
||||
def __init__(self):
|
||||
self.objects = {}
|
||||
self.lock = threading.Lock()
|
||||
self.choosing = None
|
||||
|
||||
def get(self, key):
|
||||
with self.lock:
|
||||
return (self.objects[key], None) if key in self.objects else None
|
||||
|
||||
def put(self, key, data, labels):
|
||||
with self.lock:
|
||||
self.objects[key] = data
|
||||
if self.choosing and key.startswith(service.CLAIMED) and json.loads(data)["ticket"] == 0:
|
||||
self.choosing.wait(timeout=5)
|
||||
|
||||
def delete(self, key):
|
||||
with self.lock:
|
||||
self.objects.pop(key, None)
|
||||
|
||||
def keys(self, prefix):
|
||||
with self.lock:
|
||||
return sorted(k for k in self.objects if k.startswith(prefix))
|
||||
|
||||
|
||||
def task():
|
||||
return dict(
|
||||
run="run",
|
||||
id="pin",
|
||||
task_id="run-pin",
|
||||
input="engine-queue/inputs/run/pin/source.tar",
|
||||
deadline=10,
|
||||
question_deadline=2,
|
||||
pin={"repo": "owner/repo", "pin": "commit", "units": []},
|
||||
)
|
||||
|
||||
|
||||
def test_engine_worker_racing_claimants_run_exactly_once():
|
||||
store = FakeS3()
|
||||
item = task()
|
||||
service.put(store, service.PENDING + item["task_id"], item)
|
||||
store.choosing = threading.Barrier(2)
|
||||
called = []
|
||||
|
||||
def execute(client, entry):
|
||||
called.append(entry["task_id"])
|
||||
time.sleep(0.05)
|
||||
return {**entry["pin"], "done": True, "instrument": {"bytes": 11}}
|
||||
|
||||
with ThreadPoolExecutor(2) as pool:
|
||||
futures = [pool.submit(service.Queue(store, identity).process_one, execute) for identity in ["worker-b", "worker-a"]]
|
||||
results = [f.result(timeout=10) for f in futures]
|
||||
assert sum(results) == 1
|
||||
assert called == ["run-pin"]
|
||||
result = service.read(store, service.result_key(item))
|
||||
assert result["worker"] in {"worker-a", "worker-b"}
|
||||
assert not store.keys(service.CLAIMED)
|
||||
assert not store.keys(service.PENDING)
|
||||
|
||||
|
||||
def test_engine_worker_expired_lease_recovered_without_pending():
|
||||
store = FakeS3()
|
||||
item = task()
|
||||
service.put(store, "engine-queue/claimed/dead/run-pin", dict(task=item, ticket=1, expires=99))
|
||||
queue = service.Queue(store, "new", now=lambda: 100)
|
||||
assert queue.process_one(lambda *_: {**item["pin"], "done": True})
|
||||
assert service.read(store, service.result_key(item))["worker"] == "new"
|
||||
assert not store.keys(service.CLAIMED)
|
||||
|
||||
|
||||
def test_engine_worker_late_arrival_cannot_displace_running_owner():
|
||||
store = FakeS3()
|
||||
item = task()
|
||||
first = service.Queue(store, "z-first")
|
||||
key = first.claim(item)
|
||||
assert key
|
||||
stop = threading.Event()
|
||||
later = service.Queue(store, "a-later", stop=stop)
|
||||
with ThreadPoolExecutor(1) as pool:
|
||||
future = pool.submit(later.claim, item)
|
||||
assert future.result(timeout=3) is None
|
||||
assert service.read(store, key) is not None
|
||||
service.put(store, service.result_key(item), {"done": True})
|
||||
store.delete(key)
|
||||
|
||||
|
||||
def test_engine_queue_dispatch_roundtrip_and_timeout(tmp_path):
|
||||
store = FakeS3()
|
||||
checkout, output = tmp_path / "case", tmp_path / "out"
|
||||
checkout.mkdir()
|
||||
output.mkdir()
|
||||
(checkout / "app.py").write_text("print(1)\n")
|
||||
args = SimpleNamespace(deadline=10, question_deadline=2, queue_timeout=2, image="image@sha256:abc")
|
||||
unit = {"unit": "u", "supported": True, "reason": "", "family": "injection"}
|
||||
pin = {"repo": "owner/repo", "pin": "sha", "units": [unit], "questions": []}
|
||||
|
||||
def execute(client, item):
|
||||
unpacked = tmp_path / "unpacked"
|
||||
unpacked.mkdir()
|
||||
service.unpack(client.get(item["input"])[0], unpacked)
|
||||
assert (unpacked / "app.py").read_bytes() == b"print(1)\n"
|
||||
return {
|
||||
**item["pin"],
|
||||
"done": True,
|
||||
"container_exit": 0,
|
||||
"units": [service.unit_record(unit, "asked-nothing")],
|
||||
"instrument": {"bytes": 9, "heap_mb": 2560},
|
||||
}
|
||||
|
||||
queue = service.Queue(store, "pod-1")
|
||||
dispatcher = transport.Dispatcher(args, store, pause=lambda _: queue.process_one(execute))
|
||||
completed, worker = dispatcher.execute(checkout, output, pin, "pin")
|
||||
assert completed.returncode == 0
|
||||
assert worker == "pod-1"
|
||||
result = json.loads((output / "result.json").read_text())
|
||||
assert result["executor"] == "queue"
|
||||
assert result["units"][0]["status"] == "asked-nothing"
|
||||
assert result["instrument"]["heap_mb"] == 2560
|
||||
assert dispatcher.status() == {"pending": 0, "claimed": 0, "done": 1}
|
||||
assert not store.keys("engine-queue/inputs/")
|
||||
ticks = iter([0, 3])
|
||||
dispatcher = transport.Dispatcher(args, store, clock=lambda: next(ticks))
|
||||
assert dispatcher.execute(checkout, output, pin, "other") == (None, None)
|
||||
assert not store.keys(service.PENDING)
|
||||
|
||||
|
||||
def test_engine_worker_imports_without_yaml():
|
||||
code = "import sys; sys.modules['yaml'] = None; import engine_worker_service; import engine_trace_worker"
|
||||
done = subprocess.run(
|
||||
[sys.executable, "-c", code],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env={**os.environ, "PYTHONPATH": f"{ROOT / 'src'}:{ROOT / 'benchmarks/learn'}"},
|
||||
)
|
||||
assert done.returncode == 0, done.stderr
|
||||
|
||||
|
||||
def test_engine_k8s_render_contract():
|
||||
base = ROOT / "ops/k8s/engine"
|
||||
config = yaml.safe_load((base / "kustomization.yaml").read_text())
|
||||
# Render this base's resource inclusion, namespace, image and replacement transforms offline.
|
||||
resources = [yaml.safe_load((base / name).read_text()) for name in config["resources"]]
|
||||
deployment = next(r for r in resources if r["kind"] == "Deployment")
|
||||
deployment["metadata"]["namespace"] = config["namespace"]
|
||||
image = config["images"][0]
|
||||
pod = deployment["spec"]["template"]["spec"]
|
||||
container = pod["containers"][0]
|
||||
assert container["image"] == image["name"]
|
||||
container["image"] = image["newName"] + "@" + image["digest"]
|
||||
replacement = config["replacements"][0]
|
||||
assert replacement["source"]["fieldPath"] == "spec.template.spec.containers.0.image"
|
||||
assert replacement["targets"][0]["fieldPaths"] == ["spec.template.spec.containers.0.env.0.value"]
|
||||
container["env"][0]["value"] = container["image"]
|
||||
assert "@sha256:" in container["image"]
|
||||
assert deployment["spec"]["replicas"] == 2
|
||||
assert container["resources"] == {"requests": {"memory": "3Gi", "cpu": "1"}, "limits": {"memory": "3Gi"}}
|
||||
affinity = pod["affinity"]["podAntiAffinity"]["requiredDuringSchedulingIgnoredDuringExecution"][0]
|
||||
assert affinity["topologyKey"] == "kubernetes.io/hostname"
|
||||
assert affinity["labelSelector"]["matchLabels"] == deployment["spec"]["template"]["metadata"]["labels"]
|
||||
assert pod["terminationGracePeriodSeconds"] == 900
|
||||
assert container["envFrom"] == [{"secretRef": {"name": "ousast-engine-s3"}}]
|
||||
for probe in ["readinessProbe", "livenessProbe"]:
|
||||
assert container[probe]["httpGet"] == {"path": "/healthz", "port": "health"}
|
||||
assert not pod["automountServiceAccountToken"]
|
||||
assert pod["securityContext"]["runAsNonRoot"]
|
||||
assert "hostPath" not in json.dumps(resources)
|
||||
assert "privileged" not in json.dumps(resources)
|
||||
assert all(r["kind"] != "Secret" for r in resources)
|
||||
secret = yaml.safe_load((base / "secret.example.yaml").read_text())
|
||||
assert secret["stringData"] == {}
|
||||
assert "AWS_SECRET_ACCESS_KEY" not in json.dumps(resources)
|
||||
|
||||
|
||||
def test_engine_image_and_rustfs_policy():
|
||||
dockerfile = (ROOT / "Dockerfile").read_text()
|
||||
assert "COPY benchmarks ./benchmarks" in dockerfile
|
||||
assert ".[semantic,s3]" in dockerfile
|
||||
for name in ["engine_trace_worker.py", "engine_trace_parse.py", "engine_worker_service.py"]:
|
||||
assert (ROOT / "benchmarks/learn" / name).stat().st_size > 0
|
||||
doc = (ROOT / "docs/rustfs.md").read_text().split("## Scoped account for engine workers")[1]
|
||||
policy = json.loads(doc.split("```json\n")[1].split("```")[0])
|
||||
listing, objects = policy["Statement"]
|
||||
assert listing["Action"] == "s3:ListBucket"
|
||||
assert set(listing["Condition"]["StringLike"]["s3:prefix"]) == {
|
||||
"engine-queue/",
|
||||
"engine-queue/*",
|
||||
"engine-results/",
|
||||
"engine-results/*",
|
||||
}
|
||||
assert set(objects["Action"]) == {"s3:GetObject", "s3:PutObject", "s3:DeleteObject"}
|
||||
assert set(objects["Resource"]) == {"arn:aws:s3:::sast-memory/engine-queue/*", "arn:aws:s3:::sast-memory/engine-results/*"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("hang", [False, True])
|
||||
def test_engine_worker_process_deadline_and_shipped_schema(tmp_path, monkeypatch, hang):
|
||||
store = FakeS3()
|
||||
item = task()
|
||||
item["image"] = "image@sha256:abc"
|
||||
item["deadline"] = 0.1 if hang else 5
|
||||
item["pin"]["units"] = [{"unit": "u", "supported": True, "reason": ""}]
|
||||
(tmp_path / "app.py").write_text("print(1)\n")
|
||||
store.put(item["input"], service.pack(tmp_path), {})
|
||||
monkeypatch.setenv("ENGINE_IMAGE", item["image"])
|
||||
|
||||
def fake_worker(pin, root, out, deadline, question_deadline):
|
||||
assert (root / "app.py").read_bytes() == b"print(1)\n"
|
||||
if hang:
|
||||
time.sleep(10)
|
||||
out.write_text(json.dumps({**pin, "done": True, "units": [service.unit_record(pin["units"][0], "asked-nothing")]}))
|
||||
|
||||
monkeypatch.setattr(service, "worker", fake_worker)
|
||||
result = service.execute_task(store, item)
|
||||
assert result["done"]
|
||||
assert result["units"][0]["status"] == ("timeout" if hang else "asked-nothing")
|
||||
assert result["container_exit"] == (None if hang else 0)
|
||||
assert result["seconds"] < 5
|
||||
|
||||
|
||||
def test_engine_worker_stopping_does_not_claim():
|
||||
store = FakeS3()
|
||||
item = task()
|
||||
service.put(store, service.PENDING + item["task_id"], item)
|
||||
stop = threading.Event()
|
||||
stop.set()
|
||||
assert not service.Queue(store, stop=stop).process_one(lambda *_: pytest.fail("claimed after SIGTERM"))
|
||||
assert store.keys(service.PENDING)
|
||||
assert not store.keys(service.CLAIMED)
|
||||
|
||||
|
||||
def test_engine_worker_rejects_archive_links_and_escape(tmp_path):
|
||||
import io
|
||||
import tarfile
|
||||
|
||||
for name, kind in [("../escape", tarfile.REGTYPE), ("link", tarfile.SYMTYPE)]:
|
||||
stream = io.BytesIO()
|
||||
with tarfile.open(fileobj=stream, mode="w") as archive:
|
||||
member = tarfile.TarInfo(name)
|
||||
member.type = kind
|
||||
member.linkname = "/etc/passwd"
|
||||
archive.addfile(member)
|
||||
with pytest.raises(ValueError, match="unsafe source archive"):
|
||||
service.unpack(stream.getvalue(), tmp_path)
|
||||
|
||||
|
||||
def test_engine_worker_timeout_kills_separate_session_descendant(tmp_path, monkeypatch):
|
||||
store = FakeS3()
|
||||
item = task()
|
||||
item.update(image="image@sha256:abc", deadline=0.5)
|
||||
item["pin"]["units"] = [{"unit": "u", "supported": True, "reason": ""}]
|
||||
(tmp_path / "app.py").write_text("print(1)\n")
|
||||
store.put(item["input"], service.pack(tmp_path), {})
|
||||
monkeypatch.setenv("ENGINE_IMAGE", item["image"])
|
||||
pid_file = tmp_path / "escaped-pid"
|
||||
|
||||
def fake_worker(*_):
|
||||
child = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"], start_new_session=True)
|
||||
pid_file.write_text(str(child.pid))
|
||||
time.sleep(30)
|
||||
|
||||
monkeypatch.setattr(service, "worker", fake_worker)
|
||||
result = service.execute_task(store, item)
|
||||
assert result["units"][0]["status"] == "timeout"
|
||||
assert pid_file.exists(), "instrument: fake worker never launched its descendant"
|
||||
pid = int(pid_file.read_text())
|
||||
assert not Path(f"/proc/{pid}").exists(), "escaped session survived the task deadline"
|
||||
+292
-1
@@ -534,7 +534,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 = []
|
||||
@@ -1161,3 +1169,286 @@ 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
|
||||
|
||||
|
||||
@pytest.mark.parametrize("outcome", ["success", "timeout", "failure"])
|
||||
def test_engine_queue_pin_persistence_resume_and_progress(tmp_path, monkeypatch, outcome):
|
||||
from engine_trace_queue import Dispatcher
|
||||
from engine_worker_service import Queue
|
||||
from test_engine_queue import FakeS3
|
||||
|
||||
store = FakeS3()
|
||||
args = k8s_args(tmp_path, executor="queue", queue_timeout=2, image="image@sha256:abc")
|
||||
plan = trace.make_plan([unit()], tmp_path)
|
||||
|
||||
class Inputs:
|
||||
def materialize(self, pin, target):
|
||||
target.mkdir()
|
||||
(target / "app.py").write_text("def run(): pass")
|
||||
return {"files": 1, "bytes": 15}
|
||||
|
||||
def execute(client, task):
|
||||
return {
|
||||
**task["pin"],
|
||||
"done": True,
|
||||
"container_exit": 0,
|
||||
"units": [worker.unit_record(u, "asked-nothing") for u in task["pin"]["units"]],
|
||||
}
|
||||
|
||||
queue = Queue(store, "pod-worker")
|
||||
if outcome == "timeout":
|
||||
ticks = iter([0, 3])
|
||||
args.dispatcher = Dispatcher(args, store, clock=lambda: next(ticks))
|
||||
else:
|
||||
args.dispatcher = Dispatcher(args, store, pause=lambda _: queue.process_one(execute))
|
||||
if outcome == "failure":
|
||||
|
||||
def fail(*_):
|
||||
raise RuntimeError("credential-must-not-appear")
|
||||
|
||||
monkeypatch.setattr(store, "put", fail)
|
||||
monkeypatch.setattr(trace, "prepare_questions", lambda *_: {"a": {}})
|
||||
monkeypatch.setattr(subprocess, "run", lambda *_a, **_k: pytest.fail("queue must not launch Docker/kubectl"))
|
||||
assert trace.run_plan(plan, None, {"image": args.image}, Inputs(), args) is (outcome != "success")
|
||||
result = json.loads((tmp_path / trace.pin_name(plan[0])).read_text())
|
||||
assert result["executor"] == "queue"
|
||||
assert result["units"][0]["status"] == {"success": "asked-nothing", "timeout": "timeout", "failure": "failed"}[outcome]
|
||||
assert "credential-must-not-appear" not in json.dumps(result)
|
||||
if outcome == "success":
|
||||
assert result["worker"] == "pod-worker"
|
||||
assert trace.make_plan([unit()], tmp_path) == []
|
||||
else:
|
||||
assert trace.make_plan([unit()], tmp_path, rerun_status=("failed", "timeout"))
|
||||
progress = json.loads((tmp_path / "progress.json").read_text())
|
||||
assert progress["pins_done"] == 1 and progress["pins_in_flight"] == 0
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
"""Synthetic, offline result snapshots; no analyzer, model or corpus access."""
|
||||
|
||||
import hashlib
|
||||
import importlib
|
||||
import json
|
||||
from pathlib import Path
|
||||
@@ -159,3 +160,31 @@ def test_bad_inputs_fail_loudly(audit_module, tmp_path):
|
||||
write_result(tmp_path, audit_module, [{**result(population[0], "path"), "label": 0}])
|
||||
with pytest.raises(ValueError, match="metadata mismatch"):
|
||||
run(audit_module, tmp_path, population)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("done", [False, True])
|
||||
def test_record_uses_hashes_without_repository_identity(audit_module, tmp_path, done):
|
||||
repo = "https://github.com/owner/name"
|
||||
population = [{**u, "group": "owner/name", "repo": repo, "source": "population-v1"} for u in units(1)]
|
||||
root = tmp_path / "github.com" / "owner" / "name"
|
||||
root.mkdir(parents=True)
|
||||
paths = []
|
||||
for i, u in enumerate(population):
|
||||
row = {**result(u, "path"), "source": "population-v1"}
|
||||
path = write_result(root, audit_module, [row], f"name-{1 - i}", done=done)
|
||||
record = json.loads(path.read_text())
|
||||
record.update(repo=repo, pin="owner/name")
|
||||
path.write_text(json.dumps(record))
|
||||
paths.append(path)
|
||||
report = run(audit_module, root, population)
|
||||
serialized = json.dumps(report)
|
||||
assert repo not in serialized and "owner/name" not in serialized
|
||||
assert "files" not in report["inputs"]
|
||||
inventory = sorted((p.name, p.stat().st_size) for p in paths)
|
||||
assert report["inputs"]["files_sha256"] == hashlib.sha256(json.dumps(inventory, separators=(",", ":")).encode()).hexdigest()
|
||||
assert report["inputs"]["files_read"] == 2
|
||||
assert report["inputs"]["bytes_read"] == sum(size for _, size in inventory)
|
||||
assert report["inputs"]["result_rows_read"] == (2 if done else 0)
|
||||
assert report["by_source"]["population-v1"]["rows"] == 2
|
||||
assert [r["unit"] for r in report["units"]] == [u["unit"] for u in population]
|
||||
assert [r["pair"] for r in report["units"]] == [u["pair"] for u in population]
|
||||
|
||||
@@ -150,7 +150,9 @@ def test_learn_slice_audit_reads_and_never_writes(audit_module: ModuleType, monk
|
||||
|
||||
raw = b" 10 def f(x):\n 11 execute(x)\n"
|
||||
digest = hashlib.sha256(raw).hexdigest()
|
||||
example = replace(corpus(groups=1)[0], excerpt_sha=digest, source="pairs", language="python", family="injection")
|
||||
example = replace(
|
||||
corpus(groups=1)[0], excerpt_sha=digest, source="pairs", language="python", family="injection", group="github.com/owner/name"
|
||||
)
|
||||
monkeypatch.setattr(audit_module, "load_examples", lambda store: [example])
|
||||
|
||||
class ReadOnly:
|
||||
@@ -161,6 +163,8 @@ def test_learn_slice_audit_reads_and_never_writes(audit_module: ModuleType, monk
|
||||
units = [{"family": example.family, "label": example.label, "group": example.group, "unit": example.id, "pair": None, "fold": "f"}]
|
||||
report = audit_module.audit(ReadOnly(), units, expected_rows=1, expected_pairs=0)
|
||||
assert report["example_rows_read"] == report["matched_rows"] == report["excerpts_read"] == 1
|
||||
assert "owner/name" not in json.dumps(report)
|
||||
assert report["units"][0]["unit"] == example.id
|
||||
assert report["excerpt_bytes"] == len(raw)
|
||||
assert report["coverage"]["injection"]["flow"] == 1
|
||||
|
||||
|
||||
Reference in New Issue
Block a user