mirror of
https://github.com/norandom/OpenUltraSAST.git
synced 2026-10-16 19:59:26 +00:00
plane(1.2): literals out of doctor, egress, router, generate -- the profile is handed in
doctor(profile) checks the profile's kube context, the two namespaces behind a reachable context only, the images file's runner reference in the profile's registry (a local registry over its v2 API, a remote one via docker manifest inspect) and opens the memory store; its first line reports every configured address (Req 1.2). Egress(executable, context), open_router(context, url=...), receiver_cluster_ip(context) and Receiver(..., context) take the profile's context and build none; generate.increment/render and harvest take the profile's atespace. `ousast plane doctor|run|workspaces|harvest --profile`. The test fakes (kubectl-ate, kubectl, ax) record the context they received and the tests assert it is the profile's. grep "kind-" src/openultrasast/plane: no match. Gates: pytest 2006 passed, 42 skipped, exit 0; ruff check 0; ruff format 0; mypy 0. Co-Authored-By: Claude Fable 5.1 <[email protected]> Claude-Session: https://claude.ai/code/session_01EmM1j46qgKMh3zqD1QMwEo
This commit is contained in:
co-authored by
Claude Fable 5.1
parent
8eac1b2c69
commit
77a03bad61
+23
-10
@@ -272,16 +272,21 @@ def _main(argv: list[str] | None) -> int:
|
||||
learn_labels.add_argument("--include-title", action="store_true", help="the title-tier pairs as well (a sensitivity arm)")
|
||||
learn_labels.add_argument("--no-assumed-benign", action="store_true", help="skip mining the assumed-benign history source")
|
||||
add_learn_commands(learn_sub)
|
||||
plane = subparsers.add_parser("plane", help="maintainer: the ax-backed service plane on this host's kind cluster")
|
||||
plane = subparsers.add_parser("plane", help="maintainer: the ax-backed service plane on the profile's cluster (ops/k8s/profiles/)")
|
||||
profile_help = "the PlaneProfile: a name under ops/k8s/profiles/ or a .toml path (default OUSAST_PLANE_PROFILE, else kind)"
|
||||
plane_sub = plane.add_subparsers(dest="plane_command", required=True)
|
||||
plane_run = plane_sub.add_parser("run", help="execute a Run manifest on ax; a rerun skips done tasks")
|
||||
plane_run.add_argument("run_manifest", type=Path)
|
||||
plane_run.add_argument("--workers", type=int, default=1, help="tasks in flight at once (serialize labels never overlap)")
|
||||
plane_run.add_argument("--ax", default="ax", help="path of the ax CLI")
|
||||
plane_run.add_argument("--profile", help=profile_help)
|
||||
plane_status = plane_sub.add_parser("status", help="per-task status and the token attribution table of a run")
|
||||
plane_status.add_argument("run")
|
||||
plane_status.add_argument("--units", action="store_true", help="per-unit rows from each task's units.jsonl")
|
||||
plane_sub.add_parser("doctor", help="check kind, Agent Substrate, the ax controller and the runner image")
|
||||
plane_doctor = plane_sub.add_parser(
|
||||
"doctor", help="check the profile's cluster, Agent Substrate, the ax controller, the runner image and the memory store"
|
||||
)
|
||||
plane_doctor.add_argument("--profile", help=profile_help)
|
||||
plane_remember = plane_sub.add_parser("remember", help="ingest a run's memory rows into the store (OUSAST_MEMORY)")
|
||||
plane_remember.add_argument("run")
|
||||
plane_remember.add_argument("--plane", type=Path, default=Path("plane"), help="where runs/<run>.yaml names population and split")
|
||||
@@ -319,9 +324,11 @@ def _main(argv: list[str] | None) -> int:
|
||||
)
|
||||
plane_harvest.add_argument("--ceiling", type=float, default=10.0, help="USD: each Run's task budgets sum to at most this")
|
||||
plane_harvest.add_argument("--runner-image", type=Path, help="re-pin the task templates to the image in this file first (ops/ax/up.sh)")
|
||||
plane_harvest.add_argument("--profile", help=profile_help)
|
||||
plane_ws = plane_sub.add_parser("workspaces", help="one Workspace manifest per case pin of a population file")
|
||||
plane_ws.add_argument("population", type=Path)
|
||||
plane_ws.add_argument("--out", type=Path, default=Path("plane/workspaces"))
|
||||
plane_ws.add_argument("--profile", help=profile_help)
|
||||
plane_ws.add_argument("--validation-set", type=Path, help="JSON {case: [[path, function], ...]}: emit that set's increment instead")
|
||||
plane_ws.add_argument("--candidates", type=Path, help="with --validation-set: per-case candidate records (independent-v2-modelsinks)")
|
||||
plane_ws.add_argument("--triage", type=Path, help="with --validation-set: recorded batched-check scans (<case>--vulnerable_a.jsonl)")
|
||||
@@ -1379,10 +1386,16 @@ def _learn_labels(args: argparse.Namespace) -> int:
|
||||
|
||||
def _plane(args: argparse.Namespace) -> int:
|
||||
from .plane import reconciler # lazy: the plane is a maintainer surface, not the scan path
|
||||
from .plane.profile import ProfileError, load_profile
|
||||
|
||||
try:
|
||||
profile = load_profile(getattr(args, "profile", None))
|
||||
except ProfileError as exc:
|
||||
print(f"plane: {exc}", file=sys.stderr)
|
||||
return 2
|
||||
if args.plane_command == "run":
|
||||
_plane_memory("seed", args.run_manifest)
|
||||
result = reconciler.run(args.run_manifest, workers=args.workers, ax=args.ax)
|
||||
result = reconciler.run(args.run_manifest, workers=args.workers, ax=args.ax, profile=profile)
|
||||
print(f"run finished: {result}")
|
||||
_plane_memory("ingest", args.run_manifest)
|
||||
return 0 if result == "done" else 1
|
||||
@@ -1397,14 +1410,14 @@ def _plane(args: argparse.Namespace) -> int:
|
||||
if args.plane_command == "alerts-engine":
|
||||
return _plane_alerts_engine(args)
|
||||
if args.plane_command == "harvest":
|
||||
return _plane_harvest(args)
|
||||
return _plane_harvest(args, profile.atespace)
|
||||
if args.plane_command == "doctor":
|
||||
checks = reconciler.doctor()
|
||||
checks = reconciler.doctor(profile)
|
||||
for name, ok, text in checks:
|
||||
print(f"[{'ok' if ok else 'FAIL'}] {name}: {text}")
|
||||
return 0 if all(ok for _, ok, _ in checks) else 1
|
||||
if args.validation_set is not None:
|
||||
return _plane_increment(args)
|
||||
return _plane_increment(args, profile.atespace)
|
||||
from .plane.workspaces import workspaces
|
||||
|
||||
written = workspaces(args.population, args.out)
|
||||
@@ -1412,7 +1425,7 @@ def _plane(args: argparse.Namespace) -> int:
|
||||
return 0
|
||||
|
||||
|
||||
def _plane_harvest(args: argparse.Namespace) -> int:
|
||||
def _plane_harvest(args: argparse.Namespace, atespace: str) -> int:
|
||||
from .plane.generate import repin_templates
|
||||
from .plane.harvest import write_harvest
|
||||
|
||||
@@ -1421,7 +1434,7 @@ def _plane_harvest(args: argparse.Namespace) -> int:
|
||||
command = f"ousast plane harvest --labels {args.labels.name} --name {args.name} --ceiling {args.ceiling:g}"
|
||||
files = write_harvest(
|
||||
args.labels, args.plane, name=args.name, command=command, catalog=args.catalog, cache=args.cache, templates=args.templates,
|
||||
ceiling=args.ceiling,
|
||||
ceiling=args.ceiling, atespace=atespace,
|
||||
) # fmt: skip
|
||||
for relative in sorted(k for k in files if k.startswith("runs/")):
|
||||
print(f"{args.plane / relative}: {files[relative].splitlines()[1].lstrip('# ')}")
|
||||
@@ -1502,7 +1515,7 @@ def _plane_memory(
|
||||
return 0
|
||||
|
||||
|
||||
def _plane_increment(args: argparse.Namespace) -> int:
|
||||
def _plane_increment(args: argparse.Namespace, atespace: str) -> int:
|
||||
from .plane.generate import LOOP_MANIFEST, Loop, increment
|
||||
|
||||
if args.candidates is None or args.triage is None:
|
||||
@@ -1528,7 +1541,7 @@ def _plane_increment(args: argparse.Namespace) -> int:
|
||||
command += "".join(f" --qualify-population {p}" for p in loop.populations)
|
||||
written = increment(
|
||||
args.population, args.validation_set, args.candidates, args.triage, plane=args.plane, run_name=args.run, command=command,
|
||||
repos=args.repos, runner_image=args.runner_image, loop=loop,
|
||||
repos=args.repos, runner_image=args.runner_image, loop=loop, atespace=atespace,
|
||||
) # fmt: skip
|
||||
print(f"{len(written)} manifests written under {args.plane} (Run {args.run})")
|
||||
return 0
|
||||
|
||||
@@ -1,15 +1,35 @@
|
||||
"""``ousast plane doctor``: four checks of this host's ax before a Run is submitted (ai-service-plane Req 3.5).
|
||||
"""``ousast plane doctor``: checks of the configured cluster before a Run is submitted (ai-service-plane Req 3.5,
|
||||
plane-on-kubernetes Req 1.2).
|
||||
|
||||
Everything is read from the :class:`~openultrasast.plane.profile.PlaneProfile`: the kube context is checked for
|
||||
reachability and, when reachable, the two namespaces for ready pods; the ``images`` file's runner reference is
|
||||
resolved in its registry (a local registry over its v2 API, a remote one through ``docker manifest inspect``); the
|
||||
memory store is opened, which verifies an S3 bucket; and the first line reports every configured address. Nothing
|
||||
here names a cluster, a registry or an address of its own: those are the profile's.
|
||||
|
||||
A sibling of ``reconciler.py`` (tasks.md task 0 allows either), which re-exports :func:`doctor`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import ipaddress
|
||||
import re
|
||||
import subprocess
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
|
||||
from .profile import PlaneProfile, ProfileError, load_profile
|
||||
|
||||
UP = "ops/ax/up.sh brings it up"
|
||||
RUNBOOK = "docs/deployment.md lists the prerequisites"
|
||||
_MANIFEST_TYPES = ", ".join(
|
||||
(
|
||||
"application/vnd.oci.image.index.v1+json",
|
||||
"application/vnd.oci.image.manifest.v1+json",
|
||||
"application/vnd.docker.distribution.manifest.list.v2+json",
|
||||
"application/vnd.docker.distribution.manifest.v2+json",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
# --- checks ----------------------------------------------------------------------------------------------
|
||||
@@ -34,18 +54,68 @@ def _pods_ready(namespace: str, context: str) -> tuple[bool, str]:
|
||||
return not bad, f"not ready: {', '.join(bad)}" if bad else f"{len(rows)} pods ready"
|
||||
|
||||
|
||||
def doctor() -> list[tuple[str, bool, str]]:
|
||||
"""Four checks of the host's ax: kind, Agent Substrate, the ax controller, the runner image (Req 3.5)."""
|
||||
context = "kind-" + (os.environ.get("KIND_CLUSTER_NAME") or "ousast")
|
||||
def _is_local(registry: str) -> bool:
|
||||
host = registry.rsplit(":", 1)[0] if re.search(r":\d+$", registry) else registry
|
||||
if host == "localhost":
|
||||
return True
|
||||
try:
|
||||
return ipaddress.ip_address(host.strip("[]")).is_loopback or ipaddress.ip_address(host.strip("[]")).is_private
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
|
||||
def _registry_has(image: str) -> tuple[bool, str]:
|
||||
"""Resolve a digest-pinned reference ``<registry>/<repository>@sha256:<hex>``: a local registry answers its v2
|
||||
API for the manifest by digest (plain HTTP, as kind's registry speaks); a remote one is asked through
|
||||
``docker manifest inspect``, which carries the operator's registry login."""
|
||||
name, _, digest = image.partition("@")
|
||||
registry, _, repository = name.partition("/")
|
||||
short = f"{repository}@{digest[:19]}..."
|
||||
if _is_local(registry):
|
||||
url = f"http://{registry}/v2/{repository}/manifests/{digest}"
|
||||
request = urllib.request.Request(url, method="HEAD", headers={"Accept": _MANIFEST_TYPES})
|
||||
try:
|
||||
with urllib.request.urlopen(request, timeout=20) as response: # noqa: S310 - the profile's registry
|
||||
found = 200 <= int(response.status) < 300
|
||||
except urllib.error.HTTPError as exc:
|
||||
return False, f"{short} missing from {registry} (HTTP {exc.code})"
|
||||
except (urllib.error.URLError, OSError) as exc:
|
||||
return False, f"{registry} unreachable ({str(exc)[:80]})"
|
||||
return found, f"{short} present in {registry}" if found else f"{short} missing from {registry}"
|
||||
ok, out = _sh("docker", "manifest", "inspect", image)
|
||||
return ok, f"{short} present in {registry}" if ok else f"{short} not resolved in {registry} ({out[:120]})"
|
||||
|
||||
|
||||
def doctor(profile: PlaneProfile | None = None) -> list[tuple[str, bool, str]]:
|
||||
"""The checks of the profile's cluster: ``(name, ok, text)`` per line, the configured addresses first."""
|
||||
profile = profile or load_profile()
|
||||
hint = UP if profile.name == "kind" else RUNBOOK
|
||||
context = profile.kube_context
|
||||
checks = [(f"profile {profile.name}", True, " ".join(f"{k}={v}" for k, v in profile.addresses().items()))]
|
||||
ok, out = _sh("kubectl", "--context", context, "cluster-info")
|
||||
checks = [("kind cluster", ok, "reachable" if ok else f"context {context} unreachable ({out[:80]}); {UP}")]
|
||||
for label, ns in (("agent substrate (ate-system)", "ate-system"), ("ax controller (ax-system)", "ax-system")):
|
||||
ok, out = _pods_ready(ns, context)
|
||||
checks.append((label, ok, out if ok else f"{out}; {UP}"))
|
||||
ok, out = _sh("curl", "-s", "localhost:5001/v2/_catalog")
|
||||
present = ok and "ousast-runner" in out
|
||||
text = "ousast-runner present" if present else f"missing from localhost:5001 ({out[:80]}); {UP} and loads it"
|
||||
checks.append(("runner image in kind registry", present, text))
|
||||
checks.append((f"cluster {context}", ok, "reachable" if ok else f"context {context} unreachable ({out[:80]}); {hint}"))
|
||||
if ok: # behind an unreachable context the namespaces cannot be told apart from it
|
||||
for label, ns in (("agent substrate (ate-system)", "ate-system"), ("ax controller (ax-system)", "ax-system")):
|
||||
ok, out = _pods_ready(ns, context)
|
||||
checks.append((label, ok, out if ok else f"{out}; {hint}"))
|
||||
try:
|
||||
images = profile.load_images()
|
||||
runner = images["runner"]
|
||||
except ProfileError as exc:
|
||||
checks.append((f"runner image in {profile.registry}", False, f"{exc}; {hint}"))
|
||||
except KeyError:
|
||||
checks.append((f"runner image in {profile.registry}", False, f"{profile.images} names no runner image; {hint}"))
|
||||
else:
|
||||
ok, out = _registry_has(runner)
|
||||
checks.append((f"runner image in {profile.registry}", ok, out if ok else f"{out}; {hint} and loads it"))
|
||||
from .memory import MemoryStoreError, open_store # lazy: boto3 only when the profile names an S3 store
|
||||
|
||||
try:
|
||||
store = open_store(profile.memory)
|
||||
except (MemoryStoreError, OSError, ValueError) as exc:
|
||||
checks.append((f"memory store {profile.memory}", False, str(exc)[:300]))
|
||||
else:
|
||||
checks.append((f"memory store {profile.memory}", True, f"{store.describe()} opens (an S3 bucket is verified)"))
|
||||
return checks
|
||||
|
||||
|
||||
|
||||
@@ -110,9 +110,11 @@ def rules_of(policy: Mapping[str, Any]) -> list[dict[str, Any]]:
|
||||
return sorted(out, key=lambda r: json.dumps(r, sort_keys=True))
|
||||
|
||||
|
||||
def receiver_cluster_ip(context: str | None = None) -> str | None:
|
||||
"""The ClusterIP of the ``ousast-receiver`` Service (``kubectl get svc``); None when kubectl cannot say."""
|
||||
context = context or "kind-" + (os.environ.get("KIND_CLUSTER_NAME") or "ousast")
|
||||
def receiver_cluster_ip(context: str) -> str | None:
|
||||
"""The ClusterIP of the ``ousast-receiver`` Service (``kubectl get svc`` in the profile's ``context``); None when
|
||||
kubectl cannot say."""
|
||||
if not context:
|
||||
raise RuntimeError("receiver_cluster_ip needs the profile's kube context")
|
||||
argv = ["kubectl", "--context", context, "-n", "ax-system", "get", "svc", "ousast-receiver", "-o", "jsonpath={.spec.clusterIP}"]
|
||||
try:
|
||||
proc = subprocess.run(argv, capture_output=True, text=True, check=False, timeout=20)
|
||||
@@ -122,8 +124,9 @@ def receiver_cluster_ip(context: str | None = None) -> str | None:
|
||||
return address if proc.returncode == 0 and _is_ip(address) else None
|
||||
|
||||
|
||||
def receiver_address(port: int) -> tuple[str, str | None]:
|
||||
"""(``OUSAST_ARTIFACT_URL``, ``OUSAST_ARTIFACT_DIAL``) for a receiver bound on ``port``.
|
||||
def receiver_address(port: int, context: str | None) -> tuple[str, str | None]:
|
||||
"""(``OUSAST_ARTIFACT_URL``, ``OUSAST_ARTIFACT_DIAL``) for a receiver bound on ``port``; ``context`` is the
|
||||
profile's kube context the Service's ClusterIP is looked up in.
|
||||
|
||||
``OUSAST_ARTIFACT_HOST`` set: the old direct URL ``http://<host>:<port>/`` and no dial (tests, a host the actor
|
||||
reaches without the gateway). Otherwise the receiver Service's name, dialled at ``OUSAST_ARTIFACT_DIAL`` or the
|
||||
@@ -135,7 +138,9 @@ def receiver_address(port: int) -> tuple[str, str | None]:
|
||||
return f"http://{host}:{port}/", None
|
||||
dial = os.environ.get("OUSAST_ARTIFACT_DIAL")
|
||||
if not dial:
|
||||
address = receiver_cluster_ip()
|
||||
if not context:
|
||||
raise RuntimeError("receiver_address needs the profile's kube context to look the receiver Service up")
|
||||
address = receiver_cluster_ip(context)
|
||||
dial = f"{address}:80" if address else None
|
||||
return f"http://{RECEIVER_HOST}/", dial
|
||||
|
||||
@@ -143,9 +148,11 @@ def receiver_address(port: int) -> tuple[str, str | None]:
|
||||
class Egress:
|
||||
"""``kubectl ate`` for one actor's egress policy; the executable is injectable so tests use a fake."""
|
||||
|
||||
def __init__(self, executable: str | None = None, context: str | None = None) -> None:
|
||||
def __init__(self, executable: str | None, context: str) -> None:
|
||||
if not context:
|
||||
raise ValueError("Egress needs the profile's kube context; it never builds one")
|
||||
self.executable = executable or os.environ.get("OUSAST_KUBECTL_ATE") or "kubectl-ate"
|
||||
self.context = context or "kind-" + (os.environ.get("KIND_CLUSTER_NAME") or "ousast")
|
||||
self.context = context
|
||||
|
||||
def _run(self, *args: str, stdin: str | None = None) -> subprocess.CompletedProcess[str]:
|
||||
argv = [self.executable, "--context", self.context, *args]
|
||||
@@ -214,12 +221,12 @@ class Receiver(ThreadingHTTPServer):
|
||||
|
||||
daemon_threads = True
|
||||
|
||||
def __init__(self, run_name: str, base: Path, port: int) -> None:
|
||||
def __init__(self, run_name: str, base: Path, port: int, context: str | None = None) -> None:
|
||||
super().__init__(("0.0.0.0", port), _Handler)
|
||||
self.run_name, self.base, self.running, self.delivered = run_name, base, set[str](), set[str]()
|
||||
self.inputs: dict[str, frozenset[str]] = {}
|
||||
self.lock = threading.Lock()
|
||||
self.url, self.dial = receiver_address(self.server_address[1])
|
||||
self.url, self.dial = receiver_address(self.server_address[1], context)
|
||||
|
||||
def _task(self, headers: Any) -> str | None:
|
||||
"""The ``X-Ousast-Task`` of a request naming this run and a task that is running now, else None."""
|
||||
|
||||
@@ -82,7 +82,7 @@ TRIAGE_BASIS = (
|
||||
RANGES_BASIS = "git diff -U0 <vulnerable> <fixed>, old side: benchmarks/independent/evaluate.hunks(case, 'old')"
|
||||
HEADROOM = 1.5
|
||||
CALLS_PER_HUNT = MAX_STEPS + 2 # every tool step, the answer, one empty-content retry
|
||||
ATESPACE = "default"
|
||||
ATESPACE = "default" # the default when no profile is given; ``increment`` takes ``profile.atespace``
|
||||
CASE_PATH = "/workspace/case"
|
||||
INPUTS_PATH = "/workspace/inputs"
|
||||
FIXED_PATH = "/workspace/fixed"
|
||||
@@ -236,10 +236,10 @@ def _case_record(item: CaseInputs) -> dict[str, Any]:
|
||||
return record
|
||||
|
||||
|
||||
def _workspaces(item: CaseInputs) -> tuple[dict[str, Any], dict[str, Any]]:
|
||||
def _workspaces(item: CaseInputs, atespace: str = ATESPACE) -> tuple[dict[str, Any], dict[str, Any]]:
|
||||
case = item.case
|
||||
git = [{"name": "repo", "repo": case["repo"], "dir": "repo", "depth": 1}]
|
||||
pinned = _doc("Workspace", _ax_name(case["id"], "vulnerable"), {"git": git}, ATESPACE)
|
||||
pinned = _doc("Workspace", _ax_name(case["id"], "vulnerable"), {"git": git}, atespace)
|
||||
pinned["metadata"]["annotations"] = {GIT_COMMITS_ANNOTATION: f"repo={case['vulnerable']}"}
|
||||
files = {
|
||||
"candidates.json": _json({"id": case["id"], "family": case["family"], "candidates": [list(c) for c in item.kept]}),
|
||||
@@ -249,7 +249,7 @@ def _workspaces(item: CaseInputs) -> tuple[dict[str, Any], dict[str, Any]]:
|
||||
{"candidates": [list(c) for c in item.candidates], "kept": [list(c) for c in item.kept], "triage": dict(item.triage)}
|
||||
),
|
||||
}
|
||||
inputs = _doc("Workspace", _ax_name(case["id"], "inputs"), {"files": [{"path": k, "content": v} for k, v in files.items()]}, ATESPACE)
|
||||
inputs = _doc("Workspace", _ax_name(case["id"], "inputs"), {"files": [{"path": k, "content": v} for k, v in files.items()]}, atespace)
|
||||
return pinned, inputs
|
||||
|
||||
|
||||
@@ -339,11 +339,11 @@ def _remember(item: CaseInputs, templates: Mapping[str, Task], population: str,
|
||||
return task, entry
|
||||
|
||||
|
||||
def _alerts(item: CaseInputs, templates: Mapping[str, Task]) -> tuple[dict, dict, dict]:
|
||||
def _alerts(item: CaseInputs, templates: Mapping[str, Task], atespace: str = ATESPACE) -> tuple[dict, dict, dict]:
|
||||
"""(fixed-pin Workspace, Task, Run entry) of the case's `alerts`: quick-mode rules on both pins."""
|
||||
case = item.case
|
||||
git = [{"name": "repo", "repo": case["repo"], "dir": "repo", "depth": 1}]
|
||||
fixed = _doc("Workspace", _ax_name(case["id"], "fixed"), {"git": git}, ATESPACE)
|
||||
fixed = _doc("Workspace", _ax_name(case["id"], "fixed"), {"git": git}, atespace)
|
||||
fixed["metadata"]["annotations"] = {GIT_COMMITS_ANNOTATION: f"repo={case['fixed']}"}
|
||||
bindings = [(_ax_name(case["id"], "vulnerable"), CASE_PATH), (fixed["metadata"]["name"], FIXED_PATH)]
|
||||
bindings.append((_ax_name(case["id"], "inputs"), INPUTS_PATH))
|
||||
@@ -358,10 +358,12 @@ def _alerts(item: CaseInputs, templates: Mapping[str, Task]) -> tuple[dict, dict
|
||||
return fixed, task, entry
|
||||
|
||||
|
||||
def _loop(cases: Sequence[CaseInputs], templates: Mapping[str, Task], loop: Loop) -> tuple[dict, list[dict], list[dict]]:
|
||||
def _loop(
|
||||
cases: Sequence[CaseInputs], templates: Mapping[str, Task], loop: Loop, atespace: str = ATESPACE
|
||||
) -> tuple[dict, list[dict], list[dict]]:
|
||||
"""(project Workspace, Tasks, Run entries) of the loop's singletons: snapshot -> measure -> propose -> improve."""
|
||||
git = [{"name": "repo", "repo": loop.repo, "dir": "repo", "depth": 1}]
|
||||
project = _doc("Workspace", _ax_name("openultrasast", loop.commit[:12]), {"git": git}, ATESPACE)
|
||||
project = _doc("Workspace", _ax_name("openultrasast", loop.commit[:12]), {"git": git}, atespace)
|
||||
project["metadata"]["annotations"] = {GIT_COMMITS_ANNOTATION: f"repo={loop.commit}"}
|
||||
bound = [(project["metadata"]["name"], PROJECT_PATH)]
|
||||
env = {
|
||||
@@ -399,6 +401,7 @@ def render(
|
||||
population: str = "",
|
||||
split: str = "",
|
||||
loop: Loop | None = None,
|
||||
atespace: str = ATESPACE,
|
||||
) -> dict[str, str]:
|
||||
"""Relative path under the plane root -> file text, for every generated manifest. ``population`` and ``split``
|
||||
(the population file's and the set's stems) become the Run's annotations, which the memory rows carry. Each
|
||||
@@ -411,7 +414,7 @@ def render(
|
||||
tiebreak: list[dict[str, Any]] = []
|
||||
for item in cases:
|
||||
case_id = item.case["id"]
|
||||
pinned, inputs = _workspaces(item)
|
||||
pinned, inputs = _workspaces(item, atespace)
|
||||
for doc in (pinned, inputs):
|
||||
out[f"workspaces/{doc['metadata']['name']}.yaml"] = _dump(header, [doc])
|
||||
repo = [(pinned["metadata"]["name"], CASE_PATH), (inputs["metadata"]["name"], INPUTS_PATH)]
|
||||
@@ -451,7 +454,7 @@ def render(
|
||||
entries.extend(tiebreak) # appended: the first Run's entries stay byte-identical, so a rerun keeps their state
|
||||
for item in cases: # appended after the tie-break for the same reason
|
||||
if loop is not None:
|
||||
fixed, alerts_task, alerts_entry = _alerts(item, templates)
|
||||
fixed, alerts_task, alerts_entry = _alerts(item, templates, atespace)
|
||||
out[f"workspaces/{fixed['metadata']['name']}.yaml"] = _dump(header, [fixed])
|
||||
tasks.append(alerts_task)
|
||||
entries.append(alerts_entry)
|
||||
@@ -462,7 +465,7 @@ def render(
|
||||
tasks.append(remember_task)
|
||||
entries.append(remember_entry)
|
||||
if loop is not None:
|
||||
project, loop_tasks, loop_entries = _loop(cases, templates, loop)
|
||||
project, loop_tasks, loop_entries = _loop(cases, templates, loop, atespace)
|
||||
out[f"workspaces/{project['metadata']['name']}.yaml"] = _dump(header, [project])
|
||||
tasks.extend(loop_tasks)
|
||||
entries.extend(loop_entries)
|
||||
@@ -509,16 +512,24 @@ def increment(
|
||||
repos: Path | None = None,
|
||||
runner_image: Path | None = None,
|
||||
loop: Loop | None = None,
|
||||
atespace: str | None = None,
|
||||
) -> list[Path]:
|
||||
"""Write the generated manifests under ``plane`` from the templates in ``plane/tasks`` (re-pinned first when
|
||||
``runner_image`` is given), with the improvement loop when ``loop`` is given; returns the paths."""
|
||||
``runner_image`` is given), with the improvement loop when ``loop`` is given; returns the paths. Workspaces go
|
||||
to ``atespace`` (default: the profile's)."""
|
||||
if atespace is None:
|
||||
from .profile import load_profile
|
||||
|
||||
atespace = load_profile().atespace
|
||||
if loop is not None and repos is None:
|
||||
raise ValueError("--loop needs --repos: the alerts task marks fixed-pin alerts by the fix's new-side ranges")
|
||||
if runner_image is not None:
|
||||
repin_templates(plane, runner_image)
|
||||
loaded = load_manifests([plane / "tasks" / f"{name}.yaml" for name in TEMPLATES])
|
||||
cases = read_cases(validation_set, candidates_dir, triage_dir, population, repos, fixed_side=loop is not None)
|
||||
files = render(cases, loaded.tasks, run_name, command, population=population.stem, split=validation_set.stem, loop=loop)
|
||||
files = render(
|
||||
cases, loaded.tasks, run_name, command, population=population.stem, split=validation_set.stem, loop=loop, atespace=atespace
|
||||
)
|
||||
written: list[Path] = []
|
||||
for relative, text in files.items():
|
||||
path = plane / relative
|
||||
|
||||
@@ -210,17 +210,19 @@ def _reader(cache: Path, source: str, ref: str, pin: str) -> Any:
|
||||
# --- manifests -------------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _git_workspace(unit: Unit) -> dict[str, Any]:
|
||||
doc = _doc("Workspace", _ax_name(unit.name, "pin"), {"git": [{"name": REPO, "repo": unit.repo, "dir": REPO, "depth": 1}]}, ATESPACE)
|
||||
def _git_workspace(unit: Unit, atespace: str = ATESPACE) -> dict[str, Any]:
|
||||
doc = _doc("Workspace", _ax_name(unit.name, "pin"), {"git": [{"name": REPO, "repo": unit.repo, "dir": REPO, "depth": 1}]}, atespace)
|
||||
doc["metadata"]["annotations"] = {GIT_COMMITS_ANNOTATION: f"{REPO}={unit.pin}"}
|
||||
return doc
|
||||
|
||||
|
||||
def _files_workspace(name: str, files: Mapping[str, str]) -> dict[str, Any]:
|
||||
return _doc("Workspace", name, {"files": [{"path": k, "content": v} for k, v in sorted(files.items())]}, ATESPACE)
|
||||
def _files_workspace(name: str, files: Mapping[str, str], atespace: str = ATESPACE) -> dict[str, Any]:
|
||||
return _doc("Workspace", name, {"files": [{"path": k, "content": v} for k, v in sorted(files.items())]}, atespace)
|
||||
|
||||
|
||||
def _verify_docs(unit: Unit, templates: Mapping[str, Task], usd_per_hunt: float) -> tuple[list[dict], list[dict], list[dict]]:
|
||||
def _verify_docs(
|
||||
unit: Unit, templates: Mapping[str, Task], usd_per_hunt: float, atespace: str = ATESPACE
|
||||
) -> tuple[list[dict], list[dict], list[dict]]:
|
||||
"""(Workspaces, Tasks, Run entries) of one verify unit."""
|
||||
asked = list(unit.verify or unit.candidates)
|
||||
candidates = {"family": unit.family, "candidates": [list(c) for c in asked]}
|
||||
@@ -232,15 +234,15 @@ def _verify_docs(unit: Unit, templates: Mapping[str, Task], usd_per_hunt: float)
|
||||
entries: list[dict] = []
|
||||
if unit.files:
|
||||
files = {f"{REPO}/{path}": text for path, text in unit.files} | {"inputs/candidates.json": _json(candidates)}
|
||||
workspace = _files_workspace(_ax_name(unit.name, "side"), files)
|
||||
workspace = _files_workspace(_ax_name(unit.name, "side"), files, atespace)
|
||||
workspaces.append(workspace)
|
||||
bindings = [(workspace["metadata"]["name"], CASE_PATH)]
|
||||
inputs = f"{CASE_PATH}/inputs/candidates.json"
|
||||
else:
|
||||
pinned = _git_workspace(unit)
|
||||
pinned = _git_workspace(unit, atespace)
|
||||
functions = [{"path": p, "function": fn, "line": ln} for p, fn, ln in asked]
|
||||
files_doc = _files_workspace(
|
||||
_ax_name(unit.name, "inputs"), {"candidates.json": _json(candidates), "functions.json": _json(functions)}
|
||||
_ax_name(unit.name, "inputs"), {"candidates.json": _json(candidates), "functions.json": _json(functions)}, atespace
|
||||
)
|
||||
workspaces += [pinned, files_doc]
|
||||
bindings = [(pinned["metadata"]["name"], CASE_PATH), (files_doc["metadata"]["name"], INPUTS_PATH)]
|
||||
@@ -309,7 +311,7 @@ def _roles_groups(units: Sequence[Unit]) -> list[list[Unit]]:
|
||||
|
||||
|
||||
def render_harvest(
|
||||
units: Sequence[Unit], templates: Mapping[str, Task], name: str, command: str, *, ceiling: float = 10.0
|
||||
units: Sequence[Unit], templates: Mapping[str, Task], name: str, command: str, *, ceiling: float = 10.0, atespace: str = ATESPACE
|
||||
) -> dict[str, str]:
|
||||
"""Relative path under the plane root -> file text: both Runs, their Workspaces and Tasks, and ``units.json``.
|
||||
A population file whose size is unknown is estimated at :data:`UNKNOWN_SIZE`."""
|
||||
@@ -323,7 +325,7 @@ def render_harvest(
|
||||
tasks: list[dict] = []
|
||||
entries: list[dict] = []
|
||||
for unit in verify_units:
|
||||
w, t, e = _verify_docs(unit, templates, USD_PER_HUNT[_kind(unit)] * scale)
|
||||
w, t, e = _verify_docs(unit, templates, USD_PER_HUNT[_kind(unit)] * scale, atespace)
|
||||
workspaces += w
|
||||
tasks += t
|
||||
entries += e
|
||||
@@ -333,7 +335,7 @@ def render_harvest(
|
||||
for group in _roles_groups([u for u in units if u.files and _inline(u, roles=True) <= MAX_INLINE_BYTES]):
|
||||
files = {f"{REPO}/{u.name}/{p}": t for u in group for p, t in u.files if p in u.paths}
|
||||
listing = sorted(f.removeprefix(f"{REPO}/") for f in files)
|
||||
workspace = _files_workspace(_ax_name(name, "roles", group[0].name), files | {"inputs/files.json": _json(listing)})
|
||||
workspace = _files_workspace(_ax_name(name, "roles", group[0].name), files | {"inputs/files.json": _json(listing)}, atespace)
|
||||
workspaces.append(workspace)
|
||||
estimates.append(_roles_estimate((len(text), len(text.splitlines())) for text in files.values()))
|
||||
bound = [(workspace["metadata"]["name"], CASE_PATH)]
|
||||
@@ -350,8 +352,8 @@ def render_harvest(
|
||||
same = [u for u in units if not u.files and (u.repo, u.pin) == (unit.repo, unit.pin)]
|
||||
paths = sorted({p for u in same for p in u.paths})
|
||||
known = {p: (chars, lines) for u in same for p, chars, lines in u.sizes}
|
||||
pinned = _git_workspace(unit)
|
||||
files_doc = _files_workspace(_ax_name(unit.name, "roles-in"), {"files.json": _json(paths)})
|
||||
pinned = _git_workspace(unit, atespace)
|
||||
files_doc = _files_workspace(_ax_name(unit.name, "roles-in"), {"files.json": _json(paths)}, atespace)
|
||||
workspaces += [w for w in (pinned, files_doc) if w["metadata"]["name"] not in emitted]
|
||||
estimates.append(_roles_estimate(known.get(p, UNKNOWN_SIZE) for p in paths))
|
||||
bound = [(pinned["metadata"]["name"], CASE_PATH), (files_doc["metadata"]["name"], INPUTS_PATH)]
|
||||
@@ -380,7 +382,16 @@ def _run_files(header: str, run: str, workspaces: list[dict], tasks: list[dict],
|
||||
|
||||
|
||||
def write_harvest(
|
||||
labels: Path, plane: Path, *, name: str, command: str, catalog: Path, cache: Path, templates: Path, ceiling: float = 10.0
|
||||
labels: Path,
|
||||
plane: Path,
|
||||
*,
|
||||
name: str,
|
||||
command: str,
|
||||
catalog: Path,
|
||||
cache: Path,
|
||||
templates: Path,
|
||||
ceiling: float = 10.0,
|
||||
atespace: str = ATESPACE,
|
||||
) -> dict[str, str]:
|
||||
"""Build the units, render, write under ``plane`` (with the Model manifests copied beside); returns the rendered
|
||||
files."""
|
||||
@@ -388,7 +399,7 @@ def write_harvest(
|
||||
|
||||
units = pair_units(labels, catalog) + population_units(labels, cache)
|
||||
loaded = load_manifests([templates / "tasks" / f"{n}.yaml" for n in ("repo-facts", "verify", "agree", "roles")]).tasks
|
||||
files = render_harvest(units, loaded, name, command, ceiling=ceiling)
|
||||
files = render_harvest(units, loaded, name, command, ceiling=ceiling, atespace=atespace)
|
||||
for path in sorted((templates / "models").glob("*.yaml")):
|
||||
files[f"models/{path.name}"] = path.read_text(encoding="utf-8")
|
||||
for relative, text in files.items():
|
||||
|
||||
@@ -1,9 +1,10 @@
|
||||
"""ax-backed reconciler of a ``Run``: submission, resume, start, completion by delivery (ai-service-plane Req 3, 4, 7).
|
||||
|
||||
ax executes every task on this host's kind cluster; nothing here runs a task locally or knows a prompt, a score or
|
||||
a model call. Per ready task: one rendered YAML file (the Task with its env extended, its Workspaces, the bound
|
||||
Model), ``ax apply -f``, ``ax resume task`` (retried on DeadlineExceeded/Unavailable -- a new image's first resume
|
||||
waits for Agent Substrate's golden snapshot -- for up to ``OUSAST_RESUME_TIMEOUT`` s, 900), and once ax says
|
||||
ax executes every task on the profile's cluster (``profile.py``: kube context, registry, store); nothing here runs a
|
||||
task locally or knows a prompt, a score or a model call. Per ready task: one rendered YAML file (the Task with its
|
||||
env extended, its Workspaces, the bound Model), ``ax apply -f``, ``ax resume task`` (retried on
|
||||
DeadlineExceeded/Unavailable -- a new image's first resume waits for Agent Substrate's golden snapshot -- for up to
|
||||
``OUSAST_RESUME_TIMEOUT`` s, 900), and once ax says
|
||||
``Running`` one start request through Substrate's router (``router.py``, Req 4.5) carrying the bound Model's
|
||||
credential from this process's environment, never from a file (Req 4.6). Completion is the runner's artifact
|
||||
delivery (ax has no Completed phase); ``Failed``, the task vanishing, a refused start or ``OUSAST_TASK_TIMEOUT`` s
|
||||
@@ -44,6 +45,7 @@ import yaml
|
||||
from .doctor import doctor
|
||||
from .egress import RECEIVER_PORT, Egress, Receiver, policy_for
|
||||
from .manifests import AX_API_VERSION, Manifests, Model, Run, RunTask, Task, Workspace, load_manifests
|
||||
from .profile import PlaneProfile, load_profile
|
||||
from .router import Router, credentials, open_router
|
||||
|
||||
MODEL_ANNOTATION = "openultrasast.io/model"
|
||||
@@ -340,17 +342,27 @@ def _execute(
|
||||
return ("failed", failure) if failure else outcome_of(_read_json(base / entry.name / "summary.json"))
|
||||
|
||||
|
||||
def run(run_manifest: Path, *, workers: int = 1, ax: str = "ax", results_root: Path | None = None, kubectl_ate: str | None = None) -> str:
|
||||
"""Execute a Run on ax and return its status: ``done``, ``unfinished`` or ``failed`` (Req 3.1-3.3, 2.2)."""
|
||||
def run(
|
||||
run_manifest: Path,
|
||||
*,
|
||||
workers: int = 1,
|
||||
ax: str = "ax",
|
||||
results_root: Path | None = None,
|
||||
kubectl_ate: str | None = None,
|
||||
profile: PlaneProfile | None = None,
|
||||
) -> str:
|
||||
"""Execute a Run on ax and return its status: ``done``, ``unfinished`` or ``failed`` (Req 3.1-3.3, 2.2). Every
|
||||
address comes from ``profile`` (default: the one ``OUSAST_PLANE_PROFILE`` names)."""
|
||||
profile = profile or load_profile()
|
||||
run_spec, manifests = load_run(Path(run_manifest))
|
||||
base = run_dir(run_spec.metadata.name, results_root)
|
||||
base.mkdir(parents=True, exist_ok=True)
|
||||
lock = _acquire_lock(base)
|
||||
state = _State(base / "state.json", _read_json(base / "state.json") or {"run": run_spec.metadata.name, "started": _now(), "tasks": {}})
|
||||
port = os.environ.get("OUSAST_ARTIFACT_PORT") or (0 if os.environ.get("OUSAST_ARTIFACT_HOST") else RECEIVER_PORT)
|
||||
receiver = Receiver(run_spec.metadata.name, base, int(port)) # the receiver Service targets RECEIVER_PORT
|
||||
receiver = Receiver(run_spec.metadata.name, base, int(port), profile.kube_context) # the Service targets RECEIVER_PORT
|
||||
threading.Thread(target=receiver.serve_forever, daemon=True).start()
|
||||
cli, egress, finished = Ax(ax), Egress(kubectl_ate), queue.Queue[tuple[str, str, str]]()
|
||||
cli, egress, finished = Ax(ax), Egress(kubectl_ate, profile.kube_context), queue.Queue[tuple[str, str, str]]()
|
||||
pending = [t for t in run_spec.tasks if state.status(t.name) != "done"]
|
||||
for entry in run_spec.tasks:
|
||||
state.set(entry.name, **({"status": "pending"} if entry in pending else {}))
|
||||
@@ -364,7 +376,7 @@ def run(run_manifest: Path, *, workers: int = 1, ax: str = "ax", results_root: P
|
||||
finished.put((entry.name, *_execute(cli, egress, run_spec, entry, manifests, base, receiver, router)))
|
||||
|
||||
try:
|
||||
with open_router() as router: # one port-forward to Substrate's router for the whole run
|
||||
with open_router(profile.kube_context, url=profile.router_url or None) as router: # one tunnel per run (ax-tunnel)
|
||||
while pending or running:
|
||||
blocked = {t.serialize for t in running.values() if t.serialize}
|
||||
for entry in list(pending):
|
||||
|
||||
@@ -131,18 +131,20 @@ def _drain(stream: Iterable[str], lines: queue.Queue[str | None]) -> None:
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def open_router(context: str | None = None, *, kubectl: str = "kubectl", ready_timeout: float = 30.0) -> Iterator[Router]:
|
||||
def open_router(context: str, *, url: str | None = None, kubectl: str = "kubectl", ready_timeout: float = 30.0) -> Iterator[Router]:
|
||||
"""A :class:`Router` for the duration of the block.
|
||||
|
||||
``OUSAST_ROUTER_URL`` set: that URL, nothing started. Otherwise ``kubectl port-forward svc/atenet-router :80``
|
||||
in ``ate-system`` of ``context`` (default ``kind-$KIND_CLUSTER_NAME``, ``kind-ousast``), the local port read
|
||||
from kubectl's ``Forwarding from 127.0.0.1:NNNN`` line; on exit that child process is terminated.
|
||||
``url`` (the profile's ``router_url``) or ``OUSAST_ROUTER_URL`` set: that URL, nothing started. Otherwise
|
||||
``kubectl port-forward svc/atenet-router :80`` in ``ate-system`` of the profile's ``context`` (ax-tunnel: the
|
||||
Kubernetes API is the one authenticated door), the local port read from kubectl's ``Forwarding from
|
||||
127.0.0.1:NNNN`` line; on exit that child process is terminated. No context is ever built here.
|
||||
"""
|
||||
override = os.environ.get("OUSAST_ROUTER_URL")
|
||||
override = url or os.environ.get("OUSAST_ROUTER_URL")
|
||||
if override:
|
||||
yield Router(override)
|
||||
return
|
||||
context = context or "kind-" + (os.environ.get("KIND_CLUSTER_NAME") or "ousast")
|
||||
if not context:
|
||||
raise StartError("open_router needs the profile's kube context (or router_url)")
|
||||
argv = [kubectl, "--context", context, "-n", NAMESPACE, "port-forward", SERVICE, ":80"]
|
||||
proc = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
|
||||
try:
|
||||
|
||||
Reference in New Issue
Block a user