mirror of
https://github.com/norandom/OpenUltraSAST.git
synced 2026-10-16 03:49:26 +00:00
engine: worker Deployment from the image (S3 queue with leases, dispatcher --executor queue, engine-image workflow, scoped RustFS policy); merge main
Implemented with Codex. Cluster deploy pending the k3s kubeconfig and the scoped account. Co-Authored-By: Claude Opus 5.5 <[email protected]>
This commit is contained in:
Marius Ciepluch
@@ -0,0 +1,315 @@
|
||||
"""Image-resident queue workers. Requires strongly consistent S3 GET/PUT/LIST.
|
||||
|
||||
Per-task Lamport bakery tickets use only ordinary object operations. A choosing
|
||||
marker precedes ticket selection; contenders wait for choosing peers and lower
|
||||
(ticket, worker) pairs. Late arrivals cannot displace a running owner. Leases
|
||||
assume bounded clock skew (less than MARGIN / 2) and a worker killed by its deadline.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import contextlib
|
||||
import ctypes
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
import multiprocessing
|
||||
import os
|
||||
import signal
|
||||
import socket
|
||||
import tarfile
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from pathlib import Path
|
||||
|
||||
from engine_trace_worker import unit_record, worker
|
||||
|
||||
from openultrasast.plane.memory import S3Client, s3_settings
|
||||
|
||||
PENDING = "engine-queue/pending/"
|
||||
CLAIMED = "engine-queue/claimed/"
|
||||
RESULTS = "engine-results/"
|
||||
MARGIN = 60
|
||||
LOG = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def client_from_env():
|
||||
settings = s3_settings(os.environ)
|
||||
bucket = os.environ.get("S3_BUCKET")
|
||||
if not bucket:
|
||||
raise ValueError("S3_BUCKET is required")
|
||||
return S3Client(bucket=bucket, **settings)
|
||||
|
||||
|
||||
def read(client, key):
|
||||
value = client.get(key)
|
||||
return json.loads(value[0]) if value else None
|
||||
|
||||
|
||||
def put(client, key, value):
|
||||
client.put(key, json.dumps(value).encode(), {})
|
||||
|
||||
|
||||
def result_key(task):
|
||||
return f"{RESULTS}{task['run']}/{task['id']}.json"
|
||||
|
||||
|
||||
def pack(root):
|
||||
stream = io.BytesIO()
|
||||
with tarfile.open(fileobj=stream, mode="w") as archive:
|
||||
archive.add(root, arcname=".")
|
||||
return stream.getvalue()
|
||||
|
||||
|
||||
def unpack(data, root):
|
||||
with tarfile.open(fileobj=io.BytesIO(data)) as archive:
|
||||
# Source snapshots need only regular files and directories, never links/devices.
|
||||
for member in archive.getmembers():
|
||||
target = (root / member.name).resolve()
|
||||
if not target.is_relative_to(root.resolve()) or not (member.isfile() or member.isdir()):
|
||||
raise ValueError("unsafe source archive")
|
||||
archive.extractall(root, filter="data")
|
||||
|
||||
|
||||
class Queue:
|
||||
def __init__(self, client, identity=None, *, now=time.time, pause=time.sleep, stop=None):
|
||||
self.client = client
|
||||
self.identity = identity or f"{socket.gethostname()}-{uuid.uuid4().hex}"
|
||||
self.now, self.pause = now, pause
|
||||
self.stop = stop or threading.Event()
|
||||
|
||||
def markers(self, task_id=None):
|
||||
result = []
|
||||
for key in self.client.keys(CLAIMED):
|
||||
if task_id is None or key.rsplit("/", 1)[-1] == task_id:
|
||||
marker = read(self.client, key)
|
||||
if marker:
|
||||
result.append((key, marker))
|
||||
return result
|
||||
|
||||
def tasks(self):
|
||||
tasks = {}
|
||||
for key in self.client.keys(PENDING):
|
||||
task = read(self.client, key)
|
||||
if task:
|
||||
tasks[task["task_id"]] = task
|
||||
# A dead worker may have removed pending already. Its task remains in the lease.
|
||||
for _, marker in self.markers():
|
||||
if marker["expires"] < self.now():
|
||||
task = marker["task"]
|
||||
tasks.setdefault(task["task_id"], task)
|
||||
return list(tasks.values())
|
||||
|
||||
def claim(self, task):
|
||||
if read(self.client, result_key(task)) is not None or self.stop.is_set():
|
||||
return None
|
||||
key = f"{CLAIMED}{self.identity}/{task['task_id']}"
|
||||
marker = dict(task=task, ticket=0, expires=self.now() + task["deadline"] + MARGIN)
|
||||
put(self.client, key, marker)
|
||||
try:
|
||||
live = [m for _, m in self.markers(task["task_id"]) if m["expires"] > self.now()]
|
||||
marker["ticket"] = 1 + max((m["ticket"] for m in live), default=0)
|
||||
put(self.client, key, marker)
|
||||
while not self.stop.is_set():
|
||||
# Do not start once the execution budget no longer fits within our lease.
|
||||
if self.now() >= marker["expires"] - task["deadline"] - MARGIN / 2:
|
||||
break
|
||||
peers = [(k, m) for k, m in self.markers(task["task_id"]) if k != key and m["expires"] > self.now()]
|
||||
if any(m["ticket"] == 0 for _, m in peers):
|
||||
self.pause(0.1)
|
||||
continue
|
||||
if any((m["ticket"], k) < (marker["ticket"], key) for k, m in peers):
|
||||
break # loser backs off and can try another pending task
|
||||
if read(self.client, key) != marker or read(self.client, result_key(task)) is not None:
|
||||
break
|
||||
self.client.delete(PENDING + task["task_id"])
|
||||
return key
|
||||
except Exception:
|
||||
self.client.delete(key)
|
||||
raise
|
||||
self.client.delete(key)
|
||||
return None
|
||||
|
||||
def process_one(self, execute=None):
|
||||
for task in self.tasks():
|
||||
key = self.claim(task)
|
||||
if key is None:
|
||||
continue
|
||||
try:
|
||||
lease = read(self.client, key)
|
||||
if lease is None:
|
||||
continue
|
||||
budget = min(task["deadline"], lease["expires"] - self.now() - MARGIN / 2)
|
||||
if budget <= 0:
|
||||
continue # Leave the expired lease's task available to recovery.
|
||||
result = (execute or execute_task)(self.client, {**task, "runtime_budget": budget})
|
||||
result.update(executor="queue", worker=self.identity, pod=socket.gethostname(), image=os.environ.get("ENGINE_IMAGE", ""))
|
||||
put(self.client, result_key(task), result)
|
||||
# Publish the result before removing either the input or the claim.
|
||||
self.client.delete(task["input"])
|
||||
for old_key, old in self.markers(task["task_id"]):
|
||||
if old_key == key or old["expires"] < self.now():
|
||||
self.client.delete(old_key)
|
||||
return True
|
||||
except Exception:
|
||||
# Keep the lease and task for recovery, without logging credential-bearing errors.
|
||||
LOG.error("task failed; lease retained for recovery")
|
||||
return False
|
||||
return False
|
||||
|
||||
|
||||
def _child(pin, root, out, deadline, question_deadline):
|
||||
os.setsid()
|
||||
signal.signal(signal.SIGTERM, signal.SIG_DFL)
|
||||
worker(pin, root, out, deadline, question_deadline)
|
||||
|
||||
|
||||
def adopt_orphans():
|
||||
# Linux PR_SET_CHILD_SUBREAPER: escaped Joern sessions reparent here, not
|
||||
# outside the service. There is exactly one task process tree per worker.
|
||||
if ctypes.CDLL(None, use_errno=True).prctl(36, 1, 0, 0, 0) != 0:
|
||||
raise OSError("cannot enable task descendant cleanup")
|
||||
|
||||
|
||||
def kill_children():
|
||||
"""Freeze parents before walking children, including new-session JVMs."""
|
||||
seen = set()
|
||||
|
||||
def freeze(parent):
|
||||
for entry in Path("/proc").iterdir():
|
||||
if not entry.name.isdigit() or int(entry.name) in seen:
|
||||
continue
|
||||
try:
|
||||
# comm can contain spaces or parentheses. Fields after it begin at state.
|
||||
fields = (entry / "stat").read_text().rsplit(")", 1)[1].split()
|
||||
if int(fields[1]) != parent:
|
||||
continue
|
||||
pid = int(entry.name)
|
||||
os.kill(pid, signal.SIGSTOP)
|
||||
seen.add(pid)
|
||||
freeze(pid)
|
||||
except (FileNotFoundError, ProcessLookupError, PermissionError):
|
||||
continue
|
||||
|
||||
# A child that exited during discovery can reparent descendants to this
|
||||
# subreaper. Repeat until every remaining descendant is frozen.
|
||||
while True:
|
||||
count = len(seen)
|
||||
freeze(os.getpid())
|
||||
if len(seen) == count:
|
||||
break
|
||||
for pid in seen:
|
||||
with contextlib.suppress(ProcessLookupError):
|
||||
os.kill(pid, signal.SIGKILL)
|
||||
|
||||
|
||||
def execute_task(client, task):
|
||||
adopt_orphans()
|
||||
pin = task["pin"]
|
||||
started = time.monotonic()
|
||||
budget = task.get("runtime_budget", task["deadline"])
|
||||
status, reason, code = "failed", "worker produced no result", None
|
||||
with tempfile.TemporaryDirectory(prefix="engine-task-") as scratch:
|
||||
root, out = Path(scratch) / "case", Path(scratch) / "result.json"
|
||||
root.mkdir()
|
||||
try:
|
||||
if task["deadline"] > 900 or task["deadline"] <= 0:
|
||||
raise ValueError("deadline exceeds deployment grace")
|
||||
if task.get("image") != os.environ.get("ENGINE_IMAGE"):
|
||||
raise ValueError("task image differs from deployed image")
|
||||
source = client.get(task["input"])
|
||||
if source is None:
|
||||
raise ValueError("missing source")
|
||||
unpack(source[0], root)
|
||||
# The shipped worker also checks frontend input and graph output for every language.
|
||||
files = [p for p in root.rglob("*") if p.is_file()]
|
||||
if not files:
|
||||
raise ValueError("empty source")
|
||||
LOG.info("input proof: first file bytes=%d", len(files[0].read_bytes()))
|
||||
if time.monotonic() - started >= budget:
|
||||
status, reason = "timeout", "task input deadline"
|
||||
raise TimeoutError(reason)
|
||||
process = multiprocessing.get_context("fork").Process(target=_child, args=(pin, root, out, budget, task["question_deadline"]))
|
||||
process.start()
|
||||
try:
|
||||
process.join(max(0, budget - (time.monotonic() - started)))
|
||||
if process.is_alive():
|
||||
status, reason = "timeout", "task deadline"
|
||||
else:
|
||||
reason = f"worker exit {process.exitcode}"
|
||||
code = process.exitcode
|
||||
finally:
|
||||
# Joern creates separate sessions. Kill the whole descendant tree,
|
||||
# not merely the Python wrapper's process group.
|
||||
kill_children()
|
||||
process.join()
|
||||
# Reap adopted JVMs and their launchers after the direct child.
|
||||
while True:
|
||||
try:
|
||||
os.waitpid(-1, 0)
|
||||
except ChildProcessError:
|
||||
break
|
||||
record = json.loads(out.read_text()) if out.exists() else {**pin, "units": []}
|
||||
except Exception:
|
||||
record = {**pin, "units": []}
|
||||
reason = "worker input or execution failed"
|
||||
if not record.get("done") or code != 0:
|
||||
previous = {u["unit"]: u for u in record["units"]}
|
||||
record["units"] = [
|
||||
previous[u["unit"]]
|
||||
if code is None and previous.get(u["unit"], {}).get("status") in {"path", "asked-nothing"}
|
||||
else unit_record(u, status if u["supported"] else "unsupported", reason if u["supported"] else u["reason"])
|
||||
for u in pin["units"]
|
||||
]
|
||||
record.update(done=True, container_exit=code, seconds=time.monotonic() - started)
|
||||
return record
|
||||
|
||||
|
||||
def status(client):
|
||||
return {
|
||||
"pending": len(client.keys(PENDING)),
|
||||
"claimed": len(client.keys(CLAIMED)),
|
||||
"done": len([k for k in client.keys(RESULTS) if k.endswith(".json")]),
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--port", type=int, default=8080)
|
||||
args = parser.parse_args()
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
stop = threading.Event()
|
||||
signal.signal(signal.SIGTERM, lambda *_: stop.set())
|
||||
signal.signal(signal.SIGINT, lambda *_: stop.set())
|
||||
client = client_from_env()
|
||||
queue = Queue(client, stop=stop)
|
||||
|
||||
class Health(BaseHTTPRequestHandler):
|
||||
def do_GET(self):
|
||||
self.send_response(200 if self.path == "/healthz" and not stop.is_set() else 503)
|
||||
self.end_headers()
|
||||
|
||||
def log_message(self, *_):
|
||||
pass
|
||||
|
||||
server = ThreadingHTTPServer(("0.0.0.0", args.port), Health)
|
||||
threading.Thread(target=server.serve_forever, daemon=True).start()
|
||||
try:
|
||||
while not stop.is_set():
|
||||
try:
|
||||
if not queue.process_one():
|
||||
stop.wait(2)
|
||||
except Exception:
|
||||
LOG.error("queue unavailable; retrying")
|
||||
stop.wait(5)
|
||||
finally:
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user