mirror of
https://github.com/norandom/OpenUltraSAST.git
synced 2026-10-16 11:49:28 +00:00
feat(plane): per-task egress policy; receiver by name through the gateway
Agent Substrate's egress gateway denies an actor without an EgressPolicy, so no task could deliver its artifacts or reach a model. egress.py grants each task exactly: http/80 to ousast-receiver.ax-system.svc.cluster.local, and TLS passthrough/443 to the bound Workspaces' https Git hosts and the bound Model's openultrasast.io/egress-hosts (api.deepseek.com for deepseek-flash). No IPs, no wildcards, no https (MITM) rules. The reconciler writes the policy with kubectl-ate between ax apply and resume, reads it back, waits OUSAST_EGRESS_SETTLE (11 s, the gateway's cached deny), and after ax delete task confirms the policy went with its actor (kubectl-ate has no delete verb; the store cascades). The receiver is a selector-less Service (ops/ax/receiver-service.yaml.tmpl, applied by up.sh together with the no-build agentgateway egress gateway) on port 18090 of the host. Measured live: the gateway checks the policy on Host but connects to the address the actor dialled, so the runner dials the Service's ClusterIP (OUSAST_ARTIFACT_DIAL) with the receiver's name as Host. The runner now labels its delivery with the Run's task name (OUSAST_TASK), which is what the receiver accepts; the ax Task name was refused with 403. The Receiver moved to egress.py to keep reconciler.py at 460 lines. Co-Authored-By: Claude Opus 5.5 <[email protected]> Claude-Session: https://claude.ai/code/session_01EmM1j46qgKMh3zqD1QMwEo
This commit is contained in:
co-authored by
Claude Opus 5.5
parent
0744a157bf
commit
137c6c8ed5
@@ -0,0 +1,252 @@
|
||||
"""Per-task network egress on Agent Substrate (ai-service-plane, task 0: gateway + receiver Service).
|
||||
|
||||
Every TCP connection an actor opens (DNS aside) is redirected to Substrate's egress gateway
|
||||
(``atenet-egress.ate-system``), which denies by default: an actor without an ``EgressPolicy`` gets no tunnel, and
|
||||
a refused connection is accepted and then closed. :func:`policy_for` grants a task exactly what it needs:
|
||||
|
||||
* ``http`` on port 80 to the artifact receiver's cluster name (``ousast-receiver.ax-system.svc.cluster.local``, a
|
||||
selector-less Service whose endpoint is this host's receiver, ``ops/ax/receiver-service.yaml.tmpl``). The actor
|
||||
cannot resolve cluster names, so the runner dials ``OUSAST_ARTIFACT_DIAL`` (the Service's ClusterIP, port 80)
|
||||
with the receiver's name as ``Host``. The gateway decides the policy on that ``Host`` but connects to the
|
||||
address the actor dialled (measured live: agentgateway's upstream is the CONNECT authority, not the Host), so
|
||||
the dial address must be one the gateway pod reaches: the ClusterIP, which kube-proxy maps to the host.
|
||||
* ``tls_passthrough`` on 443 for the Git hosts of the bound Workspaces' ``https`` repos (smart HTTP talks to the
|
||||
repo's host only) and for the bound Model's ``openultrasast.io/egress-hosts`` (``api.deepseek.com``).
|
||||
|
||||
Nothing else: no wildcard, no IP address (the gateway rejects them), no ``https`` rule (that one intercepts TLS and
|
||||
the actor would have to trust the gateway CA, which ax templates do not mount). :class:`Egress` writes the policy
|
||||
with ``kubectl ate create egress-policy`` after ``ax apply`` created the actor, reads it back, and before resume
|
||||
waits ``OUSAST_EGRESS_SETTLE`` s (11) because the gateway caches a missing policy as deny for 10 s. The kubectl-ate
|
||||
CLI has no delete verb: Substrate's store deletes the policy with its actor (``ON DELETE CASCADE``), so
|
||||
:meth:`Egress.delete` runs after ``ax delete task`` and confirms the policy is gone.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import ipaddress
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import tarfile
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Iterable, Mapping
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
from .manifests import Model, Task, Workspace
|
||||
|
||||
__all__ = [
|
||||
"RECEIVER_HOST",
|
||||
"RECEIVER_PORT",
|
||||
"Egress",
|
||||
"Receiver",
|
||||
"git_hosts",
|
||||
"receiver_cluster_ip",
|
||||
"policy_for",
|
||||
"receiver_address",
|
||||
"rules_of",
|
||||
]
|
||||
|
||||
RECEIVER_HOST = "ousast-receiver.ax-system.svc.cluster.local"
|
||||
RECEIVER_PORT = 18090 # the receiver Service's targetPort; ops/ax/up.sh renders the Service with the same number
|
||||
_NOT_FOUND = ("not found", "notfound")
|
||||
|
||||
|
||||
def _is_ip(host: str) -> bool:
|
||||
try:
|
||||
ipaddress.ip_address(host.strip("[]"))
|
||||
except ValueError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def git_hosts(workspaces: Iterable[Workspace]) -> tuple[list[str], list[str]]:
|
||||
"""(https hosts, http hosts) of the Workspaces' ``spec.git[].repo`` URLs; ssh and IP-address repos add nothing."""
|
||||
tls, plain = set[str](), set[str]()
|
||||
for ws in workspaces:
|
||||
for source in ws.git:
|
||||
parts = urlsplit(source.repo)
|
||||
host = (parts.hostname or "").lower()
|
||||
if (
|
||||
host
|
||||
and not _is_ip(host)
|
||||
and parts.scheme in ("https", "http")
|
||||
and parts.port in (None, 443 if parts.scheme == "https" else 80)
|
||||
):
|
||||
(tls if parts.scheme == "https" else plain).add(host)
|
||||
return sorted(tls), sorted(plain)
|
||||
|
||||
|
||||
def policy_for(task: Task, workspaces: Iterable[Workspace], model: Model | None, receiver_host: str | None) -> dict[str, Any]:
|
||||
"""The task's EgressPolicy body: the receiver over HTTP/80, Git and Model hosts as TLS passthrough on 443."""
|
||||
tls, plain = git_hosts(workspaces)
|
||||
tls = sorted(set(tls) | set(model.egress_hosts if model else ()))
|
||||
if receiver_host and not _is_ip(receiver_host): # an address cannot be allowed; only the direct test mode has one
|
||||
plain = sorted({*plain, receiver_host.lower()})
|
||||
rules: list[dict[str, Any]] = []
|
||||
if plain:
|
||||
rules.append({"http": {"hostnames": plain, "ports": {"numbers": [80]}}})
|
||||
if tls:
|
||||
rules.append({"tlsPassthrough": {"hostnames": tls, "ports": {"numbers": [443]}}})
|
||||
return {"rules": rules}
|
||||
|
||||
|
||||
def rules_of(policy: Mapping[str, Any]) -> list[dict[str, Any]]:
|
||||
"""The rules of a policy as ``kubectl ate get -o json`` prints them, normalised for comparison."""
|
||||
out = []
|
||||
for rule in policy.get("rules") or []:
|
||||
for kind, body in rule.items():
|
||||
key = "tlsPassthrough" if kind in ("tls_passthrough", "tlsPassthrough") else kind
|
||||
ports = body.get("ports") or {}
|
||||
norm_ports: dict[str, Any] = {"all": {}} if "all" in ports else {"numbers": sorted(int(n) for n in ports.get("numbers") or [])}
|
||||
out.append({key: {"hostnames": sorted(body.get("hostnames") or []), "ports": norm_ports}})
|
||||
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")
|
||||
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)
|
||||
except (OSError, subprocess.TimeoutExpired):
|
||||
return None
|
||||
address = proc.stdout.strip()
|
||||
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``.
|
||||
|
||||
``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
|
||||
Service's ClusterIP on port 80: the gateway checks the policy against ``Host`` but connects to the address the
|
||||
actor dialled (its CONNECT authority), so that address must reach the receiver from the gateway pod.
|
||||
"""
|
||||
host = os.environ.get("OUSAST_ARTIFACT_HOST")
|
||||
if host:
|
||||
return f"http://{host}:{port}/", None
|
||||
dial = os.environ.get("OUSAST_ARTIFACT_DIAL")
|
||||
if not dial:
|
||||
address = receiver_cluster_ip()
|
||||
dial = f"{address}:80" if address else None
|
||||
return f"http://{RECEIVER_HOST}/", dial
|
||||
|
||||
|
||||
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:
|
||||
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")
|
||||
|
||||
def _run(self, *args: str, stdin: str | None = None) -> subprocess.CompletedProcess[str]:
|
||||
argv = [self.executable, "--context", self.context, *args]
|
||||
return subprocess.run(argv, input=stdin, capture_output=True, text=True, check=False)
|
||||
|
||||
@staticmethod
|
||||
def _error(proc: subprocess.CompletedProcess[str]) -> str:
|
||||
return (proc.stderr.strip() or proc.stdout.strip() or f"exit {proc.returncode}")[-400:]
|
||||
|
||||
def get(self, actor: str, atespace: str) -> dict[str, Any] | None:
|
||||
"""The stored policy, None when there is none (or no actor); RuntimeError when kubectl-ate fails otherwise."""
|
||||
proc = self._run("get", "egress-policy", actor, "-a", atespace, "-o", "json")
|
||||
if proc.returncode != 0:
|
||||
if "notfound" in self._error(proc).lower().replace(" ", ""):
|
||||
return None
|
||||
raise RuntimeError(f"kubectl ate get egress-policy {atespace}/{actor} failed: {self._error(proc)}")
|
||||
if not proc.stdout.strip(): # an actor without a policy: exit 0, "has no egress policy" on stderr only
|
||||
return None
|
||||
try:
|
||||
doc = json.loads(proc.stdout)
|
||||
except ValueError as exc:
|
||||
raise RuntimeError(f"kubectl ate get egress-policy {atespace}/{actor}: unreadable output ({exc})") from None
|
||||
return doc if isinstance(doc, dict) else None
|
||||
|
||||
def apply(self, actor: str, atespace: str, policy: Mapping[str, Any], *, timeout: float | None = None) -> None:
|
||||
"""Create (or replace) the policy, retrying while the actor does not exist yet, then read it back."""
|
||||
limit = float(os.environ.get("OUSAST_EGRESS_TIMEOUT") or 60) if timeout is None else timeout
|
||||
deadline, backoff = time.monotonic() + limit, 0.5
|
||||
while True:
|
||||
existing = self.get(actor, atespace)
|
||||
if existing is None:
|
||||
proc = self._run("create", "egress-policy", actor, "-a", atespace, "-f", "-", stdin=json.dumps(dict(policy)))
|
||||
else:
|
||||
meta = {k: existing.get("metadata", {}).get(k) for k in ("name", "atespace", "uid", "version")}
|
||||
body = json.dumps({**dict(policy), "metadata": meta})
|
||||
proc = self._run("update", "egress-policy", actor, "-a", atespace, "-f", "-", stdin=body)
|
||||
if proc.returncode == 0:
|
||||
break
|
||||
error = self._error(proc)
|
||||
retry = any(word in error.lower() for word in (*_NOT_FOUND, "already exists", "conflict", "unavailable"))
|
||||
if not retry or time.monotonic() + backoff > deadline:
|
||||
raise RuntimeError(f"egress policy for {atespace}/{actor} not written: {error}")
|
||||
time.sleep(backoff)
|
||||
backoff = min(backoff * 2, 5.0)
|
||||
stored = self.get(actor, atespace)
|
||||
if stored is None or rules_of(stored) != rules_of(policy):
|
||||
raise RuntimeError(f"egress policy for {atespace}/{actor} reads back {stored!r}, not {dict(policy)!r}")
|
||||
|
||||
def delete(self, actor: str, atespace: str) -> str | None:
|
||||
"""After ``ax delete task``: None once the policy is gone with its actor, else why it is still there."""
|
||||
try:
|
||||
stored = self.get(actor, atespace)
|
||||
except RuntimeError as exc:
|
||||
return str(exc)
|
||||
return None if stored is None else f"egress policy for {atespace}/{actor} outlived its actor"
|
||||
|
||||
@staticmethod
|
||||
def settle() -> None:
|
||||
"""Wait out the gateway's cached deny (10 s) before the actor's first connection: ``OUSAST_EGRESS_SETTLE``."""
|
||||
time.sleep(float(os.environ.get("OUSAST_EGRESS_SETTLE") or 11))
|
||||
|
||||
|
||||
class Receiver(ThreadingHTTPServer):
|
||||
"""Accepts one tar per running task and extracts it under ``<run dir>/<task>/``."""
|
||||
|
||||
daemon_threads = True
|
||||
|
||||
def __init__(self, run_name: str, base: Path, port: int) -> 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.lock = threading.Lock()
|
||||
self.url, self.dial = receiver_address(self.server_address[1])
|
||||
|
||||
def accept(self, headers: Any, body: bytes) -> tuple[int, str]:
|
||||
task = headers.get("X-Ousast-Task", "")
|
||||
if headers.get("X-Ousast-Run") != self.run_name or task not in self.running:
|
||||
return 403, "unknown run or task"
|
||||
if not str(headers.get("Content-Type", "")).startswith("application/x-tar"):
|
||||
return 415, "expected application/x-tar"
|
||||
dest = (self.base / task).resolve()
|
||||
try:
|
||||
with tarfile.open(fileobj=io.BytesIO(body), mode="r:") as tar:
|
||||
members = tar.getmembers()
|
||||
for m in members:
|
||||
if m.issym() or m.islnk() or not (dest / m.name).resolve().is_relative_to(dest):
|
||||
return 400, f"rejected member {m.name}"
|
||||
dest.mkdir(parents=True, exist_ok=True)
|
||||
tar.extractall(dest, members) # every member checked above
|
||||
except tarfile.TarError as exc:
|
||||
return 400, f"bad tar: {exc}"
|
||||
with self.lock:
|
||||
self.delivered.add(task)
|
||||
return 200, "ok"
|
||||
|
||||
|
||||
class _Handler(BaseHTTPRequestHandler):
|
||||
def do_POST(self) -> None:
|
||||
server: Receiver = self.server # type: ignore[assignment]
|
||||
body = self.rfile.read(int(self.headers.get("Content-Length") or 0))
|
||||
code, text = server.accept(self.headers, body)
|
||||
self.send_response(code)
|
||||
self.send_header("Content-Type", "text/plain")
|
||||
self.end_headers()
|
||||
self.wfile.write(text.encode())
|
||||
|
||||
def log_message(self, format: str, *args: object) -> None:
|
||||
return
|
||||
@@ -24,6 +24,7 @@ import yaml
|
||||
__all__ = [
|
||||
"AX_API_VERSION",
|
||||
"AX_PROVIDERS",
|
||||
"EGRESS_HOSTS_ANNOTATION",
|
||||
"EXTENSION_PROVIDERS",
|
||||
"GIT_COMMITS_ANNOTATION",
|
||||
"PROVIDER_EXTENSION_ANNOTATION",
|
||||
@@ -55,6 +56,8 @@ AX_API_VERSION = "ax.io/v1alpha1"
|
||||
RUN_API_VERSION = "openultrasast.io/v1alpha1"
|
||||
PROVIDER_EXTENSION_ANNOTATION = "openultrasast.io/provider-extension"
|
||||
GIT_COMMITS_ANNOTATION = "openultrasast.io/git-commits"
|
||||
EGRESS_HOSTS_ANNOTATION = "openultrasast.io/egress-hosts" # the Model API hosts a bound task may reach
|
||||
_DNS_NAME = re.compile(r"^(?=.{1,253}$)[a-z0-9]([a-z0-9-]{0,61}[a-z0-9])?(\.[a-z0-9]([a-z0-9-]{0,61}[a-z0-9])?)*$")
|
||||
AX_PROVIDERS = frozenset({"google", "anthropic"})
|
||||
EXTENSION_PROVIDERS = frozenset({"deepseek"})
|
||||
|
||||
@@ -168,6 +171,15 @@ class Model:
|
||||
secret_key: SecretKey | None = None
|
||||
parameters: Mapping[str, Any] = field(default_factory=dict)
|
||||
|
||||
@property
|
||||
def egress_hosts(self) -> tuple[str, ...]:
|
||||
"""DNS names from ``metadata.annotations["openultrasast.io/egress-hosts"]`` (comma-separated)."""
|
||||
return _split_hosts(self.metadata.annotations.get(EGRESS_HOSTS_ANNOTATION, ""))
|
||||
|
||||
|
||||
def _split_hosts(text: str) -> tuple[str, ...]:
|
||||
return tuple(h.strip() for h in text.split(",") if h.strip())
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Budget:
|
||||
@@ -426,6 +438,10 @@ def _parse_model(check: _Check, metadata: Metadata, spec: Mapping[str, Any]) ->
|
||||
)
|
||||
elif provider not in AX_PROVIDERS:
|
||||
raise check.fail("spec.provider", f'"{provider}" is not one of {", ".join(sorted(AX_PROVIDERS | EXTENSION_PROVIDERS))}')
|
||||
for host in _split_hosts(metadata.annotations.get(EGRESS_HOSTS_ANNOTATION, "")):
|
||||
if not _DNS_NAME.match(host) or host.replace(".", "").isdigit():
|
||||
path = f'metadata.annotations["{EGRESS_HOSTS_ANNOTATION}"]'
|
||||
raise check.fail(path, f"{host!r} is not a lowercase DNS name (no wildcard, IP address, port or scheme)")
|
||||
secret = check.optional_mapping(spec, "spec", "secretKey", ("name", "key"))
|
||||
parameters = spec.get("parameters") or {}
|
||||
if not isinstance(parameters, dict):
|
||||
|
||||
@@ -12,40 +12,39 @@ vanishing, a refused start or ``OUSAST_TASK_TIMEOUT`` s (7200) without delivery
|
||||
ax rejects unknown fields, so rendered documents carry only ax's fields; ``openultrasast.io/model`` reaches the
|
||||
task as ``OUSAST_MODEL`` plus ``OUSAST_MODEL_PARAMS`` (Req 7.1), each Workspace's ``openultrasast.io/git-commits``
|
||||
as ``OUSAST_GIT_PINS``; ``OUSAST_INPUT_<NAME>`` points into ``/workspace/.ousast-in/<task>``. The receiver takes
|
||||
the runner's tar (``X-Ousast-Run``, ``X-Ousast-Task``) on ``OUSAST_ARTIFACT_PORT``, advertised as
|
||||
``OUSAST_ARTIFACT_HOST`` (default ``172.17.0.1``, the docker bridge). State: ``~/ousast-results/plane/<run>/``
|
||||
(``OUSAST_RESULTS``), polled every ``OUSAST_POLL_SECONDS`` (3).
|
||||
the runner's tar (``X-Ousast-Run``, ``X-Ousast-Task``) on ``OUSAST_ARTIFACT_PORT`` (18090), reached by its Service
|
||||
name through the egress gateway; between ``ax apply`` and resume the task's egress policy is written (``egress.py``).
|
||||
State: ``~/ousast-results/plane/<run>/`` (``OUSAST_RESULTS``), polled every ``OUSAST_POLL_SECONDS`` (3).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import queue
|
||||
import re
|
||||
import subprocess
|
||||
import tarfile
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable, Iterable, Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
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 .router import Router, credentials, open_router
|
||||
|
||||
MODEL_ANNOTATION = "openultrasast.io/model"
|
||||
OUTPUT_ROOT = "/workspace/.ousast-out"
|
||||
INPUTS_ROOT = "/workspace/.ousast-in"
|
||||
DEFAULT_ARTIFACT_HOST = "172.17.0.1"
|
||||
TRANSIENT = re.compile(r"deadline ?exceeded|unavailable", re.IGNORECASE) # the golden-snapshot build outlives a resume
|
||||
_ROW = "{:<28} {:<20} {:>7} {:>10} {:>10} {:>9} {:>10}"
|
||||
_COLUMNS = ("calls", "prompt", "cache_hit", "output", "usd")
|
||||
@@ -134,13 +133,15 @@ def _input_files(entry: RunTask, base: Path) -> list[dict[str, str]]:
|
||||
return files
|
||||
|
||||
|
||||
def render_task(run: Run, entry: RunTask, manifests: Manifests, base: Path, artifact_url: str) -> list[dict[str, Any]]:
|
||||
def render_task(
|
||||
run: Run, entry: RunTask, manifests: Manifests, base: Path, artifact_url: str, dial: str | None = None
|
||||
) -> list[dict[str, Any]]:
|
||||
"""The documents ``ax apply`` receives for one Run task: Task, its Model, its Workspaces, the inputs Workspace."""
|
||||
task: Task = manifests.tasks[entry.task]
|
||||
ax_name = _ax_name(run.metadata.name, entry.name)
|
||||
env = {e.name: e.value for e in task.env}
|
||||
env.update(OUSAST_RUN=run.metadata.name, OUSAST_TASK=entry.name, OUSAST_OUTPUT_DIR=f"{OUTPUT_ROOT}/{entry.name}")
|
||||
env["OUSAST_ARTIFACT_URL"] = artifact_url
|
||||
env.update(OUSAST_ARTIFACT_URL=artifact_url, **({"OUSAST_ARTIFACT_DIAL": dial} if dial else {}))
|
||||
budget = _fields(entry.budget) if entry.budget else {}
|
||||
env.update({f"OUSAST_BUDGET_{k.upper()}": str(v) for k, v in budget.items()})
|
||||
for input_name, ref in entry.inputs.items():
|
||||
@@ -214,55 +215,6 @@ class Ax:
|
||||
self._run("delete", "task", name)
|
||||
|
||||
|
||||
class Receiver(ThreadingHTTPServer):
|
||||
"""Accepts one tar per running task and extracts it under ``<run dir>/<task>/``."""
|
||||
|
||||
daemon_threads = True
|
||||
|
||||
def __init__(self, run_name: str, base: Path, port: int) -> 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.lock = threading.Lock()
|
||||
|
||||
def url(self) -> str:
|
||||
return f"http://{os.environ.get('OUSAST_ARTIFACT_HOST') or DEFAULT_ARTIFACT_HOST}:{self.server_address[1]}/"
|
||||
|
||||
def accept(self, headers: Any, body: bytes) -> tuple[int, str]:
|
||||
task = headers.get("X-Ousast-Task", "")
|
||||
if headers.get("X-Ousast-Run") != self.run_name or task not in self.running:
|
||||
return 403, "unknown run or task"
|
||||
if not str(headers.get("Content-Type", "")).startswith("application/x-tar"):
|
||||
return 415, "expected application/x-tar"
|
||||
dest = (self.base / task).resolve()
|
||||
try:
|
||||
with tarfile.open(fileobj=io.BytesIO(body), mode="r:") as tar:
|
||||
members = tar.getmembers()
|
||||
for m in members:
|
||||
if m.issym() or m.islnk() or not (dest / m.name).resolve().is_relative_to(dest):
|
||||
return 400, f"rejected member {m.name}"
|
||||
dest.mkdir(parents=True, exist_ok=True)
|
||||
tar.extractall(dest, members) # every member checked above
|
||||
except tarfile.TarError as exc:
|
||||
return 400, f"bad tar: {exc}"
|
||||
with self.lock:
|
||||
self.delivered.add(task)
|
||||
return 200, "ok"
|
||||
|
||||
|
||||
class _Handler(BaseHTTPRequestHandler):
|
||||
def do_POST(self) -> None:
|
||||
server: Receiver = self.server # type: ignore[assignment]
|
||||
body = self.rfile.read(int(self.headers.get("Content-Length") or 0))
|
||||
code, text = server.accept(self.headers, body)
|
||||
self.send_response(code)
|
||||
self.send_header("Content-Type", "text/plain")
|
||||
self.end_headers()
|
||||
self.wfile.write(text.encode())
|
||||
|
||||
def log_message(self, format: str, *args: object) -> None:
|
||||
return
|
||||
|
||||
|
||||
# --- the run -------------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
@@ -347,35 +299,45 @@ def _await(ax: Ax, name: str, delivered: Callable[[], bool], start: Callable[[],
|
||||
return None
|
||||
|
||||
|
||||
def _execute(ax: Ax, run: Run, entry: RunTask, manifests: Manifests, base: Path, receiver: Receiver, router: Router) -> tuple[str, str]:
|
||||
def _execute(
|
||||
ax: Ax, egress: Egress, run: Run, entry: RunTask, manifests: Manifests, base: Path, receiver: Receiver, router: Router
|
||||
) -> tuple[str, str]:
|
||||
ax_name, task = _ax_name(run.metadata.name, entry.name), manifests.tasks[entry.task]
|
||||
manifest = base / entry.name / "task.yaml"
|
||||
atespace, manifest = task.metadata.atespace or "default", base / entry.name / "task.yaml"
|
||||
manifest.parent.mkdir(parents=True, exist_ok=True)
|
||||
docs = render_task(run, entry, manifests, base, receiver.url())
|
||||
docs = render_task(run, entry, manifests, base, receiver.url, receiver.dial)
|
||||
manifest.write_text(yaml.safe_dump_all(docs, sort_keys=False), encoding="utf-8")
|
||||
model = manifests.models.get(task.metadata.annotations.get(MODEL_ANNOTATION) or "")
|
||||
policy = policy_for(task, [manifests.workspaces[b.name] for b in task.workspaces], model, urlsplit(receiver.url).hostname)
|
||||
try: # Req 4.6: the bound Model's credential is read here and only ever sent in the start request
|
||||
secrets = credentials(manifests.models.get(task.metadata.annotations.get(MODEL_ANNOTATION) or ""))
|
||||
start = router.starter(task.metadata.atespace or "default", ax_name, run.metadata.name, secrets)
|
||||
ax.apply(manifest)
|
||||
start = router.starter(atespace, ax_name, run.metadata.name, credentials(model))
|
||||
ax.apply(manifest) # creates the Substrate actor; its egress is denied until the policy exists
|
||||
egress.apply(ax_name, atespace, policy)
|
||||
egress.settle()
|
||||
except RuntimeError as exc: # StartError included: a Model-bound task whose secret is unset
|
||||
ax.delete(ax_name)
|
||||
return "failed", str(exc)
|
||||
try:
|
||||
failure = _await(ax, ax_name, lambda: entry.name in receiver.delivered, start)
|
||||
finally:
|
||||
ax.delete(ax_name)
|
||||
leftover = egress.delete(ax_name, atespace) # the policy goes with its actor; a survivor is reported
|
||||
if leftover:
|
||||
print(leftover, file=sys.stderr)
|
||||
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) -> str:
|
||||
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)."""
|
||||
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": {}})
|
||||
receiver = Receiver(run_spec.metadata.name, base, int(os.environ.get("OUSAST_ARTIFACT_PORT") or 0))
|
||||
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
|
||||
threading.Thread(target=receiver.serve_forever, daemon=True).start()
|
||||
cli, finished = Ax(ax), queue.Queue[tuple[str, str, str]]()
|
||||
cli, egress, finished = Ax(ax), Egress(kubectl_ate), 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 {}))
|
||||
@@ -386,7 +348,7 @@ def run(run_manifest: Path, *, workers: int = 1, ax: str = "ax", results_root: P
|
||||
if state.status(entry.name) in ("running", "suspended"):
|
||||
cli.delete(_ax_name(run_spec.metadata.name, entry.name)) # an interrupted attempt; ax must not keep it
|
||||
state.set(entry.name, status="running", started=_now(), finished=None)
|
||||
finished.put((entry.name, *_execute(cli, run_spec, entry, manifests, base, receiver, router)))
|
||||
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
|
||||
|
||||
@@ -43,6 +43,8 @@ Environment the runner reads, all optional except the two ax variables:
|
||||
OUSAST_CASE_CACHE a directory of git checkouts (``benchmarks/independent`` cache layout)
|
||||
OUSAST_GIT_PINS JSON ``{"<workspace>/<git name>": "<40-hex commit>"}`` from the reconciler
|
||||
OUSAST_ARTIFACT_URL, OUSAST_RUN where to deliver, and the run name carried in ``X-Ousast-Run``
|
||||
OUSAST_ARTIFACT_DIAL ``host:port`` to connect to instead of the URL's host, which is sent as
|
||||
``Host`` (the egress gateway decides on and resolves that name)
|
||||
OUSAST_DELIVERY_BACKOFF seconds before the second delivery attempt, doubled per retry
|
||||
|
||||
The runner exports ``OUSAST_WORKSPACE_DIR`` (the first bound workspace) when the Task's env leaves it unset,
|
||||
@@ -72,6 +74,7 @@ from dataclasses import dataclass
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
import yaml
|
||||
|
||||
@@ -632,11 +635,26 @@ def build_tar(output_dir: Path) -> bytes:
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
def deliver(url: str, payload: bytes, headers: Mapping[str, str], *, attempts: int = DELIVERY_ATTEMPTS, backoff: float = 1.0) -> int:
|
||||
"""POST ``payload`` to ``url`` with ``headers``; retry with doubling backoff; return the 2xx status."""
|
||||
last = "no attempt made"
|
||||
def deliver(
|
||||
url: str,
|
||||
payload: bytes,
|
||||
headers: Mapping[str, str],
|
||||
*,
|
||||
attempts: int = DELIVERY_ATTEMPTS,
|
||||
backoff: float = 1.0,
|
||||
dial: str | None = None,
|
||||
) -> int:
|
||||
"""POST ``payload`` to ``url`` with ``headers``; retry with doubling backoff; return the 2xx status.
|
||||
|
||||
``dial`` (``host:port``): connect there instead and send the URL's host as ``Host``, as the actor must -- it
|
||||
cannot resolve the receiver's cluster name, and the egress gateway decides on ``Host`` and resolves it itself.
|
||||
"""
|
||||
last, target, headers = "no attempt made", url, dict(headers)
|
||||
if dial:
|
||||
parts = urlsplit(url)
|
||||
target, headers["Host"] = parts._replace(netloc=dial).geturl(), parts.netloc
|
||||
for attempt in range(1, attempts + 1):
|
||||
request = urllib.request.Request(url, data=payload, method="POST", headers=dict(headers))
|
||||
request = urllib.request.Request(target, data=payload, method="POST", headers=headers)
|
||||
try:
|
||||
with urllib.request.urlopen(request, timeout=60) as response: # noqa: S310 - the reconciler's URL
|
||||
status = int(response.status)
|
||||
@@ -658,12 +676,13 @@ def _deliver_output(task: Task, output_dir: Path, environ: Mapping[str, str]) ->
|
||||
url = environ.get("OUSAST_ARTIFACT_URL")
|
||||
if not url:
|
||||
return False
|
||||
headers = {"Content-Type": "application/x-tar", "X-Ousast-Task": task.metadata.name}
|
||||
# the Run's task name (``OUSAST_TASK``) is what the receiver knows; the ax Task is named ``<run>-<task>``
|
||||
headers = {"Content-Type": "application/x-tar", "X-Ousast-Task": environ.get("OUSAST_TASK") or task.metadata.name}
|
||||
run = environ.get("OUSAST_RUN")
|
||||
if run:
|
||||
headers["X-Ousast-Run"] = run
|
||||
backoff = float(environ.get("OUSAST_DELIVERY_BACKOFF", "1.0"))
|
||||
deliver(url, build_tar(output_dir), headers, backoff=backoff)
|
||||
deliver(url, build_tar(output_dir), headers, backoff=backoff, dial=environ.get("OUSAST_ARTIFACT_DIAL") or None)
|
||||
return True
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user