mirror of
https://github.com/norandom/OpenUltraSAST.git
synced 2026-11-13 09:48:18 +00:00
feat(detection-core): migrate benchmark readiness and admission consumers
This commit is contained in:
@@ -14,6 +14,8 @@ import urllib.error
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
|
||||
from benchmarks.valid.artifact_contract import write_benchmark_artifact
|
||||
|
||||
|
||||
class TaskEntryError(RuntimeError):
|
||||
"""A task-entry diagnostic containing only locally controlled text."""
|
||||
@@ -91,68 +93,94 @@ def failure_record(pin, reason):
|
||||
}
|
||||
|
||||
|
||||
def main(*, work_root=Path("/tmp"), worker_script=None, readiness_budget=60):
|
||||
pin = {}
|
||||
failed = False
|
||||
def _read_worker_result(path, pin):
|
||||
"""Validate the engine worker's raw handoff before benchmark-owned persistence."""
|
||||
try:
|
||||
value = json.loads(path.read_text())
|
||||
except (OSError, json.JSONDecodeError) as error:
|
||||
raise TaskEntryError("worker result is unreadable") from error
|
||||
if not isinstance(value, dict) or {"schema_version", "artifact_kind", "artifact_digest"}.intersection(value):
|
||||
raise TaskEntryError("worker result has an unsupported raw schema")
|
||||
units = value.get("units")
|
||||
expected = [unit.get("unit") for unit in pin.get("units", [])]
|
||||
if value.get("done") is not True or not isinstance(units, list):
|
||||
raise TaskEntryError("worker did not write a completed result")
|
||||
actual = [unit.get("unit") for unit in units if isinstance(unit, dict)]
|
||||
if len(actual) != len(units) or actual != expected:
|
||||
raise TaskEntryError("worker result units do not match the prepared pin")
|
||||
return value
|
||||
|
||||
|
||||
def _produce_result(work_root, worker_script, readiness_budget, state):
|
||||
output = work_root / "out"
|
||||
result = output / "result.json"
|
||||
step = "prepare output"
|
||||
output.mkdir(parents=True, exist_ok=True)
|
||||
result.unlink(missing_ok=True)
|
||||
questions = work_root / "questions.json"
|
||||
state["step"] = "GET questions"
|
||||
download(os.environ["QUESTIONS_URL"], questions, readiness_budget=readiness_budget)
|
||||
state["step"] = "parse questions"
|
||||
pin = json.loads(questions.read_text())
|
||||
if not isinstance(pin, dict):
|
||||
raise ValueError("questions object must contain a prepared pin")
|
||||
state["pin"] = pin
|
||||
source, case = work_root / "source.tar", work_root / "case"
|
||||
state["step"] = "GET source"
|
||||
download(os.environ["SOURCE_URL"], source)
|
||||
state["step"] = "unpack source"
|
||||
case.mkdir(parents=True, exist_ok=True)
|
||||
unpack(source, case)
|
||||
deadline = int(os.environ.get("DEADLINE", "900"))
|
||||
command = _worker_command(worker_script, questions, case, result, deadline)
|
||||
state["step"] = "run worker"
|
||||
completed = subprocess.run(command, timeout=deadline + 5, check=False)
|
||||
if completed.returncode:
|
||||
raise TaskEntryError(f"worker exited {completed.returncode}")
|
||||
state["step"] = "read worker result"
|
||||
worker_result = _read_worker_result(result, pin) if result.is_file() else {}
|
||||
if not worker_result:
|
||||
raise TaskEntryError("worker did not write a completed result")
|
||||
write_benchmark_artifact(result, "engine-task-result", worker_result, replace=True)
|
||||
return output, result
|
||||
|
||||
|
||||
def _worker_command(worker_script, questions, case, result, deadline):
|
||||
return [
|
||||
sys.executable,
|
||||
str(worker_script or Path(__file__).with_name("engine_trace_worker.py")),
|
||||
"--worker",
|
||||
str(questions),
|
||||
"--root",
|
||||
str(case),
|
||||
"--output",
|
||||
str(result),
|
||||
"--heap-profile",
|
||||
"cluster",
|
||||
"--deadline",
|
||||
str(deadline),
|
||||
"--question-deadline",
|
||||
os.environ.get("QUESTION_DEADLINE", "120"),
|
||||
]
|
||||
|
||||
|
||||
def _failure_detail(error):
|
||||
if isinstance(error, urllib.error.HTTPError):
|
||||
return f"HTTP {error.code}"
|
||||
return str(error) if isinstance(error, TaskEntryError) else type(error).__name__
|
||||
|
||||
|
||||
def _persist_failure(work_root, state, error):
|
||||
output, result = work_root / "out", work_root / "out/result.json"
|
||||
reason = f"{state['step']}: {_failure_detail(error)}"
|
||||
try:
|
||||
output.mkdir(parents=True, exist_ok=True)
|
||||
result.unlink(missing_ok=True)
|
||||
questions = work_root / "questions.json"
|
||||
step = "GET questions"
|
||||
download(os.environ["QUESTIONS_URL"], questions, readiness_budget=readiness_budget)
|
||||
step = "parse questions"
|
||||
pin = json.loads(questions.read_text())
|
||||
if not isinstance(pin, dict):
|
||||
pin = {}
|
||||
raise ValueError("questions object must contain a prepared pin")
|
||||
source = work_root / "source.tar"
|
||||
step = "GET source"
|
||||
download(os.environ["SOURCE_URL"], source)
|
||||
step = "unpack source"
|
||||
case = work_root / "case"
|
||||
case.mkdir(parents=True, exist_ok=True)
|
||||
unpack(source, case)
|
||||
step = "configure worker"
|
||||
deadline = int(os.environ.get("DEADLINE", "900"))
|
||||
command = [
|
||||
sys.executable,
|
||||
str(worker_script or Path(__file__).with_name("engine_trace_worker.py")),
|
||||
"--worker",
|
||||
str(questions),
|
||||
"--root",
|
||||
str(case),
|
||||
"--output",
|
||||
str(result),
|
||||
"--heap-profile",
|
||||
"cluster",
|
||||
"--deadline",
|
||||
str(deadline),
|
||||
"--question-deadline",
|
||||
os.environ.get("QUESTION_DEADLINE", "120"),
|
||||
]
|
||||
step = "run worker"
|
||||
completed = subprocess.run(command, timeout=deadline + 5, check=False)
|
||||
if completed.returncode:
|
||||
raise TaskEntryError(f"worker exited {completed.returncode}")
|
||||
step = "read worker result"
|
||||
if not result.is_file() or not json.loads(result.read_text()).get("done"):
|
||||
raise TaskEntryError("worker did not write a completed result")
|
||||
except Exception as exc:
|
||||
failed = True
|
||||
# Avoid including exception text containing a presigned URL in durable output.
|
||||
if isinstance(exc, urllib.error.HTTPError):
|
||||
detail = f"HTTP {exc.code}"
|
||||
else:
|
||||
detail = str(exc) if isinstance(exc, TaskEntryError) else type(exc).__name__
|
||||
reason = f"{step}: {detail}"
|
||||
try:
|
||||
output.mkdir(parents=True, exist_ok=True)
|
||||
result.write_text(json.dumps(failure_record(pin, reason)) + "\n")
|
||||
except Exception:
|
||||
print("task entry could not write failure result", file=sys.stderr)
|
||||
write_benchmark_artifact(result, "engine-task-failure", failure_record(state["pin"], reason), replace=True)
|
||||
except Exception:
|
||||
print("task entry could not write failure result", file=sys.stderr)
|
||||
return output
|
||||
|
||||
|
||||
def _pack_result(work_root, output, readiness_budget, failed):
|
||||
step = "pack result"
|
||||
try:
|
||||
archive = work_root / "result.tar"
|
||||
@@ -166,6 +194,17 @@ def main(*, work_root=Path("/tmp"), worker_script=None, readiness_budget=60):
|
||||
except Exception as exc:
|
||||
detail = f"HTTP {exc.code}" if isinstance(exc, urllib.error.HTTPError) else type(exc).__name__
|
||||
print(f"{step}: {detail}; dispatcher deadline remains authoritative", file=sys.stderr)
|
||||
|
||||
|
||||
def main(*, work_root=Path("/tmp"), worker_script=None, readiness_budget=60):
|
||||
state = {"pin": {}, "step": "prepare output"}
|
||||
failed = False
|
||||
try:
|
||||
output, _result = _produce_result(work_root, worker_script, readiness_budget, state)
|
||||
except Exception as error:
|
||||
failed = True
|
||||
output = _persist_failure(work_root, state, error)
|
||||
_pack_result(work_root, output, readiness_budget, failed)
|
||||
return 0
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user