From cafe0f5deabf8535ec5507b9d1bc6bdce8509785 Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Fri, 10 Jul 2026 15:43:53 +0000 Subject: [PATCH 01/31] test(k8s): add Tier-1 KIND lineage smoke harness Phase-0 spike for the Kubernetes training lineage integration: a KIND harness that proves the existing roar runtime works in vanilla training pods before any roar.backends.k8s code exists. - bootstrap_k8s.sh: pinned kind/kubectl download, 1cp+2worker cluster, dist/ wheel hostPath mounts, in-cluster glaas Service->host endpoint wiring, image pre-pull, pod->glaas preflight probe; destroy_k8s.sh - hand-wrapped single-pod Job template modeling the future `roar k8s prepare` output: wheel staging, `roar run --tracer preload` command wrap, Secret-delivered fragment-session credentials, downward-API identity env - roar-unaware synthetic trainer + wrapper that exports the recorded job as an ExecutionFragment (backend=k8s, task_id contract pod_uid:container:index:attempt) and streams it via the shared fragment transport - smoke tests (k8s_e2e marker): traced reads/writes with hashes, identity metadata, and host-side merge into .roar/roar.db through the shared fragment lineage engine Verified live: 3/3 passing repeatedly (~16s warm); default gate, ruff, and mypy green. Co-Authored-By: Claude Fable 5 --- .gitignore | 3 + pyproject.toml | 1 + tests/backends/k8s/README.md | 58 +++++ tests/backends/k8s/__init__.py | 0 tests/backends/k8s/e2e/__init__.py | 0 tests/backends/k8s/e2e/conftest.py | 230 ++++++++++++++++++ tests/backends/k8s/e2e/test_k8s_smoke.py | 155 ++++++++++++ tests/backends/k8s/kind-config.yaml | 26 ++ .../k8s/manifests/job-single.yaml.tpl | 78 ++++++ tests/backends/k8s/manifests/minio.yaml | 51 ++++ tests/backends/k8s/scripts/bootstrap_k8s.sh | 187 ++++++++++++++ tests/backends/k8s/scripts/destroy_k8s.sh | 22 ++ .../backends/k8s/workloads/emit_fragments.py | 54 ++++ tests/backends/k8s/workloads/roar-wrapper.sh | 28 +++ tests/backends/k8s/workloads/train.py | 36 +++ 15 files changed, 929 insertions(+) create mode 100644 tests/backends/k8s/README.md create mode 100644 tests/backends/k8s/__init__.py create mode 100644 tests/backends/k8s/e2e/__init__.py create mode 100644 tests/backends/k8s/e2e/conftest.py create mode 100644 tests/backends/k8s/e2e/test_k8s_smoke.py create mode 100644 tests/backends/k8s/kind-config.yaml create mode 100644 tests/backends/k8s/manifests/job-single.yaml.tpl create mode 100644 tests/backends/k8s/manifests/minio.yaml create mode 100755 tests/backends/k8s/scripts/bootstrap_k8s.sh create mode 100755 tests/backends/k8s/scripts/destroy_k8s.sh create mode 100644 tests/backends/k8s/workloads/emit_fragments.py create mode 100644 tests/backends/k8s/workloads/roar-wrapper.sh create mode 100644 tests/backends/k8s/workloads/train.py diff --git a/.gitignore b/.gitignore index 35b53a8f..71837386 100644 --- a/.gitignore +++ b/.gitignore @@ -230,3 +230,6 @@ tests/benchmarks/results/ # Design docs — internal only, not for version control docs/design/ scratch/ + +# k8s e2e harness downloaded tooling +tests/backends/k8s/.tools/ diff --git a/pyproject.toml b/pyproject.toml index 0b76d8ae..7795903f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -123,6 +123,7 @@ markers = [ "diagnostic: Opt-in diagnostics or aspirational performance budgets outside the default gate", "large_pipeline: Stress-style pipeline coverage with larger DAG fixtures", "osmo_e2e: OSMO end-to-end tests requiring a Docker Compose managed KIND harness", + "k8s_e2e: Kubernetes end-to-end tests requiring the KIND harness (tests/backends/k8s)", "ray_e2e: Ray end-to-end tests requiring a running Docker cluster", "ray_contract: User-facing Ray contract tests using `roar run ray job submit ...`", "ray_diagnostic: Diagnostic Ray tests that intentionally inspect internal runtime details", diff --git a/tests/backends/k8s/README.md b/tests/backends/k8s/README.md new file mode 100644 index 00000000..3a3759ea --- /dev/null +++ b/tests/backends/k8s/README.md @@ -0,0 +1,58 @@ +# k8s Lineage E2E Harness (Tier 1) + +KIND-based harness for pressure-testing roar lineage capture in Kubernetes +training pods. This is the Phase-0/Tier-1 slice from +`design-docs/k8s-training-lineage-integration.md`: no `roar.backends.k8s` +exists yet — the fixtures hand-wrap a Job manifest the way the future +`roar k8s prepare` will, so the runtime assumptions (in-pod tracing, wheel +staging, fragment streaming, identity contract) are proven first. + +## Prerequisites + +- Docker +- A packaged wheel: `bash scripts/build_wheel_with_bins.sh` (repo root) +- Local glaas-api on `http://localhost:3001` (e.g. via pm2) +- `kind`/`kubectl` are downloaded automatically into `.tools/bin` if missing + +## Usage + +```bash +# one-time (and after wheel changes) +bash scripts/build_wheel_with_bins.sh + +# create cluster + wire glaas + preflight (add --with-minio for S3 scenarios) +bash tests/backends/k8s/scripts/bootstrap_k8s.sh + +# run the smoke tests (addopts override needed: e2e dirs are ignored by default) +pytest tests/backends/k8s/e2e -o addopts='' -m k8s_e2e -v + +# tear down +bash tests/backends/k8s/scripts/destroy_k8s.sh +``` + +## Topology + +- KIND cluster `roar-k8s-e2e`: 1 control plane + 2 workers, k8s 1.33 +- `dist/` mounted into nodes at `/roar-dist` → pods install the wheel from a + hostPath volume (no network fetch, no stale artifact URLs) +- Host-visible vs cluster-visible endpoints are modeled separately on purpose: + - glaas: `http://localhost:3001` (host) vs `http://glaas:3001` (pods, via a + Service/Endpoints pair pointing at the kind docker-network gateway) + - MinIO (optional): `http://localhost:39000` (host) vs `http://minio:9000` + (pods) + +## What the smoke test proves + +1. The packaged wheel installs and traces (preload) inside a vanilla + `python:3.12-slim` pod. +2. A roar-unaware training script's file I/O is captured with content hashes. +3. Fragments stream from the pod to glaas-api through the encrypted + fragment-session pipeline using Secret-delivered credentials. +4. Fragments carry the k8s identity contract + (`pod_uid:container:completion_index:restart_attempt` + pod/node metadata). +5. Decrypted fragments merge into a local `.roar/roar.db` via the shared + fragment lineage engine. + +Infra diagnosis lives in `scripts/bootstrap_k8s.sh` (preflight probe), not in +the tests: if the tests skip or fail, re-run bootstrap first to separate +infra failures from product failures. diff --git a/tests/backends/k8s/__init__.py b/tests/backends/k8s/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/backends/k8s/e2e/__init__.py b/tests/backends/k8s/e2e/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/backends/k8s/e2e/conftest.py b/tests/backends/k8s/e2e/conftest.py new file mode 100644 index 00000000..30dfda26 --- /dev/null +++ b/tests/backends/k8s/e2e/conftest.py @@ -0,0 +1,230 @@ +"""Fixtures for the Tier-1 KIND k8s lineage smoke tests. + +Requires the harness cluster created by +`tests/backends/k8s/scripts/bootstrap_k8s.sh` and a local glaas-api on +http://localhost:3001. Tests skip (with instructions) when either is +missing so the default gate stays green. +""" + +from __future__ import annotations + +import base64 +import json +import shutil +import string +import subprocess +import time +import urllib.parse +import urllib.request +import uuid +from collections.abc import Sequence +from pathlib import Path +from typing import Any + +import pytest +from cryptography.hazmat.primitives.ciphers.aead import AESGCM + +HARNESS_DIR = Path(__file__).resolve().parent.parent +MANIFESTS_DIR = HARNESS_DIR / "manifests" +WORKLOADS_DIR = HARNESS_DIR / "workloads" +TOOLS_BIN = HARNESS_DIR / ".tools" / "bin" + +CLUSTER_NAME = "roar-k8s-e2e" +KUBE_CONTEXT = f"kind-{CLUSTER_NAME}" +NAMESPACE = "roar-e2e" +HOST_GLAAS_URL = "http://localhost:3001" +JOB_TIMEOUT_SECONDS = 420 +BOOTSTRAP_HINT = "run: bash tests/backends/k8s/scripts/bootstrap_k8s.sh" + + +def _kubectl_bin() -> str | None: + tool = TOOLS_BIN / "kubectl" + if tool.is_file(): + return str(tool) + return shutil.which("kubectl") + + +def kubectl( + args: Sequence[str], + *, + input_text: str | None = None, + check: bool = True, +) -> subprocess.CompletedProcess[str]: + binary = _kubectl_bin() + assert binary, f"kubectl not found; {BOOTSTRAP_HINT}" + result = subprocess.run( + [binary, "--context", KUBE_CONTEXT, *args], + input=input_text, + capture_output=True, + text=True, + check=False, + ) + if check and result.returncode != 0: + raise RuntimeError( + f"kubectl {' '.join(args)} failed ({result.returncode}):\n" + f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + ) + return result + + +@pytest.fixture(scope="session") +def k8s_cluster() -> None: + if _kubectl_bin() is None: + pytest.skip(f"kubectl not available; {BOOTSTRAP_HINT}") + result = kubectl(["cluster-info"], check=False) + if result.returncode != 0: + pytest.skip(f"KIND cluster {CLUSTER_NAME} not reachable; {BOOTSTRAP_HINT}") + result = kubectl(["get", "service", "glaas", "-n", NAMESPACE], check=False) + if result.returncode != 0: + pytest.skip(f"glaas Service missing in namespace {NAMESPACE}; {BOOTSTRAP_HINT}") + + +@pytest.fixture(scope="session") +def glaas_health() -> str: + try: + with urllib.request.urlopen(f"{HOST_GLAAS_URL}/api/v1/health", timeout=5) as response: + assert response.status == 200 + except Exception: + pytest.skip(f"glaas-api not reachable at {HOST_GLAAS_URL}; start it first") + return HOST_GLAAS_URL + + +def register_fragment_session_with_glaas() -> dict[str, str]: + from roar.execution.fragments.sessions import generate_fragment_session + from roar.integrations.glaas import GlaasClient + + session = generate_fragment_session() + client = GlaasClient(base_url=HOST_GLAAS_URL, force_anonymous=True) + result, error = client.register_fragment_session( + session["session_id"], + session["token_hash"], + ttl_seconds=3600, + ) + assert error is None, f"fragment session registration failed: {error}" + assert result is not None + return session + + +def fetch_fragment_batches(session_id: str, token: str) -> list[dict[str, Any]]: + encoded_token = urllib.parse.quote(token, safe="") + url = f"{HOST_GLAAS_URL}/api/v1/fragments/sessions/{session_id}/fragments?token={encoded_token}" + with urllib.request.urlopen(url, timeout=10) as response: + payload = json.loads(response.read().decode("utf-8")) + fragments = payload.get("fragments") + if fragments is None and isinstance(payload.get("data"), dict): + fragments = payload["data"].get("fragments") + assert isinstance(fragments, list), f"Expected fragment batches from {url}, got: {payload!r}" + return [item for item in fragments if isinstance(item, dict)] + + +def decrypt_fragment_batches( + batches: Sequence[dict[str, Any]], + token: str, +) -> list[dict[str, Any]]: + key = bytes.fromhex(token) + decrypted: list[dict[str, Any]] = [] + for batch in batches: + encrypted_batch = batch.get("encrypted_batch") + if not isinstance(encrypted_batch, str) or not encrypted_batch: + continue + payload = base64.b64decode(encrypted_batch) + plaintext = AESGCM(key).decrypt(payload[:12], payload[12:], None) + decoded = json.loads(plaintext.decode("utf-8")) + if isinstance(decoded, list): + decrypted.extend(item for item in decoded if isinstance(item, dict)) + return decrypted + + +def _render_job_manifest(job_name: str, configmap_name: str, secret_name: str) -> str: + template = string.Template((MANIFESTS_DIR / "job-single.yaml.tpl").read_text(encoding="utf-8")) + return template.substitute( + job_name=job_name, + namespace=NAMESPACE, + configmap_name=configmap_name, + secret_name=secret_name, + ) + + +def _wait_for_job(job_name: str) -> tuple[bool, str]: + """Poll the Job until it completes or fails, so failures surface fast.""" + deadline = time.time() + JOB_TIMEOUT_SECONDS + while time.time() < deadline: + result = kubectl(["get", f"job/{job_name}", "-n", NAMESPACE, "-o", "json"], check=False) + if result.returncode == 0: + payload = json.loads(result.stdout) + for condition in payload.get("status", {}).get("conditions") or []: + if condition.get("status") != "True": + continue + if condition.get("type") in ("Complete", "SuccessCriteriaMet"): + return True, "" + if condition.get("type") in ("Failed", "FailureTarget"): + return False, str(condition.get("message") or "job failed") + time.sleep(5) + return False, f"timed out after {JOB_TIMEOUT_SECONDS}s" + + +def _pod_logs_for_job(job_name: str) -> str: + result = kubectl( + ["logs", "-n", NAMESPACE, "-l", f"job-name={job_name}", "--tail=200"], + check=False, + ) + return result.stdout + result.stderr + + +@pytest.fixture(scope="module") +def smoke_run(k8s_cluster: None, glaas_health: str) -> dict[str, Any]: + """Run the wrapped single-pod training Job once and return its lineage.""" + suffix = uuid.uuid4().hex[:6] + job_name = f"roar-smoke-{suffix}" + configmap_name = f"roar-smoke-workload-{suffix}" + secret_name = f"roar-smoke-session-{suffix}" + + session = register_fragment_session_with_glaas() + + kubectl( + [ + "create", + "configmap", + configmap_name, + "-n", + NAMESPACE, + f"--from-file={WORKLOADS_DIR}", + ] + ) + kubectl( + [ + "create", + "secret", + "generic", + secret_name, + "-n", + NAMESPACE, + f"--from-literal=session_id={session['session_id']}", + f"--from-literal=token={session['token']}", + ] + ) + + try: + manifest = _render_job_manifest(job_name, configmap_name, secret_name) + kubectl(["apply", "-f", "-"], input_text=manifest) + + succeeded, failure_reason = _wait_for_job(job_name) + logs = _pod_logs_for_job(job_name) + assert succeeded, f"Job {job_name} did not complete: {failure_reason}\npod logs:\n{logs}" + + batches = fetch_fragment_batches(session["session_id"], session["token"]) + fragments = decrypt_fragment_batches(batches, session["token"]) + return { + "job_name": job_name, + "session": session, + "batches": batches, + "fragments": fragments, + "logs": logs, + } + finally: + for kind_name in ( + f"job/{job_name}", + f"configmap/{configmap_name}", + f"secret/{secret_name}", + ): + kubectl(["delete", kind_name, "-n", NAMESPACE, "--ignore-not-found"], check=False) diff --git a/tests/backends/k8s/e2e/test_k8s_smoke.py b/tests/backends/k8s/e2e/test_k8s_smoke.py new file mode 100644 index 00000000..e02d82e9 --- /dev/null +++ b/tests/backends/k8s/e2e/test_k8s_smoke.py @@ -0,0 +1,155 @@ +"""Tier-1 smoke: single-pod k8s Job lineage captured and streamed to GLaaS. + +Phase-0 product path (pre-backend): a roar-unaware training script runs in +a vanilla pod under `roar run`, the recorded job is exported as an +execution fragment, streamed to the local glaas-api through the shared +fragment-session pipeline, then fetched, decrypted, and merged into a +local roar db on the host — the same loop the future `roar k8s` backend +will own end to end. +""" + +from __future__ import annotations + +import sqlite3 +from pathlib import Path +from typing import Any + +import pytest + +from .conftest import NAMESPACE + +pytestmark = [ + pytest.mark.e2e, + pytest.mark.k8s_e2e, + # The module fixture runs a full Job (image pull + wheel install + trace) + # inside the first test's budget; the default 60s gate timeout is too low. + pytest.mark.timeout(600), +] + +SMOKE_TASK_NAME = "k8s-smoke-train" + + +def _smoke_fragment(smoke_run: dict[str, Any]) -> dict[str, Any]: + fragments = [ + fragment + for fragment in smoke_run["fragments"] + if fragment.get("task_name") == SMOKE_TASK_NAME + ] + assert fragments, ( + f"No {SMOKE_TASK_NAME} fragment found in " + f"{len(smoke_run['fragments'])} decrypted fragment(s).\n" + f"pod logs:\n{smoke_run['logs']}" + ) + return fragments[-1] + + +def _refs_by_suffix(refs: list[dict[str, Any]], suffix: str) -> dict[str, Any] | None: + for ref in refs: + if str(ref.get("path", "")).endswith(suffix): + return ref + return None + + +def test_job_streams_traced_lineage_fragment(smoke_run: dict[str, Any]) -> None: + fragment = _smoke_fragment(smoke_run) + + assert fragment.get("backend") == "k8s" + assert fragment.get("exit_code") == 0 + + reads = [ref for ref in fragment.get("reads", []) if isinstance(ref, dict)] + writes = [ref for ref in fragment.get("writes", []) if isinstance(ref, dict)] + + dataset = _refs_by_suffix(reads, "dataset.csv") + assert dataset is not None, f"dataset.csv not in reads: {reads}\nlogs:\n{smoke_run['logs']}" + assert dataset.get("hash"), f"dataset.csv read has no hash: {dataset}" + + for expected in ("model.bin", "metrics.json"): + ref = _refs_by_suffix(writes, expected) + assert ref is not None, f"{expected} not in writes: {writes}\nlogs:\n{smoke_run['logs']}" + assert ref.get("hash"), f"{expected} write has no hash: {ref}" + assert int(ref.get("size") or 0) > 0, f"{expected} write has no size: {ref}" + + +def test_fragment_carries_k8s_identity_metadata(smoke_run: dict[str, Any]) -> None: + fragment = _smoke_fragment(smoke_run) + metadata = fragment.get("backend_metadata") or {} + + pod_uid = metadata.get("k8s_pod_uid") + assert pod_uid, f"missing k8s_pod_uid in backend_metadata: {metadata}" + assert metadata.get("k8s_namespace") == NAMESPACE + assert metadata.get("k8s_pod_name", "").startswith("roar-smoke-") + assert metadata.get("k8s_node_name") + + task_id = str(fragment.get("task_id") or "") + assert task_id == f"{pod_uid}:trainer:0:0", ( + f"task_id does not follow the pod-uid:container:index:attempt contract: {task_id}" + ) + + +def test_fragments_reconstitute_into_local_db(smoke_run: dict[str, Any], tmp_path: Path) -> None: + from roar.db.context import create_database_context + from roar.execution.fragments.lineage import ( + FragmentLineageBackend, + merge_execution_fragments, + ) + from roar.execution.fragments.models import ( + ExecutionFragment, + derive_fragment_identity, + ) + + k8s_test_lineage_backend = FragmentLineageBackend( + job_type="k8s_task", + command_for_fragment=lambda fragment: f"k8s_task:{fragment.task_name}", + script_for_fragment=lambda fragment: fragment.task_name or None, + execution_role_from_fragment=lambda fragment, _parent: "task", + metadata_from_fragment=lambda fragment, _parent: { + "k8s_task_id": fragment.task_id, + **(fragment.backend_metadata or {}), + }, + task_identity_from_metadata=lambda parent_job_uid, job_uid, metadata: ( + derive_fragment_identity( + "k8s", + parent_job_uid, + str(metadata.get("k8s_task_id") or ""), + job_uid, + ) + ), + ) + + project_dir = tmp_path / "project" + roar_dir = project_dir / ".roar" + roar_dir.mkdir(parents=True) + db_path = roar_dir / "roar.db" + with create_database_context(roar_dir): + pass # connect() initializes the schema + + fragments = [ + ExecutionFragment.from_dict(fragment) + for fragment in smoke_run["fragments"] + if fragment.get("task_name") == SMOKE_TASK_NAME + ] + assert fragments, "no smoke fragments available for reconstitution" + + merge_execution_fragments( + fragments=fragments, + project_dir=str(project_dir), + backend=k8s_test_lineage_backend, + ) + + conn = sqlite3.connect(db_path) + conn.row_factory = sqlite3.Row + try: + jobs = conn.execute("SELECT * FROM jobs WHERE job_type = 'k8s_task'").fetchall() + assert len(jobs) == 1, f"expected one merged k8s_task job, got {len(jobs)}" + + artifact_paths = { + str(row["path"]) + for row in conn.execute( + "SELECT path FROM job_outputs WHERE job_id = ?", + (jobs[0]["id"],), + ).fetchall() + } + assert any(path.endswith("model.bin") for path in artifact_paths), artifact_paths + assert any(path.endswith("metrics.json") for path in artifact_paths), artifact_paths + finally: + conn.close() diff --git a/tests/backends/k8s/kind-config.yaml b/tests/backends/k8s/kind-config.yaml new file mode 100644 index 00000000..b29c5a6d --- /dev/null +++ b/tests/backends/k8s/kind-config.yaml @@ -0,0 +1,26 @@ +# KIND cluster template for the roar k8s lineage e2e harness. +# +# __ROAR_DIST_DIR__ is replaced by scripts/bootstrap_k8s.sh with the absolute +# path of the repo's dist/ directory so pods can install the locally built +# roar_cli wheel from a hostPath volume (no network fetch, no stale URLs). +# +# The control-plane node maps the MinIO NodePort (30900) to host port 39000 so +# S3 scenarios can model host-visible vs cluster-visible endpoints explicitly. +kind: Cluster +apiVersion: kind.x-k8s.io/v1alpha4 +nodes: + - role: control-plane + extraPortMappings: + - containerPort: 30900 + hostPort: 39000 + protocol: TCP + - role: worker + extraMounts: + - hostPath: __ROAR_DIST_DIR__ + containerPath: /roar-dist + readOnly: true + - role: worker + extraMounts: + - hostPath: __ROAR_DIST_DIR__ + containerPath: /roar-dist + readOnly: true diff --git a/tests/backends/k8s/manifests/job-single.yaml.tpl b/tests/backends/k8s/manifests/job-single.yaml.tpl new file mode 100644 index 00000000..06e1448b --- /dev/null +++ b/tests/backends/k8s/manifests/job-single.yaml.tpl @@ -0,0 +1,78 @@ +# Single-pod training Job template for the Tier-1 smoke test. +# +# Rendered by tests/backends/k8s/e2e/conftest.py with string.Template. +# This is the hand-wrapped stand-in for what `roar k8s prepare` will +# eventually generate: runtime staged from a hostPath wheel mount, the +# command wrapped through roar, fragment-session credentials from a +# Secret (never inline env), and identity from the downward API. +apiVersion: batch/v1 +kind: Job +metadata: + name: ${job_name} + namespace: ${namespace} + labels: + app.kubernetes.io/part-of: roar-k8s-e2e +spec: + backoffLimit: 0 + ttlSecondsAfterFinished: 1800 + template: + metadata: + labels: + app.kubernetes.io/part-of: roar-k8s-e2e + spec: + restartPolicy: Never + volumes: + - name: roar-wheels + hostPath: + path: /roar-dist + type: Directory + - name: workload + configMap: + name: ${configmap_name} + - name: work + emptyDir: {} + containers: + - name: trainer + image: python:3.12-slim + command: ["bash", "/workload/roar-wrapper.sh"] + workingDir: /work + env: + - name: ROAR_SESSION_ID + valueFrom: + secretKeyRef: + name: ${secret_name} + key: session_id + - name: ROAR_FRAGMENT_TOKEN + valueFrom: + secretKeyRef: + name: ${secret_name} + key: token + - name: GLAAS_URL + value: http://glaas:3001 + - name: ROAR_NO_TELEMETRY + value: "1" + - name: ROAR_K8S_NAMESPACE + valueFrom: + fieldRef: + fieldPath: metadata.namespace + - name: ROAR_K8S_POD_NAME + valueFrom: + fieldRef: + fieldPath: metadata.name + - name: ROAR_K8S_POD_UID + valueFrom: + fieldRef: + fieldPath: metadata.uid + - name: ROAR_K8S_NODE_NAME + valueFrom: + fieldRef: + fieldPath: spec.nodeName + volumeMounts: + - name: roar-wheels + mountPath: /wheels + readOnly: true + - name: workload + mountPath: /workload + readOnly: true + - name: work + mountPath: /work diff --git a/tests/backends/k8s/manifests/minio.yaml b/tests/backends/k8s/manifests/minio.yaml new file mode 100644 index 00000000..cdfa099c --- /dev/null +++ b/tests/backends/k8s/manifests/minio.yaml @@ -0,0 +1,51 @@ +# MinIO for S3 lineage scenarios. +# +# Cluster-visible endpoint: http://minio:9000 (namespace roar-e2e) +# Host-visible endpoint: http://localhost:39000 (NodePort 30900 mapped by +# kind-config.yaml on the control-plane node) +apiVersion: apps/v1 +kind: Deployment +metadata: + name: minio + namespace: roar-e2e +spec: + replicas: 1 + selector: + matchLabels: + app: minio + template: + metadata: + labels: + app: minio + spec: + containers: + - name: minio + image: quay.io/minio/minio:latest + args: ["server", "/data"] + env: + - name: MINIO_ROOT_USER + value: minioadmin + - name: MINIO_ROOT_PASSWORD + value: minioadmin + ports: + - containerPort: 9000 + readinessProbe: + httpGet: + path: /minio/health/ready + port: 9000 + initialDelaySeconds: 2 + periodSeconds: 3 +--- +apiVersion: v1 +kind: Service +metadata: + name: minio + namespace: roar-e2e +spec: + type: NodePort + selector: + app: minio + ports: + - port: 9000 + targetPort: 9000 + nodePort: 30900 diff --git a/tests/backends/k8s/scripts/bootstrap_k8s.sh b/tests/backends/k8s/scripts/bootstrap_k8s.sh new file mode 100755 index 00000000..e9e4927e --- /dev/null +++ b/tests/backends/k8s/scripts/bootstrap_k8s.sh @@ -0,0 +1,187 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Tier-1 KIND harness bootstrap for the roar k8s lineage e2e tests. +# +# What it does: +# 1. Ensures kind/kubectl are available (downloads pinned versions into +# tests/backends/k8s/.tools/bin when missing from PATH). +# 2. Requires a packaged roar_cli wheel in /dist (built with +# scripts/build_wheel_with_bins.sh) and mounts dist/ into cluster nodes. +# 3. Creates the KIND cluster (1 control plane + 2 workers). +# 4. Wires an in-cluster `glaas` Service/Endpoints to the host glaas-api so +# pods use the cluster-visible URL (http://glaas:3001) while the host +# keeps using the host-visible URL (http://localhost:3001). +# 5. Pre-pulls the workload image and preflights pod -> glaas reachability +# with a probe pod, so infra failures are diagnosed here rather than +# inside product tests. +# 6. Optionally deploys MinIO for S3 scenarios (--with-minio). +# +# Usage: +# bash tests/backends/k8s/scripts/bootstrap_k8s.sh [--with-minio] [--skip-glaas] + +HARNESS_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +REPO_ROOT="$(cd "$HARNESS_DIR/../../.." && pwd)" +TOOLS_BIN="$HARNESS_DIR/.tools/bin" +DIST_DIR="$REPO_ROOT/dist" + +KIND_VERSION="v0.29.0" +KUBECTL_VERSION="v1.33.1" +CLUSTER_NAME="roar-k8s-e2e" +KUBE_CONTEXT="kind-${CLUSTER_NAME}" +NAMESPACE="roar-e2e" +WORKLOAD_IMAGE="docker.io/library/python:3.12-slim" +HOST_GLAAS_URL="${HOST_GLAAS_URL:-http://localhost:3001}" +CLUSTER_GLAAS_PORT=3001 + +WITH_MINIO=0 +SKIP_GLAAS=0 +for arg in "$@"; do + case "$arg" in + --with-minio) WITH_MINIO=1 ;; + --skip-glaas) SKIP_GLAAS=1 ;; + *) + echo "error: unknown flag: $arg" >&2 + exit 2 + ;; + esac +done + +case "$(uname -m)" in + x86_64 | amd64) ARCH="amd64" ;; + aarch64 | arm64) ARCH="arm64" ;; + *) + echo "error: unsupported host arch $(uname -m)" >&2 + exit 1 + ;; +esac + +mkdir -p "$TOOLS_BIN" +export PATH="$TOOLS_BIN:$PATH" + +ensure_tool() { + local name="$1" + local url="$2" + if command -v "$name" >/dev/null 2>&1; then + return + fi + echo "▶ Downloading $name into $TOOLS_BIN" + curl -fsSL -o "$TOOLS_BIN/$name" "$url" + chmod +x "$TOOLS_BIN/$name" +} + +ensure_tool kind "https://kind.sigs.k8s.io/dl/${KIND_VERSION}/kind-linux-${ARCH}" +ensure_tool kubectl "https://dl.k8s.io/release/${KUBECTL_VERSION}/bin/linux/${ARCH}/kubectl" + +kubectl_ctx() { + kubectl --context "$KUBE_CONTEXT" "$@" +} + +echo "▶ Checking for a packaged roar_cli wheel in $DIST_DIR" +shopt -s nullglob +wheels=("$DIST_DIR"/roar_cli-*.whl) +shopt -u nullglob +if ((${#wheels[@]} == 0)); then + echo "error: no roar_cli wheel in $DIST_DIR" >&2 + echo "hint: build one first: bash scripts/build_wheel_with_bins.sh" >&2 + exit 1 +fi +if ! printf '%s\n' "${wheels[@]}" | grep -Eq 'cp312|abi3'; then + echo "warning: no cp312/abi3 wheel found; the workload image ($WORKLOAD_IMAGE) runs Python 3.12" >&2 +fi +echo " found: $(basename "${wheels[-1]}")" + +if kind get clusters 2>/dev/null | grep -qx "$CLUSTER_NAME"; then + echo "▶ KIND cluster $CLUSTER_NAME already exists, reusing it" +else + echo "▶ Creating KIND cluster $CLUSTER_NAME" + rendered_config="$(mktemp)" + trap 'rm -f "$rendered_config"' EXIT + sed "s|__ROAR_DIST_DIR__|$DIST_DIR|g" "$HARNESS_DIR/kind-config.yaml" >"$rendered_config" + kind create cluster --name "$CLUSTER_NAME" --config "$rendered_config" --wait 180s +fi + +echo "▶ Ensuring namespace $NAMESPACE" +kubectl_ctx create namespace "$NAMESPACE" --dry-run=client -o yaml | kubectl_ctx apply -f - + +echo "▶ Wiring in-cluster glaas Service to the host glaas-api" +gateway_ip="$(docker network inspect kind -f '{{range .IPAM.Config}}{{.Gateway}} {{end}}' \ + | grep -Eo '[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+' | head -1)" +if [[ -z "$gateway_ip" ]]; then + echo "error: could not resolve the kind docker network gateway IP" >&2 + exit 1 +fi +host_glaas_port="${HOST_GLAAS_URL##*:}" +kubectl_ctx apply -f - </dev/null & +done +wait + +if ((SKIP_GLAAS == 0)); then + echo "▶ Checking host glaas-api at $HOST_GLAAS_URL" + if ! curl -fsS --max-time 5 "$HOST_GLAAS_URL/api/v1/health" >/dev/null; then + echo "error: glaas-api is not reachable at $HOST_GLAAS_URL" >&2 + echo "hint: start it (e.g. pm2 start glaas-api) or pass --skip-glaas" >&2 + exit 1 + fi + + echo "▶ Preflight: pod -> glaas reachability probe" + kubectl_ctx -n "$NAMESPACE" delete pod roar-glaas-probe --ignore-not-found >/dev/null + kubectl_ctx -n "$NAMESPACE" run roar-glaas-probe \ + --image="$WORKLOAD_IMAGE" --restart=Never --command -- \ + python -c "import urllib.request; r = urllib.request.urlopen('http://glaas:${CLUSTER_GLAAS_PORT}/api/v1/health', timeout=10); print('glaas reachable from pod:', r.status)" + if ! kubectl_ctx -n "$NAMESPACE" wait --for=jsonpath='{.status.phase}'=Succeeded \ + pod/roar-glaas-probe --timeout=120s >/dev/null; then + echo "error: probe pod could not reach glaas from inside the cluster" >&2 + kubectl_ctx -n "$NAMESPACE" describe pod roar-glaas-probe >&2 || true + kubectl_ctx -n "$NAMESPACE" logs roar-glaas-probe >&2 || true + echo "hint: is glaas-api listening on 0.0.0.0? is a firewall blocking ${gateway_ip}?" >&2 + exit 1 + fi + kubectl_ctx -n "$NAMESPACE" logs roar-glaas-probe + kubectl_ctx -n "$NAMESPACE" delete pod roar-glaas-probe >/dev/null +fi + +if ((WITH_MINIO == 1)); then + echo "▶ Deploying MinIO" + kubectl_ctx apply -f "$HARNESS_DIR/manifests/minio.yaml" + kubectl_ctx -n "$NAMESPACE" rollout status deployment/minio --timeout=180s +fi + +echo +echo "✓ Harness ready" +echo " cluster: $CLUSTER_NAME (context $KUBE_CONTEXT)" +echo " namespace: $NAMESPACE" +echo " host glaas URL: $HOST_GLAAS_URL" +echo " cluster glaas URL: http://glaas:${CLUSTER_GLAAS_PORT} (via ${gateway_ip})" +if ((WITH_MINIO == 1)); then + echo " MinIO (host): http://localhost:39000" + echo " MinIO (cluster): http://minio:9000" +fi +echo +echo "Run the smoke tests:" +echo " pytest tests/backends/k8s/e2e -o addopts='' -m k8s_e2e -v" diff --git a/tests/backends/k8s/scripts/destroy_k8s.sh b/tests/backends/k8s/scripts/destroy_k8s.sh new file mode 100755 index 00000000..12365d91 --- /dev/null +++ b/tests/backends/k8s/scripts/destroy_k8s.sh @@ -0,0 +1,22 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Tears down the roar k8s lineage e2e KIND cluster. + +HARNESS_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +TOOLS_BIN="$HARNESS_DIR/.tools/bin" +CLUSTER_NAME="roar-k8s-e2e" + +export PATH="$TOOLS_BIN:$PATH" + +if ! command -v kind >/dev/null 2>&1; then + echo "kind is not installed; nothing to destroy" >&2 + exit 0 +fi + +if kind get clusters 2>/dev/null | grep -qx "$CLUSTER_NAME"; then + kind delete cluster --name "$CLUSTER_NAME" + echo "✓ Deleted KIND cluster $CLUSTER_NAME" +else + echo "KIND cluster $CLUSTER_NAME does not exist" +fi diff --git a/tests/backends/k8s/workloads/emit_fragments.py b/tests/backends/k8s/workloads/emit_fragments.py new file mode 100644 index 00000000..65d9dd20 --- /dev/null +++ b/tests/backends/k8s/workloads/emit_fragments.py @@ -0,0 +1,54 @@ +"""Export the recorded in-pod roar job and stream it to GLaaS. + +Phase-0 seam: reuses the OSMO bundle exporter (backend-neutral in +practice) and the shared fragment transport. The future k8s worker +bootstrap replaces this file; the identity contract it encodes +(pod UID + container + node index + attempt) is the piece under test. +""" + +import json +import os +import sys +from pathlib import Path + +from roar.backends.osmo.export import export_osmo_lineage_bundle +from roar.execution.fragments.transport import emit_fragment_dicts + +TASK_NAME = "k8s-smoke-train" + + +def main() -> int: + pod_uid = os.environ.get("ROAR_K8S_POD_UID", "unknown-pod") + completion_index = os.environ.get("JOB_COMPLETION_INDEX", "0") + restart_attempt = "0" + task_id = f"{pod_uid}:trainer:{completion_index}:{restart_attempt}" + + bundle_path = Path("/tmp/roar-fragments.json") + export = export_osmo_lineage_bundle( + roar_dir=Path.cwd() / ".roar", + output_path=bundle_path, + task_id=task_id, + task_name=TASK_NAME, + backend_name="k8s", + ) + print(f"[emit-fragments] exported job {export.exported_job_uid} as task {task_id}") + + payload = json.loads(bundle_path.read_text(encoding="utf-8")) + fragments = payload["fragments"] + for fragment in fragments: + metadata = fragment.setdefault("backend_metadata", {}) + metadata["k8s_namespace"] = os.environ.get("ROAR_K8S_NAMESPACE") + metadata["k8s_pod_name"] = os.environ.get("ROAR_K8S_POD_NAME") + metadata["k8s_pod_uid"] = pod_uid + metadata["k8s_node_name"] = os.environ.get("ROAR_K8S_NODE_NAME") + metadata["k8s_container"] = "trainer" + metadata["k8s_completion_index"] = completion_index + metadata["k8s_restart_attempt"] = restart_attempt + + result = emit_fragment_dicts(fragments) + print(f"[emit-fragments] emit result: {result} ({len(fragments)} fragment(s))") + return 0 if result == "streamed" else 3 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/backends/k8s/workloads/roar-wrapper.sh b/tests/backends/k8s/workloads/roar-wrapper.sh new file mode 100644 index 00000000..a57d7d10 --- /dev/null +++ b/tests/backends/k8s/workloads/roar-wrapper.sh @@ -0,0 +1,28 @@ +#!/usr/bin/env bash +set -euo pipefail + +# In-pod wrapper for the Tier-1 smoke test. +# +# This is the hand-rolled stand-in for the future `roar k8s` runtime wrapper: +# stage the roar runtime from the mounted wheel, trace the (roar-unaware) +# training command, then export the recorded job as an execution fragment and +# stream it to GLaaS via the fragment-session credentials in the environment. + +echo "[roar-wrapper] installing roar-cli from /wheels" +# The wheel comes from the mounted dist/; its dependencies (blake3, click, +# ...) still come from the index. A fully hermetic runtime is the job of the +# future roar-runtime image (see the k8s integration design doc). +wheel="$(ls /wheels/roar_cli-*.whl | sort -V | tail -1)" +pip install --quiet "$wheel" +roar --version + +mkdir -p /work/project +cd /work/project + +printf 'x,y\n1.0,2.0\n2.0,3.9\n3.0,6.1\n4.0,8.2\n' >dataset.csv + +roar init -n +roar run --tracer "${ROAR_K8S_TRACER:-preload}" python /workload/train.py dataset.csv + +echo "[roar-wrapper] exporting and streaming lineage fragments" +python /workload/emit_fragments.py diff --git a/tests/backends/k8s/workloads/train.py b/tests/backends/k8s/workloads/train.py new file mode 100644 index 00000000..a8a8fad4 --- /dev/null +++ b/tests/backends/k8s/workloads/train.py @@ -0,0 +1,36 @@ +"""Roar-unaware synthetic trainer for the k8s smoke test. + +Reads a CSV dataset, "trains" a scalar weight, and writes a model +artifact plus metrics. Deliberately torch-free so Tier-1 runs stay +fast on a slim CPU image; distributed torchrun workloads arrive with +the Tier-2 harness. +""" + +import hashlib +import json +import sys +from pathlib import Path + + +def main() -> None: + dataset = Path(sys.argv[1] if len(sys.argv) > 1 else "dataset.csv") + rows = dataset.read_text(encoding="utf-8").strip().splitlines() + + weight = 0.0 + for _epoch in range(3): + for row in rows[1:]: + x_raw, y_raw = row.split(",")[:2] + x, y = float(x_raw), float(y_raw) + weight += 0.01 * (y - weight * x) * x + + digest = hashlib.blake2b(f"{weight:.9f}".encode()).digest() + Path("model.bin").write_bytes(digest * 8) + Path("metrics.json").write_text( + json.dumps({"rows": len(rows) - 1, "weight": weight}) + "\n", + encoding="utf-8", + ) + print(f"trained weight={weight:.6f} over {len(rows) - 1} rows") + + +if __name__ == "__main__": + main() From 7bde87d4f8cdabb9650e04ad10cca3007dfa9d30 Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Fri, 10 Jul 2026 17:25:36 +0000 Subject: [PATCH 02/31] feat(k8s): add Kubernetes Job execution backend Phase 1 of the k8s training lineage integration: a `roar.backends.k8s` backend that instruments plain batch/v1 Jobs submitted through `roar run kubectl apply|create -f `, gated by k8s.enabled. - submit planning: pre-registers a GLaaS fragment session, rewrites the manifest (0600 prepared file + context sidecar), returns session_id so the shared submit finalizer reconstitutes after completion; degrades to an uninstrumented submit with a warning when GLaaS is unavailable - manifest rewriter: wraps explicit container commands via a /bin/sh bootstrap that pip-installs the roar runtime and execs pod_entry, falling back to the original command on any bootstrap failure; injects the env/identity contract (Secret-backed session creds, downward-API pod identity, cluster-visible GLAAS_URL); appends the Secret document; fails actionably on missing explicit commands, generateName, or multiple Jobs - pod_entry: roar init + roar run (preload) around the original command, then exports the recorded job as an ExecutionFragment with task identity pod_uid:container:completion_index:restart_attempt and streams it via the shared fragment transport (best-effort; training exit code always wins) - host execution: kubectl submit, prepared-manifest cleanup, Job terminal-condition wait, submit job recorded with the ORIGINAL command and plan-generated parent uid so reproduce re-enters the backend and pod fragments link to the submit node - reconstituter + K8S_FRAGMENT_LINEAGE_BACKEND merge k8s_task child jobs through the shared fragment lineage engine - `roar k8s prepare` CLI for inspectable rewrites (no Secret embedded) - shared roar.execution.fragments.export extracted from the OSMO bundle exporter (OSMO delegates, behavior unchanged) - registry: k8s builtin (priority 95), config section, job-env marker Verified: product-path e2e green on the KIND harness (4 tests: rewrite, streaming, reconstitution into .roar/roar.db, secret hygiene) plus the Phase-0 runtime smoke (7 total); 23 new unit tests; default gate 1993 passed; ruff + mypy clean. Co-Authored-By: Claude Fable 5 --- docs/developer/k8s-integration.md | 88 +++++ roar/backends/k8s/__init__.py | 0 roar/backends/k8s/config.py | 131 ++++++++ roar/backends/k8s/fragment_reconstituter.py | 212 ++++++++++++ roar/backends/k8s/host_execution.py | 262 +++++++++++++++ roar/backends/k8s/lineage.py | 119 +++++++ roar/backends/k8s/manifest.py | 269 +++++++++++++++ roar/backends/k8s/plugin.py | 104 ++++++ roar/backends/k8s/pod_entry.py | 148 +++++++++ roar/backends/k8s/submit.py | 307 ++++++++++++++++++ roar/backends/osmo/export.py | 168 +--------- roar/cli/command_registry.py | 3 + roar/cli/commands/k8s.py | 111 +++++++ roar/execution/fragments/export.py | 207 ++++++++++++ roar/execution/framework/registry.py | 7 +- tests/backends/k8s/README.md | 21 +- .../backends/k8s/e2e/test_k8s_product_path.py | 282 ++++++++++++++++ tests/backends/k8s/unit/__init__.py | 0 tests/backends/k8s/unit/conftest.py | 50 +++ .../k8s/unit/test_backend_selection.py | 62 ++++ .../k8s/unit/test_manifest_rewrite.py | 151 +++++++++ .../backends/k8s/unit/test_submit_planning.py | 148 +++++++++ 22 files changed, 2689 insertions(+), 161 deletions(-) create mode 100644 docs/developer/k8s-integration.md create mode 100644 roar/backends/k8s/__init__.py create mode 100644 roar/backends/k8s/config.py create mode 100644 roar/backends/k8s/fragment_reconstituter.py create mode 100644 roar/backends/k8s/host_execution.py create mode 100644 roar/backends/k8s/lineage.py create mode 100644 roar/backends/k8s/manifest.py create mode 100644 roar/backends/k8s/plugin.py create mode 100644 roar/backends/k8s/pod_entry.py create mode 100644 roar/backends/k8s/submit.py create mode 100644 roar/cli/commands/k8s.py create mode 100644 roar/execution/fragments/export.py create mode 100644 tests/backends/k8s/e2e/test_k8s_product_path.py create mode 100644 tests/backends/k8s/unit/__init__.py create mode 100644 tests/backends/k8s/unit/conftest.py create mode 100644 tests/backends/k8s/unit/test_backend_selection.py create mode 100644 tests/backends/k8s/unit/test_manifest_rewrite.py create mode 100644 tests/backends/k8s/unit/test_submit_planning.py diff --git a/docs/developer/k8s-integration.md b/docs/developer/k8s-integration.md new file mode 100644 index 00000000..a9b16149 --- /dev/null +++ b/docs/developer/k8s-integration.md @@ -0,0 +1,88 @@ +# Kubernetes Integration (Developer) + +## 1. High-level summary + +The k8s backend instruments plain `batch/v1` Jobs submitted with +`kubectl apply|create -f ` through `roar run`. Unlike Ray (runtime +hooks) it works at the manifest layer, like OSMO — but transports lineage via +GLaaS fragment streaming, like Ray, rather than bundle files. + +Phase 1 scope (see `design-docs/k8s-training-lineage-integration.md` in the +dev meta-repo): exactly one Job per manifest, explicit container commands +only, streaming transport only. Kubeflow TrainJob/JobSet/RayJob adapters, +bundle fallback, `roar k8s attach`, and the admission-webhook injector are +later phases. + +## 2. Flow + +1. **Match** (`roar/backends/k8s/submit.py`): binary `kubectl`, verb + `apply|create`, `-f` pointing at a manifest containing exactly one + `batch/v1` Job; gated by `k8s.enabled` (default off). +2. **Plan** (same module): pre-registers a GLaaS fragment session (saves the + `.key` under `.roar/fragment-sessions/`), rewrites the manifest + (`manifest.py`), writes it 0600 under `.roar/k8s/prepared/` with a + `.context.json` sidecar, and returns the rewritten command plus + `session_id` — the framework attaches the shared submit finalizer. + GLaaS/registration failures degrade to an uninstrumented submit with a + warning; lineage never blocks training. +3. **Rewrite** (`manifest.py`): wraps each container that has an explicit + `command` with a `/bin/sh -c` script that pip-installs the roar runtime + (`k8s.runtime_install_requirement`, default pinned `roar-cli`) and execs + `python3 -m roar.backends.k8s.pod_entry "$@"`; on any bootstrap failure it + falls back to exec'ing the original command uninstrumented. Injects the + env contract (`GLAAS_URL` = cluster-visible URL, Secret-backed + `ROAR_SESSION_ID`/`ROAR_FRAGMENT_TOKEN`, downward-API identity fields) and + appends the Secret document. Containers without an explicit command are + skipped with a warning; a Job with none fails actionably (ENTRYPOINT + resolution is a later phase). +4. **Host execution** (`host_execution.py`): runs kubectl, deletes the + prepared manifest (it embeds the token Secret), polls the Job to a + terminal condition (`k8s.wait_for_completion`, default on), and records + the submit as a local job — with the **original** command and the + plan-generated `parent_job_uid`, so reproduce re-enters the backend and + pod fragments link to the submit node. +5. **In pod** (`pod_entry.py`): `roar init -n` + `roar run --tracer preload` + around the original command, then exports the recorded job as an + `ExecutionFragment` (task identity + `pod_uid:container:completion_index:restart_attempt`, parent from + `ROAR_K8S_PARENT_JOB_UID`) and streams it via the shared fragment + transport. Best-effort: the training exit code always wins. +6. **Reconstitution** (`fragment_reconstituter.py` + `lineage.py`): the + shared finalizer fetches/decrypts session batches, dedupes by task + identity, and merges `k8s_task` jobs through `merge_execution_fragments` + with `K8S_FRAGMENT_LINEAGE_BACKEND`. + +## 3. CLI + +`roar k8s prepare -f job.yaml -o prepared.yaml` runs the same rewriter for +inspection: no session is registered and no Secret is embedded; the user +creates the named Secret out of band or uses the managed path. + +## 4. Config + +Section `k8s` (registered `BackendConfigAdapter`): `enabled`, `tracer`, +`runtime_install_requirement`, `cluster_glaas_url`, `wait_for_completion`, +`wait_timeout_seconds`, `poll_interval_seconds`, +`fragment_session_ttl_seconds`. Env overrides beat config: +`ROAR_CLUSTER_PIP_REQ`, `ROAR_CLUSTER_GLAAS_URL`. + +## 5. Tests + +- Unit (default gate): `tests/backends/k8s/unit/` — matching, planning, + rewriting, identity contract. +- E2E (KIND harness, `k8s_e2e` marker): `tests/backends/k8s/e2e/` — see + `tests/backends/k8s/README.md` for the harness and the product-path vs + runtime-diagnostic split. + +## 6. Notes and caveats + +- Host-visible vs cluster-visible endpoints stay separate on purpose + (`glaas.url` vs `k8s.cluster_glaas_url`). +- The prepared manifest embeds the fragment token in a Secret document; it is + written 0600 and deleted immediately after kubectl reads it. The token also + lives in `.roar/fragment-sessions/*.key`, matching the Ray path. +- `glaas.url` has a hosted default, so the "GLaaS unconfigured" branch is + effectively unreachable via config; degradation is driven by registration + failure instead. +- The distributed adapter's worker-bootstrap hooks are inert stubs: pods are + instrumented by manifest rewriting, not `roar-worker` entrypoints. diff --git a/roar/backends/k8s/__init__.py b/roar/backends/k8s/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/roar/backends/k8s/config.py b/roar/backends/k8s/config.py new file mode 100644 index 00000000..fbad0c47 --- /dev/null +++ b/roar/backends/k8s/config.py @@ -0,0 +1,131 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field + +from roar.execution.framework.contract import BackendConfigAdapter, ConfigurableKeySpec + + +class K8sBackendConfig(BaseModel): + """Kubernetes backend configuration.""" + + model_config = ConfigDict( + strict=False, + validate_assignment=True, + extra="ignore", + revalidate_instances="never", + ) + + enabled: bool = False + tracer: str = "preload" + runtime_install_requirement: str = "" + cluster_glaas_url: str = "" + wait_for_completion: bool = True + wait_timeout_seconds: int = Field(default=30 * 60, ge=1) + poll_interval_seconds: float = Field(default=5.0, gt=0.0) + fragment_session_ttl_seconds: int = Field(default=86400, ge=60) + + +K8S_CONFIGURABLE_KEYS = { + "k8s.enabled": ConfigurableKeySpec( + value_type=bool, + default=False, + description="Enable automatic kubectl Job submit handling in roar run", + ), + "k8s.tracer": ConfigurableKeySpec( + value_type=str, + default="preload", + description="Tracer backend used inside instrumented pods (preload|ptrace|auto)", + ), + "k8s.runtime_install_requirement": ConfigurableKeySpec( + value_type=str, + default="", + description=( + "Pinned requirement or wheel URL installed inside pods to bootstrap the roar " + "runtime; defaults to the submitter's pinned roar-cli version. Wheels must " + "include packaged tracer binaries" + ), + ), + "k8s.cluster_glaas_url": ConfigurableKeySpec( + value_type=str, + default="", + description=( + "Cluster-visible GLaaS URL injected into pods when it differs from the " + "host-visible glaas.url (ROAR_CLUSTER_GLAAS_URL env always wins)" + ), + ), + "k8s.wait_for_completion": ConfigurableKeySpec( + value_type=bool, + default=True, + description="Wait for submitted Jobs to reach a terminal state before completing roar run", + ), + "k8s.wait_timeout_seconds": ConfigurableKeySpec( + value_type=int, + default=30 * 60, + description="Maximum time to wait for a submitted Job to reach a terminal state", + ), + "k8s.poll_interval_seconds": ConfigurableKeySpec( + value_type=float, + default=5.0, + description="Polling interval in seconds when waiting for Job completion", + ), + "k8s.fragment_session_ttl_seconds": ConfigurableKeySpec( + value_type=int, + default=86400, + description="TTL requested when pre-registering the GLaaS fragment session for a Job", + ), +} + +K8S_INIT_TEMPLATE = """\ +[k8s] +# Enable kubectl Job submit recognition in roar run +enabled = false +# Tracer backend used inside instrumented pods +tracer = "preload" +# Optional pinned requirement or wheel URL installed inside pods +# For roar-cli, use a packaged wheel or index source that includes bundled tracer binaries +runtime_install_requirement = "" +# Cluster-visible GLaaS URL when pods cannot reach the host-visible glaas.url +cluster_glaas_url = "" +# Wait for submitted Jobs to finish so lineage can be reconstituted immediately +wait_for_completion = true +""" + + +def normalize_k8s_backend_config(section: Mapping[str, Any] | None) -> dict[str, Any]: + return K8sBackendConfig.model_validate(dict(section or {})).model_dump() + + +def load_k8s_backend_config(start_dir: str | None = None) -> dict[str, Any]: + try: + from roar.integrations.config import load_config + + config = load_config(start_dir=start_dir) + except Exception: + return dict(K8S_BACKEND_CONFIG.default_values) + + section = config.get("k8s", {}) + if not isinstance(section, Mapping): + return dict(K8S_BACKEND_CONFIG.default_values) + return normalize_k8s_backend_config(section) + + +K8S_BACKEND_CONFIG = BackendConfigAdapter( + section_name="k8s", + default_values=K8sBackendConfig().model_dump(), + configurable_keys=K8S_CONFIGURABLE_KEYS, + init_template=K8S_INIT_TEMPLATE, + normalize_section=normalize_k8s_backend_config, +) + + +__all__ = [ + "K8S_BACKEND_CONFIG", + "K8S_CONFIGURABLE_KEYS", + "K8S_INIT_TEMPLATE", + "K8sBackendConfig", + "load_k8s_backend_config", + "normalize_k8s_backend_config", +] diff --git a/roar/backends/k8s/fragment_reconstituter.py b/roar/backends/k8s/fragment_reconstituter.py new file mode 100644 index 00000000..a0271244 --- /dev/null +++ b/roar/backends/k8s/fragment_reconstituter.py @@ -0,0 +1,212 @@ +"""Fetch, decrypt, and merge k8s fragment sessions from GLaaS.""" + +from __future__ import annotations + +import base64 +import json +import sqlite3 +import urllib.request +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from cryptography.hazmat.primitives.ciphers.aead import AESGCM + +from roar.backends.k8s.lineage import collect_k8s_fragments, resolve_active_session_context + + +def _get_logger(): + from roar.core.logging import get_logger + + return get_logger() + + +@dataclass(frozen=True) +class K8sReconstitutionResult: + jobs_merged: int = 0 + artifacts_merged: int = 0 + fragments_processed: int = 0 + + +class K8sFragmentReconstituter: + def __init__( + self, + session_id: str, + token: str, + glaas_url: str, + roar_db_path: Path, + ) -> None: + self._session_id = session_id + self._token = token + self._glaas_url = glaas_url.rstrip("/") + self._roar_db_path = roar_db_path + + def reconstitute(self) -> K8sReconstitutionResult: + batches = self._fetch_batches() + if not batches: + return K8sReconstitutionResult() + + try: + key = bytes.fromhex(self._token) + except ValueError as exc: + _get_logger().warning( + "Invalid fragment token for session %s: %s", self._session_id, exc + ) + return K8sReconstitutionResult() + + fragments: list[dict[str, Any]] = [] + for batch in batches: + fragments.extend(self._decrypt_batch(batch, key)) + if not fragments: + return K8sReconstitutionResult() + + fragments = self._deduplicate_by_task_identity(fragments) + driver_job_uid = next( + ( + str(fragment.get("parent_job_uid") or "").strip() + for fragment in fragments + if str(fragment.get("parent_job_uid") or "").strip() + ), + None, + ) + + jobs_before, artifacts_before = self._count_local_rows() + session_id, step_number = resolve_active_session_context(str(self._roar_db_path)) + try: + collect_k8s_fragments( + fragments, + project_dir=str(self._project_dir()), + driver_job_uid=driver_job_uid, + session_id=session_id, + step_number=step_number, + ) + except Exception as exc: + _get_logger().warning( + "Failed to merge reconstituted fragments for session %s: %s", + self._session_id, + exc, + ) + return K8sReconstitutionResult(fragments_processed=len(fragments)) + + jobs_after, artifacts_after = self._count_local_rows() + return K8sReconstitutionResult( + jobs_merged=max(0, jobs_after - jobs_before), + artifacts_merged=max(0, artifacts_after - artifacts_before), + fragments_processed=len(fragments), + ) + + def _fetch_batches(self) -> list[dict[str, Any]]: + request = urllib.request.Request( + url=f"{self._glaas_url}/api/v1/fragments/sessions/{self._session_id}/fragments", + headers={"x-roar-fragment-token": self._token}, + method="GET", + ) + try: + with urllib.request.urlopen(request, timeout=10) as response: + payload = json.loads(response.read().decode("utf-8")) + except Exception as exc: + _get_logger().warning( + "Failed to fetch fragments for session %s: %s", self._session_id, exc + ) + return [] + + rows = payload.get("data", {}).get("fragments", payload.get("fragments")) + if not isinstance(rows, list): + _get_logger().warning( + "Invalid fragment response for session %s: missing fragments list", + self._session_id, + ) + return [] + + batches = [item for item in rows if isinstance(item, dict)] + return sorted(batches, key=self._sequence_key) + + def _decrypt_batch(self, batch: dict[str, Any], key: bytes) -> list[dict[str, Any]]: + encrypted_batch = batch.get("encrypted_batch") + sequence = batch.get("sequence") + if not isinstance(encrypted_batch, str) or not encrypted_batch: + return [] + + try: + payload = base64.b64decode(encrypted_batch) + if len(payload) <= 12: + raise ValueError("payload too short") + plaintext = AESGCM(key).decrypt(payload[:12], payload[12:], None) + decoded = json.loads(plaintext.decode("utf-8")) + if not isinstance(decoded, list): + raise ValueError("decrypted payload is not a list") + except Exception as exc: + _get_logger().warning( + "Skipping undecryptable fragment batch for session %s sequence %s: %s", + self._session_id, + sequence, + exc, + ) + return [] + + return [item for item in decoded if isinstance(item, dict)] + + @staticmethod + def _deduplicate_by_task_identity(fragments: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Keep the last emission per task identity (batches are sequence-sorted). + + Re-emissions and retries replace earlier snapshots of the same task + rather than duplicating them; the merge engine additionally + reconciles job UIDs for anything that slips through. + """ + winners: dict[str, dict[str, Any]] = {} + for index, fragment in enumerate(fragments): + identity = str(fragment.get("task_identity") or "").strip() or f"fragment:{index}" + winners[identity] = fragment + return list(winners.values()) + + @staticmethod + def _sequence_key(batch: dict[str, Any]) -> int: + sequence = batch.get("sequence") + if sequence is None: + return 2**31 - 1 + try: + return int(sequence) + except (TypeError, ValueError): + return 2**31 - 1 + + def _project_dir(self) -> Path: + parent = self._roar_db_path.parent + if parent.name == ".roar": + return parent.parent + return parent + + def _count_local_rows(self) -> tuple[int, int]: + if not self._roar_db_path.exists(): + return 0, 0 + + conn = sqlite3.connect(self._roar_db_path) + conn.row_factory = sqlite3.Row + try: + jobs_row = conn.execute( + "SELECT COUNT(*) AS count FROM jobs WHERE job_type = 'k8s_task'" + ).fetchone() + artifacts_row = conn.execute("SELECT COUNT(*) AS count FROM artifacts").fetchone() + jobs_count = int(jobs_row["count"]) if jobs_row is not None else 0 + artifacts_count = int(artifacts_row["count"]) if artifacts_row is not None else 0 + return jobs_count, artifacts_count + except sqlite3.Error: + return 0, 0 + finally: + conn.close() + + +def create_k8s_fragment_reconstituter( + session_id: str, + token: str, + glaas_url: str, + roar_db_path: Path, +) -> K8sFragmentReconstituter: + return K8sFragmentReconstituter(session_id, token, glaas_url, roar_db_path) + + +__all__ = [ + "K8sFragmentReconstituter", + "K8sReconstitutionResult", + "create_k8s_fragment_reconstituter", +] diff --git a/roar/backends/k8s/host_execution.py b/roar/backends/k8s/host_execution.py new file mode 100644 index 00000000..c8c27b80 --- /dev/null +++ b/roar/backends/k8s/host_execution.py @@ -0,0 +1,262 @@ +"""Host-side execution for kubectl Job submits.""" + +from __future__ import annotations + +import json +import shlex +import subprocess +import sys +import time +from pathlib import Path +from typing import Any + +from roar.backends.k8s.config import load_k8s_backend_config +from roar.backends.k8s.submit import ( + K8sSubmitContext, + _find_filename_argument, + discard_submit_context, + load_submit_context, +) +from roar.core.bootstrap import bootstrap +from roar.core.models.run import RunContext, RunResult, resolve_run_config_start_dir +from roar.core.operation_metadata import build_operation_metadata_json +from roar.db.context import create_database_context +from roar.db.hashing import hash_files_blake3 +from roar.execution.recording import LocalJobRecorder, LocalRecordedArtifact +from roar.execution.runtime.errors import ExecutionSetupError + +_TERMINAL_SUCCESS_CONDITIONS = ("Complete", "SuccessCriteriaMet") +_TERMINAL_FAILURE_CONDITIONS = ("Failed", "FailureTarget") + + +def execute_k8s_job_submit(ctx: RunContext) -> RunResult: + """Submit a (possibly instrumented) Job manifest and record it locally.""" + bootstrap(ctx.roar_dir) + started_at = time.time() + config = load_k8s_backend_config(start_dir=str(resolve_run_config_start_dir(ctx))) + + filename = _find_filename_argument(ctx.command) + prepared_path = Path(filename[1]).resolve() if filename else None + submit_context = load_submit_context(prepared_path) if prepared_path else None + + try: + completed = subprocess.run( + ctx.command, + cwd=ctx.repo_root, + capture_output=True, + text=True, + check=False, + ) + except FileNotFoundError as exc: + raise ExecutionSetupError( + "Error: kubectl not found. Install kubectl or adjust PATH." + ) from exc + finally: + if submit_context is not None and prepared_path is not None: + # The prepared manifest embeds the fragment token Secret. + prepared_path.unlink(missing_ok=True) + discard_submit_context(prepared_path) + + _emit_captured_output(completed.stdout, sys.stdout) + _emit_captured_output(completed.stderr, sys.stderr) + + final_exit_code = completed.returncode + wait_payload: dict[str, Any] | None = None + if ( + completed.returncode == 0 + and submit_context is not None + and bool(config.get("wait_for_completion", True)) + ): + succeeded, wait_payload = _wait_for_job_completion( + ctx=ctx, + submit_context=submit_context, + timeout_seconds=int(config.get("wait_timeout_seconds", 30 * 60)), + poll_interval_seconds=float(config.get("poll_interval_seconds", 5.0)), + ) + if not succeeded: + final_exit_code = 1 + + recorded_command = ( + shlex.join(submit_context.original_command) + if submit_context is not None + else shlex.join(ctx.command) + ) + payload = _build_submit_payload( + ctx=ctx, + submit_context=submit_context, + exit_code=final_exit_code, + submit_stdout=completed.stdout, + wait_payload=wait_payload, + ) + duration = max(0.0, time.time() - started_at) + + with create_database_context(ctx.roar_dir) as db_ctx: + session_id = db_ctx.sessions.get_or_create_active() + recorder = LocalJobRecorder() + job_id, job_uid = recorder.record( + db_ctx, + command=recorded_command, + timestamp=started_at, + metadata=build_operation_metadata_json("k8s_submit", payload), + execution_backend=ctx.execution_backend, + execution_role=ctx.execution_role, + job_type=ctx.job_type or "run", + input_artifacts=_build_submit_input_artifacts(submit_context), + output_artifacts=[], + duration_seconds=duration, + exit_code=final_exit_code, + session_id=session_id, + job_uid=submit_context.parent_job_uid if submit_context else None, + ) + inputs = db_ctx.jobs.get_inputs(job_id) + outputs = db_ctx.jobs.get_outputs(job_id) + db_ctx.commit() + + return RunResult( + exit_code=final_exit_code, + job_id=job_id, + job_uid=job_uid, + duration=duration, + inputs=inputs, + outputs=outputs, + interrupted=False, + is_build=ctx.job_type == "build", + backend="k8s", + ) + + +def _wait_for_job_completion( + *, + ctx: RunContext, + submit_context: K8sSubmitContext, + timeout_seconds: int, + poll_interval_seconds: float, +) -> tuple[bool, dict[str, Any]]: + kubectl_binary = ctx.command[0] + global_flags = _extract_kubectl_global_flags(ctx.command) + get_command = [ + kubectl_binary, + *global_flags, + "get", + f"job/{submit_context.job_name}", + "-n", + submit_context.namespace, + "-o", + "json", + ] + + print( + f"[roar-k8s] waiting for job/{submit_context.job_name} in namespace " + f"{submit_context.namespace} (timeout {timeout_seconds}s)" + ) + deadline = time.time() + timeout_seconds + last_error = "" + while time.time() < deadline: + result = subprocess.run(get_command, capture_output=True, text=True, check=False) + if result.returncode == 0: + try: + status = json.loads(result.stdout).get("status", {}) or {} + except json.JSONDecodeError: + status = {} + for condition in status.get("conditions") or []: + if not isinstance(condition, dict) or condition.get("status") != "True": + continue + condition_type = str(condition.get("type") or "") + if condition_type in _TERMINAL_SUCCESS_CONDITIONS: + print(f"[roar-k8s] job/{submit_context.job_name} completed") + return True, {"terminal_condition": condition_type} + if condition_type in _TERMINAL_FAILURE_CONDITIONS: + message = str(condition.get("message") or "job failed") + print( + f"[roar-k8s] job/{submit_context.job_name} failed: {message}", + file=sys.stderr, + ) + return False, { + "terminal_condition": condition_type, + "message": message, + } + else: + last_error = result.stderr.strip() + time.sleep(poll_interval_seconds) + + message = f"timed out after {timeout_seconds}s" + if last_error: + message = f"{message}; last kubectl error: {last_error}" + print(f"[roar-k8s] wait for job/{submit_context.job_name} {message}", file=sys.stderr) + return False, {"terminal_condition": None, "message": message} + + +def _extract_kubectl_global_flags(command: list[str]) -> list[str]: + flags: list[str] = [] + for index, arg in enumerate(command): + if arg in ("--context", "--kubeconfig", "--cluster", "--user") and index + 1 < len(command): + flags.extend([arg, command[index + 1]]) + elif arg.startswith(("--context=", "--kubeconfig=", "--cluster=", "--user=")): + flags.append(arg) + return flags + + +def _build_submit_input_artifacts( + submit_context: K8sSubmitContext | None, +) -> list[LocalRecordedArtifact]: + if submit_context is None: + return [] + manifest_path = Path(submit_context.manifest_path) + if not manifest_path.is_file(): + return [] + + hashes = hash_files_blake3([manifest_path]) + digest = hashes.get(str(manifest_path)) + if not digest: + return [] + return [ + LocalRecordedArtifact( + path=str(manifest_path), + hashes={"blake3": digest}, + size=manifest_path.stat().st_size, + metadata=json.dumps({"k8s_submit_input": {"role": "job_manifest"}}), + ) + ] + + +def _build_submit_payload( + *, + ctx: RunContext, + submit_context: K8sSubmitContext | None, + exit_code: int, + submit_stdout: str, + wait_payload: dict[str, Any] | None, +) -> dict[str, Any]: + payload: dict[str, Any] = { + "executed_command": shlex.join(ctx.command), + "exit_code": exit_code, + "instrumented": submit_context is not None, + "submit_stdout_tail": submit_stdout.strip().splitlines()[-5:], + } + if submit_context is not None: + payload.update( + { + "job_name": submit_context.job_name, + "namespace": submit_context.namespace, + "manifest_path": submit_context.manifest_path, + "secret_name": submit_context.secret_name, + "session_id": submit_context.session_id, + "parent_job_uid": submit_context.parent_job_uid, + "wrapped_containers": submit_context.wrapped_containers, + "skipped_containers": submit_context.skipped_containers, + } + ) + if wait_payload is not None: + payload["wait"] = wait_payload + return payload + + +def _emit_captured_output(text: str, stream: Any) -> None: + if text: + stream.write(text) + if not text.endswith("\n"): + stream.write("\n") + stream.flush() + + +__all__ = ["execute_k8s_job_submit"] diff --git a/roar/backends/k8s/lineage.py b/roar/backends/k8s/lineage.py new file mode 100644 index 00000000..3bc8dd2a --- /dev/null +++ b/roar/backends/k8s/lineage.py @@ -0,0 +1,119 @@ +"""Kubernetes fragment shaping for the shared fragment lineage engine.""" + +from __future__ import annotations + +import os +import sqlite3 +from collections.abc import Mapping +from typing import Any + +from roar.execution.fragments.lineage import ( + FragmentLineageBackend, + merge_execution_fragments, +) +from roar.execution.fragments.models import ( + ExecutionFragment, + derive_fragment_identity, +) + + +def _k8s_fragment_command(fragment: ExecutionFragment) -> str: + return f"k8s_task:{fragment.task_name or 'task'}" + + +def _k8s_fragment_metadata( + fragment: ExecutionFragment, + fallback_parent_job_uid: str | None, +) -> Mapping[str, Any] | None: + metadata: dict[str, Any] = { + "k8s_task_id": fragment.task_id, + "parent_job_uid": fragment.parent_job_uid or fallback_parent_job_uid or None, + } + for key, value in (fragment.backend_metadata or {}).items(): + if value is not None: + metadata[key] = value + return metadata + + +K8S_FRAGMENT_LINEAGE_BACKEND = FragmentLineageBackend( + job_type="k8s_task", + command_for_fragment=_k8s_fragment_command, + script_for_fragment=lambda fragment: fragment.task_name or None, + execution_role_from_fragment=lambda fragment, _fallback: str( + (fragment.backend_metadata or {}).get("execution_role") or "task" + ), + metadata_from_fragment=_k8s_fragment_metadata, + task_identity_from_metadata=lambda parent_job_uid, job_uid, metadata: derive_fragment_identity( + "k8s", + parent_job_uid, + str(metadata.get("k8s_task_id") or metadata.get("task_id") or ""), + job_uid, + ), +) + + +def collect_k8s_fragments( + fragments: list[dict[str, Any]], + *, + project_dir: str, + driver_job_uid: str | None = None, + session_id: int | None = None, + step_number: int = 1, +) -> int: + """Merge k8s fragment dicts into the local DB; returns fragments merged.""" + parsed: list[ExecutionFragment] = [] + for payload in fragments: + if not isinstance(payload, dict): + continue + try: + parsed.append(ExecutionFragment.from_dict(payload)) + except Exception: + continue + + if not parsed: + return 0 + + merge_execution_fragments( + fragments=parsed, + project_dir=project_dir, + backend=K8S_FRAGMENT_LINEAGE_BACKEND, + driver_job_uid=driver_job_uid, + session_id=session_id, + step_number=step_number, + ) + return len(parsed) + + +def resolve_active_session_context(db_path: str) -> tuple[int | None, int]: + """Return ``(active_session_id, next_step_number)`` for fragment merges.""" + if not db_path or not os.path.exists(db_path): + return None, 1 + + try: + conn = sqlite3.connect(db_path) + except sqlite3.Error: + return None, 1 + + conn.row_factory = sqlite3.Row + try: + row = conn.execute("SELECT id FROM sessions WHERE is_active = 1 LIMIT 1").fetchone() + if row is None: + return None, 1 + session_id = int(row["id"]) + step_row = conn.execute( + "SELECT COALESCE(MAX(step_number), 0) AS max_step FROM jobs WHERE session_id = ?", + (session_id,), + ).fetchone() + max_step = int(step_row["max_step"]) if step_row is not None else 0 + return session_id, max_step + 1 + except sqlite3.Error: + return None, 1 + finally: + conn.close() + + +__all__ = [ + "K8S_FRAGMENT_LINEAGE_BACKEND", + "collect_k8s_fragments", + "resolve_active_session_context", +] diff --git a/roar/backends/k8s/manifest.py b/roar/backends/k8s/manifest.py new file mode 100644 index 00000000..b2f1909d --- /dev/null +++ b/roar/backends/k8s/manifest.py @@ -0,0 +1,269 @@ +"""Kubernetes Job manifest rewriting for lineage instrumentation. + +Phase 1 scope: exactly one ``batch/v1`` Job per manifest (Indexed or +plain). The rewriter wraps explicit container commands through the roar +pod entrypoint, injects the env/identity contract, and appends a Secret +carrying the fragment-session credentials so tokens never appear inline +in pod specs. +""" + +from __future__ import annotations + +import shlex +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import yaml # type: ignore[import-untyped] + +_JOB_API_VERSION = "batch/v1" +_JOB_KIND = "Job" + +# sh -c script: "$0" is the synthetic argv0, "$@" is the original +# command+args. Lineage is best-effort by design: any failure to stage +# the roar runtime falls back to running the original command +# uninstrumented rather than failing the training job. +_POD_WRAPPER_TEMPLATE = """\ +run_fallback() {{ echo "[roar-k8s] lineage runtime unavailable; running uninstrumented" >&2; exec "$@"; }} +command -v python3 >/dev/null 2>&1 || run_fallback "$@" +python3 -m pip install --quiet {requirement} || run_fallback "$@" +exec python3 -m roar.backends.k8s.pod_entry "$@" +""" + + +class K8sManifestError(ValueError): + """Raised when a manifest cannot be instrumented, with actionable detail.""" + + +@dataclass(frozen=True) +class K8sManifestRewrite: + documents: list[dict[str, Any]] + job_name: str + namespace: str + secret_name: str + wrapped_containers: list[str] = field(default_factory=list) + skipped_containers: list[str] = field(default_factory=list) + + +def load_manifest_documents(path: Path) -> list[dict[str, Any]]: + try: + raw = path.read_text(encoding="utf-8") + except OSError as exc: + raise K8sManifestError(f"cannot read manifest {path}: {exc}") from exc + + try: + documents = [doc for doc in yaml.safe_load_all(raw) if doc is not None] + except yaml.YAMLError as exc: + raise K8sManifestError(f"invalid YAML in manifest {path}: {exc}") from exc + + return [doc for doc in documents if isinstance(doc, dict)] + + +def find_job_documents(documents: list[dict[str, Any]]) -> list[dict[str, Any]]: + return [ + doc + for doc in documents + if str(doc.get("apiVersion") or "") == _JOB_API_VERSION + and str(doc.get("kind") or "") == _JOB_KIND + ] + + +def rewrite_manifest_for_lineage( + documents: list[dict[str, Any]], + *, + secret_name: str, + session_id: str | None, + fragment_token: str | None, + requirement: str, + cluster_glaas_url: str, + tracer: str, + parent_job_uid: str, + namespace_override: str | None = None, +) -> K8sManifestRewrite: + """Return a rewritten copy of ``documents`` with lineage instrumentation. + + When ``session_id``/``fragment_token`` are provided, a Secret document is + appended and referenced from the wrapped containers; otherwise the env + contract still points at ``secret_name`` (the ``roar k8s prepare`` flow, + where the user creates the Secret out of band). + """ + jobs = find_job_documents(documents) + if not jobs: + raise K8sManifestError( + "no batch/v1 Job found in manifest; Phase 1 instruments plain Jobs only" + ) + if len(jobs) > 1: + names = [str((doc.get("metadata") or {}).get("name") or "") for doc in jobs] + raise K8sManifestError( + f"manifest contains {len(jobs)} Jobs ({', '.join(names)}); " + "Phase 1 instruments exactly one Job per submit" + ) + + job_index = next(index for index, doc in enumerate(documents) if doc is jobs[0]) + rewritten_documents = [dict(doc) for doc in documents] + job = _deep_copy(jobs[0]) + rewritten_documents[job_index] = job + + metadata = job.get("metadata") + if not isinstance(metadata, dict) or not str(metadata.get("name") or "").strip(): + hint = "" + if isinstance(metadata, dict) and metadata.get("generateName"): + hint = " (generateName is not supported yet; set a fixed metadata.name)" + raise K8sManifestError(f"the Job needs an explicit metadata.name{hint}") + job_name = str(metadata["name"]).strip() + manifest_namespace = str(metadata.get("namespace") or "").strip() + namespace = namespace_override or manifest_namespace or "default" + + pod_spec = _require_dict_path(job, ("spec", "template", "spec"), job_name=job_name) + containers = pod_spec.get("containers") + if not isinstance(containers, list) or not containers: + raise K8sManifestError(f"Job {job_name} has no spec.template.spec.containers") + + wrapped: list[str] = [] + skipped: list[str] = [] + for container in containers: + if not isinstance(container, dict): + continue + container_name = str(container.get("name") or "").strip() or "" + command = container.get("command") + if not isinstance(command, list) or not command: + skipped.append(container_name) + continue + _wrap_container( + container, + job_name=job_name, + secret_name=secret_name, + requirement=requirement, + cluster_glaas_url=cluster_glaas_url, + tracer=tracer, + parent_job_uid=parent_job_uid, + ) + wrapped.append(container_name) + + if not wrapped: + raise K8sManifestError( + f"Job {job_name} has no container with an explicit command; " + "roar wraps commands it can see — set spec.template.spec.containers[].command " + "(images relying on ENTRYPOINT are not supported yet)" + ) + + if session_id and fragment_token: + secret_doc: dict[str, Any] = { + "apiVersion": "v1", + "kind": "Secret", + "metadata": { + "name": secret_name, + "labels": {"app.kubernetes.io/managed-by": "roar"}, + }, + "type": "Opaque", + "stringData": { + "session_id": session_id, + "token": fragment_token, + }, + } + if manifest_namespace: + secret_doc["metadata"]["namespace"] = manifest_namespace + rewritten_documents.insert(0, secret_doc) + + return K8sManifestRewrite( + documents=rewritten_documents, + job_name=job_name, + namespace=namespace, + secret_name=secret_name, + wrapped_containers=wrapped, + skipped_containers=skipped, + ) + + +def dump_manifest_documents(documents: list[dict[str, Any]]) -> str: + return yaml.safe_dump_all(documents, sort_keys=False) + + +def _wrap_container( + container: dict[str, Any], + *, + job_name: str, + secret_name: str, + requirement: str, + cluster_glaas_url: str, + tracer: str, + parent_job_uid: str, +) -> None: + container_name = str(container.get("name") or "").strip() or "main" + original = [str(part) for part in container.get("command", [])] + original.extend(str(part) for part in container.get("args", []) or []) + + script = _POD_WRAPPER_TEMPLATE.format(requirement=shlex.quote(requirement)) + container["command"] = ["/bin/sh", "-c", script, "roar-k8s", *original] + container.pop("args", None) + + env = container.setdefault("env", []) + if not isinstance(env, list): + raise K8sManifestError(f"container {container_name} has a non-list env block") + existing_names = { + str(entry.get("name")) for entry in env if isinstance(entry, dict) and entry.get("name") + } + + def add_value(name: str, value: str) -> None: + if name not in existing_names: + env.append({"name": name, "value": value}) + + def add_field_ref(name: str, field_path: str) -> None: + if name not in existing_names: + env.append({"name": name, "valueFrom": {"fieldRef": {"fieldPath": field_path}}}) + + def add_secret_ref(name: str, key: str) -> None: + if name not in existing_names: + env.append( + { + "name": name, + "valueFrom": {"secretKeyRef": {"name": secret_name, "key": key}}, + } + ) + + add_value("ROAR_EXECUTION_BACKEND", "k8s") + add_value("ROAR_NO_TELEMETRY", "1") + add_value("ROAR_K8S_TRACER", tracer) + add_value("GLAAS_URL", cluster_glaas_url) + add_value("ROAR_K8S_PARENT_JOB_UID", parent_job_uid) + add_value("ROAR_K8S_JOB_NAME", job_name) + add_value("ROAR_K8S_CONTAINER", container_name) + add_value("ROAR_K8S_TASK_NAME", f"{job_name}/{container_name}") + add_secret_ref("ROAR_SESSION_ID", "session_id") + add_secret_ref("ROAR_FRAGMENT_TOKEN", "token") + add_field_ref("ROAR_K8S_NAMESPACE", "metadata.namespace") + add_field_ref("ROAR_K8S_POD_NAME", "metadata.name") + add_field_ref("ROAR_K8S_POD_UID", "metadata.uid") + add_field_ref("ROAR_K8S_NODE_NAME", "spec.nodeName") + + +def _require_dict_path( + root: dict[str, Any], + path: tuple[str, ...], + *, + job_name: str, +) -> dict[str, Any]: + node: Any = root + walked: list[str] = [] + for key in path: + walked.append(key) + node = node.get(key) if isinstance(node, dict) else None + if not isinstance(node, dict): + raise K8sManifestError(f"Job {job_name} is missing {'.'.join(walked)}") + return node + + +def _deep_copy(document: dict[str, Any]) -> dict[str, Any]: + import copy + + return copy.deepcopy(document) + + +__all__ = [ + "K8sManifestError", + "K8sManifestRewrite", + "dump_manifest_documents", + "find_job_documents", + "load_manifest_documents", + "rewrite_manifest_for_lineage", +] diff --git a/roar/backends/k8s/plugin.py b/roar/backends/k8s/plugin.py new file mode 100644 index 00000000..56eab05d --- /dev/null +++ b/roar/backends/k8s/plugin.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Any + +from roar.backends.k8s.config import K8S_BACKEND_CONFIG +from roar.backends.k8s.fragment_reconstituter import create_k8s_fragment_reconstituter +from roar.backends.k8s.host_execution import execute_k8s_job_submit +from roar.backends.k8s.submit import ( + matches_kubectl_job_submit_command, + plan_kubectl_job_submit_command, +) +from roar.execution.framework.contract import ( + DistributedRuntimeAdapter, + DriverBootstrapAdapter, + ExecutionBackend, + ExecutionPolicyAdapter, + FragmentReconstitutionAdapter, + HostExecutionAdapter, + WorkerBootstrapAdapter, +) +from roar.execution.framework.registry import register_execution_backend + + +def _no_driver_proxy_fragment( + entries: Sequence[Any], + started_at: float, + ended_at: float, + exit_code: int, + environ: Mapping[str, str], +) -> dict[str, Any] | None: + del entries, started_at, ended_at, exit_code, environ + return None + + +def _no_local_merge( + fragments: list[dict[str, Any]], + project_dir: str, + driver_job_uid: str | None, +) -> None: + del fragments, project_dir, driver_job_uid + + +def _passthrough_runtime_env( + runtime_env: Mapping[str, Any] | None, + job_id: str, + source_environ: Mapping[str, str], +) -> dict[str, Any]: + del job_id, source_environ + return dict(runtime_env or {}) + + +def _worker_startup() -> None: + return None + + +def _worker_entrypoint(argv: list[str]) -> None: + raise RuntimeError( + "the k8s backend instruments pods via manifest rewriting; " + f"roar-worker entrypoints are not used (argv: {argv!r})" + ) + + +K8S_EXECUTION_BACKEND = ExecutionBackend( + name="k8s", + priority=95, + matches_command=matches_kubectl_job_submit_command, + plan_command=plan_kubectl_job_submit_command, + host_execution=HostExecutionAdapter(execute=execute_k8s_job_submit), + distributed=DistributedRuntimeAdapter( + driver_bootstrap=DriverBootstrapAdapter( + build_proxy_fragment=_no_driver_proxy_fragment, + local_merge=_no_local_merge, + should_start_local_proxy=lambda _env: False, + ), + worker_bootstrap=WorkerBootstrapAdapter( + py_executable="roar-worker", + setup_hook="roar.execution.runtime.worker_bootstrap.startup", + prepare_runtime_env=_passthrough_runtime_env, + startup=_worker_startup, + run_entrypoint=_worker_entrypoint, + ), + fragment_reconstitution=FragmentReconstitutionAdapter( + create_reconstituter=create_k8s_fragment_reconstituter, + ), + ), + policy=ExecutionPolicyAdapter( + submit_roles=("submit",), + task_roles=("task",), + job_environment_markers=("ROAR_K8S_PARENT_JOB_UID",), + ), + config=K8S_BACKEND_CONFIG, +) + + +def register() -> ExecutionBackend: + register_execution_backend(K8S_EXECUTION_BACKEND) + return K8S_EXECUTION_BACKEND + + +__all__ = [ + "K8S_EXECUTION_BACKEND", + "register", +] diff --git a/roar/backends/k8s/pod_entry.py b/roar/backends/k8s/pod_entry.py new file mode 100644 index 00000000..0a409b89 --- /dev/null +++ b/roar/backends/k8s/pod_entry.py @@ -0,0 +1,148 @@ +"""In-pod entrypoint for roar-instrumented Kubernetes Jobs. + +Invoked by the injected container wrapper as +``python3 -m roar.backends.k8s.pod_entry `` after the +roar runtime has been installed. Runs the original command under ``roar +run``, then exports the recorded job as an execution fragment stamped +with the k8s identity contract and streams it through the shared +fragment transport. + +Lineage is best-effort: the training command's exit code is always +propagated, and lineage failures only warn. +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +import tempfile +from pathlib import Path + + +def main(argv: list[str] | None = None) -> int: + command = list(sys.argv[1:] if argv is None else argv) + if command and command[0] == "--": + command = command[1:] + if not command: + print("[roar-k8s] no command to execute", file=sys.stderr) + return 2 + + exit_code = _run_traced(command) + _emit_lineage_best_effort() + return exit_code + + +def _run_traced(command: list[str]) -> int: + workdir = Path(os.environ.get("ROAR_K8S_WORKDIR") or os.getcwd()) + try: + workdir.mkdir(parents=True, exist_ok=True) + os.chdir(workdir) + except OSError as exc: + print(f"[roar-k8s] cannot use workdir {workdir}: {exc}", file=sys.stderr) + return _run_uninstrumented(command) + + if not (workdir / ".roar").is_dir(): + init = subprocess.run( + [sys.executable, "-m", "roar", "init", "-n"], + capture_output=True, + text=True, + check=False, + ) + if init.returncode != 0: + print( + f"[roar-k8s] roar init failed; running uninstrumented:\n{init.stderr}", + file=sys.stderr, + ) + return _run_uninstrumented(command) + + tracer = str(os.environ.get("ROAR_K8S_TRACER") or "preload").strip() or "preload" + run = subprocess.run( + [sys.executable, "-m", "roar", "run", "--tracer", tracer, *command], + check=False, + ) + return run.returncode + + +def _run_uninstrumented(command: list[str]) -> int: + completed = subprocess.run(command, check=False) + return completed.returncode + + +def task_identity_from_environment(environ: dict[str, str] | None = None) -> tuple[str, str]: + """Return ``(task_id, task_name)`` per the k8s identity contract. + + ``task_id`` = ``pod_uid:container:completion_index:restart_attempt`` — + never pod names (reused across operator restarts) or ranks (unstable + under elastic rendezvous). + """ + env = os.environ if environ is None else environ + pod_uid = str(env.get("ROAR_K8S_POD_UID") or "unknown-pod").strip() or "unknown-pod" + container = str(env.get("ROAR_K8S_CONTAINER") or "main").strip() or "main" + completion_index = ( + str(env.get("JOB_COMPLETION_INDEX") or "").strip() + or str(env.get("PET_NODE_RANK") or "").strip() + or "0" + ) + restart_attempt = ( + str(env.get("ROAR_K8S_RESTART_ATTEMPT") or "").strip() + or str(env.get("TORCHELASTIC_RESTART_COUNT") or "").strip() + or "0" + ) + task_id = f"{pod_uid}:{container}:{completion_index}:{restart_attempt}" + task_name = ( + str(env.get("ROAR_K8S_TASK_NAME") or "").strip() + or str(env.get("ROAR_K8S_JOB_NAME") or "").strip() + or "k8s-task" + ) + return task_id, task_name + + +def _emit_lineage_best_effort() -> None: + try: + from roar.execution.fragments.export import export_local_job_fragment_bundle + from roar.execution.fragments.transport import emit_fragment_dicts + + task_id, task_name = task_identity_from_environment() + parent_job_uid = str(os.environ.get("ROAR_K8S_PARENT_JOB_UID") or "").strip() + + with tempfile.TemporaryDirectory(prefix="roar-k8s-") as tmp: + bundle_path = Path(tmp) / "roar-fragments.json" + export_local_job_fragment_bundle( + roar_dir=Path.cwd() / ".roar", + output_path=bundle_path, + backend_name="k8s", + task_id=task_id, + task_name=task_name, + parent_job_uid=parent_job_uid, + default_task_name="k8s-task", + ) + payload = json.loads(bundle_path.read_text(encoding="utf-8")) + + fragments = [item for item in payload.get("fragments", []) if isinstance(item, dict)] + completion_index = task_id.split(":")[2] if task_id.count(":") >= 3 else "0" + restart_attempt = task_id.split(":")[3] if task_id.count(":") >= 3 else "0" + for fragment in fragments: + metadata = fragment.setdefault("backend_metadata", {}) + metadata.update( + { + "k8s_namespace": os.environ.get("ROAR_K8S_NAMESPACE"), + "k8s_pod_name": os.environ.get("ROAR_K8S_POD_NAME"), + "k8s_pod_uid": os.environ.get("ROAR_K8S_POD_UID"), + "k8s_node_name": os.environ.get("ROAR_K8S_NODE_NAME"), + "k8s_container": os.environ.get("ROAR_K8S_CONTAINER"), + "k8s_job_name": os.environ.get("ROAR_K8S_JOB_NAME"), + "k8s_completion_index": completion_index, + "k8s_restart_attempt": restart_attempt, + } + ) + + result = emit_fragment_dicts(fragments) + print(f"[roar-k8s] lineage emit: {result} ({len(fragments)} fragment(s))") + except Exception as exc: + print(f"[roar-k8s] warning: lineage emit failed: {exc}", file=sys.stderr) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/roar/backends/k8s/submit.py b/roar/backends/k8s/submit.py new file mode 100644 index 00000000..62b189c8 --- /dev/null +++ b/roar/backends/k8s/submit.py @@ -0,0 +1,307 @@ +"""kubectl Job submit planning through the execution backend framework.""" + +from __future__ import annotations + +import json +import os +import secrets +import sys +from dataclasses import asdict, dataclass +from pathlib import Path + +from roar.backends.k8s.config import load_k8s_backend_config +from roar.backends.k8s.manifest import ( + K8sManifestError, + dump_manifest_documents, + find_job_documents, + load_manifest_documents, + rewrite_manifest_for_lineage, +) +from roar.execution.fragments.sessions import generate_fragment_session, save_fragment_session +from roar.execution.framework.contract import ExecutionCommandPlan + +_KUBECTL_VERBS = ("apply", "create") +SUBMIT_CONTEXT_SUFFIX = ".context.json" + + +@dataclass(frozen=True) +class K8sSubmitContext: + """Sidecar context linking a prepared manifest back to its submit.""" + + original_command: list[str] + manifest_path: str + prepared_path: str + job_name: str + namespace: str + secret_name: str + session_id: str | None + parent_job_uid: str + wrapped_containers: list[str] + skipped_containers: list[str] + + +def matches_kubectl_job_submit_command(command: list[str]) -> bool: + if len(command) < 4: + return False + if Path(command[0]).name.lower() != "kubectl": + return False + if command[1].lower() not in _KUBECTL_VERBS: + return False + if not _k8s_backend_enabled(): + return False + + filename = _find_filename_argument(command) + if filename is None: + return False + manifest_path = Path(filename[1]) + if not manifest_path.is_file(): + return False + + try: + documents = load_manifest_documents(manifest_path) + except K8sManifestError: + return False + return len(find_job_documents(documents)) == 1 + + +def plan_kubectl_job_submit_command(command: list[str]) -> ExecutionCommandPlan: + if not matches_kubectl_job_submit_command(command): + return ExecutionCommandPlan(backend_name="k8s", command=list(command)) + + filename = _find_filename_argument(command) + assert filename is not None + _filename_index, manifest_arg = filename + manifest_path = Path(manifest_arg).resolve() + + start_dir = os.environ.get("ROAR_PROJECT_DIR") or os.getcwd() + config = load_k8s_backend_config(start_dir=start_dir) + roar_dir = Path(os.getcwd()) / ".roar" + + glaas_url = _resolve_glaas_url() + if glaas_url is None: + _warn( + "GLaaS is not configured; submitting the Job without lineage instrumentation " + "(set glaas.url to enable fragment streaming)" + ) + return ExecutionCommandPlan( + backend_name="k8s", + command=list(command), + execution_role="submit", + ) + + session = generate_fragment_session() + try: + _register_fragment_session( + glaas_url, + session["session_id"], + session["token_hash"], + ttl=int(config.get("fragment_session_ttl_seconds", 86400)), + ) + except Exception as exc: + _warn(f"fragment session pre-registration failed ({exc}); submitting uninstrumented") + return ExecutionCommandPlan( + backend_name="k8s", + command=list(command), + execution_role="submit", + ) + + parent_job_uid = secrets.token_hex(4) + secret_name = f"roar-fragment-{session['session_id'][:8]}" + documents = load_manifest_documents(manifest_path) + rewrite = rewrite_manifest_for_lineage( + documents, + secret_name=secret_name, + session_id=session["session_id"], + fragment_token=session["token"], + requirement=resolve_runtime_requirement(config), + cluster_glaas_url=_resolve_cluster_glaas_url(config, glaas_url), + tracer=str(config.get("tracer") or "preload"), + parent_job_uid=parent_job_uid, + namespace_override=_find_namespace_argument(command), + ) + if rewrite.skipped_containers: + _warn( + "containers without an explicit command were left uninstrumented: " + + ", ".join(rewrite.skipped_containers) + ) + + save_fragment_session(roar_dir, session) + + prepared_path = _write_prepared_manifest( + roar_dir, + rewrite.job_name, + parent_job_uid, + dump_manifest_documents(rewrite.documents), + ) + submit_context = K8sSubmitContext( + original_command=list(command), + manifest_path=str(manifest_path), + prepared_path=str(prepared_path), + job_name=rewrite.job_name, + namespace=rewrite.namespace, + secret_name=secret_name, + session_id=session["session_id"], + parent_job_uid=parent_job_uid, + wrapped_containers=list(rewrite.wrapped_containers), + skipped_containers=list(rewrite.skipped_containers), + ) + _write_submit_context(prepared_path, submit_context) + + rewritten_command = _replace_filename_argument(command, str(prepared_path)) + return ExecutionCommandPlan( + backend_name="k8s", + command=rewritten_command, + execution_role="submit", + session_id=session["session_id"], + ) + + +def resolve_runtime_requirement(config: dict) -> str: + import importlib.metadata as importlib_metadata + + override = os.environ.get("ROAR_CLUSTER_PIP_REQ", "").strip() + if override: + return override + + configured = str(config.get("runtime_install_requirement") or "").strip() + if configured: + return configured + + try: + version = importlib_metadata.version("roar-cli") + return f"roar-cli=={version}" + except Exception: + return "roar-cli" + + +def load_submit_context(prepared_path: Path) -> K8sSubmitContext | None: + context_path = prepared_path.with_name(prepared_path.name + SUBMIT_CONTEXT_SUFFIX) + if not context_path.is_file(): + return None + try: + payload = json.loads(context_path.read_text(encoding="utf-8")) + return K8sSubmitContext(**payload) + except Exception: + return None + + +def discard_submit_context(prepared_path: Path) -> None: + context_path = prepared_path.with_name(prepared_path.name + SUBMIT_CONTEXT_SUFFIX) + context_path.unlink(missing_ok=True) + + +def _write_prepared_manifest( + roar_dir: Path, + job_name: str, + parent_job_uid: str, + contents: str, +) -> Path: + prepared_dir = roar_dir / "k8s" / "prepared" + prepared_dir.mkdir(parents=True, exist_ok=True) + prepared_path = prepared_dir / f"{job_name}-{parent_job_uid}.yaml" + # The prepared manifest embeds the fragment-session token in the Secret + # document; keep it private and delete it right after kubectl reads it. + fd = os.open(prepared_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + with os.fdopen(fd, "w", encoding="utf-8") as handle: + handle.write(contents) + return prepared_path + + +def _write_submit_context(prepared_path: Path, context: K8sSubmitContext) -> None: + context_path = prepared_path.with_name(prepared_path.name + SUBMIT_CONTEXT_SUFFIX) + context_path.write_text(json.dumps(asdict(context), indent=2) + "\n", encoding="utf-8") + + +def _k8s_backend_enabled() -> bool: + start_dir = os.environ.get("ROAR_PROJECT_DIR") or os.getcwd() + return bool(load_k8s_backend_config(start_dir=start_dir).get("enabled", False)) + + +def _find_filename_argument(command: list[str]) -> tuple[int, str] | None: + for index, arg in enumerate(command[2:], start=2): + if arg in ("-f", "--filename"): + if index + 1 < len(command): + return index + 1, command[index + 1] + return None + if arg.startswith("--filename="): + return index, arg.split("=", 1)[1] + if arg.startswith("-f="): + return index, arg.split("=", 1)[1] + return None + + +def _replace_filename_argument(command: list[str], prepared_path: str) -> list[str]: + rewritten = list(command) + location = _find_filename_argument(command) + assert location is not None + index, _value = location + arg = rewritten[index] + if arg.startswith("--filename="): + rewritten[index] = f"--filename={prepared_path}" + elif arg.startswith("-f="): + rewritten[index] = f"-f={prepared_path}" + else: + rewritten[index] = prepared_path + return rewritten + + +def _find_namespace_argument(command: list[str]) -> str | None: + for index, arg in enumerate(command): + if arg in ("-n", "--namespace") and index + 1 < len(command): + return command[index + 1] + if arg.startswith("--namespace="): + return arg.split("=", 1)[1] + return None + + +def _resolve_glaas_url() -> str | None: + from roar.integrations.glaas import get_glaas_url + + url = get_glaas_url() + if url is None: + return None + text = str(url).strip() + return text or None + + +def _resolve_cluster_glaas_url(config: dict, host_glaas_url: str) -> str: + env_override = os.environ.get("ROAR_CLUSTER_GLAAS_URL", "").strip() + if env_override: + return env_override + configured = str(config.get("cluster_glaas_url") or "").strip() + if configured: + return configured + return host_glaas_url + + +def _register_fragment_session( + glaas_url: str, + session_id: str, + token_hash: str, + ttl: int = 86400, +) -> None: + from roar.integrations.glaas import GlaasClient + + client = GlaasClient(base_url=glaas_url) + _result, error = client.register_fragment_session( + session_id=session_id, + token_hash=token_hash, + ttl_seconds=ttl, + ) + if error: + raise RuntimeError(f"failed to pre-register fragment session {session_id}: {error}") + + +def _warn(message: str) -> None: + print(f"[roar-k8s] warning: {message}", file=sys.stderr) + + +__all__ = [ + "SUBMIT_CONTEXT_SUFFIX", + "K8sSubmitContext", + "discard_submit_context", + "load_submit_context", + "matches_kubectl_job_submit_command", + "plan_kubectl_job_submit_command", + "resolve_runtime_requirement", +] diff --git a/roar/backends/osmo/export.py b/roar/backends/osmo/export.py index 19172a79..ef76480d 100644 --- a/roar/backends/osmo/export.py +++ b/roar/backends/osmo/export.py @@ -1,12 +1,9 @@ from __future__ import annotations -import json from dataclasses import dataclass from pathlib import Path -from typing import Any -from roar.db.context import create_database_context -from roar.execution.fragments.models import ArtifactRef, ExecutionFragment +from roar.execution.fragments.export import export_local_job_fragment_bundle @dataclass(frozen=True) @@ -27,161 +24,24 @@ def export_osmo_lineage_bundle( task_name: str | None = None, backend_name: str = "osmo", ) -> OsmoLineageBundleExport: - project_dir = str(roar_dir.parent.resolve()) - - with create_database_context(roar_dir) as db_ctx: - selected_job = _select_export_job(db_ctx, job_uid=job_uid) - if selected_job is None: - raise ValueError("no local Roar jobs are available to export") - - selected_job_id = int(selected_job["id"]) - selected_job_uid = str(selected_job.get("job_uid") or "").strip() - resolved_task_id = str(task_id or selected_job_uid).strip() - if not resolved_task_id: - raise ValueError("task_id is required when exporting a job without a persisted job_uid") - - resolved_task_name = _resolve_task_name(selected_job, task_name) - reads = [ - _build_artifact_ref(item, project_dir=project_dir) - for item in db_ctx.jobs.get_inputs(selected_job_id) - ] - writes = [ - _build_artifact_ref(item, project_dir=project_dir) - for item in db_ctx.jobs.get_outputs(selected_job_id) - ] - - started_at = float(selected_job.get("timestamp") or 0.0) - duration = max(0.0, float(selected_job.get("duration_seconds") or 0.0)) - fragment = ExecutionFragment( - job_uid=selected_job_uid or resolved_task_id, - parent_job_uid="", - task_id=resolved_task_id, - worker_id="", - node_id="", - actor_id=None, - task_name=resolved_task_name, - started_at=started_at, - ended_at=started_at + duration, - exit_code=int(selected_job.get("exit_code") or 0), - backend=str(backend_name or "osmo").strip() or "osmo", - reads=reads, - writes=writes, - backend_metadata={ - "execution_role": "task", - "source_job_uid": selected_job_uid or None, - "source_execution_backend": str(selected_job.get("execution_backend") or "").strip() - or None, - "source_execution_role": str(selected_job.get("execution_role") or "").strip() or None, - "source_job_type": str(selected_job.get("job_type") or "").strip() or None, - }, + export = export_local_job_fragment_bundle( + roar_dir=roar_dir, + output_path=output_path, + backend_name=str(backend_name or "osmo").strip() or "osmo", + job_uid=job_uid, + task_id=task_id, + task_name=task_name, + default_task_name="osmo-task", ) - - payload = { - "fragments": [fragment.to_dict()], - "metadata": { - "exported_job_uid": selected_job_uid or None, - "task_id": resolved_task_id, - "task_name": resolved_task_name, - }, - } - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") - return OsmoLineageBundleExport( - output_path=str(output_path), - exported_job_uid=selected_job_uid, - fragment_count=1, - task_id=resolved_task_id, - task_name=resolved_task_name, - ) - - -def _select_export_job(db_ctx, *, job_uid: str | None) -> dict[str, Any] | None: - if job_uid: - selected = db_ctx.jobs.get_by_uid(job_uid) - if selected is None: - raise ValueError(f"local Roar job not found for job UID {job_uid!r}") - return selected - - jobs = db_ctx.jobs.get_recent(limit=1) - if not jobs: - return None - return jobs[0] - - -def _resolve_task_name(job: dict[str, Any], requested_task_name: str | None) -> str: - explicit = str(requested_task_name or "").strip() - if explicit: - return explicit - - script = str(job.get("script") or "").strip() - if script: - return script - - command = str(job.get("command") or "").strip() - if not command: - return "osmo-task" - - first_token = command.split(" ", 1)[0].strip() - return first_token or "osmo-task" - - -def _build_artifact_ref(payload: dict[str, Any], *, project_dir: str) -> ArtifactRef: - hash_value, hash_algorithm = _select_primary_hash(payload.get("hashes")) - path = _normalize_bundle_path(str(payload.get("path") or ""), project_dir=project_dir) - return ArtifactRef( - path=path, - hash=hash_value, - hash_algorithm=hash_algorithm, - size=int(payload.get("size") or 0), - capture_method="python", + output_path=export.output_path, + exported_job_uid=export.exported_job_uid, + fragment_count=export.fragment_count, + task_id=export.task_id, + task_name=export.task_name, ) -def _select_primary_hash(hashes: Any) -> tuple[str | None, str]: - if isinstance(hashes, list): - blake3_row = next( - ( - item - for item in hashes - if isinstance(item, dict) and str(item.get("algorithm") or "").strip() == "blake3" - ), - None, - ) - if isinstance(blake3_row, dict): - digest = str(blake3_row.get("digest") or "").strip() - if digest: - return digest, "blake3" - - for item in hashes: - if not isinstance(item, dict): - continue - algorithm = str(item.get("algorithm") or "").strip() - digest = str(item.get("digest") or "").strip() - if algorithm and digest: - return digest, algorithm - - return None, "blake3" - - -def _normalize_bundle_path(path: str, *, project_dir: str) -> str: - normalized = str(path or "").strip() - if not normalized or "://" in normalized: - return normalized - - candidate = Path(normalized) - project_root = Path(project_dir) - if not candidate.is_absolute(): - candidate = (project_root / candidate).resolve(strict=False) - - try: - relative = candidate.resolve(strict=False).relative_to(project_root) - except ValueError: - return str(candidate) - - return "${ROAR_PROJECT_DIR}/" + relative.as_posix() - - __all__ = [ "OsmoLineageBundleExport", "export_osmo_lineage_bundle", diff --git a/roar/cli/command_registry.py b/roar/cli/command_registry.py index 4a855c41..5676183c 100644 --- a/roar/cli/command_registry.py +++ b/roar/cli/command_registry.py @@ -257,6 +257,9 @@ class CommandSpec: "GLaaS / TReqs Account", ), CommandSpec("osmo", "roar.cli.commands.osmo", "osmo", "Manage OSMO workflow attachment"), + CommandSpec( + "k8s", "roar.cli.commands.k8s", "k8s", "Prepare Kubernetes Jobs for lineage capture" + ), ) diff --git a/roar/cli/commands/k8s.py b/roar/cli/commands/k8s.py new file mode 100644 index 00000000..3582f1f2 --- /dev/null +++ b/roar/cli/commands/k8s.py @@ -0,0 +1,111 @@ +from __future__ import annotations + +import secrets +from pathlib import Path + +import click + +from ...backends.k8s.config import load_k8s_backend_config +from ...backends.k8s.manifest import ( + K8sManifestError, + dump_manifest_documents, + load_manifest_documents, + rewrite_manifest_for_lineage, +) +from ...backends.k8s.submit import resolve_runtime_requirement + + +@click.group("k8s", invoke_without_command=True) +@click.pass_context +def k8s(ctx: click.Context) -> None: + """Prepare Kubernetes Jobs for lineage capture.""" + if ctx.invoked_subcommand is None: + click.echo(ctx.get_help()) + + +@k8s.command("prepare") +@click.option( + "-f", + "--filename", + "manifest_path", + type=click.Path(path_type=Path, dir_okay=False, exists=True), + required=True, + help="Job manifest to instrument", +) +@click.option( + "-o", + "--output", + "output_path", + type=click.Path(path_type=Path, dir_okay=False), + required=True, + help="Where to write the instrumented manifest", +) +@click.option( + "--secret-name", + default="roar-fragment-session", + show_default=True, + help="Fragment-session Secret the pods will read credentials from", +) +@click.option( + "--requirement", + default=None, + help="Override the roar runtime install requirement (defaults to k8s.runtime_install_requirement or the pinned roar-cli version)", +) +@click.option( + "--cluster-glaas-url", + default=None, + help="Cluster-visible GLaaS URL injected into pods (defaults to k8s.cluster_glaas_url)", +) +def k8s_prepare( + manifest_path: Path, + output_path: Path, + secret_name: str, + requirement: str | None, + cluster_glaas_url: str | None, +) -> None: + """Rewrite a Job manifest with lineage instrumentation for inspection. + + Unlike the managed `roar run kubectl apply -f ...` path, no fragment + session is registered and no Secret is embedded: create the Secret named + by --secret-name (keys: session_id, token) before applying the output. + """ + config = load_k8s_backend_config() + resolved_requirement = requirement or resolve_runtime_requirement(config) + resolved_cluster_glaas = ( + cluster_glaas_url or str(config.get("cluster_glaas_url") or "") or "http://glaas:3001" + ) + + try: + documents = load_manifest_documents(manifest_path) + rewrite = rewrite_manifest_for_lineage( + documents, + secret_name=secret_name, + session_id=None, + fragment_token=None, + requirement=resolved_requirement, + cluster_glaas_url=resolved_cluster_glaas, + tracer=str(config.get("tracer") or "preload"), + parent_job_uid=secrets.token_hex(4), + ) + except K8sManifestError as exc: + raise click.ClickException(str(exc)) from exc + + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text(dump_manifest_documents(rewrite.documents), encoding="utf-8") + + click.echo(f"Prepared manifest written to {output_path}") + click.echo(f" job: {rewrite.job_name} (namespace {rewrite.namespace})") + click.echo(f" wrapped containers: {', '.join(rewrite.wrapped_containers)}") + if rewrite.skipped_containers: + click.echo(f" skipped (no explicit command): {', '.join(rewrite.skipped_containers)}") + click.echo(f" runtime install: {resolved_requirement}") + click.echo() + click.echo("Before applying, create the fragment-session Secret:") + click.echo( + f" kubectl create secret generic {secret_name} " + "--from-literal=session_id= --from-literal=token=" + ) + click.echo("Or use the managed path instead: roar run kubectl apply -f ") + + +__all__ = ["k8s"] diff --git a/roar/execution/fragments/export.py b/roar/execution/fragments/export.py new file mode 100644 index 00000000..ec506bda --- /dev/null +++ b/roar/execution/fragments/export.py @@ -0,0 +1,207 @@ +"""Shared export of a locally recorded job as an execution-fragment bundle. + +Backend-neutral: OSMO and Kubernetes runtime wrappers both export the +job that `roar run` recorded inside a remote task/pod as a single +``ExecutionFragment`` bundle, either returned as a file (OSMO datasets) +or streamed to GLaaS (k8s fragment sessions). +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from roar.db.context import create_database_context +from roar.execution.fragments.models import ArtifactRef, ExecutionFragment + + +@dataclass(frozen=True) +class LocalJobFragmentExport: + output_path: str + exported_job_uid: str + fragment_count: int + task_id: str + task_name: str + + +def export_local_job_fragment_bundle( + *, + roar_dir: Path, + output_path: Path, + backend_name: str, + job_uid: str | None = None, + task_id: str | None = None, + task_name: str | None = None, + parent_job_uid: str = "", + default_task_name: str = "task", +) -> LocalJobFragmentExport: + project_dir = str(roar_dir.parent.resolve()) + + with create_database_context(roar_dir) as db_ctx: + selected_job = _select_export_job(db_ctx, job_uid=job_uid) + if selected_job is None: + raise ValueError("no local Roar jobs are available to export") + + selected_job_id = int(selected_job["id"]) + selected_job_uid = str(selected_job.get("job_uid") or "").strip() + resolved_task_id = str(task_id or selected_job_uid).strip() + if not resolved_task_id: + raise ValueError("task_id is required when exporting a job without a persisted job_uid") + + resolved_task_name = _resolve_task_name( + selected_job, + task_name, + default_task_name=default_task_name, + ) + reads = [ + _build_artifact_ref(item, project_dir=project_dir) + for item in db_ctx.jobs.get_inputs(selected_job_id) + ] + writes = [ + _build_artifact_ref(item, project_dir=project_dir) + for item in db_ctx.jobs.get_outputs(selected_job_id) + ] + + started_at = float(selected_job.get("timestamp") or 0.0) + duration = max(0.0, float(selected_job.get("duration_seconds") or 0.0)) + fragment = ExecutionFragment( + job_uid=selected_job_uid or resolved_task_id, + parent_job_uid=parent_job_uid, + task_id=resolved_task_id, + worker_id="", + node_id="", + actor_id=None, + task_name=resolved_task_name, + started_at=started_at, + ended_at=started_at + duration, + exit_code=int(selected_job.get("exit_code") or 0), + backend=str(backend_name or "").strip() or "local", + reads=reads, + writes=writes, + backend_metadata={ + "execution_role": "task", + "source_job_uid": selected_job_uid or None, + "source_execution_backend": str(selected_job.get("execution_backend") or "").strip() + or None, + "source_execution_role": str(selected_job.get("execution_role") or "").strip() or None, + "source_job_type": str(selected_job.get("job_type") or "").strip() or None, + }, + ) + + payload = { + "fragments": [fragment.to_dict()], + "metadata": { + "exported_job_uid": selected_job_uid or None, + "task_id": resolved_task_id, + "task_name": resolved_task_name, + }, + } + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + return LocalJobFragmentExport( + output_path=str(output_path), + exported_job_uid=selected_job_uid, + fragment_count=1, + task_id=resolved_task_id, + task_name=resolved_task_name, + ) + + +def _select_export_job(db_ctx, *, job_uid: str | None) -> dict[str, Any] | None: + if job_uid: + selected = db_ctx.jobs.get_by_uid(job_uid) + if selected is None: + raise ValueError(f"local Roar job not found for job UID {job_uid!r}") + return selected + + jobs = db_ctx.jobs.get_recent(limit=1) + if not jobs: + return None + return jobs[0] + + +def _resolve_task_name( + job: dict[str, Any], + requested_task_name: str | None, + *, + default_task_name: str, +) -> str: + explicit = str(requested_task_name or "").strip() + if explicit: + return explicit + + script = str(job.get("script") or "").strip() + if script: + return script + + command = str(job.get("command") or "").strip() + if not command: + return default_task_name + + first_token = command.split(" ", 1)[0].strip() + return first_token or default_task_name + + +def _build_artifact_ref(payload: dict[str, Any], *, project_dir: str) -> ArtifactRef: + hash_value, hash_algorithm = _select_primary_hash(payload.get("hashes")) + path = _normalize_bundle_path(str(payload.get("path") or ""), project_dir=project_dir) + return ArtifactRef( + path=path, + hash=hash_value, + hash_algorithm=hash_algorithm, + size=int(payload.get("size") or 0), + capture_method="python", + ) + + +def _select_primary_hash(hashes: Any) -> tuple[str | None, str]: + if isinstance(hashes, list): + blake3_row = next( + ( + item + for item in hashes + if isinstance(item, dict) and str(item.get("algorithm") or "").strip() == "blake3" + ), + None, + ) + if isinstance(blake3_row, dict): + digest = str(blake3_row.get("digest") or "").strip() + if digest: + return digest, "blake3" + + for item in hashes: + if not isinstance(item, dict): + continue + algorithm = str(item.get("algorithm") or "").strip() + digest = str(item.get("digest") or "").strip() + if algorithm and digest: + return digest, algorithm + + return None, "blake3" + + +def _normalize_bundle_path(path: str, *, project_dir: str) -> str: + normalized = str(path or "").strip() + if not normalized or "://" in normalized: + return normalized + + candidate = Path(normalized) + project_root = Path(project_dir) + if not candidate.is_absolute(): + candidate = (project_root / candidate).resolve(strict=False) + + try: + relative = candidate.resolve(strict=False).relative_to(project_root) + except ValueError: + return str(candidate) + + return "${ROAR_PROJECT_DIR}/" + relative.as_posix() + + +__all__ = [ + "LocalJobFragmentExport", + "export_local_job_fragment_bundle", +] diff --git a/roar/execution/framework/registry.py b/roar/execution/framework/registry.py index 5c0d584b..b3ddaf2b 100644 --- a/roar/execution/framework/registry.py +++ b/roar/execution/framework/registry.py @@ -14,6 +14,7 @@ _ENTRYPOINT_GROUP = "roar.execution_backends" _BUILTIN_EXECUTION_BACKEND_MODULES = ( "roar.backends.ray.plugin", + "roar.backends.k8s.plugin", "roar.backends.osmo.plugin", "roar.backends.local.plugin", ) @@ -31,7 +32,9 @@ # and triggering discovery would import every backend plugin module — # ~300ms of cost in exchange for checking a handful of stable env-var # names. -_BUILTIN_JOB_ENVIRONMENT_MARKERS: frozenset[str] = frozenset({"RAY_JOB_ID"}) +_BUILTIN_JOB_ENVIRONMENT_MARKERS: frozenset[str] = frozenset( + {"RAY_JOB_ID", "ROAR_K8S_PARENT_JOB_UID"} +) # Top-level TOML section names owned by built-in backends. The config # loader checks against this to decide whether a `config_get("X.Y")` @@ -41,7 +44,7 @@ # would pay ~300ms loading every backend plugin just to confirm "hints" # isn't a backend-namespaced key. Kept in sync with each builtin # backend's ``BackendConfigAdapter.section_name`` by a consistency test. -_BUILTIN_BACKEND_CONFIG_SECTIONS: frozenset[str] = frozenset({"ray", "osmo"}) +_BUILTIN_BACKEND_CONFIG_SECTIONS: frozenset[str] = frozenset({"ray", "osmo", "k8s"}) _registered_execution_backends: list[ExecutionBackend] = [] _execution_backends_discovered = False _execution_backends_discovering = False diff --git a/tests/backends/k8s/README.md b/tests/backends/k8s/README.md index 3a3759ea..7f73d88a 100644 --- a/tests/backends/k8s/README.md +++ b/tests/backends/k8s/README.md @@ -1,11 +1,22 @@ # k8s Lineage E2E Harness (Tier 1) KIND-based harness for pressure-testing roar lineage capture in Kubernetes -training pods. This is the Phase-0/Tier-1 slice from -`design-docs/k8s-training-lineage-integration.md`: no `roar.backends.k8s` -exists yet — the fixtures hand-wrap a Job manifest the way the future -`roar k8s prepare` will, so the runtime assumptions (in-pod tracing, wheel -staging, fragment streaming, identity contract) are proven first. +training pods (`design-docs/k8s-training-lineage-integration.md`). + +Two test layers share this harness: + +- `e2e/test_k8s_product_path.py` — the Phase-1 product path through the real + `roar.backends.k8s` backend: `roar run kubectl apply -f job.yaml` with a + roar-unaware manifest, plan-time rewriting, Secret-delivered credentials, + wheel served to pods over HTTP, and shared-finalizer reconstitution into + the submitting project's `.roar/roar.db`. This is the confidence test. +- `e2e/test_k8s_smoke.py` — the Phase-0 runtime diagnostic: fixtures + hand-wrap the manifest (no backend involved) to isolate the runtime pieces + (in-pod tracing, fragment streaming, identity contract) when the product + path breaks. + +Unit tests for the backend (manifest rewriting, command matching, planning) +live in `unit/` and run in the default gate — no cluster needed. ## Prerequisites diff --git a/tests/backends/k8s/e2e/test_k8s_product_path.py b/tests/backends/k8s/e2e/test_k8s_product_path.py new file mode 100644 index 00000000..15129aae --- /dev/null +++ b/tests/backends/k8s/e2e/test_k8s_product_path.py @@ -0,0 +1,282 @@ +"""Product-path e2e: `roar run kubectl apply -f job.yaml` on the KIND harness. + +Exercises the real `roar.backends.k8s` backend end to end: manifest +rewriting at plan time, Secret-delivered fragment credentials, in-pod +runtime bootstrap from a host-served wheel, fragment streaming to the +local glaas-api, and shared-finalizer reconstitution into the +submitting project's `.roar/roar.db`. +""" + +from __future__ import annotations + +import functools +import http.server +import os +import sqlite3 +import subprocess +import sys +import threading +import uuid +from pathlib import Path +from typing import Any + +import pytest + +from .conftest import HARNESS_DIR, KUBE_CONTEXT, NAMESPACE, TOOLS_BIN, kubectl + +pytestmark = [ + pytest.mark.e2e, + pytest.mark.k8s_e2e, + # The module fixture runs a full instrumented Job (wheel download + + # pip install + trace) inside the first test's budget. + pytest.mark.timeout(900), +] + +REPO_ROOT = HARNESS_DIR.parent.parent.parent +DIST_DIR = REPO_ROOT / "dist" + +JOB_MANIFEST_TEMPLATE = """\ +apiVersion: batch/v1 +kind: Job +metadata: + name: {job_name} + namespace: {namespace} + labels: + app.kubernetes.io/part-of: roar-k8s-e2e +spec: + backoffLimit: 0 + ttlSecondsAfterFinished: 1800 + template: + spec: + restartPolicy: Never + volumes: + - name: work + emptyDir: {{}} + containers: + - name: trainer + image: python:3.12-slim + workingDir: /work + volumeMounts: + - name: work + mountPath: /work + command: + - bash + - -c + - >- + printf 'x,y\\n1.0,2.0\\n2.0,3.9\\n3.0,6.1\\n' > dataset.csv && + python -c "import hashlib, json; + rows = open('dataset.csv').read().strip().splitlines(); + open('model.bin', 'wb').write(hashlib.blake2b(str(len(rows)).encode()).digest() * 4); + json.dump({{'rows': len(rows) - 1}}, open('metrics.json', 'w'))" +""" + + +def _kind_gateway_ip() -> str: + result = subprocess.run( + [ + "docker", + "network", + "inspect", + "kind", + "-f", + "{{range .IPAM.Config}}{{.Gateway}} {{end}}", + ], + capture_output=True, + text=True, + check=True, + ) + for token in result.stdout.split(): + if token.count(".") == 3: + return token + raise RuntimeError(f"no IPv4 gateway found on the kind docker network: {result.stdout!r}") + + +@pytest.fixture(scope="module") +def wheel_server() -> Any: + wheels = sorted(DIST_DIR.glob("roar_cli-*.whl")) + if not wheels: + pytest.skip(f"no roar_cli wheel in {DIST_DIR}; run scripts/build_wheel_with_bins.sh") + + handler = functools.partial( + http.server.SimpleHTTPRequestHandler, + directory=str(DIST_DIR), + ) + server = http.server.ThreadingHTTPServer(("0.0.0.0", 0), handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield { + "url": f"http://{_kind_gateway_ip()}:{server.server_address[1]}/{wheels[-1].name}", + } + finally: + server.shutdown() + + +@pytest.fixture(scope="module") +def product_run( + k8s_cluster: None, + glaas_health: str, + wheel_server: dict[str, str], + tmp_path_factory: pytest.TempPathFactory, +) -> dict[str, Any]: + """Run the full product path once and return project + output state.""" + project_dir = tmp_path_factory.mktemp("k8s-product") + roar_dir = project_dir / ".roar" + roar_dir.mkdir() + (roar_dir / "config.toml").write_text( + "\n".join( + [ + "[glaas]", + 'url = "http://localhost:3001"', + "", + "[k8s]", + "enabled = true", + f'runtime_install_requirement = "{wheel_server["url"]}"', + 'cluster_glaas_url = "http://glaas:3001"', + "wait_for_completion = true", + "wait_timeout_seconds = 300", + "poll_interval_seconds = 2.0", + ] + ) + + "\n", + encoding="utf-8", + ) + + job_name = f"roar-product-{uuid.uuid4().hex[:6]}" + manifest_path = project_dir / "job.yaml" + manifest_path.write_text( + JOB_MANIFEST_TEMPLATE.format(job_name=job_name, namespace=NAMESPACE), + encoding="utf-8", + ) + + env = dict(os.environ) + env["PATH"] = f"{TOOLS_BIN}{os.pathsep}{env.get('PATH', '')}" + env.pop("GLAAS_URL", None) + env.pop("ROAR_PROJECT_DIR", None) + + try: + completed = subprocess.run( + [ + sys.executable, + "-m", + "roar", + "run", + "kubectl", + "apply", + "--context", + KUBE_CONTEXT, + "-f", + "job.yaml", + ], + cwd=project_dir, + env=env, + capture_output=True, + text=True, + check=False, + timeout=600, + ) + pod_logs = kubectl( + ["logs", "-n", NAMESPACE, "-l", f"job-name={job_name}", "--tail=100"], + check=False, + ) + return { + "project_dir": project_dir, + "job_name": job_name, + "manifest_path": manifest_path, + "exit_code": completed.returncode, + "stdout": completed.stdout, + "stderr": completed.stderr, + "pod_logs": pod_logs.stdout + pod_logs.stderr, + } + finally: + kubectl(["delete", f"job/{job_name}", "-n", NAMESPACE, "--ignore-not-found"], check=False) + kubectl( + [ + "delete", + "secret", + "-n", + NAMESPACE, + "-l", + "app.kubernetes.io/managed-by=roar", + "--ignore-not-found", + ], + check=False, + ) + + +def _describe(run: dict[str, Any]) -> str: + return ( + f"exit={run['exit_code']}\nstdout:\n{run['stdout']}\nstderr:\n{run['stderr']}\n" + f"pod logs:\n{run.get('pod_logs', '')}" + ) + + +def _query(run: dict[str, Any], sql: str, params: tuple = ()) -> list[sqlite3.Row]: + db_path = Path(run["project_dir"]) / ".roar" / "roar.db" + assert db_path.is_file(), f"no roar.db produced\n{_describe(run)}" + conn = sqlite3.connect(db_path) + conn.row_factory = sqlite3.Row + try: + return conn.execute(sql, params).fetchall() + finally: + conn.close() + + +def test_product_run_succeeds_and_reconstitutes(product_run: dict[str, Any]) -> None: + assert product_run["exit_code"] == 0, _describe(product_run) + combined = product_run["stdout"] + product_run["stderr"] + assert "lineage reconstituted" in combined, _describe(product_run) + + +def test_submit_job_recorded_with_original_command(product_run: dict[str, Any]) -> None: + rows = _query( + product_run, + "SELECT command, execution_backend, execution_role, exit_code FROM jobs " + "WHERE execution_backend = 'k8s' AND execution_role = 'submit'", + ) + assert len(rows) == 1, _describe(product_run) + submit = rows[0] + assert submit["exit_code"] == 0 + # The recorded command must be the user's original submit (reproduce + # re-enters the backend through it), not the rewritten prepared path. + assert "-f job.yaml" in submit["command"] + assert ".roar/k8s/prepared" not in submit["command"] + + +def test_pod_lineage_merged_as_k8s_task(product_run: dict[str, Any]) -> None: + rows = _query( + product_run, + "SELECT id, job_uid, command FROM jobs WHERE job_type = 'k8s_task'", + ) + assert len(rows) == 1, _describe(product_run) + task = rows[0] + assert task["command"].startswith("k8s_task:") + + output_paths = { + str(row["path"]) + for row in _query( + product_run, + "SELECT path FROM job_outputs WHERE job_id = ?", + (task["id"],), + ) + } + assert any(path.endswith("model.bin") for path in output_paths), output_paths + assert any(path.endswith("metrics.json") for path in output_paths), output_paths + + input_paths = { + str(row["path"]) + for row in _query( + product_run, + "SELECT path FROM job_inputs WHERE job_id = ?", + (task["id"],), + ) + } + assert any(path.endswith("dataset.csv") for path in input_paths), input_paths + + +def test_prepared_manifest_with_secret_is_cleaned_up(product_run: dict[str, Any]) -> None: + prepared_dir = Path(product_run["project_dir"]) / ".roar" / "k8s" / "prepared" + leftovers = list(prepared_dir.glob("*.yaml")) if prepared_dir.is_dir() else [] + assert not leftovers, ( + f"prepared manifests (which embed the token Secret) left behind: {leftovers}" + ) diff --git a/tests/backends/k8s/unit/__init__.py b/tests/backends/k8s/unit/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/backends/k8s/unit/conftest.py b/tests/backends/k8s/unit/conftest.py new file mode 100644 index 00000000..0ce5d77a --- /dev/null +++ b/tests/backends/k8s/unit/conftest.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest +import yaml # type: ignore[import-untyped] + +SINGLE_JOB_MANIFEST = { + "apiVersion": "batch/v1", + "kind": "Job", + "metadata": {"name": "train-demo", "namespace": "ml"}, + "spec": { + "backoffLimit": 0, + "template": { + "spec": { + "restartPolicy": "Never", + "containers": [ + { + "name": "trainer", + "image": "python:3.12-slim", + "command": ["python", "train.py"], + "args": ["--epochs", "3"], + "env": [{"name": "USER_SETTING", "value": "keep"}], + } + ], + } + }, + }, +} + + +@pytest.fixture() +def k8s_project(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + """A temp project with the k8s backend enabled and cwd pointed at it.""" + roar_dir = tmp_path / ".roar" + roar_dir.mkdir() + (roar_dir / "config.toml").write_text( + '[k8s]\nenabled = true\n\n[glaas]\nurl = "http://localhost:3001"\n', + encoding="utf-8", + ) + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("ROAR_PROJECT_DIR", str(tmp_path)) + return tmp_path + + +@pytest.fixture() +def job_manifest_path(k8s_project: Path) -> Path: + manifest = k8s_project / "job.yaml" + manifest.write_text(yaml.safe_dump(SINGLE_JOB_MANIFEST, sort_keys=False), encoding="utf-8") + return manifest diff --git a/tests/backends/k8s/unit/test_backend_selection.py b/tests/backends/k8s/unit/test_backend_selection.py new file mode 100644 index 00000000..a0e1bca5 --- /dev/null +++ b/tests/backends/k8s/unit/test_backend_selection.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +from roar.execution.framework.planning import plan_execution_command +from roar.execution.framework.registry import get_execution_backend + + +def test_planner_routes_kubectl_apply_to_k8s_backend( + job_manifest_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + # Keep planning offline: glaas.url falls back to the hosted default, and + # an instrumented plan would otherwise register a session over the network. + monkeypatch.setattr("roar.backends.k8s.submit._resolve_glaas_url", lambda: None) + config_path = job_manifest_path.parent / ".roar" / "config.toml" + config_path.write_text("[k8s]\nenabled = true\n", encoding="utf-8") + + plan = plan_execution_command(["kubectl", "apply", "-f", str(job_manifest_path)]) + assert plan.backend_name == "k8s" + assert plan.execution_role == "submit" + + +def test_planner_falls_back_to_local_when_disabled( + job_manifest_path: Path, +) -> None: + config_path = job_manifest_path.parent / ".roar" / "config.toml" + config_path.write_text("[k8s]\nenabled = false\n", encoding="utf-8") + + plan = plan_execution_command(["kubectl", "apply", "-f", str(job_manifest_path)]) + assert plan.backend_name == "local" + + +def test_k8s_backend_registered_with_expected_adapters() -> None: + backend = get_execution_backend("k8s") + assert backend.priority == 95 + assert backend.distributed is not None + assert backend.distributed.fragment_reconstitution is not None + assert backend.config is not None + assert backend.config.section_name == "k8s" + assert backend.policy is not None + assert "ROAR_K8S_PARENT_JOB_UID" in backend.policy.job_environment_markers + + +def test_pod_entry_task_identity_contract(monkeypatch: pytest.MonkeyPatch) -> None: + from roar.backends.k8s.pod_entry import task_identity_from_environment + + environ = { + "ROAR_K8S_POD_UID": "abc-123", + "ROAR_K8S_CONTAINER": "trainer", + "JOB_COMPLETION_INDEX": "2", + "TORCHELASTIC_RESTART_COUNT": "1", + "ROAR_K8S_TASK_NAME": "train-demo/trainer", + } + task_id, task_name = task_identity_from_environment(environ) + assert task_id == "abc-123:trainer:2:1" + assert task_name == "train-demo/trainer" + + task_id, task_name = task_identity_from_environment({}) + assert task_id == "unknown-pod:main:0:0" + assert task_name == "k8s-task" diff --git a/tests/backends/k8s/unit/test_manifest_rewrite.py b/tests/backends/k8s/unit/test_manifest_rewrite.py new file mode 100644 index 00000000..fc21d575 --- /dev/null +++ b/tests/backends/k8s/unit/test_manifest_rewrite.py @@ -0,0 +1,151 @@ +from __future__ import annotations + +import copy +from typing import Any + +import pytest + +from roar.backends.k8s.manifest import ( + K8sManifestError, + rewrite_manifest_for_lineage, +) + +from .conftest import SINGLE_JOB_MANIFEST + + +def _rewrite(documents: list[dict[str, Any]], **overrides: Any): + kwargs: dict[str, Any] = { + "secret_name": "roar-fragment-abc", + "session_id": "session-1", + "fragment_token": "ff" * 32, + "requirement": "roar-cli==0.3.7", + "cluster_glaas_url": "http://glaas:3001", + "tracer": "preload", + "parent_job_uid": "cafe0123", + } + kwargs.update(overrides) + return rewrite_manifest_for_lineage(documents, **kwargs) + + +def _job_container(rewrite) -> dict[str, Any]: + job = next(doc for doc in rewrite.documents if doc.get("kind") == "Job") + return job["spec"]["template"]["spec"]["containers"][0] + + +def test_wraps_command_preserving_original_as_positional_args() -> None: + rewrite = _rewrite([copy.deepcopy(SINGLE_JOB_MANIFEST)]) + container = _job_container(rewrite) + + assert container["command"][:2] == ["/bin/sh", "-c"] + script = container["command"][2] + assert "python3 -m pip install --quiet roar-cli==0.3.7" in script + assert "roar.backends.k8s.pod_entry" in script + assert "run_fallback" in script + # argv0 sentinel then the original command + args, args key removed + assert container["command"][3:] == ["roar-k8s", "python", "train.py", "--epochs", "3"] + assert "args" not in container + assert rewrite.wrapped_containers == ["trainer"] + + +def test_injects_env_contract_without_clobbering_user_env() -> None: + rewrite = _rewrite([copy.deepcopy(SINGLE_JOB_MANIFEST)]) + env = {entry["name"]: entry for entry in _job_container(rewrite)["env"]} + + assert env["USER_SETTING"]["value"] == "keep" + assert env["ROAR_EXECUTION_BACKEND"]["value"] == "k8s" + assert env["GLAAS_URL"]["value"] == "http://glaas:3001" + assert env["ROAR_K8S_PARENT_JOB_UID"]["value"] == "cafe0123" + assert env["ROAR_K8S_TASK_NAME"]["value"] == "train-demo/trainer" + assert env["ROAR_SESSION_ID"]["valueFrom"]["secretKeyRef"] == { + "name": "roar-fragment-abc", + "key": "session_id", + } + assert env["ROAR_FRAGMENT_TOKEN"]["valueFrom"]["secretKeyRef"]["key"] == "token" + assert env["ROAR_K8S_POD_UID"]["valueFrom"]["fieldRef"]["fieldPath"] == "metadata.uid" + + +def test_user_env_wins_over_injected_contract() -> None: + manifest = copy.deepcopy(SINGLE_JOB_MANIFEST) + manifest["spec"]["template"]["spec"]["containers"][0]["env"].append( + {"name": "GLAAS_URL", "value": "http://user-override:9999"} + ) + rewrite = _rewrite([manifest]) + env_entries = [ + entry for entry in _job_container(rewrite)["env"] if entry["name"] == "GLAAS_URL" + ] + assert env_entries == [{"name": "GLAAS_URL", "value": "http://user-override:9999"}] + + +def test_appends_secret_document_with_token() -> None: + rewrite = _rewrite([copy.deepcopy(SINGLE_JOB_MANIFEST)]) + secret = next(doc for doc in rewrite.documents if doc.get("kind") == "Secret") + + assert secret["metadata"]["name"] == "roar-fragment-abc" + assert secret["metadata"]["namespace"] == "ml" + assert secret["stringData"]["session_id"] == "session-1" + assert secret["stringData"]["token"] == "ff" * 32 + assert rewrite.namespace == "ml" + + +def test_prepare_mode_omits_secret_document() -> None: + rewrite = _rewrite( + [copy.deepcopy(SINGLE_JOB_MANIFEST)], + session_id=None, + fragment_token=None, + ) + assert not any(doc.get("kind") == "Secret" for doc in rewrite.documents) + env = {entry["name"]: entry for entry in _job_container(rewrite)["env"]} + assert env["ROAR_SESSION_ID"]["valueFrom"]["secretKeyRef"]["name"] == "roar-fragment-abc" + + +def test_container_without_command_fails_actionably() -> None: + manifest = copy.deepcopy(SINGLE_JOB_MANIFEST) + del manifest["spec"]["template"]["spec"]["containers"][0]["command"] + + with pytest.raises(K8sManifestError, match="explicit command"): + _rewrite([manifest]) + + +def test_sidecar_without_command_is_skipped_not_fatal() -> None: + manifest = copy.deepcopy(SINGLE_JOB_MANIFEST) + manifest["spec"]["template"]["spec"]["containers"].append( + {"name": "log-sidecar", "image": "busybox"} + ) + rewrite = _rewrite([manifest]) + assert rewrite.wrapped_containers == ["trainer"] + assert rewrite.skipped_containers == ["log-sidecar"] + + +def test_multiple_jobs_rejected() -> None: + first = copy.deepcopy(SINGLE_JOB_MANIFEST) + second = copy.deepcopy(SINGLE_JOB_MANIFEST) + second["metadata"]["name"] = "train-demo-2" + + with pytest.raises(K8sManifestError, match="exactly one Job"): + _rewrite([first, second]) + + +def test_generate_name_rejected_with_hint() -> None: + manifest = copy.deepcopy(SINGLE_JOB_MANIFEST) + del manifest["metadata"]["name"] + manifest["metadata"]["generateName"] = "train-" + + with pytest.raises(K8sManifestError, match="generateName"): + _rewrite([manifest]) + + +def test_non_job_documents_pass_through_unchanged() -> None: + config_map = { + "apiVersion": "v1", + "kind": "ConfigMap", + "metadata": {"name": "settings"}, + "data": {"key": "value"}, + } + rewrite = _rewrite([copy.deepcopy(config_map), copy.deepcopy(SINGLE_JOB_MANIFEST)]) + passed_through = next(doc for doc in rewrite.documents if doc.get("kind") == "ConfigMap") + assert passed_through == config_map + + +def test_namespace_override_wins() -> None: + rewrite = _rewrite([copy.deepcopy(SINGLE_JOB_MANIFEST)], namespace_override="flagged") + assert rewrite.namespace == "flagged" diff --git a/tests/backends/k8s/unit/test_submit_planning.py b/tests/backends/k8s/unit/test_submit_planning.py new file mode 100644 index 00000000..6dec5384 --- /dev/null +++ b/tests/backends/k8s/unit/test_submit_planning.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +import yaml # type: ignore[import-untyped] + +from roar.backends.k8s.submit import ( + load_submit_context, + matches_kubectl_job_submit_command, + plan_kubectl_job_submit_command, + resolve_runtime_requirement, +) + + +def test_matches_kubectl_apply_with_job_manifest(job_manifest_path: Path) -> None: + assert matches_kubectl_job_submit_command(["kubectl", "apply", "-f", str(job_manifest_path)]) + assert matches_kubectl_job_submit_command( + ["kubectl", "create", f"--filename={job_manifest_path}", "-n", "ml"] + ) + + +def test_does_not_match_when_disabled(job_manifest_path: Path) -> None: + config_path = job_manifest_path.parent / ".roar" / "config.toml" + config_path.write_text("[k8s]\nenabled = false\n", encoding="utf-8") + assert not matches_kubectl_job_submit_command( + ["kubectl", "apply", "-f", str(job_manifest_path)] + ) + + +def test_does_not_match_non_job_manifests(k8s_project: Path) -> None: + manifest = k8s_project / "deploy.yaml" + manifest.write_text( + yaml.safe_dump( + { + "apiVersion": "apps/v1", + "kind": "Deployment", + "metadata": {"name": "web"}, + } + ), + encoding="utf-8", + ) + assert not matches_kubectl_job_submit_command(["kubectl", "apply", "-f", str(manifest)]) + + +def test_does_not_match_other_kubectl_commands(job_manifest_path: Path) -> None: + assert not matches_kubectl_job_submit_command(["kubectl", "get", "pods", "-A"]) + assert not matches_kubectl_job_submit_command( + ["kubectl", "delete", "-f", str(job_manifest_path)] + ) + + +def test_plan_without_glaas_returns_original_command( + job_manifest_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + # glaas.url has a hosted default, so "unconfigured" is modeled at the + # resolver seam rather than via env/config manipulation. + monkeypatch.setattr("roar.backends.k8s.submit._resolve_glaas_url", lambda: None) + + command = ["kubectl", "apply", "-f", str(job_manifest_path)] + plan = plan_kubectl_job_submit_command(command) + + assert plan.backend_name == "k8s" + assert plan.command == command + assert plan.session_id is None + assert plan.execution_role == "submit" + + +def test_plan_rewrites_manifest_and_persists_context( + job_manifest_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + registered: dict[str, str] = {} + + def fake_register(glaas_url: str, session_id: str, token_hash: str, ttl: int = 86400) -> None: + registered.update( + {"glaas_url": glaas_url, "session_id": session_id, "token_hash": token_hash} + ) + + monkeypatch.setattr( + "roar.backends.k8s.submit._register_fragment_session", + fake_register, + ) + + command = ["kubectl", "apply", "--context", "kind-test", "-f", str(job_manifest_path)] + plan = plan_kubectl_job_submit_command(command) + + assert plan.backend_name == "k8s" + assert plan.session_id == registered["session_id"] + assert registered["glaas_url"] == "http://localhost:3001" + + prepared_path = Path(plan.command[plan.command.index("-f") + 1]) + assert prepared_path.is_file() + assert prepared_path.parent == job_manifest_path.parent / ".roar" / "k8s" / "prepared" + # Context flags are preserved verbatim. + assert plan.command[:4] == ["kubectl", "apply", "--context", "kind-test"] + + documents = list(yaml.safe_load_all(prepared_path.read_text(encoding="utf-8"))) + kinds = [doc.get("kind") for doc in documents if doc] + assert kinds.count("Secret") == 1 + assert kinds.count("Job") == 1 + + context = load_submit_context(prepared_path) + assert context is not None + assert context.original_command == command + assert context.job_name == "train-demo" + assert context.namespace == "ml" + assert context.session_id == plan.session_id + assert len(context.parent_job_uid) == 8 + assert context.wrapped_containers == ["trainer"] + + key_path = job_manifest_path.parent / ".roar" / "fragment-sessions" / f"{plan.session_id}.key" + assert key_path.is_file() + key_payload = json.loads(key_path.read_text(encoding="utf-8")) + assert key_payload["session_id"] == plan.session_id + + +def test_plan_degrades_to_uninstrumented_when_registration_fails( + job_manifest_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + def failing_register(*args: object, **kwargs: object) -> None: + raise RuntimeError("glaas is down") + + monkeypatch.setattr( + "roar.backends.k8s.submit._register_fragment_session", + failing_register, + ) + + command = ["kubectl", "apply", "-f", str(job_manifest_path)] + plan = plan_kubectl_job_submit_command(command) + + assert plan.command == command + assert plan.session_id is None + + +def test_resolve_runtime_requirement_precedence(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("ROAR_CLUSTER_PIP_REQ", "http://host/wheel.whl") + assert ( + resolve_runtime_requirement({"runtime_install_requirement": "roar-cli==1.0"}) + == "http://host/wheel.whl" + ) + + monkeypatch.delenv("ROAR_CLUSTER_PIP_REQ") + assert ( + resolve_runtime_requirement({"runtime_install_requirement": "roar-cli==1.0"}) + == "roar-cli==1.0" + ) + assert resolve_runtime_requirement({}).startswith("roar-cli") From 89cc0d73de8d6f64590867b6f55eaa2fd48e9892 Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Fri, 10 Jul 2026 19:07:05 +0000 Subject: [PATCH 03/31] feat(k8s): add distributed workload adapters and roar k8s attach MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 2 of the k8s training lineage integration on top of the Job backend: operator kind adapters, multi-pod distributed capture, and an attach flow for fire-and-forget submits. - manifest rewriter generalized to a WORKLOAD_KINDS adapter registry: batch/v1 Job, jobset.x-k8s.io JobSet (nested replicatedJobs pod templates), kubeflow.org/v1 PyTorchJob (per-role replica specs), and trainer.kubeflow.org TrainJob (wraps the spec.trainer command/env override; requires an explicit trainer command). Task names carry the replica role; terminal wait conditions unioned across kinds in the new workload_wait module - pod identity node-index chain extended: JOB_COMPLETION_INDEX -> PET_NODE_RANK -> pod-level RANK (PyTorchJob v1) - roar k8s attach WORKLOAD: recovers the parent job uid and Secret name from the cluster object's own env contract, resolves credentials (local key -> cluster Secret -> --session-file), optionally waits, records a local attach job under the recovered parent uid, and reconstitutes streamed fragments — the CI flow - harness: bootstrap --with-jobset installs the JobSet controller (v0.12.0); new distributed e2e covers a 2-pod Indexed Job (shared fragment session, distinct completion-index identities, child-process capture, cross-pod artifact edge with topological step ordering), attach from a fresh project via the cluster Secret, and the same worker as a JobSet through the real controller Verified live on KIND: 12/12 e2e (5 new distributed + 4 product path + 3 smoke); 21 new unit tests (65 total k8s-related); default gate 2003 passed; ruff + mypy clean. Co-Authored-By: Claude Fable 5 --- docs/developer/k8s-integration.md | 63 ++- roar/backends/k8s/attach.py | 357 ++++++++++++ roar/backends/k8s/host_execution.py | 90 +-- roar/backends/k8s/manifest.py | 376 ++++++++++--- roar/backends/k8s/pod_entry.py | 5 + roar/backends/k8s/submit.py | 8 +- roar/backends/k8s/workload_wait.py | 127 +++++ roar/cli/commands/k8s.py | 75 ++- tests/backends/k8s/README.md | 12 +- .../backends/k8s/e2e/test_k8s_distributed.py | 531 ++++++++++++++++++ tests/backends/k8s/scripts/bootstrap_k8s.sh | 15 +- tests/backends/k8s/unit/test_attach.py | 75 +++ .../k8s/unit/test_manifest_rewrite.py | 2 +- .../k8s/unit/test_workload_adapters.py | 195 +++++++ 14 files changed, 1734 insertions(+), 197 deletions(-) create mode 100644 roar/backends/k8s/attach.py create mode 100644 roar/backends/k8s/workload_wait.py create mode 100644 tests/backends/k8s/e2e/test_k8s_distributed.py create mode 100644 tests/backends/k8s/unit/test_attach.py create mode 100644 tests/backends/k8s/unit/test_workload_adapters.py diff --git a/docs/developer/k8s-integration.md b/docs/developer/k8s-integration.md index a9b16149..ac543a42 100644 --- a/docs/developer/k8s-integration.md +++ b/docs/developer/k8s-integration.md @@ -2,22 +2,30 @@ ## 1. High-level summary -The k8s backend instruments plain `batch/v1` Jobs submitted with +The k8s backend instruments Kubernetes training workloads submitted with `kubectl apply|create -f ` through `roar run`. Unlike Ray (runtime hooks) it works at the manifest layer, like OSMO — but transports lineage via GLaaS fragment streaming, like Ray, rather than bundle files. -Phase 1 scope (see `design-docs/k8s-training-lineage-integration.md` in the -dev meta-repo): exactly one Job per manifest, explicit container commands -only, streaming transport only. Kubeflow TrainJob/JobSet/RayJob adapters, -bundle fallback, `roar k8s attach`, and the admission-webhook injector are -later phases. +Supported workload kinds (exactly one per manifest; other documents pass +through): `batch/v1` Job (plain or Indexed), `jobset.x-k8s.io` JobSet, +`kubeflow.org/v1` PyTorchJob, and `trainer.kubeflow.org` TrainJob. All pods +of a workload share one fragment session and one parent job uid; per-pod +identity comes from the downward API plus the completion-index/node-rank +env chain (`JOB_COMPLETION_INDEX` → `PET_NODE_RANK` → pod-level `RANK`). + +Current phase status (see `design-docs/k8s-training-lineage-integration.md` +in the dev meta-repo): Phases 1–2 implemented (submit wrapping, operator +adapters, multi-pod capture, `roar k8s attach`). Remaining: bundle-mode +fallback for GLaaS-less pods, boto3/fsspec hooks via sitecustomize, +mount-map path rewriting, RayJob delegation, and the admission-webhook +injector. ## 2. Flow 1. **Match** (`roar/backends/k8s/submit.py`): binary `kubectl`, verb `apply|create`, `-f` pointing at a manifest containing exactly one - `batch/v1` Job; gated by `k8s.enabled` (default off). + supported workload; gated by `k8s.enabled` (default off). 2. **Plan** (same module): pre-registers a GLaaS fragment session (saves the `.key` under `.roar/fragment-sessions/`), rewrites the manifest (`manifest.py`), writes it 0600 under `.roar/k8s/prepared/` with a @@ -25,20 +33,25 @@ later phases. `session_id` — the framework attaches the shared submit finalizer. GLaaS/registration failures degrade to an uninstrumented submit with a warning; lineage never blocks training. -3. **Rewrite** (`manifest.py`): wraps each container that has an explicit - `command` with a `/bin/sh -c` script that pip-installs the roar runtime - (`k8s.runtime_install_requirement`, default pinned `roar-cli`) and execs - `python3 -m roar.backends.k8s.pod_entry "$@"`; on any bootstrap failure it - falls back to exec'ing the original command uninstrumented. Injects the - env contract (`GLAAS_URL` = cluster-visible URL, Secret-backed - `ROAR_SESSION_ID`/`ROAR_FRAGMENT_TOKEN`, downward-API identity fields) and - appends the Secret document. Containers without an explicit command are - skipped with a warning; a Job with none fails actionably (ENTRYPOINT - resolution is a later phase). -4. **Host execution** (`host_execution.py`): runs kubectl, deletes the - prepared manifest (it embeds the token Secret), polls the Job to a - terminal condition (`k8s.wait_for_completion`, default on), and records - the submit as a local job — with the **original** command and the +3. **Rewrite** (`manifest.py`): a `WORKLOAD_KINDS` adapter registry locates + pod templates per kind (Job: `spec.template.spec`; JobSet: + `spec.replicatedJobs[*].template.spec.template.spec`; PyTorchJob: + `spec.pytorchReplicaSpecs.{Role}.template.spec`; TrainJob has no inline + template — its `spec.trainer.command/env` override is wrapped instead). + Each container with an explicit `command` gets a `/bin/sh -c` script that + pip-installs the roar runtime (`k8s.runtime_install_requirement`, default + pinned `roar-cli`) and execs `python3 -m roar.backends.k8s.pod_entry "$@"`; + on any bootstrap failure it falls back to exec'ing the original command + uninstrumented. Injects the env contract (`GLAAS_URL` = cluster-visible + URL, Secret-backed `ROAR_SESSION_ID`/`ROAR_FRAGMENT_TOKEN`, downward-API + identity fields) and appends the Secret document. Containers without an + explicit command are skipped with a warning; a workload with none fails + actionably (ENTRYPOINT resolution is a later phase). +4. **Host execution** (`host_execution.py` + `workload_wait.py`): runs + kubectl, deletes the prepared manifest (it embeds the token Secret), + polls the workload to a terminal condition (`k8s.wait_for_completion`, + default on; condition types unioned across kinds), and records the + submit as a local job — with the **original** command and the plan-generated `parent_job_uid`, so reproduce re-enters the backend and pod fragments link to the submit node. 5. **In pod** (`pod_entry.py`): `roar init -n` + `roar run --tracer preload` @@ -58,6 +71,14 @@ later phases. inspection: no session is registered and no Secret is embedded; the user creates the named Secret out of band or uses the managed path. +`roar k8s attach WORKLOAD [-n ns] [--context ctx] [--wait/--no-wait] +[--session-file key]` (`attach.py`) recovers lineage from an +already-submitted workload — the CI/fire-and-forget flow. It reads the +parent job uid and Secret name from the cluster object's own env contract, +resolves credentials (local `.key` → cluster Secret → `--session-file`), +optionally waits for completion, records a local `attach` job under the +recovered parent uid, and reconstitutes the streamed fragments. + ## 4. Config Section `k8s` (registered `BackendConfigAdapter`): `enabled`, `tracer`, diff --git a/roar/backends/k8s/attach.py b/roar/backends/k8s/attach.py new file mode 100644 index 00000000..18afa71b --- /dev/null +++ b/roar/backends/k8s/attach.py @@ -0,0 +1,357 @@ +"""Attach local roar lineage to an already-submitted Kubernetes workload. + +The CI/fire-and-forget flow: someone (or something) submitted an +instrumented workload; ``roar k8s attach`` recovers the fragment-session +credentials and parent job identity from the cluster object itself, waits +for completion if asked, records a local attach job, and reconstitutes +the streamed fragments into the local ``.roar/roar.db``. +""" + +from __future__ import annotations + +import base64 +import json +import shlex +import subprocess +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from roar.backends.k8s.config import load_k8s_backend_config +from roar.backends.k8s.fragment_reconstituter import create_k8s_fragment_reconstituter +from roar.backends.k8s.manifest import ( + WORKLOAD_KINDS, + WorkloadKind, + workload_kind_for_document, +) +from roar.backends.k8s.workload_wait import ( + get_workload_document, + terminal_condition, + wait_for_workload_completion, +) +from roar.core.bootstrap import bootstrap +from roar.core.operation_metadata import build_operation_metadata_json +from roar.db.context import create_database_context +from roar.execution.fragments.sessions import load_fragment_session +from roar.execution.recording import LocalJobRecorder + + +class K8sAttachError(RuntimeError): + """Raised when attach cannot recover lineage, with actionable detail.""" + + +@dataclass(frozen=True) +class K8sAttachResult: + exit_code: int + workload: str + session_id: str + jobs_merged: int + artifacts_merged: int + fragments_processed: int + + +_KIND_ALIASES = { + "job": "jobs.batch", + "jobs": "jobs.batch", + "jobset": "jobsets.jobset.x-k8s.io", + "jobsets": "jobsets.jobset.x-k8s.io", + "pytorchjob": "pytorchjobs.kubeflow.org", + "pytorchjobs": "pytorchjobs.kubeflow.org", + "trainjob": "trainjobs.trainer.kubeflow.org", + "trainjobs": "trainjobs.trainer.kubeflow.org", +} + + +def attach_k8s_workload( + *, + roar_dir: Path, + repo_root: str, + workload: str, + namespace: str, + kubectl_binary: str = "kubectl", + global_flags: list[str] | None = None, + wait: bool | None = None, + session_file: Path | None = None, +) -> K8sAttachResult: + bootstrap(roar_dir) + started_at = time.time() + flags = list(global_flags or []) + config = load_k8s_backend_config(start_dir=repo_root) + + name, document, workload_kind = _fetch_workload( + workload=workload, + namespace=namespace, + kubectl_binary=kubectl_binary, + global_flags=flags, + ) + kubectl_resource = workload_kind.kubectl_resource + + env_entries = _instrumented_env_entries(document, workload_kind) + parent_job_uid = _env_value(env_entries, "ROAR_K8S_PARENT_JOB_UID") + secret_name = _session_secret_name(env_entries) + if not parent_job_uid or not secret_name: + raise K8sAttachError( + f"{kubectl_resource}/{name} was not instrumented by roar " + "(no fragment-session env found); submit through " + "`roar run kubectl apply -f ...` or `roar k8s prepare`" + ) + + session_id, token = _resolve_session_credentials( + roar_dir=roar_dir, + session_file=session_file, + secret_name=secret_name, + namespace=namespace, + kubectl_binary=kubectl_binary, + global_flags=flags, + ) + + should_wait = bool(config.get("wait_for_completion", True)) if wait is None else wait + succeeded, state = terminal_condition(document) + wait_payload: dict[str, Any] | None = None + if succeeded is None and should_wait: + wait_ok, wait_payload = wait_for_workload_completion( + kubectl_binary=kubectl_binary, + global_flags=flags, + kubectl_resource=kubectl_resource, + name=name, + namespace=namespace, + timeout_seconds=int(config.get("wait_timeout_seconds", 30 * 60)), + poll_interval_seconds=float(config.get("poll_interval_seconds", 5.0)), + ) + succeeded = wait_ok + + exit_code = 0 if succeeded in (True, None) else 1 + payload = { + "workload_kind": workload_kind.kind, + "kubectl_resource": kubectl_resource, + "name": name, + "namespace": namespace, + "secret_name": secret_name, + "session_id": session_id, + "parent_job_uid": parent_job_uid, + "terminal_state": state or None, + "wait": wait_payload, + } + + _ensure_local_attach_job( + roar_dir=roar_dir, + parent_job_uid=parent_job_uid, + command=shlex.join(["roar", "k8s", "attach", f"{kubectl_resource}/{name}"]), + started_at=started_at, + exit_code=exit_code, + payload=payload, + ) + + from roar.integrations.glaas import get_glaas_url + + glaas_url = get_glaas_url() + if not glaas_url: + raise K8sAttachError("GLaaS is not configured; set glaas.url to fetch fragments") + + result = create_k8s_fragment_reconstituter( + session_id, + token, + str(glaas_url), + roar_dir / "roar.db", + ).reconstitute() + + return K8sAttachResult( + exit_code=exit_code, + workload=f"{kubectl_resource}/{name}", + session_id=session_id, + jobs_merged=result.jobs_merged, + artifacts_merged=result.artifacts_merged, + fragments_processed=result.fragments_processed, + ) + + +def _fetch_workload( + *, + workload: str, + namespace: str, + kubectl_binary: str, + global_flags: list[str], +) -> tuple[str, dict[str, Any], WorkloadKind]: + name = workload + candidates: list[str] + if "/" in workload: + kind_part, name = workload.split("/", 1) + resource = _KIND_ALIASES.get(kind_part.lower().split(".")[0]) + if resource is None: + raise K8sAttachError( + f"unknown workload kind {kind_part!r}; " + f"expected one of {', '.join(sorted(set(_KIND_ALIASES)))}" + ) + candidates = [resource] + else: + candidates = [kind.kubectl_resource for kind in WORKLOAD_KINDS] + + errors: list[str] = [] + for resource in candidates: + document, error = get_workload_document( + kubectl_binary=kubectl_binary, + global_flags=global_flags, + kubectl_resource=resource, + name=name, + namespace=namespace, + ) + if document is None: + if error: + errors.append(f"{resource}: {error}") + continue + workload_kind = workload_kind_for_document(document) + if workload_kind is None: + errors.append(f"{resource}: unsupported workload document") + continue + return name, document, workload_kind + + detail = "; ".join(errors) if errors else "not found under any supported workload kind" + raise K8sAttachError(f"cannot fetch workload {workload!r} in namespace {namespace}: {detail}") + + +def _instrumented_env_entries( + document: dict[str, Any], + workload_kind: WorkloadKind, +) -> list[dict[str, Any]]: + if workload_kind.locate_pod_specs is None: + trainer = document.get("spec", {}).get("trainer") + env = trainer.get("env") if isinstance(trainer, dict) else None + return [entry for entry in env or [] if isinstance(entry, dict)] + + for ref in workload_kind.locate_pod_specs(document): + containers = ref.spec.get("containers") + if not isinstance(containers, list): + continue + for container in containers: + if not isinstance(container, dict): + continue + env = [entry for entry in container.get("env") or [] if isinstance(entry, dict)] + if any(entry.get("name") == "ROAR_K8S_PARENT_JOB_UID" for entry in env): + return env + return [] + + +def _env_value(env_entries: list[dict[str, Any]], name: str) -> str: + for entry in env_entries: + if entry.get("name") == name: + return str(entry.get("value") or "").strip() + return "" + + +def _session_secret_name(env_entries: list[dict[str, Any]]) -> str: + for entry in env_entries: + if entry.get("name") != "ROAR_SESSION_ID": + continue + secret_ref = (entry.get("valueFrom") or {}).get("secretKeyRef") or {} + return str(secret_ref.get("name") or "").strip() + return "" + + +def _resolve_session_credentials( + *, + roar_dir: Path, + session_file: Path | None, + secret_name: str, + namespace: str, + kubectl_binary: str, + global_flags: list[str], +) -> tuple[str, str]: + if session_file is not None: + payload = json.loads(session_file.read_text(encoding="utf-8")) + session_id = str(payload.get("session_id") or "").strip() + token = str(payload.get("token") or "").strip() + if not session_id or not token: + raise K8sAttachError(f"session file {session_file} is missing session_id/token") + return session_id, token + + # Prefer a locally saved key (attach on the submitting machine), then + # fall back to reading the cluster Secret (attach from anywhere with RBAC). + session_id_hint = secret_name.removeprefix("roar-fragment-") + local = _find_local_session(roar_dir, session_id_hint) + if local is not None: + return local + + command = [ + kubectl_binary, + *global_flags, + "get", + "secret", + secret_name, + "-n", + namespace, + "-o", + "json", + ] + result = subprocess.run(command, capture_output=True, text=True, check=False) + if result.returncode != 0: + raise K8sAttachError( + f"cannot read Secret {secret_name} in namespace {namespace} " + f"({result.stderr.strip()}); pass --session-file if you have the local key" + ) + data = json.loads(result.stdout).get("data") or {} + try: + session_id = base64.b64decode(str(data.get("session_id") or "")).decode("utf-8").strip() + token = base64.b64decode(str(data.get("token") or "")).decode("utf-8").strip() + except Exception as exc: + raise K8sAttachError(f"Secret {secret_name} has invalid session data: {exc}") from exc + if not session_id or not token: + raise K8sAttachError(f"Secret {secret_name} is missing session_id/token keys") + return session_id, token + + +def _find_local_session(roar_dir: Path, session_id_prefix: str) -> tuple[str, str] | None: + if not session_id_prefix: + return None + sessions_dir = roar_dir / "fragment-sessions" + if not sessions_dir.is_dir(): + return None + for key_path in sorted(sessions_dir.glob(f"{session_id_prefix}*.key")): + try: + payload = load_fragment_session(roar_dir, key_path.stem) + except Exception: + continue + session_id = str(payload.get("session_id") or "").strip() + token = str(payload.get("token") or "").strip() + if session_id and token: + return session_id, token + return None + + +def _ensure_local_attach_job( + *, + roar_dir: Path, + parent_job_uid: str, + command: str, + started_at: float, + exit_code: int, + payload: dict[str, Any], +) -> None: + with create_database_context(roar_dir) as db_ctx: + existing = db_ctx.jobs.get_by_uid(parent_job_uid) + if existing is not None: + return + session_id = db_ctx.sessions.get_or_create_active() + LocalJobRecorder().record( + db_ctx, + command=command, + timestamp=started_at, + metadata=build_operation_metadata_json("k8s_attach", payload), + execution_backend="k8s", + execution_role="attach", + job_type="run", + input_artifacts=[], + output_artifacts=[], + duration_seconds=max(0.0, time.time() - started_at), + exit_code=exit_code, + session_id=session_id, + job_uid=parent_job_uid, + ) + db_ctx.commit() + + +__all__ = [ + "K8sAttachError", + "K8sAttachResult", + "attach_k8s_workload", +] diff --git a/roar/backends/k8s/host_execution.py b/roar/backends/k8s/host_execution.py index c8c27b80..2cbaca62 100644 --- a/roar/backends/k8s/host_execution.py +++ b/roar/backends/k8s/host_execution.py @@ -1,4 +1,4 @@ -"""Host-side execution for kubectl Job submits.""" +"""Host-side execution for kubectl workload submits.""" from __future__ import annotations @@ -17,6 +17,10 @@ discard_submit_context, load_submit_context, ) +from roar.backends.k8s.workload_wait import ( + extract_kubectl_global_flags, + wait_for_workload_completion, +) from roar.core.bootstrap import bootstrap from roar.core.models.run import RunContext, RunResult, resolve_run_config_start_dir from roar.core.operation_metadata import build_operation_metadata_json @@ -25,9 +29,6 @@ from roar.execution.recording import LocalJobRecorder, LocalRecordedArtifact from roar.execution.runtime.errors import ExecutionSetupError -_TERMINAL_SUCCESS_CONDITIONS = ("Complete", "SuccessCriteriaMet") -_TERMINAL_FAILURE_CONDITIONS = ("Failed", "FailureTarget") - def execute_k8s_job_submit(ctx: RunContext) -> RunResult: """Submit a (possibly instrumented) Job manifest and record it locally.""" @@ -67,9 +68,12 @@ def execute_k8s_job_submit(ctx: RunContext) -> RunResult: and submit_context is not None and bool(config.get("wait_for_completion", True)) ): - succeeded, wait_payload = _wait_for_job_completion( - ctx=ctx, - submit_context=submit_context, + succeeded, wait_payload = wait_for_workload_completion( + kubectl_binary=ctx.command[0], + global_flags=extract_kubectl_global_flags(ctx.command), + kubectl_resource=submit_context.kubectl_resource, + name=submit_context.job_name, + namespace=submit_context.namespace, timeout_seconds=int(config.get("wait_timeout_seconds", 30 * 60)), poll_interval_seconds=float(config.get("poll_interval_seconds", 5.0)), ) @@ -125,77 +129,6 @@ def execute_k8s_job_submit(ctx: RunContext) -> RunResult: ) -def _wait_for_job_completion( - *, - ctx: RunContext, - submit_context: K8sSubmitContext, - timeout_seconds: int, - poll_interval_seconds: float, -) -> tuple[bool, dict[str, Any]]: - kubectl_binary = ctx.command[0] - global_flags = _extract_kubectl_global_flags(ctx.command) - get_command = [ - kubectl_binary, - *global_flags, - "get", - f"job/{submit_context.job_name}", - "-n", - submit_context.namespace, - "-o", - "json", - ] - - print( - f"[roar-k8s] waiting for job/{submit_context.job_name} in namespace " - f"{submit_context.namespace} (timeout {timeout_seconds}s)" - ) - deadline = time.time() + timeout_seconds - last_error = "" - while time.time() < deadline: - result = subprocess.run(get_command, capture_output=True, text=True, check=False) - if result.returncode == 0: - try: - status = json.loads(result.stdout).get("status", {}) or {} - except json.JSONDecodeError: - status = {} - for condition in status.get("conditions") or []: - if not isinstance(condition, dict) or condition.get("status") != "True": - continue - condition_type = str(condition.get("type") or "") - if condition_type in _TERMINAL_SUCCESS_CONDITIONS: - print(f"[roar-k8s] job/{submit_context.job_name} completed") - return True, {"terminal_condition": condition_type} - if condition_type in _TERMINAL_FAILURE_CONDITIONS: - message = str(condition.get("message") or "job failed") - print( - f"[roar-k8s] job/{submit_context.job_name} failed: {message}", - file=sys.stderr, - ) - return False, { - "terminal_condition": condition_type, - "message": message, - } - else: - last_error = result.stderr.strip() - time.sleep(poll_interval_seconds) - - message = f"timed out after {timeout_seconds}s" - if last_error: - message = f"{message}; last kubectl error: {last_error}" - print(f"[roar-k8s] wait for job/{submit_context.job_name} {message}", file=sys.stderr) - return False, {"terminal_condition": None, "message": message} - - -def _extract_kubectl_global_flags(command: list[str]) -> list[str]: - flags: list[str] = [] - for index, arg in enumerate(command): - if arg in ("--context", "--kubeconfig", "--cluster", "--user") and index + 1 < len(command): - flags.extend([arg, command[index + 1]]) - elif arg.startswith(("--context=", "--kubeconfig=", "--cluster=", "--user=")): - flags.append(arg) - return flags - - def _build_submit_input_artifacts( submit_context: K8sSubmitContext | None, ) -> list[LocalRecordedArtifact]: @@ -236,6 +169,7 @@ def _build_submit_payload( if submit_context is not None: payload.update( { + "workload_kind": submit_context.workload_kind, "job_name": submit_context.job_name, "namespace": submit_context.namespace, "manifest_path": submit_context.manifest_path, diff --git a/roar/backends/k8s/manifest.py b/roar/backends/k8s/manifest.py index b2f1909d..14b7ed3d 100644 --- a/roar/backends/k8s/manifest.py +++ b/roar/backends/k8s/manifest.py @@ -1,24 +1,27 @@ -"""Kubernetes Job manifest rewriting for lineage instrumentation. +"""Kubernetes workload manifest rewriting for lineage instrumentation. -Phase 1 scope: exactly one ``batch/v1`` Job per manifest (Indexed or -plain). The rewriter wraps explicit container commands through the roar -pod entrypoint, injects the env/identity contract, and appends a Secret +Phase 2 scope: exactly one supported training workload per manifest — +``batch/v1`` Job (plain or Indexed), ``jobset.x-k8s.io`` JobSet, +``kubeflow.org/v1`` PyTorchJob, or ``trainer.kubeflow.org`` TrainJob. + +The rewriter wraps explicit container commands through the roar pod +entrypoint, injects the env/identity contract, and appends a Secret carrying the fragment-session credentials so tokens never appear inline -in pod specs. +in pod specs. All pods of a workload share one fragment session and one +parent job uid; per-pod identity comes from the downward API plus the +completion-index/node-rank env at runtime. """ from __future__ import annotations import shlex +from collections.abc import Callable from dataclasses import dataclass, field from pathlib import Path from typing import Any import yaml # type: ignore[import-untyped] -_JOB_API_VERSION = "batch/v1" -_JOB_KIND = "Job" - # sh -c script: "$0" is the synthetic argv0, "$@" is the original # command+args. Lineage is best-effort by design: any failure to stage # the roar runtime falls back to running the original command @@ -30,14 +33,104 @@ exec python3 -m roar.backends.k8s.pod_entry "$@" """ +# Terminal workload conditions, unioned across kinds: Job uses +# Complete/SuccessCriteriaMet + Failed/FailureTarget, JobSet uses +# Completed/Failed, PyTorchJob v1 uses Succeeded/Failed, TrainJob uses +# Complete/Failed. +WORKLOAD_SUCCESS_CONDITIONS = ("Complete", "Completed", "Succeeded", "SuccessCriteriaMet") +WORKLOAD_FAILURE_CONDITIONS = ("Failed", "FailureTarget") + class K8sManifestError(ValueError): """Raised when a manifest cannot be instrumented, with actionable detail.""" +@dataclass(frozen=True) +class PodSpecRef: + """A mutable reference to one pod spec inside a workload document.""" + + role: str + spec: dict[str, Any] + + +@dataclass(frozen=True) +class WorkloadKind: + kind: str + api_group: str + kubectl_resource: str + # Returns mutable pod-spec references inside the (copied) document. + # None marks trainer-override style workloads (TrainJob), which have + # no inline pod template. + locate_pod_specs: Callable[[dict[str, Any]], list[PodSpecRef]] | None + + +def _job_pod_specs(doc: dict[str, Any]) -> list[PodSpecRef]: + spec = _dict_at(doc, ("spec", "template", "spec")) + return [PodSpecRef(role="", spec=spec)] if spec is not None else [] + + +def _jobset_pod_specs(doc: dict[str, Any]) -> list[PodSpecRef]: + refs: list[PodSpecRef] = [] + replicated_jobs = _dict_get(doc, "spec").get("replicatedJobs") + if not isinstance(replicated_jobs, list): + return refs + for replicated in replicated_jobs: + if not isinstance(replicated, dict): + continue + role = str(replicated.get("name") or "").strip() + spec = _dict_at(replicated, ("template", "spec", "template", "spec")) + if spec is not None: + refs.append(PodSpecRef(role=role, spec=spec)) + return refs + + +def _pytorchjob_pod_specs(doc: dict[str, Any]) -> list[PodSpecRef]: + refs: list[PodSpecRef] = [] + replica_specs = _dict_get(doc, "spec").get("pytorchReplicaSpecs") + if not isinstance(replica_specs, dict): + return refs + for role, replica in replica_specs.items(): + if not isinstance(replica, dict): + continue + spec = _dict_at(replica, ("template", "spec")) + if spec is not None: + refs.append(PodSpecRef(role=str(role), spec=spec)) + return refs + + +WORKLOAD_KINDS: tuple[WorkloadKind, ...] = ( + WorkloadKind( + kind="Job", + api_group="batch", + kubectl_resource="jobs.batch", + locate_pod_specs=_job_pod_specs, + ), + WorkloadKind( + kind="JobSet", + api_group="jobset.x-k8s.io", + kubectl_resource="jobsets.jobset.x-k8s.io", + locate_pod_specs=_jobset_pod_specs, + ), + WorkloadKind( + kind="PyTorchJob", + api_group="kubeflow.org", + kubectl_resource="pytorchjobs.kubeflow.org", + locate_pod_specs=_pytorchjob_pod_specs, + ), + WorkloadKind( + kind="TrainJob", + api_group="trainer.kubeflow.org", + kubectl_resource="trainjobs.trainer.kubeflow.org", + locate_pod_specs=None, + ), +) + + @dataclass(frozen=True) class K8sManifestRewrite: documents: list[dict[str, Any]] + workload_kind: str + kubectl_resource: str job_name: str namespace: str secret_name: str @@ -59,13 +152,27 @@ def load_manifest_documents(path: Path) -> list[dict[str, Any]]: return [doc for doc in documents if isinstance(doc, dict)] -def find_job_documents(documents: list[dict[str, Any]]) -> list[dict[str, Any]]: - return [ - doc - for doc in documents - if str(doc.get("apiVersion") or "") == _JOB_API_VERSION - and str(doc.get("kind") or "") == _JOB_KIND - ] +def workload_kind_for_document(doc: dict[str, Any]) -> WorkloadKind | None: + api_version = str(doc.get("apiVersion") or "") + group = api_version.split("/", 1)[0] if "/" in api_version else "" + if api_version == "batch/v1": + group = "batch" + kind = str(doc.get("kind") or "") + for workload in WORKLOAD_KINDS: + if workload.kind == kind and workload.api_group == group: + return workload + return None + + +def find_workload_documents( + documents: list[dict[str, Any]], +) -> list[tuple[dict[str, Any], WorkloadKind]]: + found: list[tuple[dict[str, Any], WorkloadKind]] = [] + for doc in documents: + workload = workload_kind_for_document(doc) + if workload is not None: + found.append((doc, workload)) + return found def rewrite_manifest_for_lineage( @@ -87,63 +194,59 @@ def rewrite_manifest_for_lineage( contract still points at ``secret_name`` (the ``roar k8s prepare`` flow, where the user creates the Secret out of band). """ - jobs = find_job_documents(documents) - if not jobs: + workloads = find_workload_documents(documents) + if not workloads: + supported = ", ".join(w.kind for w in WORKLOAD_KINDS) + raise K8sManifestError(f"no supported training workload found (supported: {supported})") + if len(workloads) > 1: + names = [ + f"{workload.kind}/{(doc.get('metadata') or {}).get('name') or ''}" + for doc, workload in workloads + ] raise K8sManifestError( - "no batch/v1 Job found in manifest; Phase 1 instruments plain Jobs only" - ) - if len(jobs) > 1: - names = [str((doc.get("metadata") or {}).get("name") or "") for doc in jobs] - raise K8sManifestError( - f"manifest contains {len(jobs)} Jobs ({', '.join(names)}); " - "Phase 1 instruments exactly one Job per submit" + f"manifest contains {len(workloads)} workloads ({', '.join(names)}); " + "roar instruments exactly one workload per submit" ) - job_index = next(index for index, doc in enumerate(documents) if doc is jobs[0]) + source_doc, workload = workloads[0] + doc_index = next(index for index, doc in enumerate(documents) if doc is source_doc) rewritten_documents = [dict(doc) for doc in documents] - job = _deep_copy(jobs[0]) - rewritten_documents[job_index] = job + doc = _deep_copy(source_doc) + rewritten_documents[doc_index] = doc - metadata = job.get("metadata") + metadata = doc.get("metadata") if not isinstance(metadata, dict) or not str(metadata.get("name") or "").strip(): hint = "" if isinstance(metadata, dict) and metadata.get("generateName"): hint = " (generateName is not supported yet; set a fixed metadata.name)" - raise K8sManifestError(f"the Job needs an explicit metadata.name{hint}") - job_name = str(metadata["name"]).strip() + raise K8sManifestError(f"the {workload.kind} needs an explicit metadata.name{hint}") + workload_name = str(metadata["name"]).strip() manifest_namespace = str(metadata.get("namespace") or "").strip() namespace = namespace_override or manifest_namespace or "default" - pod_spec = _require_dict_path(job, ("spec", "template", "spec"), job_name=job_name) - containers = pod_spec.get("containers") - if not isinstance(containers, list) or not containers: - raise K8sManifestError(f"Job {job_name} has no spec.template.spec.containers") + contract = _EnvContract( + secret_name=secret_name, + requirement=requirement, + cluster_glaas_url=cluster_glaas_url, + tracer=tracer, + parent_job_uid=parent_job_uid, + workload_name=workload_name, + ) - wrapped: list[str] = [] - skipped: list[str] = [] - for container in containers: - if not isinstance(container, dict): - continue - container_name = str(container.get("name") or "").strip() or "" - command = container.get("command") - if not isinstance(command, list) or not command: - skipped.append(container_name) - continue - _wrap_container( - container, - job_name=job_name, - secret_name=secret_name, - requirement=requirement, - cluster_glaas_url=cluster_glaas_url, - tracer=tracer, - parent_job_uid=parent_job_uid, + if workload.locate_pod_specs is None: + wrapped, skipped = _rewrite_trainjob(doc, workload_name=workload_name, contract=contract) + else: + wrapped, skipped = _rewrite_pod_specs( + workload.locate_pod_specs(doc), + workload=workload, + workload_name=workload_name, + contract=contract, ) - wrapped.append(container_name) if not wrapped: raise K8sManifestError( - f"Job {job_name} has no container with an explicit command; " - "roar wraps commands it can see — set spec.template.spec.containers[].command " + f"{workload.kind} {workload_name} has no container with an explicit command; " + "roar wraps commands it can see — set an explicit command " "(images relying on ENTRYPOINT are not supported yet)" ) @@ -167,7 +270,9 @@ def rewrite_manifest_for_lineage( return K8sManifestRewrite( documents=rewritten_documents, - job_name=job_name, + workload_kind=workload.kind, + kubectl_resource=workload.kubectl_resource, + job_name=workload_name, namespace=namespace, secret_name=secret_name, wrapped_containers=wrapped, @@ -179,25 +284,112 @@ def dump_manifest_documents(documents: list[dict[str, Any]]) -> str: return yaml.safe_dump_all(documents, sort_keys=False) -def _wrap_container( - container: dict[str, Any], +@dataclass(frozen=True) +class _EnvContract: + secret_name: str + requirement: str + cluster_glaas_url: str + tracer: str + parent_job_uid: str + workload_name: str + + +def _rewrite_pod_specs( + pod_specs: list[PodSpecRef], *, - job_name: str, - secret_name: str, - requirement: str, - cluster_glaas_url: str, - tracer: str, - parent_job_uid: str, -) -> None: - container_name = str(container.get("name") or "").strip() or "main" - original = [str(part) for part in container.get("command", [])] - original.extend(str(part) for part in container.get("args", []) or []) + workload: WorkloadKind, + workload_name: str, + contract: _EnvContract, +) -> tuple[list[str], list[str]]: + if not pod_specs: + raise K8sManifestError( + f"{workload.kind} {workload_name} has no pod templates roar can instrument" + ) + wrapped: list[str] = [] + skipped: list[str] = [] + for ref in pod_specs: + containers = ref.spec.get("containers") + if not isinstance(containers, list) or not containers: + raise K8sManifestError( + f"{workload.kind} {workload_name} pod template " + f"{ref.role or ''} has no containers" + ) + for container in containers: + if not isinstance(container, dict): + continue + container_name = str(container.get("name") or "").strip() or "" + label = f"{ref.role}/{container_name}" if ref.role else container_name + command = container.get("command") + if not isinstance(command, list) or not command: + skipped.append(label) + continue + original = [str(part) for part in command] + original.extend(str(part) for part in container.get("args", []) or []) + container["command"] = _wrapped_command(original, requirement=contract.requirement) + container.pop("args", None) + _inject_env_contract( + container.setdefault("env", []), + contract=contract, + container_name=container_name, + role=ref.role, + ) + wrapped.append(label) + return wrapped, skipped + + +def _rewrite_trainjob( + doc: dict[str, Any], + *, + workload_name: str, + contract: _EnvContract, +) -> tuple[list[str], list[str]]: + """TrainJob has no inline pod template; wrap the trainer override. + + The trainer command/args/env override the runtime blueprint's ``node`` + container. Runtimes whose image entrypoint is torchrun (driven by + operator-injected ``PET_*`` env) expose no command roar can see, so an + explicit ``spec.trainer.command`` is required. + """ + trainer = _dict_get(doc, "spec").get("trainer") + if not isinstance(trainer, dict): + raise K8sManifestError( + f"TrainJob {workload_name} has no spec.trainer override; " + "roar needs spec.trainer.command to wrap (runtime-image entrypoints " + "are not visible at submit time)" + ) + command = trainer.get("command") + if not isinstance(command, list) or not command: + raise K8sManifestError( + f"TrainJob {workload_name} has no spec.trainer.command; " + "set an explicit command to instrument (the runtime image's " + "torchrun entrypoint is not visible at submit time)" + ) + + original = [str(part) for part in command] + original.extend(str(part) for part in trainer.get("args", []) or []) + trainer["command"] = _wrapped_command(original, requirement=contract.requirement) + trainer.pop("args", None) + + env = trainer.setdefault("env", []) + if not isinstance(env, list): + raise K8sManifestError(f"TrainJob {workload_name} has a non-list spec.trainer.env") + _inject_env_contract(env, contract=contract, container_name="node", role="") + return ["node"], [] + + +def _wrapped_command(original: list[str], *, requirement: str) -> list[str]: script = _POD_WRAPPER_TEMPLATE.format(requirement=shlex.quote(requirement)) - container["command"] = ["/bin/sh", "-c", script, "roar-k8s", *original] - container.pop("args", None) + return ["/bin/sh", "-c", script, "roar-k8s", *original] - env = container.setdefault("env", []) + +def _inject_env_contract( + env: list[Any], + *, + contract: _EnvContract, + container_name: str, + role: str, +) -> None: if not isinstance(env, list): raise K8sManifestError(f"container {container_name} has a non-list env block") existing_names = { @@ -217,18 +409,23 @@ def add_secret_ref(name: str, key: str) -> None: env.append( { "name": name, - "valueFrom": {"secretKeyRef": {"name": secret_name, "key": key}}, + "valueFrom": {"secretKeyRef": {"name": contract.secret_name, "key": key}}, } ) + task_name = contract.workload_name + if role: + task_name = f"{task_name}/{role}" + task_name = f"{task_name}/{container_name}" + add_value("ROAR_EXECUTION_BACKEND", "k8s") add_value("ROAR_NO_TELEMETRY", "1") - add_value("ROAR_K8S_TRACER", tracer) - add_value("GLAAS_URL", cluster_glaas_url) - add_value("ROAR_K8S_PARENT_JOB_UID", parent_job_uid) - add_value("ROAR_K8S_JOB_NAME", job_name) + add_value("ROAR_K8S_TRACER", contract.tracer) + add_value("GLAAS_URL", contract.cluster_glaas_url) + add_value("ROAR_K8S_PARENT_JOB_UID", contract.parent_job_uid) + add_value("ROAR_K8S_JOB_NAME", contract.workload_name) add_value("ROAR_K8S_CONTAINER", container_name) - add_value("ROAR_K8S_TASK_NAME", f"{job_name}/{container_name}") + add_value("ROAR_K8S_TASK_NAME", task_name) add_secret_ref("ROAR_SESSION_ID", "session_id") add_secret_ref("ROAR_FRAGMENT_TOKEN", "token") add_field_ref("ROAR_K8S_NAMESPACE", "metadata.namespace") @@ -237,20 +434,16 @@ def add_secret_ref(name: str, key: str) -> None: add_field_ref("ROAR_K8S_NODE_NAME", "spec.nodeName") -def _require_dict_path( - root: dict[str, Any], - path: tuple[str, ...], - *, - job_name: str, -) -> dict[str, Any]: +def _dict_get(doc: dict[str, Any], key: str) -> dict[str, Any]: + value = doc.get(key) + return value if isinstance(value, dict) else {} + + +def _dict_at(root: dict[str, Any], path: tuple[str, ...]) -> dict[str, Any] | None: node: Any = root - walked: list[str] = [] for key in path: - walked.append(key) node = node.get(key) if isinstance(node, dict) else None - if not isinstance(node, dict): - raise K8sManifestError(f"Job {job_name} is missing {'.'.join(walked)}") - return node + return node if isinstance(node, dict) else None def _deep_copy(document: dict[str, Any]) -> dict[str, Any]: @@ -260,10 +453,15 @@ def _deep_copy(document: dict[str, Any]) -> dict[str, Any]: __all__ = [ + "WORKLOAD_FAILURE_CONDITIONS", + "WORKLOAD_KINDS", + "WORKLOAD_SUCCESS_CONDITIONS", "K8sManifestError", "K8sManifestRewrite", + "WorkloadKind", "dump_manifest_documents", - "find_job_documents", + "find_workload_documents", "load_manifest_documents", "rewrite_manifest_for_lineage", + "workload_kind_for_document", ] diff --git a/roar/backends/k8s/pod_entry.py b/roar/backends/k8s/pod_entry.py index 0a409b89..c5489f11 100644 --- a/roar/backends/k8s/pod_entry.py +++ b/roar/backends/k8s/pod_entry.py @@ -80,9 +80,14 @@ def task_identity_from_environment(environ: dict[str, str] | None = None) -> tup env = os.environ if environ is None else environ pod_uid = str(env.get("ROAR_K8S_POD_UID") or "unknown-pod").strip() or "unknown-pod" container = str(env.get("ROAR_K8S_CONTAINER") or "main").strip() or "main" + # Node-index chain: Indexed Job/JobSet completion index, then torchrun's + # PET_NODE_RANK, then PyTorchJob v1's operator-injected pod-level RANK + # (stable node rank there — process-level RANK is never read because + # pod_entry runs above the launcher). completion_index = ( str(env.get("JOB_COMPLETION_INDEX") or "").strip() or str(env.get("PET_NODE_RANK") or "").strip() + or str(env.get("RANK") or "").strip() or "0" ) restart_attempt = ( diff --git a/roar/backends/k8s/submit.py b/roar/backends/k8s/submit.py index 62b189c8..416466c7 100644 --- a/roar/backends/k8s/submit.py +++ b/roar/backends/k8s/submit.py @@ -13,7 +13,7 @@ from roar.backends.k8s.manifest import ( K8sManifestError, dump_manifest_documents, - find_job_documents, + find_workload_documents, load_manifest_documents, rewrite_manifest_for_lineage, ) @@ -31,6 +31,8 @@ class K8sSubmitContext: original_command: list[str] manifest_path: str prepared_path: str + workload_kind: str + kubectl_resource: str job_name: str namespace: str secret_name: str @@ -61,7 +63,7 @@ def matches_kubectl_job_submit_command(command: list[str]) -> bool: documents = load_manifest_documents(manifest_path) except K8sManifestError: return False - return len(find_job_documents(documents)) == 1 + return len(find_workload_documents(documents)) == 1 def plan_kubectl_job_submit_command(command: list[str]) -> ExecutionCommandPlan: @@ -137,6 +139,8 @@ def plan_kubectl_job_submit_command(command: list[str]) -> ExecutionCommandPlan: original_command=list(command), manifest_path=str(manifest_path), prepared_path=str(prepared_path), + workload_kind=rewrite.workload_kind, + kubectl_resource=rewrite.kubectl_resource, job_name=rewrite.job_name, namespace=rewrite.namespace, secret_name=secret_name, diff --git a/roar/backends/k8s/workload_wait.py b/roar/backends/k8s/workload_wait.py new file mode 100644 index 00000000..7f1cd18c --- /dev/null +++ b/roar/backends/k8s/workload_wait.py @@ -0,0 +1,127 @@ +"""Shared kubectl helpers for workload status polling and retrieval.""" + +from __future__ import annotations + +import json +import subprocess +import sys +import time +from typing import Any + +from roar.backends.k8s.manifest import ( + WORKLOAD_FAILURE_CONDITIONS, + WORKLOAD_SUCCESS_CONDITIONS, +) + + +def extract_kubectl_global_flags(command: list[str]) -> list[str]: + """Pull connection flags out of a kubectl command for reuse in follow-ups.""" + flags: list[str] = [] + for index, arg in enumerate(command): + if arg in ("--context", "--kubeconfig", "--cluster", "--user") and index + 1 < len(command): + flags.extend([arg, command[index + 1]]) + elif arg.startswith(("--context=", "--kubeconfig=", "--cluster=", "--user=")): + flags.append(arg) + return flags + + +def get_workload_document( + *, + kubectl_binary: str, + global_flags: list[str], + kubectl_resource: str, + name: str, + namespace: str, +) -> tuple[dict[str, Any] | None, str]: + """Fetch a workload object as JSON; returns (document, error).""" + command = [ + kubectl_binary, + *global_flags, + "get", + f"{kubectl_resource}/{name}", + "-n", + namespace, + "-o", + "json", + ] + result = subprocess.run(command, capture_output=True, text=True, check=False) + if result.returncode != 0: + return None, result.stderr.strip() + try: + payload = json.loads(result.stdout) + except json.JSONDecodeError as exc: + return None, f"invalid kubectl JSON output: {exc}" + return (payload, "") if isinstance(payload, dict) else (None, "unexpected kubectl output") + + +def terminal_condition(document: dict[str, Any]) -> tuple[bool | None, str]: + """Return (succeeded, message) from workload status conditions. + + ``succeeded`` is None while the workload is still running. Condition + type names are unioned across Job/JobSet/PyTorchJob/TrainJob. + """ + status = document.get("status") + conditions = status.get("conditions") if isinstance(status, dict) else None + for condition in conditions or []: + if not isinstance(condition, dict) or condition.get("status") != "True": + continue + condition_type = str(condition.get("type") or "") + if condition_type in WORKLOAD_SUCCESS_CONDITIONS: + return True, condition_type + if condition_type in WORKLOAD_FAILURE_CONDITIONS: + return False, str(condition.get("message") or condition_type) + return None, "" + + +def wait_for_workload_completion( + *, + kubectl_binary: str, + global_flags: list[str], + kubectl_resource: str, + name: str, + namespace: str, + timeout_seconds: int, + poll_interval_seconds: float, +) -> tuple[bool, dict[str, Any]]: + print( + f"[roar-k8s] waiting for {kubectl_resource}/{name} in namespace " + f"{namespace} (timeout {timeout_seconds}s)" + ) + deadline = time.time() + timeout_seconds + last_error = "" + while time.time() < deadline: + document, error = get_workload_document( + kubectl_binary=kubectl_binary, + global_flags=global_flags, + kubectl_resource=kubectl_resource, + name=name, + namespace=namespace, + ) + if document is not None: + succeeded, message = terminal_condition(document) + if succeeded is True: + print(f"[roar-k8s] {kubectl_resource}/{name} completed") + return True, {"terminal_condition": message} + if succeeded is False: + print( + f"[roar-k8s] {kubectl_resource}/{name} failed: {message}", + file=sys.stderr, + ) + return False, {"terminal_condition": "Failed", "message": message} + elif error: + last_error = error + time.sleep(poll_interval_seconds) + + message = f"timed out after {timeout_seconds}s" + if last_error: + message = f"{message}; last kubectl error: {last_error}" + print(f"[roar-k8s] wait for {kubectl_resource}/{name} {message}", file=sys.stderr) + return False, {"terminal_condition": None, "message": message} + + +__all__ = [ + "extract_kubectl_global_flags", + "get_workload_document", + "terminal_condition", + "wait_for_workload_completion", +] diff --git a/roar/cli/commands/k8s.py b/roar/cli/commands/k8s.py index 3582f1f2..a37b1af7 100644 --- a/roar/cli/commands/k8s.py +++ b/roar/cli/commands/k8s.py @@ -5,6 +5,7 @@ import click +from ...backends.k8s.attach import K8sAttachError, attach_k8s_workload from ...backends.k8s.config import load_k8s_backend_config from ...backends.k8s.manifest import ( K8sManifestError, @@ -13,16 +14,83 @@ rewrite_manifest_for_lineage, ) from ...backends.k8s.submit import resolve_runtime_requirement +from ..context import RoarContext +from ..decorators import require_init @click.group("k8s", invoke_without_command=True) @click.pass_context def k8s(ctx: click.Context) -> None: - """Prepare Kubernetes Jobs for lineage capture.""" + """Prepare and attach Kubernetes workloads for lineage capture.""" if ctx.invoked_subcommand is None: click.echo(ctx.get_help()) +@k8s.command("attach") +@click.argument("workload") +@click.option( + "-n", + "--namespace", + default="default", + show_default=True, + help="Namespace of the workload", +) +@click.option( + "--context", + "kube_context", + default=None, + help="kubectl context to use", +) +@click.option( + "--wait/--no-wait", + default=None, + help="Override k8s.wait_for_completion for this attach", +) +@click.option( + "--session-file", + type=click.Path(path_type=Path, dir_okay=False, exists=True), + default=None, + help="Fragment-session .key file when the cluster Secret is not readable", +) +@click.pass_obj +@require_init +def k8s_attach( + ctx: RoarContext, + workload: str, + namespace: str, + kube_context: str | None, + wait: bool | None, + session_file: Path | None, +) -> None: + """Reconstitute lineage from an already-submitted workload. + + WORKLOAD is a name or KIND/NAME (job, jobset, pytorchjob, trainjob) that + was instrumented at submit time. Credentials come from a locally saved + fragment-session key, the cluster Secret, or --session-file. + """ + global_flags = ["--context", kube_context] if kube_context else [] + try: + result = attach_k8s_workload( + roar_dir=ctx.roar_dir, + repo_root=str(ctx.repo_root or ctx.roar_dir.parent), + workload=workload, + namespace=namespace, + global_flags=global_flags, + wait=wait, + session_file=session_file, + ) + except K8sAttachError as exc: + raise click.ClickException(str(exc)) from exc + + click.echo( + f"[roar] lineage reconstituted from {result.workload}: " + f"{result.jobs_merged} jobs, {result.artifacts_merged} artifacts " + f"({result.fragments_processed} fragments)" + ) + if result.exit_code != 0: + raise SystemExit(result.exit_code) + + @k8s.command("prepare") @click.option( "-f", @@ -94,7 +162,10 @@ def k8s_prepare( output_path.write_text(dump_manifest_documents(rewrite.documents), encoding="utf-8") click.echo(f"Prepared manifest written to {output_path}") - click.echo(f" job: {rewrite.job_name} (namespace {rewrite.namespace})") + click.echo( + f" workload: {rewrite.workload_kind}/{rewrite.job_name} " + f"(namespace {rewrite.namespace})" + ) click.echo(f" wrapped containers: {', '.join(rewrite.wrapped_containers)}") if rewrite.skipped_containers: click.echo(f" skipped (no explicit command): {', '.join(rewrite.skipped_containers)}") diff --git a/tests/backends/k8s/README.md b/tests/backends/k8s/README.md index 7f73d88a..527b3e38 100644 --- a/tests/backends/k8s/README.md +++ b/tests/backends/k8s/README.md @@ -3,13 +3,18 @@ KIND-based harness for pressure-testing roar lineage capture in Kubernetes training pods (`design-docs/k8s-training-lineage-integration.md`). -Two test layers share this harness: +Three test layers share this harness: - `e2e/test_k8s_product_path.py` — the Phase-1 product path through the real `roar.backends.k8s` backend: `roar run kubectl apply -f job.yaml` with a roar-unaware manifest, plan-time rewriting, Secret-delivered credentials, wheel served to pods over HTTP, and shared-finalizer reconstitution into the submitting project's `.roar/roar.db`. This is the confidence test. +- `e2e/test_k8s_distributed.py` — Phase-2 coverage: two-pod Indexed Job with + completion-index identity, child-process capture, and a cross-pod artifact + edge over a shared volume; `roar k8s attach` from a fresh project via the + cluster Secret; and a JobSet run through the real controller (skipped + unless bootstrapped with `--with-jobset`). - `e2e/test_k8s_smoke.py` — the Phase-0 runtime diagnostic: fixtures hand-wrap the manifest (no backend involved) to isolate the runtime pieces (in-pod tracing, fragment streaming, identity contract) when the product @@ -31,8 +36,9 @@ live in `unit/` and run in the default gate — no cluster needed. # one-time (and after wheel changes) bash scripts/build_wheel_with_bins.sh -# create cluster + wire glaas + preflight (add --with-minio for S3 scenarios) -bash tests/backends/k8s/scripts/bootstrap_k8s.sh +# create cluster + wire glaas + preflight +# (--with-minio for S3 scenarios, --with-jobset for the JobSet operator e2e) +bash tests/backends/k8s/scripts/bootstrap_k8s.sh --with-jobset # run the smoke tests (addopts override needed: e2e dirs are ignored by default) pytest tests/backends/k8s/e2e -o addopts='' -m k8s_e2e -v diff --git a/tests/backends/k8s/e2e/test_k8s_distributed.py b/tests/backends/k8s/e2e/test_k8s_distributed.py new file mode 100644 index 00000000..8d8d8ca4 --- /dev/null +++ b/tests/backends/k8s/e2e/test_k8s_distributed.py @@ -0,0 +1,531 @@ +"""Phase-2 e2e: multi-pod distributed capture, attach, and JobSet. + +Runs against the KIND harness through the real product path +(`roar run kubectl apply -f ...`). Covers: + +- Indexed Job with two pods sharing one fragment session: completion-index + identity, child-process capture, and a cross-pod artifact edge through a + shared volume. +- `roar k8s attach` from a fresh project (cluster-Secret credential path) + after a no-wait submit. +- JobSet through the same worker (skipped unless the JobSet controller is + installed; `bootstrap_k8s.sh --with-jobset`). +""" + +from __future__ import annotations + +import json +import os +import sqlite3 +import subprocess +import sys +import time +import uuid +from pathlib import Path +from typing import Any + +import pytest + +from .conftest import KUBE_CONTEXT, NAMESPACE, TOOLS_BIN, kubectl +from .test_k8s_product_path import wheel_server # noqa: F401 (module fixture reuse) + +pytestmark = [ + pytest.mark.e2e, + pytest.mark.k8s_e2e, + pytest.mark.timeout(900), +] + +# Both pods pin to one worker node so a hostPath volume can model the +# shared filesystem (RWX PVCs aren't available on the default KIND +# storage class); the cross-pod edge itself is content-hash based. +PINNED_NODE = "roar-k8s-e2e-worker" + +WORKER_SCRIPT = """\ +import json +import os +import subprocess +import sys +import time +from pathlib import Path + +index = os.environ.get("JOB_COMPLETION_INDEX", "0") +shared = Path("/shared") +if index == "0": + subprocess.run( + [sys.executable, "-c", "open('child-stats.json', 'w').write('loss 0.1')"], + check=True, + ) + weights = b"weights:" + b"0" * 64 + tmp = shared / "weights.bin.tmp" + tmp.write_bytes(weights) + tmp.rename(shared / "weights.bin") + Path("rank0.json").write_text(json.dumps({"rank": 0})) +else: + deadline = time.time() + 180 + while not (shared / "weights.bin").exists(): + if time.time() > deadline: + raise SystemExit("timed out waiting for rank 0 weights") + time.sleep(2) + data = (shared / "weights.bin").read_bytes() + Path("eval.json").write_text(json.dumps({"rank": 1, "bytes": len(data)})) +print(f"worker {index} done") +""" + +INDEXED_JOB_TEMPLATE = """\ +apiVersion: v1 +kind: ConfigMap +metadata: + name: {name}-worker + namespace: {namespace} +data: + worker.py: | +{worker_script} +--- +apiVersion: batch/v1 +kind: Job +metadata: + name: {name} + namespace: {namespace} + labels: + app.kubernetes.io/part-of: roar-k8s-e2e +spec: + backoffLimit: 0 + completions: 2 + parallelism: 2 + completionMode: Indexed + ttlSecondsAfterFinished: 1800 + template: + spec: + restartPolicy: Never + nodeSelector: + kubernetes.io/hostname: {pinned_node} + volumes: + - name: work + emptyDir: {{}} + - name: shared + hostPath: + path: /tmp/roar-e2e-shared-{name} + type: DirectoryOrCreate + - name: workload + configMap: + name: {name}-worker + containers: + - name: trainer + image: python:3.12-slim + workingDir: /work + volumeMounts: + - name: work + mountPath: /work + - name: shared + mountPath: /shared + - name: workload + mountPath: /workload + readOnly: true + command: ["python", "/workload/worker.py"] +""" + +JOBSET_TEMPLATE = """\ +apiVersion: v1 +kind: ConfigMap +metadata: + name: {name}-worker + namespace: {namespace} +data: + worker.py: | +{worker_script} +--- +apiVersion: jobset.x-k8s.io/v1alpha2 +kind: JobSet +metadata: + name: {name} + namespace: {namespace} + labels: + app.kubernetes.io/part-of: roar-k8s-e2e +spec: + replicatedJobs: + - name: workers + replicas: 1 + template: + spec: + backoffLimit: 0 + completions: 2 + parallelism: 2 + completionMode: Indexed + template: + spec: + restartPolicy: Never + nodeSelector: + kubernetes.io/hostname: {pinned_node} + volumes: + - name: work + emptyDir: {{}} + - name: shared + hostPath: + path: /tmp/roar-e2e-shared-{name} + type: DirectoryOrCreate + - name: workload + configMap: + name: {name}-worker + containers: + - name: trainer + image: python:3.12-slim + workingDir: /work + volumeMounts: + - name: work + mountPath: /work + - name: shared + mountPath: /shared + - name: workload + mountPath: /workload + readOnly: true + command: ["python", "/workload/worker.py"] +""" + +SINGLE_JOB_TEMPLATE = """\ +apiVersion: batch/v1 +kind: Job +metadata: + name: {name} + namespace: {namespace} + labels: + app.kubernetes.io/part-of: roar-k8s-e2e +spec: + backoffLimit: 0 + ttlSecondsAfterFinished: 1800 + template: + spec: + restartPolicy: Never + volumes: + - name: work + emptyDir: {{}} + containers: + - name: trainer + image: python:3.12-slim + workingDir: /work + volumeMounts: + - name: work + mountPath: /work + command: + - python + - -c + - >- + import json; + open('data.csv', 'w').write('x\\n1\\n2\\n'); + rows = open('data.csv').read().strip().splitlines(); + json.dump({{'rows': len(rows) - 1}}, open('summary.json', 'w')) +""" + + +def _indent_script(script: str, spaces: int = 4) -> str: + pad = " " * spaces + return "\n".join(f"{pad}{line}" if line else "" for line in script.splitlines()) + + +def _write_project(project_dir: Path, wheel_url: str, *, wait: bool = True) -> None: + roar_dir = project_dir / ".roar" + roar_dir.mkdir(parents=True, exist_ok=True) + (roar_dir / "config.toml").write_text( + "\n".join( + [ + "[glaas]", + 'url = "http://localhost:3001"', + "", + "[k8s]", + "enabled = true", + f'runtime_install_requirement = "{wheel_url}"', + 'cluster_glaas_url = "http://glaas:3001"', + f"wait_for_completion = {'true' if wait else 'false'}", + "wait_timeout_seconds = 420", + "poll_interval_seconds = 2.0", + ] + ) + + "\n", + encoding="utf-8", + ) + + +def _roar_env() -> dict[str, str]: + env = dict(os.environ) + env["PATH"] = f"{TOOLS_BIN}{os.pathsep}{env.get('PATH', '')}" + env.pop("GLAAS_URL", None) + env.pop("ROAR_PROJECT_DIR", None) + return env + + +def _roar(args: list[str], *, cwd: Path) -> subprocess.CompletedProcess[str]: + return subprocess.run( + [sys.executable, "-m", "roar", *args], + cwd=cwd, + env=_roar_env(), + capture_output=True, + text=True, + check=False, + timeout=700, + ) + + +def _submit(manifest_name: str, *, cwd: Path) -> subprocess.CompletedProcess[str]: + return _roar( + ["run", "kubectl", "apply", "--context", KUBE_CONTEXT, "-f", manifest_name], + cwd=cwd, + ) + + +def _query(project_dir: Path, sql: str, params: tuple = ()) -> list[sqlite3.Row]: + conn = sqlite3.connect(project_dir / ".roar" / "roar.db") + conn.row_factory = sqlite3.Row + try: + return conn.execute(sql, params).fetchall() + finally: + conn.close() + + +def _pod_logs(selector: str) -> str: + result = kubectl( + ["logs", "-n", NAMESPACE, "-l", selector, "--tail=100", "--prefix"], + check=False, + ) + return result.stdout + result.stderr + + +def _cleanup(job_name: str, *, resource: str = "job") -> None: + kubectl( + ["delete", f"{resource}/{job_name}", "-n", NAMESPACE, "--ignore-not-found"], + check=False, + ) + kubectl( + ["delete", f"configmap/{job_name}-worker", "-n", NAMESPACE, "--ignore-not-found"], + check=False, + ) + kubectl( + [ + "delete", + "secret", + "-n", + NAMESPACE, + "-l", + "app.kubernetes.io/managed-by=roar", + "--ignore-not-found", + ], + check=False, + ) + + +@pytest.fixture(scope="module") +def indexed_run( + k8s_cluster: None, + glaas_health: str, + wheel_server: dict[str, str], # noqa: F811 + tmp_path_factory: pytest.TempPathFactory, +) -> dict[str, Any]: + project_dir = tmp_path_factory.mktemp("k8s-indexed") + _write_project(project_dir, wheel_server["url"]) + + job_name = f"roar-dist-{uuid.uuid4().hex[:6]}" + (project_dir / "job.yaml").write_text( + INDEXED_JOB_TEMPLATE.format( + name=job_name, + namespace=NAMESPACE, + pinned_node=PINNED_NODE, + worker_script=_indent_script(WORKER_SCRIPT), + ), + encoding="utf-8", + ) + + try: + completed = _submit("job.yaml", cwd=project_dir) + return { + "project_dir": project_dir, + "job_name": job_name, + "exit_code": completed.returncode, + "stdout": completed.stdout, + "stderr": completed.stderr, + "pod_logs": _pod_logs(f"job-name={job_name}"), + } + finally: + _cleanup(job_name) + + +def _describe(run: dict[str, Any]) -> str: + return ( + f"exit={run['exit_code']}\nstdout:\n{run['stdout']}\nstderr:\n{run['stderr']}\n" + f"pod logs:\n{run.get('pod_logs', '')}" + ) + + +def test_indexed_job_captures_both_pods(indexed_run: dict[str, Any]) -> None: + assert indexed_run["exit_code"] == 0, _describe(indexed_run) + + tasks = _query( + indexed_run["project_dir"], + "SELECT id, metadata FROM jobs WHERE job_type = 'k8s_task' ORDER BY id", + ) + assert len(tasks) == 2, _describe(indexed_run) + + all_metadata = " ".join(str(row["metadata"] or "") for row in tasks) + assert ":trainer:0:0" in all_metadata, all_metadata + assert ":trainer:1:0" in all_metadata, all_metadata + + +def test_child_process_io_is_captured(indexed_run: dict[str, Any]) -> None: + outputs = _query( + indexed_run["project_dir"], + "SELECT o.path FROM job_outputs o JOIN jobs j ON j.id = o.job_id " + "WHERE j.job_type = 'k8s_task'", + ) + paths = {str(row["path"]) for row in outputs} + assert any(path.endswith("child-stats.json") for path in paths), ( + f"child process write missing from lineage: {paths}\n{_describe(indexed_run)}" + ) + + +def test_cross_pod_artifact_edge_connects(indexed_run: dict[str, Any]) -> None: + rows = _query( + indexed_run["project_dir"], + "SELECT o.artifact_id AS writer_artifact, i.artifact_id AS reader_artifact, " + "o.job_id AS writer_job, i.job_id AS reader_job " + "FROM job_outputs o JOIN job_inputs i ON i.artifact_id = o.artifact_id " + "WHERE o.path LIKE '%weights.bin' AND i.path LIKE '%weights.bin' " + "AND o.job_id != i.job_id", + ) + assert rows, ( + "no cross-pod edge: weights.bin written by rank 0 was not linked to " + f"rank 1's read\n{_describe(indexed_run)}" + ) + + steps = _query( + indexed_run["project_dir"], + "SELECT id, step_number FROM jobs WHERE id IN (?, ?)", + (rows[0]["writer_job"], rows[0]["reader_job"]), + ) + step_by_id = {row["id"]: row["step_number"] for row in steps} + assert step_by_id[rows[0]["reader_job"]] > step_by_id[rows[0]["writer_job"]], ( + f"reader should be downstream of writer: {step_by_id}" + ) + + +def test_attach_from_fresh_project_via_cluster_secret( + k8s_cluster: None, + glaas_health: str, + wheel_server: dict[str, str], # noqa: F811 + tmp_path_factory: pytest.TempPathFactory, +) -> None: + submit_dir = tmp_path_factory.mktemp("k8s-submit") + _write_project(submit_dir, wheel_server["url"], wait=False) + + job_name = f"roar-attach-{uuid.uuid4().hex[:6]}" + (submit_dir / "job.yaml").write_text( + SINGLE_JOB_TEMPLATE.format(name=job_name, namespace=NAMESPACE), + encoding="utf-8", + ) + + try: + submitted = _submit("job.yaml", cwd=submit_dir) + assert submitted.returncode == 0, submitted.stdout + submitted.stderr + + # The no-wait submit returns immediately; wait for the job out of band + # like CI would. + deadline = time.time() + 300 + while time.time() < deadline: + status = kubectl( + ["get", f"job/{job_name}", "-n", NAMESPACE, "-o", "json"], + check=False, + ) + if status.returncode == 0: + conditions = json.loads(status.stdout).get("status", {}).get("conditions") or [] + if any( + c.get("status") == "True" and c.get("type") in ("Complete", "Failed") + for c in conditions + ): + break + time.sleep(3) + + attach_dir = tmp_path_factory.mktemp("k8s-attach") + _write_project(attach_dir, wheel_server["url"]) + attached = _roar( + [ + "k8s", + "attach", + f"job/{job_name}", + "-n", + NAMESPACE, + "--context", + KUBE_CONTEXT, + ], + cwd=attach_dir, + ) + logs = _pod_logs(f"job-name={job_name}") + assert attached.returncode == 0, ( + f"attach failed:\n{attached.stdout}\n{attached.stderr}\npod logs:\n{logs}" + ) + assert "lineage reconstituted" in attached.stdout, attached.stdout + + tasks = _query( + attach_dir, + "SELECT j.id FROM jobs j WHERE j.job_type = 'k8s_task'", + ) + assert len(tasks) == 1, f"{attached.stdout}\n{attached.stderr}\npod logs:\n{logs}" + + attach_jobs = _query( + attach_dir, + "SELECT job_uid, execution_role FROM jobs " + "WHERE execution_backend = 'k8s' AND execution_role = 'attach'", + ) + assert len(attach_jobs) == 1 + parent_uid = attach_jobs[0]["job_uid"] + + linked = _query( + attach_dir, + "SELECT metadata FROM jobs WHERE job_type = 'k8s_task' AND metadata LIKE ?", + (f"%{parent_uid}%",), + ) + assert linked, "k8s_task should reference the attach job's parent uid" + finally: + _cleanup(job_name) + + +def _jobset_controller_available() -> bool: + result = kubectl(["get", "crd", "jobsets.jobset.x-k8s.io"], check=False) + return result.returncode == 0 + + +def test_jobset_workload_captures_both_pods( + k8s_cluster: None, + glaas_health: str, + wheel_server: dict[str, str], # noqa: F811 + tmp_path_factory: pytest.TempPathFactory, +) -> None: + if not _jobset_controller_available(): + pytest.skip("JobSet controller not installed; run bootstrap_k8s.sh --with-jobset") + + project_dir = tmp_path_factory.mktemp("k8s-jobset") + _write_project(project_dir, wheel_server["url"]) + + name = f"roar-js-{uuid.uuid4().hex[:6]}" + (project_dir / "jobset.yaml").write_text( + JOBSET_TEMPLATE.format( + name=name, + namespace=NAMESPACE, + pinned_node=PINNED_NODE, + worker_script=_indent_script(WORKER_SCRIPT), + ), + encoding="utf-8", + ) + + try: + completed = _submit("jobset.yaml", cwd=project_dir) + logs = _pod_logs(f"jobset.sigs.k8s.io/jobset-name={name}") + assert completed.returncode == 0, ( + f"{completed.stdout}\n{completed.stderr}\npod logs:\n{logs}" + ) + + tasks = _query( + project_dir, + "SELECT metadata FROM jobs WHERE job_type = 'k8s_task'", + ) + assert len(tasks) == 2, f"{completed.stdout}\n{completed.stderr}\npod logs:\n{logs}" + all_metadata = " ".join(str(row["metadata"] or "") for row in tasks) + assert ":trainer:0:0" in all_metadata + assert ":trainer:1:0" in all_metadata + finally: + _cleanup(name, resource="jobset") diff --git a/tests/backends/k8s/scripts/bootstrap_k8s.sh b/tests/backends/k8s/scripts/bootstrap_k8s.sh index e9e4927e..68044435 100755 --- a/tests/backends/k8s/scripts/bootstrap_k8s.sh +++ b/tests/backends/k8s/scripts/bootstrap_k8s.sh @@ -16,9 +16,11 @@ set -euo pipefail # with a probe pod, so infra failures are diagnosed here rather than # inside product tests. # 6. Optionally deploys MinIO for S3 scenarios (--with-minio). +# 7. Optionally installs the JobSet controller (--with-jobset) for the +# Tier-2 distributed operator e2e tests. # # Usage: -# bash tests/backends/k8s/scripts/bootstrap_k8s.sh [--with-minio] [--skip-glaas] +# bash tests/backends/k8s/scripts/bootstrap_k8s.sh [--with-minio] [--with-jobset] [--skip-glaas] HARNESS_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" REPO_ROOT="$(cd "$HARNESS_DIR/../../.." && pwd)" @@ -34,11 +36,15 @@ WORKLOAD_IMAGE="docker.io/library/python:3.12-slim" HOST_GLAAS_URL="${HOST_GLAAS_URL:-http://localhost:3001}" CLUSTER_GLAAS_PORT=3001 +JOBSET_VERSION="v0.12.0" + WITH_MINIO=0 +WITH_JOBSET=0 SKIP_GLAAS=0 for arg in "$@"; do case "$arg" in --with-minio) WITH_MINIO=1 ;; + --with-jobset) WITH_JOBSET=1 ;; --skip-glaas) SKIP_GLAAS=1 ;; *) echo "error: unknown flag: $arg" >&2 @@ -172,6 +178,13 @@ if ((WITH_MINIO == 1)); then kubectl_ctx -n "$NAMESPACE" rollout status deployment/minio --timeout=180s fi +if ((WITH_JOBSET == 1)); then + echo "▶ Installing JobSet controller ${JOBSET_VERSION}" + kubectl_ctx apply --server-side \ + -f "https://github.com/kubernetes-sigs/jobset/releases/download/${JOBSET_VERSION}/manifests.yaml" + kubectl_ctx -n jobset-system rollout status deployment/jobset-controller-manager --timeout=300s +fi + echo echo "✓ Harness ready" echo " cluster: $CLUSTER_NAME (context $KUBE_CONTEXT)" diff --git a/tests/backends/k8s/unit/test_attach.py b/tests/backends/k8s/unit/test_attach.py new file mode 100644 index 00000000..e008683d --- /dev/null +++ b/tests/backends/k8s/unit/test_attach.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +import copy +from typing import Any + +from roar.backends.k8s.attach import ( + _instrumented_env_entries, + _session_secret_name, +) +from roar.backends.k8s.manifest import ( + rewrite_manifest_for_lineage, + workload_kind_for_document, +) +from roar.backends.k8s.workload_wait import terminal_condition + +from .conftest import SINGLE_JOB_MANIFEST +from .test_workload_adapters import JOBSET_MANIFEST, TRAINJOB_MANIFEST + + +def _rewritten_workload(manifest: dict[str, Any]) -> dict[str, Any]: + rewrite = rewrite_manifest_for_lineage( + [copy.deepcopy(manifest)], + secret_name="roar-fragment-deadbeef", + session_id="session-1", + fragment_token="ff" * 32, + requirement="roar-cli", + cluster_glaas_url="http://glaas:3001", + tracer="preload", + parent_job_uid="cafe0123", + ) + return next(doc for doc in rewrite.documents if doc.get("kind") != "Secret") + + +def test_attach_recovers_identity_from_rewritten_job() -> None: + doc = _rewritten_workload(SINGLE_JOB_MANIFEST) + workload_kind = workload_kind_for_document(doc) + assert workload_kind is not None + + env = _instrumented_env_entries(doc, workload_kind) + values = {entry["name"]: entry for entry in env} + assert values["ROAR_K8S_PARENT_JOB_UID"]["value"] == "cafe0123" + assert _session_secret_name(env) == "roar-fragment-deadbeef" + + +def test_attach_recovers_identity_from_jobset_and_trainjob() -> None: + for manifest in (JOBSET_MANIFEST, TRAINJOB_MANIFEST): + doc = _rewritten_workload(manifest) + workload_kind = workload_kind_for_document(doc) + assert workload_kind is not None + env = _instrumented_env_entries(doc, workload_kind) + assert _session_secret_name(env) == "roar-fragment-deadbeef" + + +def test_attach_reports_uninstrumented_workload() -> None: + doc = copy.deepcopy(SINGLE_JOB_MANIFEST) + workload_kind = workload_kind_for_document(doc) + assert workload_kind is not None + env = _instrumented_env_entries(doc, workload_kind) + assert env == [] + assert _session_secret_name(env) == "" + + +def test_terminal_condition_union_across_kinds() -> None: + def doc_with(condition_type: str) -> dict[str, Any]: + return {"status": {"conditions": [{"type": condition_type, "status": "True"}]}} + + for success in ("Complete", "Completed", "Succeeded", "SuccessCriteriaMet"): + succeeded, _message = terminal_condition(doc_with(success)) + assert succeeded is True, success + for failure in ("Failed", "FailureTarget"): + succeeded, _message = terminal_condition(doc_with(failure)) + assert succeeded is False, failure + + succeeded, _message = terminal_condition({"status": {"conditions": []}}) + assert succeeded is None diff --git a/tests/backends/k8s/unit/test_manifest_rewrite.py b/tests/backends/k8s/unit/test_manifest_rewrite.py index fc21d575..c70c034a 100644 --- a/tests/backends/k8s/unit/test_manifest_rewrite.py +++ b/tests/backends/k8s/unit/test_manifest_rewrite.py @@ -121,7 +121,7 @@ def test_multiple_jobs_rejected() -> None: second = copy.deepcopy(SINGLE_JOB_MANIFEST) second["metadata"]["name"] = "train-demo-2" - with pytest.raises(K8sManifestError, match="exactly one Job"): + with pytest.raises(K8sManifestError, match="exactly one workload"): _rewrite([first, second]) diff --git a/tests/backends/k8s/unit/test_workload_adapters.py b/tests/backends/k8s/unit/test_workload_adapters.py new file mode 100644 index 00000000..455860e6 --- /dev/null +++ b/tests/backends/k8s/unit/test_workload_adapters.py @@ -0,0 +1,195 @@ +from __future__ import annotations + +import copy +from typing import Any + +import pytest + +from roar.backends.k8s.manifest import ( + K8sManifestError, + find_workload_documents, + rewrite_manifest_for_lineage, +) + +JOBSET_MANIFEST = { + "apiVersion": "jobset.x-k8s.io/v1alpha2", + "kind": "JobSet", + "metadata": {"name": "dist-train", "namespace": "ml"}, + "spec": { + "replicatedJobs": [ + { + "name": "workers", + "replicas": 1, + "template": { + "spec": { + "completions": 2, + "parallelism": 2, + "completionMode": "Indexed", + "template": { + "spec": { + "restartPolicy": "Never", + "containers": [ + { + "name": "trainer", + "image": "pytorch/pytorch:latest", + "command": ["torchrun", "train.py"], + "args": ["--epochs", "3"], + } + ], + } + }, + } + }, + } + ] + }, +} + +PYTORCHJOB_MANIFEST = { + "apiVersion": "kubeflow.org/v1", + "kind": "PyTorchJob", + "metadata": {"name": "pt-train"}, + "spec": { + "pytorchReplicaSpecs": { + "Master": { + "replicas": 1, + "template": { + "spec": { + "containers": [ + { + "name": "pytorch", + "image": "pytorch/pytorch:latest", + "command": ["python", "train.py"], + } + ] + } + }, + }, + "Worker": { + "replicas": 2, + "template": { + "spec": { + "containers": [ + { + "name": "pytorch", + "image": "pytorch/pytorch:latest", + "command": ["python", "train.py"], + } + ] + } + }, + }, + } + }, +} + +TRAINJOB_MANIFEST = { + "apiVersion": "trainer.kubeflow.org/v1alpha1", + "kind": "TrainJob", + "metadata": {"name": "tj-train", "namespace": "ml"}, + "spec": { + "runtimeRef": {"name": "torch-distributed"}, + "trainer": { + "numNodes": 2, + "command": ["torchrun", "train.py"], + "args": ["--epochs", "3"], + "env": [{"name": "USER_SETTING", "value": "keep"}], + }, + }, +} + + +def _rewrite(documents: list[dict[str, Any]], **overrides: Any): + kwargs: dict[str, Any] = { + "secret_name": "roar-fragment-abc", + "session_id": "session-1", + "fragment_token": "ff" * 32, + "requirement": "roar-cli==0.3.7", + "cluster_glaas_url": "http://glaas:3001", + "tracer": "preload", + "parent_job_uid": "cafe0123", + } + kwargs.update(overrides) + return rewrite_manifest_for_lineage(documents, **kwargs) + + +def test_find_workload_documents_recognizes_all_kinds() -> None: + documents = [ + copy.deepcopy(JOBSET_MANIFEST), + copy.deepcopy(PYTORCHJOB_MANIFEST), + copy.deepcopy(TRAINJOB_MANIFEST), + ] + found = find_workload_documents(documents) + assert [workload.kind for _doc, workload in found] == ["JobSet", "PyTorchJob", "TrainJob"] + + +def test_jobset_rewrite_wraps_nested_pod_template() -> None: + rewrite = _rewrite([copy.deepcopy(JOBSET_MANIFEST)]) + + assert rewrite.workload_kind == "JobSet" + assert rewrite.kubectl_resource == "jobsets.jobset.x-k8s.io" + assert rewrite.job_name == "dist-train" + assert rewrite.wrapped_containers == ["workers/trainer"] + + jobset = next(doc for doc in rewrite.documents if doc.get("kind") == "JobSet") + container = jobset["spec"]["replicatedJobs"][0]["template"]["spec"]["template"]["spec"][ + "containers" + ][0] + assert container["command"][:2] == ["/bin/sh", "-c"] + assert container["command"][3:] == ["roar-k8s", "torchrun", "train.py", "--epochs", "3"] + assert "args" not in container + + env = {entry["name"]: entry for entry in container["env"]} + assert env["ROAR_K8S_TASK_NAME"]["value"] == "dist-train/workers/trainer" + assert env["ROAR_SESSION_ID"]["valueFrom"]["secretKeyRef"]["name"] == "roar-fragment-abc" + + secret = next(doc for doc in rewrite.documents if doc.get("kind") == "Secret") + assert secret["metadata"]["namespace"] == "ml" + + +def test_pytorchjob_rewrite_wraps_all_replica_roles() -> None: + rewrite = _rewrite([copy.deepcopy(PYTORCHJOB_MANIFEST)]) + + assert rewrite.workload_kind == "PyTorchJob" + assert rewrite.kubectl_resource == "pytorchjobs.kubeflow.org" + assert sorted(rewrite.wrapped_containers) == ["Master/pytorch", "Worker/pytorch"] + + doc = next(d for d in rewrite.documents if d.get("kind") == "PyTorchJob") + for role in ("Master", "Worker"): + container = doc["spec"]["pytorchReplicaSpecs"][role]["template"]["spec"]["containers"][0] + assert container["command"][:2] == ["/bin/sh", "-c"] + env = {entry["name"]: entry for entry in container["env"]} + assert env["ROAR_K8S_TASK_NAME"]["value"] == f"pt-train/{role}/pytorch" + assert env["ROAR_K8S_PARENT_JOB_UID"]["value"] == "cafe0123" + + +def test_trainjob_rewrite_wraps_trainer_override() -> None: + rewrite = _rewrite([copy.deepcopy(TRAINJOB_MANIFEST)]) + + assert rewrite.workload_kind == "TrainJob" + assert rewrite.kubectl_resource == "trainjobs.trainer.kubeflow.org" + assert rewrite.wrapped_containers == ["node"] + + doc = next(d for d in rewrite.documents if d.get("kind") == "TrainJob") + trainer = doc["spec"]["trainer"] + assert trainer["command"][:2] == ["/bin/sh", "-c"] + assert trainer["command"][3:] == ["roar-k8s", "torchrun", "train.py", "--epochs", "3"] + assert "args" not in trainer + + env = {entry["name"]: entry for entry in trainer["env"]} + assert env["USER_SETTING"]["value"] == "keep" + assert env["ROAR_K8S_CONTAINER"]["value"] == "node" + assert env["ROAR_K8S_TASK_NAME"]["value"] == "tj-train/node" + + +def test_trainjob_without_command_fails_actionably() -> None: + manifest = copy.deepcopy(TRAINJOB_MANIFEST) + del manifest["spec"]["trainer"]["command"] + + with pytest.raises(K8sManifestError, match=r"spec\.trainer\.command"): + _rewrite([manifest]) + + +def test_multiple_workload_kinds_rejected() -> None: + with pytest.raises(K8sManifestError, match="exactly one workload"): + _rewrite([copy.deepcopy(JOBSET_MANIFEST), copy.deepcopy(PYTORCHJOB_MANIFEST)]) From c15c67cd4c31e53434e46f1d92d921a288979f2c Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Fri, 10 Jul 2026 19:34:50 +0000 Subject: [PATCH 04/31] feat(k8s): add bundle-mode fallback and in-pod object-store I/O capture MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Completes the remaining Phase-2 transport/capture slices. Bundle fallback (GLaaS-less pods): - k8s.bundle_dir config names a mounted shared volume, injected as ROAR_K8S_BUNDLE_DIR; pod_entry probes GLaaS reachability upfront and writes roar-fragments-.json bundles there when streaming is unavailable (the streamer swallows per-batch POST failures and reports "streamed" regardless — hence the probe; surfacing streamer failure counts is a noted follow-up). Training never fails on lineage-transport outages - roar k8s ingest-bundles DIR merges a host-visible copy of the bundle volume through the shared fragment lineage engine (bundles.py) Object-store I/O hooks (direct S3 is HTTP, invisible to the tracer): - the k8s backend now ships a RuntimeImportAdapter; the existing roar_inject.pth/sitecustomize dispatch (activated by ROAR_WRAP=1, which pod_entry sets for the traced tree) patches botocore.client.BaseClient._make_api_call plus aiobotocore's async variant (covers s3fs/fsspec) to append S3 data-op events to ROAR_K8S_OBJECT_IO_FILE (object_io.py); hooks only record after the real call succeeds, never raise into user code, and no-op when the env is unset - pod_entry folds the events into the exported fragment as s3:// refs with etag hashes, deduped last-wins per (mode, path) Verified live on KIND: 14/14 e2e — new bundle-fallback test (black-hole cluster GLaaS URL, bundle pulled off the node with docker cp, ingested; note /var/tmp hostPath because KIND nodes mount /tmp as tmpfs which docker cp cannot read) and MinIO S3 test (host-staged dataset, pod get->put captured as job inputs/outputs with etag artifact hashes). 10 new unit tests; default gate 2013 passed; ruff + mypy clean. Co-Authored-By: Claude Fable 5 --- docs/developer/k8s-integration.md | 30 +- roar/backends/k8s/bundles.py | 112 ++++++ roar/backends/k8s/config.py | 10 + roar/backends/k8s/manifest.py | 5 + roar/backends/k8s/object_io.py | 185 ++++++++++ roar/backends/k8s/plugin.py | 10 + roar/backends/k8s/pod_entry.py | 71 +++- roar/backends/k8s/submit.py | 1 + roar/cli/commands/k8s.py | 25 ++ tests/backends/k8s/README.md | 7 +- .../backends/k8s/e2e/test_k8s_fallback_s3.py | 347 ++++++++++++++++++ tests/backends/k8s/unit/test_bundles.py | 84 +++++ tests/backends/k8s/unit/test_object_io.py | 127 +++++++ 13 files changed, 1006 insertions(+), 8 deletions(-) create mode 100644 roar/backends/k8s/bundles.py create mode 100644 roar/backends/k8s/object_io.py create mode 100644 tests/backends/k8s/e2e/test_k8s_fallback_s3.py create mode 100644 tests/backends/k8s/unit/test_bundles.py create mode 100644 tests/backends/k8s/unit/test_object_io.py diff --git a/docs/developer/k8s-integration.md b/docs/developer/k8s-integration.md index ac543a42..c9362ca6 100644 --- a/docs/developer/k8s-integration.md +++ b/docs/developer/k8s-integration.md @@ -16,11 +16,31 @@ env chain (`JOB_COMPLETION_INDEX` → `PET_NODE_RANK` → pod-level `RANK`). Current phase status (see `design-docs/k8s-training-lineage-integration.md` in the dev meta-repo): Phases 1–2 implemented (submit wrapping, operator -adapters, multi-pod capture, `roar k8s attach`). Remaining: bundle-mode -fallback for GLaaS-less pods, boto3/fsspec hooks via sitecustomize, -mount-map path rewriting, RayJob delegation, and the admission-webhook +adapters, multi-pod capture, `roar k8s attach`, bundle-mode fallback, +object-store I/O hooks). Remaining: mount-map path rewriting, RayJob +delegation, kill-pod retry chaos coverage, and the admission-webhook injector. +Two capture channels feed each pod's fragment: + +- **File I/O**: the preload tracer under `roar run` (process tree included). +- **Object-store I/O** (`object_io.py`): direct S3 access is HTTP, invisible + to the tracer. The k8s backend's `RuntimeImportAdapter` (dispatched by + `roar_inject.pth`/sitecustomize inside every `ROAR_WRAP` child) patches + `botocore.client.BaseClient._make_api_call` (and aiobotocore's async + variant, covering s3fs/fsspec) to append S3 data-op events to + `ROAR_K8S_OBJECT_IO_FILE`; `pod_entry` folds them into the fragment as + `s3://` refs with etag hashes. Hooks no-op outside pods (env unset) and + never raise into user code. + +Transport is streaming-first with a bundle fallback: when `k8s.bundle_dir` +names a mounted shared volume and GLaaS is unreachable from the pod (probe +or non-streamed emit), `pod_entry` writes `roar-fragments-.json` there +instead; `roar k8s ingest-bundles ` merges a host-visible copy later. +Note the fragment streamer swallows per-batch POST failures (reports +"streamed" regardless), which is why the fallback needs its own probe — +surfacing streamer failure counts is an open follow-up. + ## 2. Flow 1. **Match** (`roar/backends/k8s/submit.py`): binary `kubectl`, verb @@ -82,8 +102,8 @@ recovered parent uid, and reconstitutes the streamed fragments. ## 4. Config Section `k8s` (registered `BackendConfigAdapter`): `enabled`, `tracer`, -`runtime_install_requirement`, `cluster_glaas_url`, `wait_for_completion`, -`wait_timeout_seconds`, `poll_interval_seconds`, +`runtime_install_requirement`, `cluster_glaas_url`, `bundle_dir`, +`wait_for_completion`, `wait_timeout_seconds`, `poll_interval_seconds`, `fragment_session_ttl_seconds`. Env overrides beat config: `ROAR_CLUSTER_PIP_REQ`, `ROAR_CLUSTER_GLAAS_URL`. diff --git a/roar/backends/k8s/bundles.py b/roar/backends/k8s/bundles.py new file mode 100644 index 00000000..e6d1373e --- /dev/null +++ b/roar/backends/k8s/bundles.py @@ -0,0 +1,112 @@ +"""Bundle-mode fragment fallback for GLaaS-less pods. + +When pods cannot reach GLaaS, the pod entrypoint writes its execution +fragment as ``roar-fragments-.json`` into ``ROAR_K8S_BUNDLE_DIR`` +(a mounted shared volume declared via ``k8s.bundle_dir``). Someone with +access to that volume later runs ``roar k8s ingest-bundles `` to +merge the bundles into the local lineage DB — the OSMO bundle pattern, +k8s-shaped. +""" + +from __future__ import annotations + +import json +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from roar.backends.k8s.lineage import collect_k8s_fragments, resolve_active_session_context + +BUNDLE_FILENAME_PREFIX = "roar-fragments-" + + +class K8sBundleError(RuntimeError): + """Raised when bundle ingestion cannot proceed, with actionable detail.""" + + +@dataclass(frozen=True) +class K8sBundleIngestResult: + bundles_ingested: int + fragments_merged: int + bundle_paths: list[str] + + +def bundle_filename_for_pod(pod_name: str) -> str: + safe = re.sub(r"[^A-Za-z0-9_.-]", "-", pod_name.strip() or "pod") + return f"{BUNDLE_FILENAME_PREFIX}{safe}.json" + + +def write_fragment_bundle(bundle_dir: Path, pod_name: str, fragments: list[dict[str, Any]]) -> Path: + bundle_dir.mkdir(parents=True, exist_ok=True) + target = bundle_dir / bundle_filename_for_pod(pod_name) + target.write_text( + json.dumps({"fragments": fragments}, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + return target + + +def discover_fragment_bundles(directory: Path) -> list[Path]: + if not directory.is_dir(): + return [] + return sorted(directory.glob(f"{BUNDLE_FILENAME_PREFIX}*.json")) + + +def ingest_fragment_bundles( + *, + roar_dir: Path, + directory: Path, +) -> K8sBundleIngestResult: + bundles = discover_fragment_bundles(directory) + if not bundles: + raise K8sBundleError( + f"no {BUNDLE_FILENAME_PREFIX}*.json bundles found in {directory}; " + "point at the directory pods wrote via k8s.bundle_dir" + ) + + fragments: list[dict[str, Any]] = [] + for bundle_path in bundles: + try: + payload = json.loads(bundle_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise K8sBundleError(f"cannot read bundle {bundle_path}: {exc}") from exc + items = payload.get("fragments") if isinstance(payload, dict) else None + fragments.extend(item for item in items or [] if isinstance(item, dict)) + + if not fragments: + raise K8sBundleError(f"bundles in {directory} contain no fragments") + + db_path = roar_dir / "roar.db" + session_id, step_number = resolve_active_session_context(str(db_path)) + driver_job_uid = next( + ( + str(fragment.get("parent_job_uid") or "").strip() + for fragment in fragments + if str(fragment.get("parent_job_uid") or "").strip() + ), + None, + ) + merged = collect_k8s_fragments( + fragments, + project_dir=str(roar_dir.parent), + driver_job_uid=driver_job_uid, + session_id=session_id, + step_number=step_number, + ) + return K8sBundleIngestResult( + bundles_ingested=len(bundles), + fragments_merged=merged, + bundle_paths=[str(path) for path in bundles], + ) + + +__all__ = [ + "BUNDLE_FILENAME_PREFIX", + "K8sBundleError", + "K8sBundleIngestResult", + "bundle_filename_for_pod", + "discover_fragment_bundles", + "ingest_fragment_bundles", + "write_fragment_bundle", +] diff --git a/roar/backends/k8s/config.py b/roar/backends/k8s/config.py index fbad0c47..295c709f 100644 --- a/roar/backends/k8s/config.py +++ b/roar/backends/k8s/config.py @@ -22,6 +22,7 @@ class K8sBackendConfig(BaseModel): tracer: str = "preload" runtime_install_requirement: str = "" cluster_glaas_url: str = "" + bundle_dir: str = "" wait_for_completion: bool = True wait_timeout_seconds: int = Field(default=30 * 60, ge=1) poll_interval_seconds: float = Field(default=5.0, gt=0.0) @@ -56,6 +57,15 @@ class K8sBackendConfig(BaseModel): "host-visible glaas.url (ROAR_CLUSTER_GLAAS_URL env always wins)" ), ), + "k8s.bundle_dir": ConfigurableKeySpec( + value_type=str, + default="", + description=( + "In-pod directory (a mounted shared volume) where pods write " + "roar-fragments-.json bundles when GLaaS streaming is unavailable; " + "ingest later with `roar k8s ingest-bundles`" + ), + ), "k8s.wait_for_completion": ConfigurableKeySpec( value_type=bool, default=True, diff --git a/roar/backends/k8s/manifest.py b/roar/backends/k8s/manifest.py index 14b7ed3d..913e74c8 100644 --- a/roar/backends/k8s/manifest.py +++ b/roar/backends/k8s/manifest.py @@ -185,6 +185,7 @@ def rewrite_manifest_for_lineage( cluster_glaas_url: str, tracer: str, parent_job_uid: str, + bundle_dir: str = "", namespace_override: str | None = None, ) -> K8sManifestRewrite: """Return a rewritten copy of ``documents`` with lineage instrumentation. @@ -231,6 +232,7 @@ def rewrite_manifest_for_lineage( tracer=tracer, parent_job_uid=parent_job_uid, workload_name=workload_name, + bundle_dir=bundle_dir, ) if workload.locate_pod_specs is None: @@ -292,6 +294,7 @@ class _EnvContract: tracer: str parent_job_uid: str workload_name: str + bundle_dir: str = "" def _rewrite_pod_specs( @@ -421,6 +424,8 @@ def add_secret_ref(name: str, key: str) -> None: add_value("ROAR_EXECUTION_BACKEND", "k8s") add_value("ROAR_NO_TELEMETRY", "1") add_value("ROAR_K8S_TRACER", contract.tracer) + if contract.bundle_dir: + add_value("ROAR_K8S_BUNDLE_DIR", contract.bundle_dir) add_value("GLAAS_URL", contract.cluster_glaas_url) add_value("ROAR_K8S_PARENT_JOB_UID", contract.parent_job_uid) add_value("ROAR_K8S_JOB_NAME", contract.workload_name) diff --git a/roar/backends/k8s/object_io.py b/roar/backends/k8s/object_io.py new file mode 100644 index 00000000..bd5383fe --- /dev/null +++ b/roar/backends/k8s/object_io.py @@ -0,0 +1,185 @@ +"""In-process object-store I/O capture for k8s pods. + +Direct S3 access (boto3/botocore, and s3fs/fsspec via aiobotocore) is +HTTP, not file I/O — the syscall tracer never sees it. These hooks are +installed by the sitecustomize import dispatch (the k8s backend's +``RuntimeImportAdapter``) inside every ``ROAR_WRAP``-instrumented Python +process and append one JSON line per successful S3 data operation to +``ROAR_K8S_OBJECT_IO_FILE``. The pod entrypoint folds the events into +the exported execution fragment after the traced command exits. + +Best-effort by construction: hooks only record after the real call +succeeds, never raise into user code, and no-op entirely when the +events-file env is unset. +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any + +OBJECT_IO_FILE_ENV = "ROAR_K8S_OBJECT_IO_FILE" + +_S3_READ_OPS = frozenset({"GetObject"}) +_S3_WRITE_OPS = frozenset({"PutObject", "CompleteMultipartUpload", "CopyObject"}) +_PATCH_MARKER = "_roar_k8s_object_io_patched" + + +def patch_imported_module(module_name: str, module: Any) -> None: + """Runtime-import dispatch target for the k8s backend.""" + if module_name == "botocore.client": + _patch_sync_client(module) + elif module_name == "aiobotocore.client": + _patch_async_client(module) + + +def _patch_sync_client(module: Any) -> None: + base_client = getattr(module, "BaseClient", None) + if base_client is None or getattr(base_client, _PATCH_MARKER, False): + return + original = base_client._make_api_call + + def _make_api_call(self: Any, operation_name: str, api_params: dict[str, Any]) -> Any: + response = original(self, operation_name, api_params) + _record_s3_event(self, operation_name, api_params, response) + return response + + base_client._make_api_call = _make_api_call + setattr(base_client, _PATCH_MARKER, True) + + +def _patch_async_client(module: Any) -> None: + base_client = getattr(module, "AioBaseClient", None) + if base_client is None or getattr(base_client, _PATCH_MARKER, False): + return + original = base_client._make_api_call + + async def _make_api_call(self: Any, operation_name: str, api_params: dict[str, Any]) -> Any: + response = await original(self, operation_name, api_params) + _record_s3_event(self, operation_name, api_params, response) + return response + + base_client._make_api_call = _make_api_call + setattr(base_client, _PATCH_MARKER, True) + + +def _record_s3_event( + client: Any, + operation_name: str, + api_params: Any, + response: Any, +) -> None: + try: + events_file = os.environ.get(OBJECT_IO_FILE_ENV, "").strip() + if not events_file: + return + if operation_name not in _S3_READ_OPS and operation_name not in _S3_WRITE_OPS: + return + service = getattr( + getattr(getattr(client, "meta", None), "service_model", None), "service_name", "" + ) + if service != "s3": + return + if not isinstance(api_params, dict): + return + bucket = str(api_params.get("Bucket") or "").strip() + key = str(api_params.get("Key") or "").strip() + if not bucket or not key: + return + + mode = "read" if operation_name in _S3_READ_OPS else "write" + event = { + "mode": mode, + "path": f"s3://{bucket}/{key}", + "operation": operation_name, + "etag": _normalize_etag(response), + "size": _event_size(mode, api_params, response), + } + with open(events_file, "a", encoding="utf-8") as handle: + handle.write(json.dumps(event, separators=(",", ":")) + "\n") + except Exception: + return + + +def _normalize_etag(response: Any) -> str | None: + if not isinstance(response, dict): + return None + etag = response.get("ETag") + if not isinstance(etag, str): + return None + normalized = etag.strip().strip('"').strip() + return normalized or None + + +def _event_size(mode: str, api_params: dict[str, Any], response: Any) -> int: + if mode == "read" and isinstance(response, dict): + content_length = response.get("ContentLength") + if isinstance(content_length, int) and content_length >= 0: + return content_length + return 0 + + body = api_params.get("Body") + if isinstance(body, (bytes, bytearray)): + return len(body) + if isinstance(body, str): + return len(body.encode("utf-8")) + try: + return max(0, int(getattr(body, "seekable", lambda: False)() and _stream_size(body))) + except Exception: + return 0 + + +def _stream_size(body: Any) -> int: + position = body.tell() + body.seek(0, os.SEEK_END) + size = body.tell() + body.seek(position) + return size + + +def load_object_io_refs(events_path: Path) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + """Read and deduplicate recorded events into fragment artifact refs. + + Last event per (mode, path) wins so re-reads/re-writes carry the most + recent etag/size. + """ + if not events_path.is_file(): + return [], [] + + winners: dict[tuple[str, str], dict[str, Any]] = {} + for line in events_path.read_text(encoding="utf-8").splitlines(): + line = line.strip() + if not line: + continue + try: + event = json.loads(line) + except json.JSONDecodeError: + continue + mode = str(event.get("mode") or "") + path = str(event.get("path") or "") + if mode not in ("read", "write") or not path.startswith("s3://"): + continue + winners[(mode, path)] = event + + reads: list[dict[str, Any]] = [] + writes: list[dict[str, Any]] = [] + for (mode, path), event in sorted(winners.items()): + etag = event.get("etag") + ref = { + "path": path, + "hash": etag if isinstance(etag, str) and etag else None, + "hash_algorithm": "etag" if etag else "", + "size": int(event.get("size") or 0), + "capture_method": "python", + } + (reads if mode == "read" else writes).append(ref) + return reads, writes + + +__all__ = [ + "OBJECT_IO_FILE_ENV", + "load_object_io_refs", + "patch_imported_module", +] diff --git a/roar/backends/k8s/plugin.py b/roar/backends/k8s/plugin.py index 56eab05d..8eeb1522 100644 --- a/roar/backends/k8s/plugin.py +++ b/roar/backends/k8s/plugin.py @@ -6,6 +6,7 @@ from roar.backends.k8s.config import K8S_BACKEND_CONFIG from roar.backends.k8s.fragment_reconstituter import create_k8s_fragment_reconstituter from roar.backends.k8s.host_execution import execute_k8s_job_submit +from roar.backends.k8s.object_io import patch_imported_module from roar.backends.k8s.submit import ( matches_kubectl_job_submit_command, plan_kubectl_job_submit_command, @@ -17,6 +18,7 @@ ExecutionPolicyAdapter, FragmentReconstitutionAdapter, HostExecutionAdapter, + RuntimeImportAdapter, WorkerBootstrapAdapter, ) from roar.execution.framework.registry import register_execution_backend @@ -83,6 +85,14 @@ def _worker_entrypoint(argv: list[str]) -> None: fragment_reconstitution=FragmentReconstitutionAdapter( create_reconstituter=create_k8s_fragment_reconstituter, ), + # Sitecustomize import dispatch: installs the object-store I/O + # hooks inside ROAR_WRAP-instrumented processes when botocore or + # aiobotocore is imported. The hooks no-op unless + # ROAR_K8S_OBJECT_IO_FILE is set (i.e. outside k8s pods). + runtime_import=RuntimeImportAdapter( + module_prefixes=("botocore", "aiobotocore"), + patch_module=patch_imported_module, + ), ), policy=ExecutionPolicyAdapter( submit_roles=("submit",), diff --git a/roar/backends/k8s/pod_entry.py b/roar/backends/k8s/pod_entry.py index c5489f11..ed8e908f 100644 --- a/roar/backends/k8s/pod_entry.py +++ b/roar/backends/k8s/pod_entry.py @@ -58,13 +58,24 @@ def _run_traced(command: list[str]) -> int: return _run_uninstrumented(command) tracer = str(os.environ.get("ROAR_K8S_TRACER") or "preload").strip() or "preload" + child_env = dict(os.environ) + # Activate the roar_inject.pth sitecustomize in every Python child so + # the k8s backend's runtime-import hooks (botocore/aiobotocore object + # I/O capture) install themselves; events land next to the local db. + child_env.setdefault("ROAR_WRAP", "1") + child_env.setdefault("ROAR_K8S_OBJECT_IO_FILE", str(_object_io_events_path())) run = subprocess.run( [sys.executable, "-m", "roar", "run", "--tracer", tracer, *command], + env=child_env, check=False, ) return run.returncode +def _object_io_events_path() -> Path: + return Path.cwd() / ".roar" / "k8s-object-io.jsonl" + + def _run_uninstrumented(command: list[str]) -> int: completed = subprocess.run(command, check=False) return completed.returncode @@ -107,7 +118,6 @@ def task_identity_from_environment(environ: dict[str, str] | None = None) -> tup def _emit_lineage_best_effort() -> None: try: from roar.execution.fragments.export import export_local_job_fragment_bundle - from roar.execution.fragments.transport import emit_fragment_dicts task_id, task_name = task_identity_from_environment() parent_job_uid = str(os.environ.get("ROAR_K8S_PARENT_JOB_UID") or "").strip() @@ -126,6 +136,7 @@ def _emit_lineage_best_effort() -> None: payload = json.loads(bundle_path.read_text(encoding="utf-8")) fragments = [item for item in payload.get("fragments", []) if isinstance(item, dict)] + _augment_with_object_io(fragments) completion_index = task_id.split(":")[2] if task_id.count(":") >= 3 else "0" restart_attempt = task_id.split(":")[3] if task_id.count(":") >= 3 else "0" for fragment in fragments: @@ -143,11 +154,67 @@ def _emit_lineage_best_effort() -> None: } ) - result = emit_fragment_dicts(fragments) + result = _emit_or_bundle(fragments) print(f"[roar-k8s] lineage emit: {result} ({len(fragments)} fragment(s))") except Exception as exc: print(f"[roar-k8s] warning: lineage emit failed: {exc}", file=sys.stderr) +def _augment_with_object_io(fragments: list[dict]) -> None: + """Fold captured S3 events into the fragment's read/write refs.""" + from roar.backends.k8s.object_io import load_object_io_refs + + events_path = Path(os.environ.get("ROAR_K8S_OBJECT_IO_FILE") or _object_io_events_path()) + reads, writes = load_object_io_refs(events_path) + if not reads and not writes: + return + for fragment in fragments: + fragment.setdefault("reads", []).extend(reads) + fragment.setdefault("writes", []).extend(writes) + + +def _emit_or_bundle(fragments: list[dict]) -> str: + """Stream fragments; fall back to a bundle file when GLaaS is unreachable. + + The streamer itself swallows per-batch POST failures, so a quick + reachability probe decides upfront; a non-"streamed" emit result also + falls back when a bundle directory is declared. + """ + from roar.execution.fragments.transport import emit_fragment_dicts + + bundle_dir = str(os.environ.get("ROAR_K8S_BUNDLE_DIR") or "").strip() + + if bundle_dir and not _glaas_reachable(): + return _write_bundle(bundle_dir, fragments) + + result = emit_fragment_dicts(fragments) + if result != "streamed" and bundle_dir: + return _write_bundle(bundle_dir, fragments) + return result + + +def _glaas_reachable() -> bool: + import urllib.request + + glaas_url = str(os.environ.get("GLAAS_URL") or "").strip() + if not glaas_url: + return False + try: + request = urllib.request.Request(f"{glaas_url.rstrip('/')}/api/v1/health") + with urllib.request.urlopen(request, timeout=3) as response: + return response.status == 200 + except Exception: + return False + + +def _write_bundle(bundle_dir: str, fragments: list[dict]) -> str: + from roar.backends.k8s.bundles import write_fragment_bundle + + pod_name = str(os.environ.get("ROAR_K8S_POD_NAME") or "pod").strip() or "pod" + target = write_fragment_bundle(Path(bundle_dir), pod_name, fragments) + print(f"[roar-k8s] GLaaS unavailable; wrote fragment bundle to {target}") + return "bundled" + + if __name__ == "__main__": sys.exit(main()) diff --git a/roar/backends/k8s/submit.py b/roar/backends/k8s/submit.py index 416466c7..aa990194 100644 --- a/roar/backends/k8s/submit.py +++ b/roar/backends/k8s/submit.py @@ -119,6 +119,7 @@ def plan_kubectl_job_submit_command(command: list[str]) -> ExecutionCommandPlan: cluster_glaas_url=_resolve_cluster_glaas_url(config, glaas_url), tracer=str(config.get("tracer") or "preload"), parent_job_uid=parent_job_uid, + bundle_dir=str(config.get("bundle_dir") or ""), namespace_override=_find_namespace_argument(command), ) if rewrite.skipped_containers: diff --git a/roar/cli/commands/k8s.py b/roar/cli/commands/k8s.py index a37b1af7..bde8d564 100644 --- a/roar/cli/commands/k8s.py +++ b/roar/cli/commands/k8s.py @@ -6,6 +6,7 @@ import click from ...backends.k8s.attach import K8sAttachError, attach_k8s_workload +from ...backends.k8s.bundles import K8sBundleError, ingest_fragment_bundles from ...backends.k8s.config import load_k8s_backend_config from ...backends.k8s.manifest import ( K8sManifestError, @@ -91,6 +92,30 @@ def k8s_attach( raise SystemExit(result.exit_code) +@k8s.command("ingest-bundles") +@click.argument( + "directory", + type=click.Path(path_type=Path, file_okay=False, exists=True), +) +@click.pass_obj +@require_init +def k8s_ingest_bundles(ctx: RoarContext, directory: Path) -> None: + """Merge roar-fragments-*.json bundles written by GLaaS-less pods. + + DIRECTORY is a host-visible copy of the shared volume pods wrote to + (declared via k8s.bundle_dir at submit time). + """ + try: + result = ingest_fragment_bundles(roar_dir=ctx.roar_dir, directory=directory) + except K8sBundleError as exc: + raise click.ClickException(str(exc)) from exc + + click.echo( + f"[roar] ingested {result.bundles_ingested} bundle(s): " + f"{result.fragments_merged} fragment(s) merged" + ) + + @k8s.command("prepare") @click.option( "-f", diff --git a/tests/backends/k8s/README.md b/tests/backends/k8s/README.md index 527b3e38..b823786d 100644 --- a/tests/backends/k8s/README.md +++ b/tests/backends/k8s/README.md @@ -15,6 +15,11 @@ Three test layers share this harness: edge over a shared volume; `roar k8s attach` from a fresh project via the cluster Secret; and a JobSet run through the real controller (skipped unless bootstrapped with `--with-jobset`). +- `e2e/test_k8s_fallback_s3.py` — bundle-mode fallback (black-hole cluster + GLaaS URL, bundle written to a shared volume, pulled off the node and + merged with `roar k8s ingest-bundles`) and in-pod S3 capture (MinIO via + `--with-minio`: boto3 get/put recorded as `s3://` lineage refs with etag + hashes). - `e2e/test_k8s_smoke.py` — the Phase-0 runtime diagnostic: fixtures hand-wrap the manifest (no backend involved) to isolate the runtime pieces (in-pod tracing, fragment streaming, identity contract) when the product @@ -38,7 +43,7 @@ bash scripts/build_wheel_with_bins.sh # create cluster + wire glaas + preflight # (--with-minio for S3 scenarios, --with-jobset for the JobSet operator e2e) -bash tests/backends/k8s/scripts/bootstrap_k8s.sh --with-jobset +bash tests/backends/k8s/scripts/bootstrap_k8s.sh --with-jobset --with-minio # run the smoke tests (addopts override needed: e2e dirs are ignored by default) pytest tests/backends/k8s/e2e -o addopts='' -m k8s_e2e -v diff --git a/tests/backends/k8s/e2e/test_k8s_fallback_s3.py b/tests/backends/k8s/e2e/test_k8s_fallback_s3.py new file mode 100644 index 00000000..c88ec036 --- /dev/null +++ b/tests/backends/k8s/e2e/test_k8s_fallback_s3.py @@ -0,0 +1,347 @@ +"""Phase-2 e2e: bundle-mode fallback and in-pod S3 capture. + +- Bundle fallback: a pod whose cluster-visible GLaaS URL is a black hole + writes its fragment bundle to a shared volume (``k8s.bundle_dir``); the + bundle is pulled off the node and merged with ``roar k8s ingest-bundles``. +- S3 hooks: a pod reads a dataset from MinIO and writes a model back via + boto3; the object I/O appears in lineage as ``s3://`` refs with etag + hashes alongside the tracer-captured local files. Requires + ``bootstrap_k8s.sh --with-minio``. +""" + +from __future__ import annotations + +import sqlite3 +import subprocess +import uuid +from pathlib import Path + +import pytest + +from .conftest import NAMESPACE, kubectl +from .test_k8s_distributed import ( + PINNED_NODE, + _cleanup, + _describe, + _pod_logs, + _roar, + _submit, + _write_project, +) +from .test_k8s_product_path import wheel_server # noqa: F401 (module fixture reuse) + +pytestmark = [ + pytest.mark.e2e, + pytest.mark.k8s_e2e, + pytest.mark.timeout(900), +] + +MINIO_HOST_URL = "http://localhost:39000" +MINIO_CLUSTER_URL = "http://minio:9000" + +BUNDLE_JOB_TEMPLATE = """\ +apiVersion: batch/v1 +kind: Job +metadata: + name: {name} + namespace: {namespace} + labels: + app.kubernetes.io/part-of: roar-k8s-e2e +spec: + backoffLimit: 0 + ttlSecondsAfterFinished: 1800 + template: + spec: + restartPolicy: Never + nodeSelector: + kubernetes.io/hostname: {pinned_node} + volumes: + - name: work + emptyDir: {{}} + - name: bundles + hostPath: + path: /var/tmp/roar-e2e-bundles-{name} + type: DirectoryOrCreate + containers: + - name: trainer + image: python:3.12-slim + workingDir: /work + volumeMounts: + - name: work + mountPath: /work + - name: bundles + mountPath: /bundles + command: + - python + - -c + - >- + open('data.csv', 'w').write('x\\n1\\n2\\n'); + rows = open('data.csv').read().strip().splitlines(); + open('model.bin', 'wb').write(str(len(rows)).encode() * 8) +""" + +S3_WORKER_SCRIPT = """\ +import os + +import boto3 +from botocore.client import Config + +endpoint = os.environ.get("AWS_ENDPOINT_URL", "http://minio:9000") +bucket = os.environ["ROAR_E2E_BUCKET"] +s3 = boto3.client( + "s3", + endpoint_url=endpoint, + config=Config(s3={"addressing_style": "path"}), +) + +obj = s3.get_object(Bucket=bucket, Key="datasets/train.csv") +data = obj["Body"].read() +model = data * 3 +s3.put_object(Bucket=bucket, Key="models/model.bin", Body=model) +with open("metrics.json", "w") as handle: + handle.write('{"rows": %d}' % len(data.splitlines())) +print("s3 train done") +""" + +S3_JOB_TEMPLATE = """\ +apiVersion: v1 +kind: ConfigMap +metadata: + name: {name}-worker + namespace: {namespace} +data: + s3_train.py: | +{worker_script} +--- +apiVersion: batch/v1 +kind: Job +metadata: + name: {name} + namespace: {namespace} + labels: + app.kubernetes.io/part-of: roar-k8s-e2e +spec: + backoffLimit: 0 + ttlSecondsAfterFinished: 1800 + template: + spec: + restartPolicy: Never + volumes: + - name: work + emptyDir: {{}} + - name: workload + configMap: + name: {name}-worker + containers: + - name: trainer + image: python:3.12-slim + workingDir: /work + volumeMounts: + - name: work + mountPath: /work + - name: workload + mountPath: /workload + readOnly: true + env: + - name: AWS_ACCESS_KEY_ID + value: minioadmin + - name: AWS_SECRET_ACCESS_KEY + value: minioadmin + - name: AWS_DEFAULT_REGION + value: us-east-1 + - name: AWS_ENDPOINT_URL + value: {minio_cluster_url} + - name: ROAR_E2E_BUCKET + value: {bucket} + command: + - bash + - -c + - pip install --quiet boto3 && python /workload/s3_train.py +""" + + +def _indent_script(script: str, spaces: int = 4) -> str: + pad = " " * spaces + return "\n".join(f"{pad}{line}" if line else "" for line in script.splitlines()) + + +def _query(project_dir: Path, sql: str, params: tuple = ()) -> list[sqlite3.Row]: + conn = sqlite3.connect(project_dir / ".roar" / "roar.db") + conn.row_factory = sqlite3.Row + try: + return conn.execute(sql, params).fetchall() + finally: + conn.close() + + +def test_bundle_fallback_when_glaas_unreachable( + k8s_cluster: None, + glaas_health: str, + wheel_server: dict[str, str], # noqa: F811 + tmp_path_factory: pytest.TempPathFactory, +) -> None: + project_dir = tmp_path_factory.mktemp("k8s-bundle") + _write_project(project_dir, wheel_server["url"]) + config_path = project_dir / ".roar" / "config.toml" + config = config_path.read_text(encoding="utf-8").replace( + 'cluster_glaas_url = "http://glaas:3001"', + 'cluster_glaas_url = "http://roar-blackhole.invalid:9"\nbundle_dir = "/bundles"', + ) + config_path.write_text(config, encoding="utf-8") + + job_name = f"roar-bundle-{uuid.uuid4().hex[:6]}" + (project_dir / "job.yaml").write_text( + BUNDLE_JOB_TEMPLATE.format( + name=job_name, + namespace=NAMESPACE, + pinned_node=PINNED_NODE, + ), + encoding="utf-8", + ) + + # /var/tmp, not /tmp: KIND nodes mount /tmp as tmpfs, which docker cp + # cannot read from. + node_bundle_dir = f"/var/tmp/roar-e2e-bundles-{job_name}" + try: + completed = _submit("job.yaml", cwd=project_dir) + logs = _pod_logs(f"job-name={job_name}") + run = { + "exit_code": completed.returncode, + "stdout": completed.stdout, + "stderr": completed.stderr, + "pod_logs": logs, + } + # Lineage is best-effort: the unreachable GLaaS must not fail the job. + assert completed.returncode == 0, _describe(run) + assert "wrote fragment bundle" in logs, _describe(run) + + # No fragments could stream, so nothing reconstituted yet. + assert not _query(project_dir, "SELECT id FROM jobs WHERE job_type = 'k8s_task'") + + # Pull the bundle off the node the way an operator would pull it + # off a PVC, then ingest it. + copied = tmp_path_factory.mktemp("bundles-copy") + subprocess.run( + ["docker", "cp", f"{PINNED_NODE}:{node_bundle_dir}/.", str(copied)], + check=True, + capture_output=True, + ) + ingested = _roar(["k8s", "ingest-bundles", str(copied)], cwd=project_dir) + assert ingested.returncode == 0, ingested.stdout + ingested.stderr + assert "1 fragment(s) merged" in ingested.stdout, ingested.stdout + + tasks = _query(project_dir, "SELECT id FROM jobs WHERE job_type = 'k8s_task'") + assert len(tasks) == 1 + outputs = { + str(row["path"]) + for row in _query( + project_dir, + "SELECT path FROM job_outputs WHERE job_id = ?", + (tasks[0]["id"],), + ) + } + assert any(path.endswith("model.bin") for path in outputs), outputs + finally: + _cleanup(job_name) + subprocess.run( + ["docker", "exec", PINNED_NODE, "rm", "-rf", node_bundle_dir], + check=False, + capture_output=True, + ) + + +def _minio_available() -> bool: + result = kubectl(["get", "deployment/minio", "-n", NAMESPACE], check=False) + return result.returncode == 0 + + +def test_s3_object_io_captured_in_lineage( + k8s_cluster: None, + glaas_health: str, + wheel_server: dict[str, str], # noqa: F811 + tmp_path_factory: pytest.TempPathFactory, +) -> None: + if not _minio_available(): + pytest.skip("MinIO not deployed; run bootstrap_k8s.sh --with-minio") + + import boto3 + from botocore.client import Config + + suffix = uuid.uuid4().hex[:6] + bucket = f"roar-e2e-{suffix}" + dataset_body = b"x,y\n1.0,2.0\n2.0,3.9\n3.0,6.1\n" + + # Host-visible endpoint: stage the dataset the way a data pipeline + # already would have (also exercises the host/cluster endpoint split). + host_s3 = boto3.client( + "s3", + endpoint_url=MINIO_HOST_URL, + aws_access_key_id="minioadmin", + aws_secret_access_key="minioadmin", + region_name="us-east-1", + config=Config(s3={"addressing_style": "path"}), + ) + host_s3.create_bucket(Bucket=bucket) + host_s3.put_object(Bucket=bucket, Key="datasets/train.csv", Body=dataset_body) + + project_dir = tmp_path_factory.mktemp("k8s-s3") + _write_project(project_dir, wheel_server["url"]) + + job_name = f"roar-s3-{suffix}" + (project_dir / "job.yaml").write_text( + S3_JOB_TEMPLATE.format( + name=job_name, + namespace=NAMESPACE, + worker_script=_indent_script(S3_WORKER_SCRIPT), + minio_cluster_url=MINIO_CLUSTER_URL, + bucket=bucket, + ), + encoding="utf-8", + ) + + try: + completed = _submit("job.yaml", cwd=project_dir) + logs = _pod_logs(f"job-name={job_name}") + run = { + "exit_code": completed.returncode, + "stdout": completed.stdout, + "stderr": completed.stderr, + "pod_logs": logs, + } + assert completed.returncode == 0, _describe(run) + + tasks = _query(project_dir, "SELECT id FROM jobs WHERE job_type = 'k8s_task'") + assert len(tasks) == 1, _describe(run) + task_id = tasks[0]["id"] + + input_paths = { + str(row["path"]) + for row in _query( + project_dir, + "SELECT path FROM job_inputs WHERE job_id = ?", + (task_id,), + ) + } + output_paths = { + str(row["path"]) + for row in _query( + project_dir, + "SELECT path FROM job_outputs WHERE job_id = ?", + (task_id,), + ) + } + assert f"s3://{bucket}/datasets/train.csv" in input_paths, ( + f"S3 read missing: {input_paths}\n{_describe(run)}" + ) + assert f"s3://{bucket}/models/model.bin" in output_paths, ( + f"S3 write missing: {output_paths}\n{_describe(run)}" + ) + assert any(path.endswith("metrics.json") for path in output_paths), output_paths + + etag_rows = _query( + project_dir, + "SELECT COUNT(*) AS count FROM artifact_hashes WHERE algorithm = 'etag'", + ) + assert int(etag_rows[0]["count"]) >= 2, "expected etag hashes for the S3 artifacts" + finally: + _cleanup(job_name) diff --git a/tests/backends/k8s/unit/test_bundles.py b/tests/backends/k8s/unit/test_bundles.py new file mode 100644 index 00000000..0f4c3708 --- /dev/null +++ b/tests/backends/k8s/unit/test_bundles.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +import sqlite3 +from pathlib import Path + +import pytest + +from roar.backends.k8s.bundles import ( + K8sBundleError, + bundle_filename_for_pod, + discover_fragment_bundles, + ingest_fragment_bundles, + write_fragment_bundle, +) + + +def _fragment(pod: str, index: str) -> dict: + return { + "job_uid": f"{pod}-job", + "parent_job_uid": "cafe0123", + "task_id": f"{pod}-uid:trainer:{index}:0", + "worker_id": "", + "node_id": "", + "actor_id": None, + "task_name": "bundle-train/trainer", + "started_at": 1000.0, + "ended_at": 1010.0, + "exit_code": 0, + "backend": "k8s", + "reads": [], + "writes": [ + { + "path": f"/work/out-{index}.bin", + "hash": f"digest-{index}", + "hash_algorithm": "blake3", + "size": 10, + "capture_method": "python", + } + ], + "backend_metadata": {"execution_role": "task"}, + } + + +def test_bundle_filename_sanitizes_pod_names() -> None: + assert bundle_filename_for_pod("train-0/pod x") == "roar-fragments-train-0-pod-x.json" + assert bundle_filename_for_pod("") == "roar-fragments-pod.json" + + +def test_write_and_discover_bundles(tmp_path: Path) -> None: + first = write_fragment_bundle(tmp_path / "bundles", "pod-0", [_fragment("pod-0", "0")]) + second = write_fragment_bundle(tmp_path / "bundles", "pod-1", [_fragment("pod-1", "1")]) + assert discover_fragment_bundles(tmp_path / "bundles") == [first, second] + + +def test_ingest_merges_bundles_into_db(tmp_path: Path) -> None: + from roar.db.context import create_database_context + + project_dir = tmp_path / "project" + roar_dir = project_dir / ".roar" + roar_dir.mkdir(parents=True) + with create_database_context(roar_dir): + pass # initialize schema + + bundle_dir = tmp_path / "bundles" + write_fragment_bundle(bundle_dir, "pod-0", [_fragment("pod-0", "0")]) + write_fragment_bundle(bundle_dir, "pod-1", [_fragment("pod-1", "1")]) + + result = ingest_fragment_bundles(roar_dir=roar_dir, directory=bundle_dir) + assert result.bundles_ingested == 2 + assert result.fragments_merged == 2 + + conn = sqlite3.connect(roar_dir / "roar.db") + try: + count = conn.execute("SELECT COUNT(*) FROM jobs WHERE job_type = 'k8s_task'").fetchone() + assert count[0] == 2 + finally: + conn.close() + + +def test_ingest_empty_directory_fails_actionably(tmp_path: Path) -> None: + empty = tmp_path / "empty" + empty.mkdir() + with pytest.raises(K8sBundleError, match="no roar-fragments-"): + ingest_fragment_bundles(roar_dir=tmp_path / ".roar", directory=empty) diff --git a/tests/backends/k8s/unit/test_object_io.py b/tests/backends/k8s/unit/test_object_io.py new file mode 100644 index 00000000..e62b0ac0 --- /dev/null +++ b/tests/backends/k8s/unit/test_object_io.py @@ -0,0 +1,127 @@ +from __future__ import annotations + +import json +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest + +from roar.backends.k8s.object_io import ( + OBJECT_IO_FILE_ENV, + load_object_io_refs, + patch_imported_module, +) + + +class _FakeBaseClient: + def __init__(self) -> None: + self.meta = SimpleNamespace(service_model=SimpleNamespace(service_name="s3")) + + def _make_api_call(self, operation_name: str, api_params: dict[str, Any]) -> dict[str, Any]: + if operation_name == "GetObject": + return {"ETag": '"abc123"', "ContentLength": 42} + return {"ETag": '"def456"'} + + +def _patched_client_module() -> Any: + module = SimpleNamespace(BaseClient=type("BaseClient", (_FakeBaseClient,), {})) + patch_imported_module("botocore.client", module) + return module + + +def test_hooks_record_s3_reads_and_writes(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + events_file = tmp_path / "events.jsonl" + monkeypatch.setenv(OBJECT_IO_FILE_ENV, str(events_file)) + + module = _patched_client_module() + client = module.BaseClient() + client._make_api_call("GetObject", {"Bucket": "data", "Key": "train/input.csv"}) + client._make_api_call("PutObject", {"Bucket": "models", "Key": "run/model.pt", "Body": b"xy"}) + client._make_api_call("ListObjectsV2", {"Bucket": "data"}) # control op, ignored + + events = [json.loads(line) for line in events_file.read_text().splitlines()] + assert len(events) == 2 + read, write = events + assert read == { + "mode": "read", + "path": "s3://data/train/input.csv", + "operation": "GetObject", + "etag": "abc123", + "size": 42, + } + assert write["mode"] == "write" + assert write["path"] == "s3://models/run/model.pt" + assert write["etag"] == "def456" + assert write["size"] == 2 + + +def test_hooks_noop_without_events_file(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv(OBJECT_IO_FILE_ENV, raising=False) + module = _patched_client_module() + result = module.BaseClient()._make_api_call("GetObject", {"Bucket": "b", "Key": "k"}) + assert result["ContentLength"] == 42 # call passes through untouched + + +def test_patch_is_idempotent(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + events_file = tmp_path / "events.jsonl" + monkeypatch.setenv(OBJECT_IO_FILE_ENV, str(events_file)) + + module = _patched_client_module() + patch_imported_module("botocore.client", module) # second patch must not double-wrap + module.BaseClient()._make_api_call("GetObject", {"Bucket": "b", "Key": "k"}) + + events = events_file.read_text().splitlines() + assert len(events) == 1 + + +def test_non_s3_services_are_ignored(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + events_file = tmp_path / "events.jsonl" + monkeypatch.setenv(OBJECT_IO_FILE_ENV, str(events_file)) + + module = _patched_client_module() + client = module.BaseClient() + client.meta = SimpleNamespace(service_model=SimpleNamespace(service_name="dynamodb")) + client._make_api_call("GetObject", {"Bucket": "b", "Key": "k"}) + assert not events_file.exists() + + +def test_load_object_io_refs_dedupes_last_wins(tmp_path: Path) -> None: + events_file = tmp_path / "events.jsonl" + events_file.write_text( + "\n".join( + [ + json.dumps({"mode": "read", "path": "s3://d/a", "etag": "old", "size": 1}), + json.dumps({"mode": "read", "path": "s3://d/a", "etag": "new", "size": 2}), + json.dumps({"mode": "write", "path": "s3://m/out", "etag": None, "size": 9}), + "not json", + json.dumps({"mode": "bogus", "path": "s3://x/y"}), + ] + ), + encoding="utf-8", + ) + + reads, writes = load_object_io_refs(events_file) + assert reads == [ + { + "path": "s3://d/a", + "hash": "new", + "hash_algorithm": "etag", + "size": 2, + "capture_method": "python", + } + ] + assert writes == [ + { + "path": "s3://m/out", + "hash": None, + "hash_algorithm": "", + "size": 9, + "capture_method": "python", + } + ] + + +def test_load_object_io_refs_missing_file(tmp_path: Path) -> None: + reads, writes = load_object_io_refs(tmp_path / "absent.jsonl") + assert reads == [] and writes == [] From 9f1aa7458cc09f605980825038d7efb7d92e52c3 Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Mon, 13 Jul 2026 15:08:32 +0000 Subject: [PATCH 05/31] feat(k8s): capture ranged S3 reads as byte_ranges through fragments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes the ranged-read fidelity gap that was previously the roar-proxy's main advantage over the in-process hooks (the proxy itself is now explicitly a Phase-3 opt-in webhook-injected sidecar for non-Python S3 clients — see the design doc's resolved open question 5). - object_io: GetObject Range headers are parsed (fully-specified bytes=start-end pairs only) and recorded on events; the loader accumulates deduped ranges per (mode, path) across events while etag/size stay last-wins - shared fragment path (additive): ArtifactRef gains an optional byte_ranges field with validating hydration, and the fragment merge engine now carries it into job_inputs/job_outputs.byte_ranges — previously fragments dropped ranges entirely (Ray/OSMO fragments gain the capability for free) Verified live on KIND: MinIO e2e extended with a ranged GetObject, asserting byte_ranges [[0,9]] on the s3:// input row; 14/14 e2e, 3 new unit tests, default gate 2015 passed, ruff + mypy clean. Co-Authored-By: Claude Fable 5 --- docs/developer/k8s-integration.md | 13 +++++- roar/backends/k8s/object_io.py | 43 +++++++++++++++++-- roar/execution/fragments/lineage.py | 18 +++++--- roar/execution/fragments/models.py | 18 ++++++++ .../backends/k8s/e2e/test_k8s_fallback_s3.py | 11 +++++ tests/backends/k8s/unit/test_bundles.py | 20 ++++++++- tests/backends/k8s/unit/test_object_io.py | 33 ++++++++++++++ 7 files changed, 143 insertions(+), 13 deletions(-) diff --git a/docs/developer/k8s-integration.md b/docs/developer/k8s-integration.md index c9362ca6..5ad346cd 100644 --- a/docs/developer/k8s-integration.md +++ b/docs/developer/k8s-integration.md @@ -30,8 +30,17 @@ Two capture channels feed each pod's fragment: `botocore.client.BaseClient._make_api_call` (and aiobotocore's async variant, covering s3fs/fsspec) to append S3 data-op events to `ROAR_K8S_OBJECT_IO_FILE`; `pod_entry` folds them into the fragment as - `s3://` refs with etag hashes. Hooks no-op outside pods (env unset) and - never raise into user code. + `s3://` refs with etag hashes. Ranged ``GetObject`` calls record their + ``Range`` header as ``byte_ranges`` (ranges accumulate per object across + events and flow through `ArtifactRef.byte_ranges` into + `job_inputs/job_outputs.byte_ranges`). Hooks no-op outside pods (env + unset) and never raise into user code. + + The `roar-proxy` S3 sidecar is deliberately not part of the CLI-side + backend: the hooks win on attribution and avoid `AWS_ENDPOINT_URL` + rewiring (which explicit-`endpoint_url` clients bypass). The proxy joins + in the Phase-3 webhook injector as an opt-in sidecar for non-Python S3 + clients (see the design doc's Phase 3). Transport is streaming-first with a bundle fallback: when `k8s.bundle_dir` names a mounted shared volume and GLaaS is unreachable from the pod (probe diff --git a/roar/backends/k8s/object_io.py b/roar/backends/k8s/object_io.py index bd5383fe..80d8c2ce 100644 --- a/roar/backends/k8s/object_io.py +++ b/roar/backends/k8s/object_io.py @@ -90,19 +90,45 @@ def _record_s3_event( return mode = "read" if operation_name in _S3_READ_OPS else "write" - event = { + event: dict[str, Any] = { "mode": mode, "path": f"s3://{bucket}/{key}", "operation": operation_name, "etag": _normalize_etag(response), "size": _event_size(mode, api_params, response), } + byte_ranges = _parse_range_header(api_params.get("Range")) + if byte_ranges: + event["byte_ranges"] = byte_ranges with open(events_file, "a", encoding="utf-8") as handle: handle.write(json.dumps(event, separators=(",", ":")) + "\n") except Exception: return +def _parse_range_header(value: Any) -> list[list[int]]: + """Parse an HTTP Range header into [[start, end], ...] pairs. + + Only fully-specified ``bytes=start-end`` ranges are recorded; + open-ended and suffix forms are skipped (the whole-object size is + already captured separately). + """ + if not isinstance(value, str): + return [] + header = value.strip() + if not header.lower().startswith("bytes="): + return [] + ranges: list[list[int]] = [] + for part in header[len("bytes=") :].split(","): + start_text, separator, end_text = part.strip().partition("-") + if not separator or not start_text.isdigit() or not end_text.isdigit(): + continue + start, end = int(start_text), int(end_text) + if end >= start: + ranges.append([start, end]) + return ranges + + def _normalize_etag(response: Any) -> str | None: if not isinstance(response, dict): return None @@ -142,13 +168,15 @@ def _stream_size(body: Any) -> int: def load_object_io_refs(events_path: Path) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: """Read and deduplicate recorded events into fragment artifact refs. - Last event per (mode, path) wins so re-reads/re-writes carry the most - recent etag/size. + Last event per (mode, path) wins for etag/size, while byte ranges + accumulate across all events so repeated ranged reads of one object + keep every range touched. """ if not events_path.is_file(): return [], [] winners: dict[tuple[str, str], dict[str, Any]] = {} + ranges_by_key: dict[tuple[str, str], list[list[int]]] = {} for line in events_path.read_text(encoding="utf-8").splitlines(): line = line.strip() if not line: @@ -162,18 +190,25 @@ def load_object_io_refs(events_path: Path) -> tuple[list[dict[str, Any]], list[d if mode not in ("read", "write") or not path.startswith("s3://"): continue winners[(mode, path)] = event + for byte_range in event.get("byte_ranges") or []: + accumulated = ranges_by_key.setdefault((mode, path), []) + if byte_range not in accumulated: + accumulated.append(byte_range) reads: list[dict[str, Any]] = [] writes: list[dict[str, Any]] = [] for (mode, path), event in sorted(winners.items()): etag = event.get("etag") - ref = { + ref: dict[str, Any] = { "path": path, "hash": etag if isinstance(etag, str) and etag else None, "hash_algorithm": "etag" if etag else "", "size": int(event.get("size") or 0), "capture_method": "python", } + byte_ranges = ranges_by_key.get((mode, path)) + if byte_ranges: + ref["byte_ranges"] = byte_ranges (reads if mode == "read" else writes).append(ref) return reads, writes diff --git a/roar/execution/fragments/lineage.py b/roar/execution/fragments/lineage.py index 535d1c9f..d1078831 100644 --- a/roar/execution/fragments/lineage.py +++ b/roar/execution/fragments/lineage.py @@ -116,10 +116,10 @@ def merge_execution_fragments( conn.execute( """ INSERT OR IGNORE INTO job_outputs - (job_id, artifact_id, path) - VALUES (?, ?, ?) + (job_id, artifact_id, path, byte_ranges) + VALUES (?, ?, ?, ?) """, - (job_id, artifact_id, ref.path), + (job_id, artifact_id, ref.path, _byte_ranges_json(ref)), ) for ref in fragment.reads: @@ -134,10 +134,10 @@ def merge_execution_fragments( conn.execute( """ INSERT OR IGNORE INTO job_inputs - (job_id, artifact_id, path) - VALUES (?, ?, ?) + (job_id, artifact_id, path, byte_ranges) + VALUES (?, ?, ?, ?) """, - (job_id, artifact_id, ref.path), + (job_id, artifact_id, ref.path, _byte_ranges_json(ref)), ) conn.commit() @@ -149,6 +149,12 @@ def merge_execution_fragments( _refresh_fragment_job_system_labels(project_dir=project_dir, job_ids=touched_job_ids) +def _byte_ranges_json(ref: ArtifactRef) -> str | None: + if not ref.byte_ranges: + return None + return json.dumps(ref.byte_ranges, separators=(",", ":")) + + def _refresh_fragment_job_system_labels(*, project_dir: str, job_ids: set[int]) -> None: if not job_ids: return diff --git a/roar/execution/fragments/models.py b/roar/execution/fragments/models.py index 3a4352ee..ef7781ce 100644 --- a/roar/execution/fragments/models.py +++ b/roar/execution/fragments/models.py @@ -16,6 +16,9 @@ class ArtifactRef: hash_algorithm: str size: int capture_method: str + # Optional [[start, end], ...] byte ranges for ranged object-store I/O; + # carried through to job_inputs/job_outputs.byte_ranges on merge. + byte_ranges: list[list[int]] | None = None @dataclass @@ -291,4 +294,19 @@ def _artifact_ref_from_mapping(item: Mapping[str, Any]) -> ArtifactRef: hash_algorithm=str(item.get("hash_algorithm") or ""), size=_normalize_size(item.get("size")), capture_method=str(item.get("capture_method") or ""), + byte_ranges=_normalize_byte_ranges(item.get("byte_ranges")), ) + + +def _normalize_byte_ranges(value: Any) -> list[list[int]] | None: + if not isinstance(value, list): + return None + ranges: list[list[int]] = [] + for item in value: + if ( + isinstance(item, (list, tuple)) + and len(item) == 2 + and all(isinstance(bound, int) and bound >= 0 for bound in item) + ): + ranges.append([int(item[0]), int(item[1])]) + return ranges or None diff --git a/tests/backends/k8s/e2e/test_k8s_fallback_s3.py b/tests/backends/k8s/e2e/test_k8s_fallback_s3.py index c88ec036..3b32220a 100644 --- a/tests/backends/k8s/e2e/test_k8s_fallback_s3.py +++ b/tests/backends/k8s/e2e/test_k8s_fallback_s3.py @@ -94,6 +94,8 @@ config=Config(s3={"addressing_style": "path"}), ) +header = s3.get_object(Bucket=bucket, Key="datasets/train.csv", Range="bytes=0-9") +header["Body"].read() obj = s3.get_object(Bucket=bucket, Key="datasets/train.csv") data = obj["Body"].read() model = data * 3 @@ -343,5 +345,14 @@ def test_s3_object_io_captured_in_lineage( "SELECT COUNT(*) AS count FROM artifact_hashes WHERE algorithm = 'etag'", ) assert int(etag_rows[0]["count"]) >= 2, "expected etag hashes for the S3 artifacts" + + ranged = _query( + project_dir, + "SELECT byte_ranges FROM job_inputs WHERE path = ? AND job_id = ?", + (f"s3://{bucket}/datasets/train.csv", task_id), + ) + assert ranged and ranged[0]["byte_ranges"] == "[[0,9]]", ( + f"ranged read not captured: {[dict(row) for row in ranged]}\n{_describe(run)}" + ) finally: _cleanup(job_name) diff --git a/tests/backends/k8s/unit/test_bundles.py b/tests/backends/k8s/unit/test_bundles.py index 0f4c3708..2e04adab 100644 --- a/tests/backends/k8s/unit/test_bundles.py +++ b/tests/backends/k8s/unit/test_bundles.py @@ -62,7 +62,18 @@ def test_ingest_merges_bundles_into_db(tmp_path: Path) -> None: pass # initialize schema bundle_dir = tmp_path / "bundles" - write_fragment_bundle(bundle_dir, "pod-0", [_fragment("pod-0", "0")]) + ranged_fragment = _fragment("pod-0", "0") + ranged_fragment["reads"] = [ + { + "path": "s3://data/shard.bin", + "hash": "etag-1", + "hash_algorithm": "etag", + "size": 4096, + "capture_method": "python", + "byte_ranges": [[0, 1023], [2048, 4095]], + } + ] + write_fragment_bundle(bundle_dir, "pod-0", [ranged_fragment]) write_fragment_bundle(bundle_dir, "pod-1", [_fragment("pod-1", "1")]) result = ingest_fragment_bundles(roar_dir=roar_dir, directory=bundle_dir) @@ -70,9 +81,16 @@ def test_ingest_merges_bundles_into_db(tmp_path: Path) -> None: assert result.fragments_merged == 2 conn = sqlite3.connect(roar_dir / "roar.db") + conn.row_factory = sqlite3.Row try: count = conn.execute("SELECT COUNT(*) FROM jobs WHERE job_type = 'k8s_task'").fetchone() assert count[0] == 2 + + ranged = conn.execute( + "SELECT byte_ranges FROM job_inputs WHERE path = 's3://data/shard.bin'" + ).fetchone() + assert ranged is not None + assert ranged["byte_ranges"] == "[[0,1023],[2048,4095]]" finally: conn.close() diff --git a/tests/backends/k8s/unit/test_object_io.py b/tests/backends/k8s/unit/test_object_io.py index e62b0ac0..16105e6e 100644 --- a/tests/backends/k8s/unit/test_object_io.py +++ b/tests/backends/k8s/unit/test_object_io.py @@ -125,3 +125,36 @@ def test_load_object_io_refs_dedupes_last_wins(tmp_path: Path) -> None: def test_load_object_io_refs_missing_file(tmp_path: Path) -> None: reads, writes = load_object_io_refs(tmp_path / "absent.jsonl") assert reads == [] and writes == [] + + +def test_ranged_reads_record_and_accumulate( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + events_file = tmp_path / "events.jsonl" + monkeypatch.setenv(OBJECT_IO_FILE_ENV, str(events_file)) + + module = _patched_client_module() + client = module.BaseClient() + client._make_api_call("GetObject", {"Bucket": "d", "Key": "shard.bin", "Range": "bytes=0-1023"}) + client._make_api_call( + "GetObject", {"Bucket": "d", "Key": "shard.bin", "Range": "bytes=2048-4095"} + ) + client._make_api_call("GetObject", {"Bucket": "d", "Key": "shard.bin", "Range": "bytes=0-1023"}) + + reads, _writes = load_object_io_refs(events_file) + assert len(reads) == 1 + assert reads[0]["path"] == "s3://d/shard.bin" + # Ranges accumulate (deduped); etag/size come from the last event. + assert reads[0]["byte_ranges"] == [[0, 1023], [2048, 4095]] + + +def test_range_header_parsing_edge_cases(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + from roar.backends.k8s.object_io import _parse_range_header + + assert _parse_range_header("bytes=0-99") == [[0, 99]] + assert _parse_range_header("bytes=0-99,200-299") == [[0, 99], [200, 299]] + assert _parse_range_header("bytes=500-") == [] # open-ended skipped + assert _parse_range_header("bytes=-500") == [] # suffix skipped + assert _parse_range_header("items=0-9") == [] + assert _parse_range_header(None) == [] + assert _parse_range_header("bytes=99-0") == [] # inverted skipped From 3e4c71161cf669117a7c7969ce21d7f434a2eff4 Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Mon, 13 Jul 2026 15:38:41 +0000 Subject: [PATCH 06/31] feat(k8s): add mount-map rewriting and retry-chaos coverage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Completes Phase 2 of the k8s training lineage integration. Mount-map rewriting (mount_map.py): - per-container mount maps derived at rewrite time from inline CSI volumes the pod spec can name (GCS FUSE, Mountpoint-for-S3 -> gs://bucket / s3://bucket URIs, subPath-aware) plus an explicit [k8s.mount_map] config table for PVC-backed FUSE drivers whose bucket lives in the cluster-side PV; PVC mounts get pvc://claim identity tags and are never rewritten (their cross-pod edges connect through content hashes) - injected as ROAR_K8S_MOUNT_MAP; pod_entry stamps the map into fragment metadata so capture stays raw and the mapping is auditable; collect_k8s_fragments rewrites ref paths at reconstitution with longest-prefix wins (streaming, bundle, and attach paths all covered) Retry chaos e2e: - a backoffLimit-1 Job whose first pod fails mid-run yields two attempt-distinct k8s_task jobs keyed by pod UID — one exit 1 with partial outputs, one exit 0 with final outputs, never conflated — proving the identity contract under real k8s pod recreation Verified live on KIND: 16/16 e2e (chaos + mount-map new; mount-map via config-declared map over a hostPath volume, asserting rewritten s3:// inputs/outputs and no raw mount paths leaking); 8 new unit tests; default gate 2021 passed; ruff + mypy clean. Co-Authored-By: Claude Fable 5 --- docs/developer/k8s-integration.md | 31 +- roar/backends/k8s/config.py | 10 + roar/backends/k8s/lineage.py | 3 + roar/backends/k8s/manifest.py | 20 +- roar/backends/k8s/mount_map.py | 182 ++++++++++ roar/backends/k8s/pod_entry.py | 7 + roar/backends/k8s/submit.py | 8 + tests/backends/k8s/README.md | 5 + .../backends/k8s/e2e/test_k8s_chaos_mounts.py | 317 ++++++++++++++++++ tests/backends/k8s/unit/test_mount_map.py | 103 ++++++ 10 files changed, 678 insertions(+), 8 deletions(-) create mode 100644 roar/backends/k8s/mount_map.py create mode 100644 tests/backends/k8s/e2e/test_k8s_chaos_mounts.py create mode 100644 tests/backends/k8s/unit/test_mount_map.py diff --git a/docs/developer/k8s-integration.md b/docs/developer/k8s-integration.md index 5ad346cd..a70f55e8 100644 --- a/docs/developer/k8s-integration.md +++ b/docs/developer/k8s-integration.md @@ -15,11 +15,27 @@ identity comes from the downward API plus the completion-index/node-rank env chain (`JOB_COMPLETION_INDEX` → `PET_NODE_RANK` → pod-level `RANK`). Current phase status (see `design-docs/k8s-training-lineage-integration.md` -in the dev meta-repo): Phases 1–2 implemented (submit wrapping, operator -adapters, multi-pod capture, `roar k8s attach`, bundle-mode fallback, -object-store I/O hooks). Remaining: mount-map path rewriting, RayJob -delegation, kill-pod retry chaos coverage, and the admission-webhook -injector. +in the dev meta-repo): Phases 1–2 fully implemented (submit wrapping, +operator adapters, multi-pod capture, `roar k8s attach`, bundle-mode +fallback, object-store I/O hooks, mount-map rewriting, retry-chaos +coverage). Remaining: RayJob delegation and the Phase-3 admission-webhook +injector (incl. the opt-in proxy sidecar). + +**Mounted storage** (`mount_map.py`): FUSE CSI mounts surface object I/O as +local file syscalls under a mount path. The rewriter derives a per-container +mount map (inline CSI volumes it can see — GCS FUSE, Mountpoint-for-S3 — +plus explicit `[k8s.mount_map]` config for PVC-backed drivers whose bucket +lives in the cluster-side PV) and injects it as `ROAR_K8S_MOUNT_MAP`; +`pod_entry` stamps it into fragment metadata, and reconstitution rewrites +ref paths (`/data/foo` → `gs://bucket/foo`, longest prefix wins). Capture +stays raw in the fragment; the mapping used is auditable in metadata. PVC +mounts get a `pvc://claim` identity tag with no rewrite — their cross-pod +edges connect through content hashes. + +**Retry semantics**: Job retries produce attempt-distinct lineage — each +attempt's fragment is keyed by pod UID inside the task identity, so a +failed first attempt and its successful retry land as separate `k8s_task` +jobs with their own outputs (covered by the chaos e2e). Two capture channels feed each pod's fragment: @@ -113,8 +129,9 @@ recovered parent uid, and reconstitutes the streamed fragments. Section `k8s` (registered `BackendConfigAdapter`): `enabled`, `tracer`, `runtime_install_requirement`, `cluster_glaas_url`, `bundle_dir`, `wait_for_completion`, `wait_timeout_seconds`, `poll_interval_seconds`, -`fragment_session_ttl_seconds`. Env overrides beat config: -`ROAR_CLUSTER_PIP_REQ`, `ROAR_CLUSTER_GLAAS_URL`. +`fragment_session_ttl_seconds`, plus the `[k8s.mount_map]` table +(config-file-only; maps mount paths to object-store URIs). Env overrides +beat config: `ROAR_CLUSTER_PIP_REQ`, `ROAR_CLUSTER_GLAAS_URL`. ## 5. Tests diff --git a/roar/backends/k8s/config.py b/roar/backends/k8s/config.py index 295c709f..3d979c56 100644 --- a/roar/backends/k8s/config.py +++ b/roar/backends/k8s/config.py @@ -23,6 +23,10 @@ class K8sBackendConfig(BaseModel): runtime_install_requirement: str = "" cluster_glaas_url: str = "" bundle_dir: str = "" + # Explicit mount-path -> object-store URI mapping for mounted storage + # whose remote identity isn't visible in the pod spec (PVC-backed FUSE + # CSI). TOML table, e.g. [k8s.mount_map] "/data" = "gs://bucket/prefix". + mount_map: dict[str, str] = Field(default_factory=dict) wait_for_completion: bool = True wait_timeout_seconds: int = Field(default=30 * 60, ge=1) poll_interval_seconds: float = Field(default=5.0, gt=0.0) @@ -101,6 +105,12 @@ class K8sBackendConfig(BaseModel): cluster_glaas_url = "" # Wait for submitted Jobs to finish so lineage can be reconstituted immediately wait_for_completion = true + +# Optional: map mounted storage paths to their object-store URIs when the +# pod spec can't reveal them (PVC-backed FUSE CSI drivers). Inline CSI +# volumes (GCS FUSE, Mountpoint-for-S3) are detected automatically. +# [k8s.mount_map] +# "/data" = "gs://my-bucket/datasets" """ diff --git a/roar/backends/k8s/lineage.py b/roar/backends/k8s/lineage.py index 3bc8dd2a..5f1658d6 100644 --- a/roar/backends/k8s/lineage.py +++ b/roar/backends/k8s/lineage.py @@ -61,11 +61,14 @@ def collect_k8s_fragments( step_number: int = 1, ) -> int: """Merge k8s fragment dicts into the local DB; returns fragments merged.""" + from roar.backends.k8s.mount_map import rewrite_fragment_paths + parsed: list[ExecutionFragment] = [] for payload in fragments: if not isinstance(payload, dict): continue try: + rewrite_fragment_paths(payload) parsed.append(ExecutionFragment.from_dict(payload)) except Exception: continue diff --git a/roar/backends/k8s/manifest.py b/roar/backends/k8s/manifest.py index 913e74c8..63a678ab 100644 --- a/roar/backends/k8s/manifest.py +++ b/roar/backends/k8s/manifest.py @@ -22,6 +22,8 @@ import yaml # type: ignore[import-untyped] +from roar.backends.k8s.mount_map import build_container_mount_map, dump_mount_map + # sh -c script: "$0" is the synthetic argv0, "$@" is the original # command+args. Lineage is best-effort by design: any failure to stage # the roar runtime falls back to running the original command @@ -186,6 +188,7 @@ def rewrite_manifest_for_lineage( tracer: str, parent_job_uid: str, bundle_dir: str = "", + mount_map: dict[str, str] | None = None, namespace_override: str | None = None, ) -> K8sManifestRewrite: """Return a rewritten copy of ``documents`` with lineage instrumentation. @@ -233,6 +236,7 @@ def rewrite_manifest_for_lineage( parent_job_uid=parent_job_uid, workload_name=workload_name, bundle_dir=bundle_dir, + config_mount_map=dict(mount_map or {}), ) if workload.locate_pod_specs is None: @@ -295,6 +299,7 @@ class _EnvContract: parent_job_uid: str workload_name: str bundle_dir: str = "" + config_mount_map: dict[str, str] = field(default_factory=dict) def _rewrite_pod_specs( @@ -336,6 +341,9 @@ def _rewrite_pod_specs( contract=contract, container_name=container_name, role=ref.role, + mount_map_entries=build_container_mount_map( + ref.spec, container, contract.config_mount_map + ), ) wrapped.append(label) return wrapped, skipped @@ -377,7 +385,14 @@ def _rewrite_trainjob( env = trainer.setdefault("env", []) if not isinstance(env, list): raise K8sManifestError(f"TrainJob {workload_name} has a non-list spec.trainer.env") - _inject_env_contract(env, contract=contract, container_name="node", role="") + _inject_env_contract( + env, + contract=contract, + container_name="node", + role="", + # TrainJob has no visible pod spec; only explicit config entries apply. + mount_map_entries=build_container_mount_map({}, {}, contract.config_mount_map), + ) return ["node"], [] @@ -392,6 +407,7 @@ def _inject_env_contract( contract: _EnvContract, container_name: str, role: str, + mount_map_entries: list[dict[str, str]] | None = None, ) -> None: if not isinstance(env, list): raise K8sManifestError(f"container {container_name} has a non-list env block") @@ -426,6 +442,8 @@ def add_secret_ref(name: str, key: str) -> None: add_value("ROAR_K8S_TRACER", contract.tracer) if contract.bundle_dir: add_value("ROAR_K8S_BUNDLE_DIR", contract.bundle_dir) + if mount_map_entries: + add_value("ROAR_K8S_MOUNT_MAP", dump_mount_map(mount_map_entries)) add_value("GLAAS_URL", contract.cluster_glaas_url) add_value("ROAR_K8S_PARENT_JOB_UID", contract.parent_job_uid) add_value("ROAR_K8S_JOB_NAME", contract.workload_name) diff --git a/roar/backends/k8s/mount_map.py b/roar/backends/k8s/mount_map.py new file mode 100644 index 00000000..e4ef0a1d --- /dev/null +++ b/roar/backends/k8s/mount_map.py @@ -0,0 +1,182 @@ +"""Mount-map derivation and path rewriting for mounted object storage. + +FUSE CSI mounts (GCS FUSE, Mountpoint-for-S3) surface object-store I/O +as local file syscalls under a mount path — the tracer captures them, +but at `/data/foo` instead of `gs://bucket/foo`. The rewriter derives a +mount map per container at submit time (inline CSI volumes it can see, +plus explicit `k8s.mount_map` config for PVC-backed drivers whose +bucket lives in the cluster-side PV), injects it as +``ROAR_K8S_MOUNT_MAP``, and the pod entrypoint stamps it into fragment +metadata. Reconstitution rewrites ref paths — capture stays raw in the +fragment; the transformation happens at merge time and the mapping used +is preserved in metadata. + +PVC mounts get a ``volume`` identity tag (no rewrite): their cross-pod +edges already connect through content hashes. +""" + +from __future__ import annotations + +import json +from typing import Any + +MOUNT_MAP_ENV = "ROAR_K8S_MOUNT_MAP" + +# Inline-CSI drivers whose pod spec carries enough to name the remote URI. +_CSI_URI_SCHEMES = { + "gcsfuse.csi.storage.gke.io": "gs", + "s3.csi.aws.com": "s3", +} + + +def build_container_mount_map( + pod_spec: dict[str, Any], + container: dict[str, Any], + config_mount_map: dict[str, str] | None = None, +) -> list[dict[str, str]]: + """Return mount-map entries for one container's volume mounts. + + Entry shapes: + - ``{"mount_path": "/data", "uri": "gs://bucket"}`` — rewrite target + - ``{"mount_path": "/ckpt", "volume": "pvc://claim"}`` — identity tag + Explicit ``k8s.mount_map`` config entries win over derived ones. + """ + volumes_by_name: dict[str, dict[str, Any]] = {} + for volume in pod_spec.get("volumes") or []: + if isinstance(volume, dict) and volume.get("name"): + volumes_by_name[str(volume["name"])] = volume + + entries: list[dict[str, str]] = [] + configured = dict(config_mount_map or {}) + + for mount in container.get("volumeMounts") or []: + if not isinstance(mount, dict): + continue + mount_path = str(mount.get("mountPath") or "").rstrip("/") + if not mount_path: + continue + if mount_path in configured: + continue # explicit config wins; added below + + volume = volumes_by_name.get(str(mount.get("name") or "")) + if not isinstance(volume, dict): + continue + + uri = _inline_csi_uri(volume) + if uri: + sub_path = str(mount.get("subPath") or "").strip("/") + if sub_path: + uri = f"{uri}/{sub_path}" + entries.append({"mount_path": mount_path, "uri": uri}) + continue + + claim = volume.get("persistentVolumeClaim") + if isinstance(claim, dict) and claim.get("claimName"): + entries.append({"mount_path": mount_path, "volume": f"pvc://{claim['claimName']}"}) + + for mount_path, uri in configured.items(): + normalized = str(mount_path or "").rstrip("/") + target = str(uri or "").rstrip("/") + if normalized and target: + entries.append({"mount_path": normalized, "uri": target}) + + return entries + + +def _inline_csi_uri(volume: dict[str, Any]) -> str | None: + csi = volume.get("csi") + if not isinstance(csi, dict): + return None + scheme = _CSI_URI_SCHEMES.get(str(csi.get("driver") or "")) + if scheme is None: + return None + attributes = csi.get("volumeAttributes") + bucket = "" + if isinstance(attributes, dict): + bucket = str(attributes.get("bucketName") or "").strip() + if not bucket: + return None + return f"{scheme}://{bucket}" + + +def dump_mount_map(entries: list[dict[str, str]]) -> str: + return json.dumps(entries, separators=(",", ":")) + + +def parse_mount_map(raw: str | None) -> list[dict[str, str]]: + if not raw: + return [] + try: + payload = json.loads(raw) + except json.JSONDecodeError: + return [] + if not isinstance(payload, list): + return [] + entries: list[dict[str, str]] = [] + for item in payload: + if not isinstance(item, dict): + continue + mount_path = str(item.get("mount_path") or "").rstrip("/") + if not mount_path: + continue + entry = {"mount_path": mount_path} + if item.get("uri"): + entry["uri"] = str(item["uri"]).rstrip("/") + if item.get("volume"): + entry["volume"] = str(item["volume"]) + entries.append(entry) + return entries + + +def rewrite_fragment_paths(fragment: dict[str, Any]) -> None: + """Rewrite mounted-object paths in a fragment dict, in place. + + Uses the ``k8s_mount_map`` recorded in the fragment's backend + metadata; longest mount-path prefix wins. Entries without a ``uri`` + (PVC identity tags) never rewrite. + """ + metadata = fragment.get("backend_metadata") + if not isinstance(metadata, dict): + return + entries = [ + entry + for entry in metadata.get("k8s_mount_map") or [] + if isinstance(entry, dict) and entry.get("uri") and entry.get("mount_path") + ] + if not entries: + return + entries.sort(key=lambda entry: len(str(entry["mount_path"])), reverse=True) + + for list_key in ("reads", "writes"): + refs = fragment.get(list_key) + if not isinstance(refs, list): + continue + for ref in refs: + if not isinstance(ref, dict): + continue + path = str(ref.get("path") or "") + rewritten = _rewrite_path(path, entries) + if rewritten is not None: + ref["path"] = rewritten + + +def _rewrite_path(path: str, entries: list[dict[str, Any]]) -> str | None: + if not path.startswith("/"): + return None + for entry in entries: + mount_path = str(entry["mount_path"]) + uri = str(entry["uri"]) + if path == mount_path: + return uri + if path.startswith(mount_path + "/"): + return uri + path[len(mount_path) :] + return None + + +__all__ = [ + "MOUNT_MAP_ENV", + "build_container_mount_map", + "dump_mount_map", + "parse_mount_map", + "rewrite_fragment_paths", +] diff --git a/roar/backends/k8s/pod_entry.py b/roar/backends/k8s/pod_entry.py index ed8e908f..3e1a12d2 100644 --- a/roar/backends/k8s/pod_entry.py +++ b/roar/backends/k8s/pod_entry.py @@ -135,12 +135,19 @@ def _emit_lineage_best_effort() -> None: ) payload = json.loads(bundle_path.read_text(encoding="utf-8")) + from roar.backends.k8s.mount_map import MOUNT_MAP_ENV, parse_mount_map + fragments = [item for item in payload.get("fragments", []) if isinstance(item, dict)] _augment_with_object_io(fragments) + mount_map = parse_mount_map(os.environ.get(MOUNT_MAP_ENV)) completion_index = task_id.split(":")[2] if task_id.count(":") >= 3 else "0" restart_attempt = task_id.split(":")[3] if task_id.count(":") >= 3 else "0" for fragment in fragments: metadata = fragment.setdefault("backend_metadata", {}) + if mount_map: + # Recorded raw; reconstitution applies the rewrite so the + # mapping used stays auditable next to the captured paths. + metadata["k8s_mount_map"] = mount_map metadata.update( { "k8s_namespace": os.environ.get("ROAR_K8S_NAMESPACE"), diff --git a/roar/backends/k8s/submit.py b/roar/backends/k8s/submit.py index aa990194..4f79777a 100644 --- a/roar/backends/k8s/submit.py +++ b/roar/backends/k8s/submit.py @@ -120,6 +120,7 @@ def plan_kubectl_job_submit_command(command: list[str]) -> ExecutionCommandPlan: tracer=str(config.get("tracer") or "preload"), parent_job_uid=parent_job_uid, bundle_dir=str(config.get("bundle_dir") or ""), + mount_map=_config_mount_map(config), namespace_override=_find_namespace_argument(command), ) if rewrite.skipped_containers: @@ -259,6 +260,13 @@ def _find_namespace_argument(command: list[str]) -> str | None: return None +def _config_mount_map(config: dict) -> dict[str, str]: + raw = config.get("mount_map") + if not isinstance(raw, dict): + return {} + return {str(key): str(value) for key, value in raw.items() if key and value} + + def _resolve_glaas_url() -> str | None: from roar.integrations.glaas import get_glaas_url diff --git a/tests/backends/k8s/README.md b/tests/backends/k8s/README.md index b823786d..165d6b06 100644 --- a/tests/backends/k8s/README.md +++ b/tests/backends/k8s/README.md @@ -20,6 +20,11 @@ Three test layers share this harness: merged with `roar k8s ingest-bundles`) and in-pod S3 capture (MinIO via `--with-minio`: boto3 get/put recorded as `s3://` lineage refs with etag hashes). +- `e2e/test_k8s_chaos_mounts.py` — retry chaos (backoffLimit-1 Job whose + first pod fails mid-run: both attempts land as distinct, non-conflated + `k8s_task` jobs keyed by pod UID) and mount-map rewriting (a hostPath + volume standing in for a FUSE mount, declared via `[k8s.mount_map]`, + rewrites to `s3://` URIs at reconstitution). - `e2e/test_k8s_smoke.py` — the Phase-0 runtime diagnostic: fixtures hand-wrap the manifest (no backend involved) to isolate the runtime pieces (in-pod tracing, fragment streaming, identity contract) when the product diff --git a/tests/backends/k8s/e2e/test_k8s_chaos_mounts.py b/tests/backends/k8s/e2e/test_k8s_chaos_mounts.py new file mode 100644 index 00000000..2b5968bb --- /dev/null +++ b/tests/backends/k8s/e2e/test_k8s_chaos_mounts.py @@ -0,0 +1,317 @@ +"""Phase-2 e2e: retry chaos and mount-map rewriting. + +- Chaos: a Job whose first pod fails mid-run (backoffLimit 1) must yield + attempt-distinct, non-conflated lineage — both attempts recorded as + separate k8s_task jobs keyed by pod UID, with the retry succeeding. +- Mount map: a hostPath volume standing in for a FUSE mount, declared via + ``[k8s.mount_map]`` config, has its file I/O rewritten to object-store + URIs at reconstitution. +""" + +from __future__ import annotations + +import json +import sqlite3 +import subprocess +import uuid +from pathlib import Path + +import pytest + +from .conftest import NAMESPACE +from .test_k8s_distributed import ( + PINNED_NODE, + _cleanup, + _describe, + _indent_script, + _pod_logs, + _submit, + _write_project, +) +from .test_k8s_product_path import wheel_server # noqa: F401 (module fixture reuse) + +pytestmark = [ + pytest.mark.e2e, + pytest.mark.k8s_e2e, + pytest.mark.timeout(900), +] + +FLAKY_SCRIPT = """\ +import json +import sys +from pathlib import Path + +marker = Path("/state/attempt-1-done") +if not marker.exists(): + marker.write_text("first attempt failed here") + Path("partial.bin").write_bytes(b"partial-progress") + print("attempt 1 failing deliberately") + sys.exit(1) + +Path("dataset.csv").write_text("x\\n1\\n2\\n") +rows = Path("dataset.csv").read_text().strip().splitlines() +Path("model.bin").write_bytes(str(len(rows)).encode() * 8) +print("attempt 2 succeeded") +""" + +CHAOS_JOB_TEMPLATE = """\ +apiVersion: v1 +kind: ConfigMap +metadata: + name: {name}-worker + namespace: {namespace} +data: + flaky_train.py: | +{worker_script} +--- +apiVersion: batch/v1 +kind: Job +metadata: + name: {name} + namespace: {namespace} + labels: + app.kubernetes.io/part-of: roar-k8s-e2e +spec: + backoffLimit: 1 + ttlSecondsAfterFinished: 1800 + template: + spec: + restartPolicy: Never + nodeSelector: + kubernetes.io/hostname: {pinned_node} + volumes: + - name: work + emptyDir: {{}} + - name: state + hostPath: + path: /var/tmp/roar-e2e-state-{name} + type: DirectoryOrCreate + - name: workload + configMap: + name: {name}-worker + containers: + - name: trainer + image: python:3.12-slim + workingDir: /work + volumeMounts: + - name: work + mountPath: /work + - name: state + mountPath: /state + - name: workload + mountPath: /workload + readOnly: true + command: ["python", "/workload/flaky_train.py"] +""" + +MOUNT_JOB_TEMPLATE = """\ +apiVersion: batch/v1 +kind: Job +metadata: + name: {name} + namespace: {namespace} + labels: + app.kubernetes.io/part-of: roar-k8s-e2e +spec: + backoffLimit: 0 + ttlSecondsAfterFinished: 1800 + template: + spec: + restartPolicy: Never + nodeSelector: + kubernetes.io/hostname: {pinned_node} + volumes: + - name: work + emptyDir: {{}} + - name: data + hostPath: + path: /var/tmp/roar-e2e-mount-{name} + type: DirectoryOrCreate + containers: + - name: trainer + image: python:3.12-slim + workingDir: /work + volumeMounts: + - name: work + mountPath: /work + - name: data + mountPath: /data + command: + - python + - -c + - >- + from pathlib import Path; + data = Path('/data/input.csv').read_text(); + Path('/data/derived').mkdir(exist_ok=True); + Path('/data/derived/copy.csv').write_text(data); + Path('model.bin').write_bytes(data.encode() * 2) +""" + + +def _query(project_dir: Path, sql: str, params: tuple = ()) -> list[sqlite3.Row]: + conn = sqlite3.connect(project_dir / ".roar" / "roar.db") + conn.row_factory = sqlite3.Row + try: + return conn.execute(sql, params).fetchall() + finally: + conn.close() + + +def _node_exec(*args: str) -> subprocess.CompletedProcess[str]: + return subprocess.run( + ["docker", "exec", PINNED_NODE, *args], + capture_output=True, + text=True, + check=False, + ) + + +def test_pod_retry_produces_attempt_distinct_lineage( + k8s_cluster: None, + glaas_health: str, + wheel_server: dict[str, str], # noqa: F811 + tmp_path_factory: pytest.TempPathFactory, +) -> None: + project_dir = tmp_path_factory.mktemp("k8s-chaos") + _write_project(project_dir, wheel_server["url"]) + + job_name = f"roar-chaos-{uuid.uuid4().hex[:6]}" + (project_dir / "job.yaml").write_text( + CHAOS_JOB_TEMPLATE.format( + name=job_name, + namespace=NAMESPACE, + pinned_node=PINNED_NODE, + worker_script=_indent_script(FLAKY_SCRIPT), + ), + encoding="utf-8", + ) + + state_dir = f"/var/tmp/roar-e2e-state-{job_name}" + try: + completed = _submit("job.yaml", cwd=project_dir) + logs = _pod_logs(f"job-name={job_name}") + run = { + "exit_code": completed.returncode, + "stdout": completed.stdout, + "stderr": completed.stderr, + "pod_logs": logs, + } + # The retry succeeded, so the submit as a whole succeeds. + assert completed.returncode == 0, _describe(run) + + tasks = _query( + project_dir, + "SELECT id, exit_code, metadata FROM jobs WHERE job_type = 'k8s_task' ORDER BY id", + ) + assert len(tasks) == 2, ( + f"expected both attempts as distinct tasks, got {len(tasks)}\n{_describe(run)}" + ) + + pod_uids = set() + for row in tasks: + metadata = json.loads(row["metadata"] or "{}") + pod_uids.add(str(metadata.get("k8s_pod_uid") or "")) + assert len(pod_uids) == 2 and "" not in pod_uids, ( + f"attempts must be keyed by distinct pod UIDs: {pod_uids}\n{_describe(run)}" + ) + + exit_codes = sorted(int(row["exit_code"] or 0) for row in tasks) + assert exit_codes == [0, 1], ( + f"expected one failed and one successful attempt, got {exit_codes}" + ) + + outputs_by_task: dict[int, set[str]] = {} + for row in tasks: + outputs_by_task[int(row["id"])] = { + str(item["path"]) + for item in _query( + project_dir, + "SELECT path FROM job_outputs WHERE job_id = ?", + (row["id"],), + ) + } + all_outputs = set().union(*outputs_by_task.values()) + assert any(path.endswith("partial.bin") for path in all_outputs), all_outputs + assert any(path.endswith("model.bin") for path in all_outputs), all_outputs + # The failed attempt's partial output must not be conflated into + # the successful attempt's job. + assert not any( + {p for p in paths if p.endswith("partial.bin")} + and {p for p in paths if p.endswith("model.bin")} + for paths in outputs_by_task.values() + ), outputs_by_task + finally: + _cleanup(job_name) + _node_exec("rm", "-rf", state_dir) + + +def test_mounted_storage_paths_rewrite_to_object_uris( + k8s_cluster: None, + glaas_health: str, + wheel_server: dict[str, str], # noqa: F811 + tmp_path_factory: pytest.TempPathFactory, +) -> None: + project_dir = tmp_path_factory.mktemp("k8s-mounts") + _write_project(project_dir, wheel_server["url"]) + config_path = project_dir / ".roar" / "config.toml" + config_path.write_text( + config_path.read_text(encoding="utf-8") + + '\n[k8s.mount_map]\n"/data" = "s3://mounted-bucket/datasets"\n', + encoding="utf-8", + ) + + job_name = f"roar-mount-{uuid.uuid4().hex[:6]}" + (project_dir / "job.yaml").write_text( + MOUNT_JOB_TEMPLATE.format( + name=job_name, + namespace=NAMESPACE, + pinned_node=PINNED_NODE, + ), + encoding="utf-8", + ) + + mount_dir = f"/var/tmp/roar-e2e-mount-{job_name}" + seeded = _node_exec( + "sh", "-c", f"mkdir -p {mount_dir} && printf 'x,y\\n1,2\\n' > {mount_dir}/input.csv" + ) + assert seeded.returncode == 0, seeded.stderr + + try: + completed = _submit("job.yaml", cwd=project_dir) + logs = _pod_logs(f"job-name={job_name}") + run = { + "exit_code": completed.returncode, + "stdout": completed.stdout, + "stderr": completed.stderr, + "pod_logs": logs, + } + assert completed.returncode == 0, _describe(run) + + tasks = _query(project_dir, "SELECT id FROM jobs WHERE job_type = 'k8s_task'") + assert len(tasks) == 1, _describe(run) + task_id = tasks[0]["id"] + + input_paths = { + str(row["path"]) + for row in _query( + project_dir, "SELECT path FROM job_inputs WHERE job_id = ?", (task_id,) + ) + } + output_paths = { + str(row["path"]) + for row in _query( + project_dir, "SELECT path FROM job_outputs WHERE job_id = ?", (task_id,) + ) + } + assert "s3://mounted-bucket/datasets/input.csv" in input_paths, ( + f"mounted read not rewritten: {input_paths}\n{_describe(run)}" + ) + assert "s3://mounted-bucket/datasets/derived/copy.csv" in output_paths, ( + f"mounted write not rewritten: {output_paths}\n{_describe(run)}" + ) + assert any(path.endswith("model.bin") for path in output_paths), output_paths + # No raw mount paths may leak through. + assert not any(path.startswith("/data/") for path in input_paths | output_paths) + finally: + _cleanup(job_name) + _node_exec("rm", "-rf", mount_dir) diff --git a/tests/backends/k8s/unit/test_mount_map.py b/tests/backends/k8s/unit/test_mount_map.py new file mode 100644 index 00000000..3f1c5ae2 --- /dev/null +++ b/tests/backends/k8s/unit/test_mount_map.py @@ -0,0 +1,103 @@ +from __future__ import annotations + +from typing import Any + +from roar.backends.k8s.mount_map import ( + build_container_mount_map, + dump_mount_map, + parse_mount_map, + rewrite_fragment_paths, +) + +POD_SPEC = { + "volumes": [ + { + "name": "gcs-data", + "csi": { + "driver": "gcsfuse.csi.storage.gke.io", + "volumeAttributes": {"bucketName": "training-data"}, + }, + }, + { + "name": "s3-data", + "csi": { + "driver": "s3.csi.aws.com", + "volumeAttributes": {"bucketName": "shards"}, + }, + }, + {"name": "ckpt", "persistentVolumeClaim": {"claimName": "shared-ckpt"}}, + {"name": "scratch", "emptyDir": {}}, + ] +} + +CONTAINER = { + "name": "trainer", + "volumeMounts": [ + {"name": "gcs-data", "mountPath": "/data"}, + {"name": "s3-data", "mountPath": "/shards", "subPath": "run-42"}, + {"name": "ckpt", "mountPath": "/ckpt"}, + {"name": "scratch", "mountPath": "/scratch"}, + ], +} + + +def test_derives_inline_csi_and_pvc_entries() -> None: + entries = build_container_mount_map(POD_SPEC, CONTAINER) + assert entries == [ + {"mount_path": "/data", "uri": "gs://training-data"}, + {"mount_path": "/shards", "uri": "s3://shards/run-42"}, + {"mount_path": "/ckpt", "volume": "pvc://shared-ckpt"}, + ] + + +def test_config_entries_win_over_derived() -> None: + entries = build_container_mount_map( + POD_SPEC, + CONTAINER, + {"/data": "gs://training-data/curated/", "/extra": "s3://other"}, + ) + by_path = {entry["mount_path"]: entry for entry in entries} + assert by_path["/data"] == {"mount_path": "/data", "uri": "gs://training-data/curated"} + assert by_path["/extra"] == {"mount_path": "/extra", "uri": "s3://other"} + + +def test_mount_map_round_trips_through_env() -> None: + entries = build_container_mount_map(POD_SPEC, CONTAINER) + assert parse_mount_map(dump_mount_map(entries)) == entries + assert parse_mount_map(None) == [] + assert parse_mount_map("not json") == [] + + +def _fragment_with_paths(read_path: str, write_path: str) -> dict[str, Any]: + return { + "backend_metadata": { + "k8s_mount_map": [ + {"mount_path": "/data", "uri": "gs://training-data"}, + {"mount_path": "/data/nested", "uri": "s3://nested-bucket"}, + {"mount_path": "/ckpt", "volume": "pvc://shared-ckpt"}, + ] + }, + "reads": [{"path": read_path}], + "writes": [{"path": write_path}], + } + + +def test_rewrite_uses_longest_prefix_and_skips_volume_tags() -> None: + fragment = _fragment_with_paths("/data/nested/shard-0.bin", "/ckpt/step-100.pt") + rewrite_fragment_paths(fragment) + assert fragment["reads"][0]["path"] == "s3://nested-bucket/shard-0.bin" + # PVC identity tags never rewrite: local path + content hash stay canonical. + assert fragment["writes"][0]["path"] == "/ckpt/step-100.pt" + + +def test_rewrite_exact_mount_path_and_untouched_paths() -> None: + fragment = _fragment_with_paths("/data", "/work/model.bin") + rewrite_fragment_paths(fragment) + assert fragment["reads"][0]["path"] == "gs://training-data" + assert fragment["writes"][0]["path"] == "/work/model.bin" + + +def test_rewrite_noops_without_mount_map() -> None: + fragment = {"backend_metadata": {}, "reads": [{"path": "/data/x"}], "writes": []} + rewrite_fragment_paths(fragment) + assert fragment["reads"][0]["path"] == "/data/x" From 9a4d034fa964504339c7376c4e170477eb7024d7 Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Tue, 14 Jul 2026 14:35:10 +0000 Subject: [PATCH 07/31] feat(k8s): validate operators live and delegate RayJob to the Ray backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Live operator validation: - bootstrap --with-kubeflow installs training-operator v1 (v1.9.2) and trainer v2 (v2.2.1) + runtimes; --with-kuberay installs KubeRay v1.4.0 - PyTorchJob e2e proves the per-role adapter against the real controller (Master+Worker, RANK-chain node indices 0/1); TrainJob e2e runs the real TrainJob->JobSet pipeline via a slim-image clone of the shipped torch-distributed ClusterTrainingRuntime RayJob delegation (rayjob.py, rewrite_style="rayjob"): - entrypoint wrapped through the Ray driver entrypoint; runtimeEnvYAML merged with the roar pip requirement, worker setup hook, and the Ray env contract (ROAR_JOB_ID = the k8s parent uid so fragments link to the submit node); credentials only as Secret refs in the RayCluster pod templates; status.jobStatus terminal handling; reconstitution and attach delegate to the Ray backend's reconstituter (fragments are Ray TaskFragments merging as ray_task jobs) - two live-found env-contract fixes: the merged local-proxy redirect (AWS_ENDPOINT_URL/ROAR_PROXY_PORT) is stripped (no proxy runs in the pods; user S3 traffic would hit a dead localhost port), and ROAR_WRAP stays off (the roar_inject.pth fires in Ray pip virtualenvs before site-packages are importable and the sitecustomize ABI-repair path blows the 60s worker registration timeout — observed as an infinite supervisor start/kill/retry loop) Registry robustness (shared, found by the live smoke): backend plugin imports that fail during early sitecustomize discovery were cached forever ("unknown execution backend: ray" in worker setup hooks even though deps import fine moments later); get_execution_backend now retries skipped builtin imports on a lookup miss. Known gap documented in the e2e and docs: per-task I/O fidelity inside KubeRay workers under setup-hook-only bootstrap (fragments arrive as ray_task shells without file refs) — a Ray-backend follow-up; the live smoke asserts the delegation transport scope. Verified live on KIND: 19/19 e2e (PyTorchJob + TrainJob + RayJob new); 10 new unit tests; default gate 2029 passed; ruff + mypy clean. Co-Authored-By: Claude Fable 5 --- docs/developer/k8s-integration.md | 21 +- roar/backends/k8s/attach.py | 53 ++- roar/backends/k8s/manifest.py | 30 +- roar/backends/k8s/rayjob.py | 205 ++++++++++++ roar/backends/k8s/submit.py | 10 + roar/backends/k8s/workload_wait.py | 8 +- roar/execution/framework/registry.py | 34 ++ tests/backends/k8s/README.md | 12 +- .../backends/k8s/e2e/test_k8s_distributed.py | 9 +- tests/backends/k8s/e2e/test_k8s_operators.py | 315 ++++++++++++++++++ tests/backends/k8s/e2e/test_k8s_rayjob.py | 229 +++++++++++++ tests/backends/k8s/scripts/bootstrap_k8s.sh | 31 ++ tests/backends/k8s/unit/test_rayjob.py | 165 +++++++++ .../backends/k8s/unit/test_registry_retry.py | 33 ++ 14 files changed, 1138 insertions(+), 17 deletions(-) create mode 100644 roar/backends/k8s/rayjob.py create mode 100644 tests/backends/k8s/e2e/test_k8s_operators.py create mode 100644 tests/backends/k8s/e2e/test_k8s_rayjob.py create mode 100644 tests/backends/k8s/unit/test_rayjob.py create mode 100644 tests/backends/k8s/unit/test_registry_retry.py diff --git a/docs/developer/k8s-integration.md b/docs/developer/k8s-integration.md index a70f55e8..f1f17ab9 100644 --- a/docs/developer/k8s-integration.md +++ b/docs/developer/k8s-integration.md @@ -15,12 +15,27 @@ identity comes from the downward API plus the completion-index/node-rank env chain (`JOB_COMPLETION_INDEX` → `PET_NODE_RANK` → pod-level `RANK`). Current phase status (see `design-docs/k8s-training-lineage-integration.md` -in the dev meta-repo): Phases 1–2 fully implemented (submit wrapping, -operator adapters, multi-pod capture, `roar k8s attach`, bundle-mode +in the dev meta-repo): Phases 1–2 fully implemented and live-validated +(submit wrapping, operator adapters incl. real training-operator v1 and +trainer v2 controllers, multi-pod capture, `roar k8s attach`, bundle-mode fallback, object-store I/O hooks, mount-map rewriting, retry-chaos -coverage). Remaining: RayJob delegation and the Phase-3 admission-webhook +coverage, RayJob delegation). Remaining: the Phase-3 admission-webhook injector (incl. the opt-in proxy sidecar). +**RayJob delegation** (`rayjob.py`): KubeRay overwrites container commands +with `ray start --block`, so command-wrapping can't see user code. Instead +the RayJob rewrite reuses the Ray backend's runtime surface: the +entrypoint is wrapped through the Ray driver entrypoint, `runtimeEnvYAML` +gains the roar pip requirement + worker setup hook + Ray env contract +(`ROAR_EXECUTION_BACKEND=ray`, `ROAR_JOB_ID` = the k8s parent uid so Ray +fragments link to the recorded submit job; node agents/proxy stay off per +the proxy decision), and fragment-session credentials go into the +RayCluster pod templates as Secret refs — never inline in the CR. RayJob +signals completion via `status.jobStatus` (not conditions), and +reconstitution/attach delegate to the **Ray** backend's reconstituter, +since the streamed fragments are Ray `TaskFragment` payloads that merge +as `ray_task` jobs. + **Mounted storage** (`mount_map.py`): FUSE CSI mounts surface object I/O as local file syscalls under a mount path. The rewriter derives a per-container mount map (inline CSI volumes it can see — GCS FUSE, Mountpoint-for-S3 — diff --git a/roar/backends/k8s/attach.py b/roar/backends/k8s/attach.py index 18afa71b..80fcf710 100644 --- a/roar/backends/k8s/attach.py +++ b/roar/backends/k8s/attach.py @@ -87,9 +87,7 @@ def attach_k8s_workload( ) kubectl_resource = workload_kind.kubectl_resource - env_entries = _instrumented_env_entries(document, workload_kind) - parent_job_uid = _env_value(env_entries, "ROAR_K8S_PARENT_JOB_UID") - secret_name = _session_secret_name(env_entries) + parent_job_uid, secret_name = _recover_identity(document, workload_kind) if not parent_job_uid or not secret_name: raise K8sAttachError( f"{kubectl_resource}/{name} was not instrumented by roar " @@ -149,7 +147,7 @@ def attach_k8s_workload( if not glaas_url: raise K8sAttachError("GLaaS is not configured; set glaas.url to fetch fragments") - result = create_k8s_fragment_reconstituter( + result = _reconstituter_factory(workload_kind)( session_id, token, str(glaas_url), @@ -210,6 +208,44 @@ def _fetch_workload( raise K8sAttachError(f"cannot fetch workload {workload!r} in namespace {namespace}: {detail}") +def _recover_identity( + document: dict[str, Any], + workload_kind: WorkloadKind, +) -> tuple[str, str]: + """Return (parent_job_uid, secret_name) from the cluster object.""" + env_entries = _instrumented_env_entries(document, workload_kind) + secret_name = _session_secret_name(env_entries) + + if workload_kind.kind == "RayJob": + # RayJob carries the parent uid in the Ray env contract + # (runtimeEnvYAML env_vars.ROAR_JOB_ID), not pod env. + import yaml # type: ignore[import-untyped] + + raw = document.get("spec", {}).get("runtimeEnvYAML") + try: + runtime_env = yaml.safe_load(raw) if isinstance(raw, str) else None + except yaml.YAMLError: + runtime_env = None + env_vars = runtime_env.get("env_vars") if isinstance(runtime_env, dict) else None + parent = str((env_vars or {}).get("ROAR_JOB_ID") or "").strip() + return parent, secret_name + + return _env_value(env_entries, "ROAR_K8S_PARENT_JOB_UID"), secret_name + + +def _reconstituter_factory(workload_kind: WorkloadKind): + if workload_kind.kind == "RayJob": + # RayJob fragments are Ray TaskFragments; use the Ray reconstituter. + from roar.execution.framework.registry import get_execution_backend + + backend = get_execution_backend("ray") + distributed = backend.distributed + if distributed is None or distributed.fragment_reconstitution is None: + raise K8sAttachError("ray backend has no fragment reconstitution adapter") + return distributed.fragment_reconstitution.create_reconstituter + return create_k8s_fragment_reconstituter + + def _instrumented_env_entries( document: dict[str, Any], workload_kind: WorkloadKind, @@ -219,6 +255,13 @@ def _instrumented_env_entries( env = trainer.get("env") if isinstance(trainer, dict) else None return [entry for entry in env or [] if isinstance(entry, dict)] + def _is_instrumented(entry: dict[str, Any]) -> bool: + if entry.get("name") == "ROAR_K8S_PARENT_JOB_UID": + return True + return entry.get("name") == "ROAR_SESSION_ID" and bool( + (entry.get("valueFrom") or {}).get("secretKeyRef") + ) + for ref in workload_kind.locate_pod_specs(document): containers = ref.spec.get("containers") if not isinstance(containers, list): @@ -227,7 +270,7 @@ def _instrumented_env_entries( if not isinstance(container, dict): continue env = [entry for entry in container.get("env") or [] if isinstance(entry, dict)] - if any(entry.get("name") == "ROAR_K8S_PARENT_JOB_UID" for entry in env): + if any(_is_instrumented(entry) for entry in env): return env return [] diff --git a/roar/backends/k8s/manifest.py b/roar/backends/k8s/manifest.py index 63a678ab..7f2a1687 100644 --- a/roar/backends/k8s/manifest.py +++ b/roar/backends/k8s/manifest.py @@ -61,9 +61,12 @@ class WorkloadKind: api_group: str kubectl_resource: str # Returns mutable pod-spec references inside the (copied) document. - # None marks trainer-override style workloads (TrainJob), which have - # no inline pod template. + # None marks workloads with no inline pod template (TrainJob). locate_pod_specs: Callable[[dict[str, Any]], list[PodSpecRef]] | None + # How instrumentation is applied: command-wrapping of pod-spec + # containers (default), the TrainJob trainer override, or delegation + # to the Ray backend's runtime surface for RayJob. + rewrite_style: str = "pod_specs" def _job_pod_specs(doc: dict[str, Any]) -> list[PodSpecRef]: @@ -86,6 +89,12 @@ def _jobset_pod_specs(doc: dict[str, Any]) -> list[PodSpecRef]: return refs +def _rayjob_pod_specs(doc: dict[str, Any]) -> list[PodSpecRef]: + from roar.backends.k8s.rayjob import rayjob_pod_specs + + return rayjob_pod_specs(doc) + + def _pytorchjob_pod_specs(doc: dict[str, Any]) -> list[PodSpecRef]: refs: list[PodSpecRef] = [] replica_specs = _dict_get(doc, "spec").get("pytorchReplicaSpecs") @@ -124,6 +133,14 @@ def _pytorchjob_pod_specs(doc: dict[str, Any]) -> list[PodSpecRef]: api_group="trainer.kubeflow.org", kubectl_resource="trainjobs.trainer.kubeflow.org", locate_pod_specs=None, + rewrite_style="trainer_override", + ), + WorkloadKind( + kind="RayJob", + api_group="ray.io", + kubectl_resource="rayjobs.ray.io", + locate_pod_specs=_rayjob_pod_specs, + rewrite_style="rayjob", ), ) @@ -239,9 +256,16 @@ def rewrite_manifest_for_lineage( config_mount_map=dict(mount_map or {}), ) - if workload.locate_pod_specs is None: + if workload.rewrite_style == "trainer_override": wrapped, skipped = _rewrite_trainjob(doc, workload_name=workload_name, contract=contract) + elif workload.rewrite_style == "rayjob": + from roar.backends.k8s.rayjob import rewrite_rayjob_for_lineage + + wrapped, skipped = rewrite_rayjob_for_lineage( + doc, workload_name=workload_name, contract=contract + ) else: + assert workload.locate_pod_specs is not None wrapped, skipped = _rewrite_pod_specs( workload.locate_pod_specs(doc), workload=workload, diff --git a/roar/backends/k8s/rayjob.py b/roar/backends/k8s/rayjob.py new file mode 100644 index 00000000..8b626736 --- /dev/null +++ b/roar/backends/k8s/rayjob.py @@ -0,0 +1,205 @@ +"""RayJob delegation: k8s manifest plumbing, Ray backend runtime semantics. + +KubeRay overwrites container commands with ``ray start --block`` and runs +user code in Ray worker actors, so the generic command-wrapping adapter +cannot see it. Instead, a RayJob rewrite reuses the Ray backend's proven +instrumentation surface: + +- ``spec.entrypoint`` is wrapped through the Ray driver entrypoint, +- ``spec.runtimeEnvYAML`` gains the roar pip requirement, the worker + setup hook, and the Ray env contract (``ROAR_EXECUTION_BACKEND=ray``), +- fragment-session credentials go into the RayCluster pod templates as + Secret refs (never plaintext in the CR), + +and reconstitution is delegated to the Ray backend's fragment +reconstituter, since the streamed fragments are Ray ``TaskFragment`` +payloads. +""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING, Any + +import yaml # type: ignore[import-untyped] + +from roar.backends.ray.env_contract import merge_worker_bootstrap_env +from roar.backends.ray.submit_context import ( + build_submit_instrumentation_context, + build_submit_source_environ, +) +from roar.execution.framework.contract import ROAR_EXECUTION_BACKEND_ENV + +if TYPE_CHECKING: + from roar.backends.k8s.manifest import _EnvContract + +_DRIVER_ENTRYPOINT_PREFIX = "python -m roar.execution.runtime.driver_entrypoint -- " +_WORKER_SETUP_HOOK = "roar.execution.runtime.worker_bootstrap.startup" + +RAYJOB_SUCCESS_STATUSES = frozenset({"SUCCEEDED"}) +RAYJOB_FAILURE_STATUSES = frozenset({"FAILED", "STOPPED"}) + + +def rewrite_rayjob_for_lineage( + doc: dict[str, Any], + *, + workload_name: str, + contract: _EnvContract, +) -> tuple[list[str], list[str]]: + """Instrument a RayJob document in place; returns (wrapped, skipped).""" + from roar.backends.k8s.manifest import K8sManifestError + + spec = doc.get("spec") + if not isinstance(spec, dict): + raise K8sManifestError(f"RayJob {workload_name} has no spec") + + entrypoint = str(spec.get("entrypoint") or "").strip() + if not entrypoint: + raise K8sManifestError( + f"RayJob {workload_name} has no spec.entrypoint; roar instruments " + "job-mode RayJobs (interactive mode is not supported)" + ) + if _DRIVER_ENTRYPOINT_PREFIX.strip() not in entrypoint: + spec["entrypoint"] = _DRIVER_ENTRYPOINT_PREFIX + entrypoint + + spec["runtimeEnvYAML"] = _merged_runtime_env_yaml( + spec.get("runtimeEnvYAML"), + workload_name=workload_name, + contract=contract, + ) + + wrapped = ["entrypoint"] + for ref in rayjob_pod_specs(doc): + containers = ref.spec.get("containers") + if not isinstance(containers, list): + continue + for container in containers: + if isinstance(container, dict): + _inject_secret_env(container, secret_name=contract.secret_name) + wrapped.append(f"{ref.role}/pods") + return wrapped, [] + + +def rayjob_pod_specs(doc: dict[str, Any]): + """Head/worker pod-spec refs (used for Secret env and attach recovery).""" + from roar.backends.k8s.manifest import PodSpecRef, _dict_at, _dict_get + + refs = [] + cluster_spec = _dict_get(doc.get("spec") or {}, "rayClusterSpec") + head_spec = _dict_at(cluster_spec, ("headGroupSpec", "template", "spec")) + if head_spec is not None: + refs.append(PodSpecRef(role="head", spec=head_spec)) + for group in cluster_spec.get("workerGroupSpecs") or []: + if not isinstance(group, dict): + continue + group_spec = _dict_at(group, ("template", "spec")) + if group_spec is not None: + refs.append(PodSpecRef(role=str(group.get("groupName") or "workers"), spec=group_spec)) + return refs + + +def _merged_runtime_env_yaml( + raw: Any, + *, + workload_name: str, + contract: _EnvContract, +) -> str: + from roar.backends.k8s.manifest import K8sManifestError + + runtime_env: dict[str, Any] = {} + if isinstance(raw, str) and raw.strip(): + try: + loaded = yaml.safe_load(raw) + except yaml.YAMLError as exc: + raise K8sManifestError( + f"RayJob {workload_name} has invalid runtimeEnvYAML: {exc}" + ) from exc + if isinstance(loaded, dict): + runtime_env = loaded + + pip = runtime_env.get("pip") + packages: list[str] = [] + if isinstance(pip, dict): + packages = [str(item) for item in pip.get("packages") or []] + elif isinstance(pip, list): + packages = [str(item) for item in pip] + if contract.requirement and contract.requirement not in packages: + packages.append(contract.requirement) + runtime_env["pip"] = packages + + runtime_env["worker_process_setup_hook"] = _WORKER_SETUP_HOOK + + # The Ray env contract, keyed to the k8s parent uid so Ray fragments + # attach to the recorded submit job. + context = build_submit_instrumentation_context( + os.environ, + cwd=os.getcwd(), + host_glaas_url=None, + job_id=contract.parent_job_uid, + ) + env_vars = merge_worker_bootstrap_env( + dict(runtime_env.get("env_vars") or {}), + build_submit_source_environ(context), + job_id=context.job_id, + overwrite_existing=True, + ) + env_vars[ROAR_EXECUTION_BACKEND_ENV] = "ray" + env_vars["GLAAS_URL"] = contract.cluster_glaas_url + # Per the proxy decision, node agents/proxy sidecars stay off for + # RayJob delegation v1; in-process hooks are the capture surface. + env_vars["ROAR_RAY_NODE_AGENTS"] = "0" + # No proxy runs in the pods, so the merged local-proxy redirect would + # point user S3 traffic at a dead localhost port — strip it. + env_vars.pop("AWS_ENDPOINT_URL", None) + env_vars.pop("ROAR_PROXY_PORT", None) + # ROAR_WRAP is deliberately NOT set: in Ray pip virtualenvs the + # roar_inject.pth fires at worker interpreter startup before the + # virtualenv's site-packages are importable, and the sitecustomize + # ABI-repair path then blows the worker registration timeout + # (observed live: supervisor start -> hang -> kill -> retry forever). + # worker_process_setup_hook is the capture surface for RayJob v1. + env_vars.pop("ROAR_WRAP", None) + # Credentials come from the pod-level Secret refs, never the CR. + env_vars.pop("ROAR_SESSION_ID", None) + env_vars.pop("ROAR_FRAGMENT_TOKEN", None) + runtime_env["env_vars"] = env_vars + + return yaml.safe_dump(runtime_env, sort_keys=False) + + +def _inject_secret_env(container: dict[str, Any], *, secret_name: str) -> None: + env = container.setdefault("env", []) + if not isinstance(env, list): + return + existing = { + str(entry.get("name")) for entry in env if isinstance(entry, dict) and entry.get("name") + } + for name, key in (("ROAR_SESSION_ID", "session_id"), ("ROAR_FRAGMENT_TOKEN", "token")): + if name not in existing: + env.append( + { + "name": name, + "valueFrom": {"secretKeyRef": {"name": secret_name, "key": key}}, + } + ) + + +def rayjob_terminal_status(document: dict[str, Any]) -> tuple[bool | None, str]: + """RayJob signals completion via status.jobStatus, not conditions.""" + status = document.get("status") + job_status = str((status or {}).get("jobStatus") or "").strip().upper() + if job_status in RAYJOB_SUCCESS_STATUSES: + return True, job_status + if job_status in RAYJOB_FAILURE_STATUSES: + message = str((status or {}).get("message") or job_status) + return False, message + return None, "" + + +__all__ = [ + "RAYJOB_FAILURE_STATUSES", + "RAYJOB_SUCCESS_STATUSES", + "rayjob_pod_specs", + "rayjob_terminal_status", + "rewrite_rayjob_for_lineage", +] diff --git a/roar/backends/k8s/submit.py b/roar/backends/k8s/submit.py index 4f79777a..699d09b9 100644 --- a/roar/backends/k8s/submit.py +++ b/roar/backends/k8s/submit.py @@ -154,11 +154,21 @@ def plan_kubectl_job_submit_command(command: list[str]) -> ExecutionCommandPlan: _write_submit_context(prepared_path, submit_context) rewritten_command = _replace_filename_argument(command, str(prepared_path)) + + # RayJob pods stream Ray TaskFragments, so reconstitution is delegated + # to the Ray backend's reconstituter rather than the k8s one. + finalize_run = None + if rewrite.workload_kind == "RayJob": + from roar.execution.fragments.reconstitution import build_submit_finalizer + + finalize_run = build_submit_finalizer("ray", session["session_id"]) + return ExecutionCommandPlan( backend_name="k8s", command=rewritten_command, execution_role="submit", session_id=session["session_id"], + finalize_run=finalize_run, ) diff --git a/roar/backends/k8s/workload_wait.py b/roar/backends/k8s/workload_wait.py index 7f1cd18c..e09c14d5 100644 --- a/roar/backends/k8s/workload_wait.py +++ b/roar/backends/k8s/workload_wait.py @@ -58,8 +58,14 @@ def terminal_condition(document: dict[str, Any]) -> tuple[bool | None, str]: """Return (succeeded, message) from workload status conditions. ``succeeded`` is None while the workload is still running. Condition - type names are unioned across Job/JobSet/PyTorchJob/TrainJob. + type names are unioned across Job/JobSet/PyTorchJob/TrainJob; RayJob + signals through ``status.jobStatus`` instead of conditions. """ + if str(document.get("kind") or "") == "RayJob": + from roar.backends.k8s.rayjob import rayjob_terminal_status + + return rayjob_terminal_status(document) + status = document.get("status") conditions = status.get("conditions") if isinstance(status, dict) else None for condition in conditions or []: diff --git a/roar/execution/framework/registry.py b/roar/execution/framework/registry.py index b3ddaf2b..f083041e 100644 --- a/roar/execution/framework/registry.py +++ b/roar/execution/framework/registry.py @@ -74,9 +74,43 @@ def get_execution_backend(name: str) -> ExecutionBackend: for backend in _registered_execution_backends: if backend.name == normalized_name: return backend + + # A miss with skipped builtin imports may be a transient early-import + # failure, not a permanent one: sitecustomize (ROAR_WRAP=1) can trigger + # discovery at interpreter startup before the runtime environment's + # site-packages are fully importable (observed in Ray pip virtualenvs, + # where roar's deps resolve fine moments later in the worker setup + # hook). Retry the skipped imports once per lookup before giving up. + if _skipped_builtin_backend_imports: + _retry_skipped_builtin_backend_imports() + for backend in _registered_execution_backends: + if backend.name == normalized_name: + return backend + raise LookupError(f"unknown execution backend: {normalized_name or ''}") +def _retry_skipped_builtin_backend_imports() -> None: + for module_name in list(_skipped_builtin_backend_imports): + if module_name not in _BUILTIN_EXECUTION_BACKEND_MODULES: + continue + try: + module = importlib.import_module(module_name) + except ImportError as exc: + _skipped_builtin_backend_imports[module_name] = str(exc) + continue + register = getattr(module, "register", None) + if not callable(register): + continue + try: + _register_entrypoint_payload(register) + except Exception as exc: + _skipped_builtin_backend_imports[f"{module_name}:register"] = str(exc) + continue + _skipped_builtin_backend_imports.pop(module_name, None) + _skipped_builtin_backend_imports.pop(f"{module_name}:register", None) + + def match_execution_backend_for_module(module_name: str) -> ExecutionBackend | None: _ensure_execution_backends_discovered() normalized_name = str(module_name or "").strip() diff --git a/tests/backends/k8s/README.md b/tests/backends/k8s/README.md index 165d6b06..7d041c44 100644 --- a/tests/backends/k8s/README.md +++ b/tests/backends/k8s/README.md @@ -25,6 +25,13 @@ Three test layers share this harness: `k8s_task` jobs keyed by pod UID) and mount-map rewriting (a hostPath volume standing in for a FUSE mount, declared via `[k8s.mount_map]`, rewrites to `s3://` URIs at reconstitution). +- `e2e/test_k8s_operators.py` — live Kubeflow operator validation: + PyTorchJob v1 (Master+Worker through the real training-operator) and + TrainJob v2 (real TrainJob→JobSet pipeline via a slim-image clone of the + shipped torch-distributed runtime). Skips without `--with-kubeflow`. +- `e2e/test_k8s_rayjob.py` — live RayJob delegation smoke on KubeRay (the + heaviest test: multi-GB Ray image + per-job pip env). Skips without + `--with-kuberay`. - `e2e/test_k8s_smoke.py` — the Phase-0 runtime diagnostic: fixtures hand-wrap the manifest (no backend involved) to isolate the runtime pieces (in-pod tracing, fragment streaming, identity contract) when the product @@ -47,8 +54,9 @@ live in `unit/` and run in the default gate — no cluster needed. bash scripts/build_wheel_with_bins.sh # create cluster + wire glaas + preflight -# (--with-minio for S3 scenarios, --with-jobset for the JobSet operator e2e) -bash tests/backends/k8s/scripts/bootstrap_k8s.sh --with-jobset --with-minio +# (--with-minio: S3 scenarios; --with-jobset: JobSet e2e; --with-kubeflow: +# PyTorchJob/TrainJob e2e; --with-kuberay: RayJob delegation e2e) +bash tests/backends/k8s/scripts/bootstrap_k8s.sh --with-jobset --with-minio --with-kubeflow --with-kuberay # run the smoke tests (addopts override needed: e2e dirs are ignored by default) pytest tests/backends/k8s/e2e -o addopts='' -m k8s_e2e -v diff --git a/tests/backends/k8s/e2e/test_k8s_distributed.py b/tests/backends/k8s/e2e/test_k8s_distributed.py index 8d8d8ca4..22660617 100644 --- a/tests/backends/k8s/e2e/test_k8s_distributed.py +++ b/tests/backends/k8s/e2e/test_k8s_distributed.py @@ -252,7 +252,7 @@ def _roar_env() -> dict[str, str]: return env -def _roar(args: list[str], *, cwd: Path) -> subprocess.CompletedProcess[str]: +def _roar(args: list[str], *, cwd: Path, timeout: int = 700) -> subprocess.CompletedProcess[str]: return subprocess.run( [sys.executable, "-m", "roar", *args], cwd=cwd, @@ -260,14 +260,17 @@ def _roar(args: list[str], *, cwd: Path) -> subprocess.CompletedProcess[str]: capture_output=True, text=True, check=False, - timeout=700, + timeout=timeout, ) -def _submit(manifest_name: str, *, cwd: Path) -> subprocess.CompletedProcess[str]: +def _submit( + manifest_name: str, *, cwd: Path, timeout: int = 700 +) -> subprocess.CompletedProcess[str]: return _roar( ["run", "kubectl", "apply", "--context", KUBE_CONTEXT, "-f", manifest_name], cwd=cwd, + timeout=timeout, ) diff --git a/tests/backends/k8s/e2e/test_k8s_operators.py b/tests/backends/k8s/e2e/test_k8s_operators.py new file mode 100644 index 00000000..6847eb9a --- /dev/null +++ b/tests/backends/k8s/e2e/test_k8s_operators.py @@ -0,0 +1,315 @@ +"""Live operator validation: PyTorchJob (training-operator v1) and TrainJob (trainer v2). + +Requires ``bootstrap_k8s.sh --with-kubeflow``; tests skip when the CRDs are +absent. Workloads use slim Python images — the operators' pod/env wiring is +what's under test (per-role identity, RANK/PET_NODE_RANK index chains, real +terminal conditions), not torch itself. The TrainJob test clones the shipped +``torch-distributed`` ClusterTrainingRuntime with a slim image so the real +TrainJob -> JobSet -> pod pipeline runs without multi-GB pulls. +""" + +from __future__ import annotations + +import json +import sqlite3 +import uuid +from pathlib import Path +from typing import Any + +import pytest + +from .conftest import NAMESPACE, kubectl +from .test_k8s_distributed import ( + PINNED_NODE, + _cleanup, + _describe, + _indent_script, + _pod_logs, + _submit, + _write_project, +) +from .test_k8s_product_path import wheel_server # noqa: F401 (module fixture reuse) + +pytestmark = [ + pytest.mark.e2e, + pytest.mark.k8s_e2e, + pytest.mark.timeout(900), +] + +PT_WORKER_SCRIPT = """\ +import json +import os +import time +from pathlib import Path + +rank = os.environ.get("RANK") or os.environ.get("PET_NODE_RANK") or "0" +Path(f"out-rank-{rank}.json").write_text(json.dumps({"rank": rank})) +if rank == "0": + deadline = time.time() + 180 + while not Path("/state/worker-done").exists(): + if time.time() > deadline: + raise SystemExit("timed out waiting for worker") + time.sleep(2) + time.sleep(2) +else: + Path("/state/worker-done").write_text("done") +print(f"pytorch replica rank {rank} done") +""" + +PYTORCHJOB_TEMPLATE = """\ +apiVersion: v1 +kind: ConfigMap +metadata: + name: {name}-worker + namespace: {namespace} +data: + pt_worker.py: | +{worker_script} +--- +apiVersion: kubeflow.org/v1 +kind: PyTorchJob +metadata: + name: {name} + namespace: {namespace} + labels: + app.kubernetes.io/part-of: roar-k8s-e2e +spec: + pytorchReplicaSpecs: + Master: + replicas: 1 + restartPolicy: Never + template: + spec: + nodeSelector: + kubernetes.io/hostname: {pinned_node} + volumes: + - name: work + emptyDir: {{}} + - name: state + hostPath: + path: /var/tmp/roar-e2e-pt-{name} + type: DirectoryOrCreate + - name: workload + configMap: + name: {name}-worker + containers: + - name: pytorch + image: python:3.12-slim + workingDir: /work + volumeMounts: + - name: work + mountPath: /work + - name: state + mountPath: /state + - name: workload + mountPath: /workload + readOnly: true + command: ["python", "/workload/pt_worker.py"] + Worker: + replicas: 1 + restartPolicy: Never + template: + spec: + nodeSelector: + kubernetes.io/hostname: {pinned_node} + volumes: + - name: work + emptyDir: {{}} + - name: state + hostPath: + path: /var/tmp/roar-e2e-pt-{name} + type: DirectoryOrCreate + - name: workload + configMap: + name: {name}-worker + containers: + - name: pytorch + image: python:3.12-slim + workingDir: /work + volumeMounts: + - name: work + mountPath: /work + - name: state + mountPath: /state + - name: workload + mountPath: /workload + readOnly: true + command: ["python", "/workload/pt_worker.py"] +""" + +TRAINJOB_TEMPLATE = """\ +apiVersion: trainer.kubeflow.org/v1alpha1 +kind: TrainJob +metadata: + name: {name} + namespace: {namespace} + labels: + app.kubernetes.io/part-of: roar-k8s-e2e +spec: + runtimeRef: + name: {runtime} + trainer: + numNodes: 2 + command: + - python + - -c + - >- + import json, os; + rank = os.environ.get('PET_NODE_RANK') or os.environ.get('JOB_COMPLETION_INDEX') or '0'; + open('node-' + rank + '.json', 'w').write(json.dumps({{'rank': rank}})); + open('model.bin', 'wb').write(rank.encode() * 8) +""" + + +def _query(project_dir: Path, sql: str, params: tuple = ()) -> list[sqlite3.Row]: + conn = sqlite3.connect(project_dir / ".roar" / "roar.db") + conn.row_factory = sqlite3.Row + try: + return conn.execute(sql, params).fetchall() + finally: + conn.close() + + +def _crd_available(crd: str) -> bool: + return kubectl(["get", "crd", crd], check=False).returncode == 0 + + +def _task_rows(project_dir: Path) -> list[dict[str, Any]]: + rows = _query( + project_dir, + "SELECT id, metadata FROM jobs WHERE job_type = 'k8s_task' ORDER BY id", + ) + return [{"id": int(row["id"]), "metadata": json.loads(row["metadata"] or "{}")} for row in rows] + + +def test_pytorchjob_live_captures_master_and_worker( + k8s_cluster: None, + glaas_health: str, + wheel_server: dict[str, str], # noqa: F811 + tmp_path_factory: pytest.TempPathFactory, +) -> None: + if not _crd_available("pytorchjobs.kubeflow.org"): + pytest.skip("training-operator v1 not installed; run bootstrap_k8s.sh --with-kubeflow") + + project_dir = tmp_path_factory.mktemp("k8s-ptjob") + _write_project(project_dir, wheel_server["url"]) + + name = f"roar-pt-{uuid.uuid4().hex[:6]}" + (project_dir / "job.yaml").write_text( + PYTORCHJOB_TEMPLATE.format( + name=name, + namespace=NAMESPACE, + pinned_node=PINNED_NODE, + worker_script=_indent_script(PT_WORKER_SCRIPT), + ), + encoding="utf-8", + ) + + try: + completed = _submit("job.yaml", cwd=project_dir) + logs = _pod_logs(f"training.kubeflow.org/job-name={name}") + run = { + "exit_code": completed.returncode, + "stdout": completed.stdout, + "stderr": completed.stderr, + "pod_logs": logs, + } + assert completed.returncode == 0, _describe(run) + + tasks = _task_rows(project_dir) + assert len(tasks) == 2, _describe(run) + + task_names = {str(task["metadata"].get("k8s_task_id") or "") for task in tasks} + indices = {task_id.split(":")[2] for task_id in task_names if task_id.count(":") >= 3} + assert indices == {"0", "1"}, ( + f"expected RANK-derived node indices 0/1, got {task_names}\n{_describe(run)}" + ) + + # Per-role task names flow from the replica-spec adapter. + all_metadata = json.dumps([task["metadata"] for task in tasks]) + assert f"{name}/Master/pytorch" in all_metadata or "Master" in all_metadata + assert "Worker" in all_metadata, all_metadata + finally: + _cleanup(name, resource="pytorchjob") + import subprocess + + subprocess.run( + ["docker", "exec", PINNED_NODE, "rm", "-rf", f"/var/tmp/roar-e2e-pt-{name}"], + check=False, + capture_output=True, + ) + + +@pytest.fixture() +def slim_training_runtime() -> Any: + """Clone the shipped torch-distributed runtime with a slim image.""" + runtime_name = "roar-e2e-slim" + result = kubectl( + ["get", "clustertrainingruntime", "torch-distributed", "-o", "json"], + check=False, + ) + if result.returncode != 0: + pytest.skip("trainer v2 runtimes not installed; run bootstrap_k8s.sh --with-kubeflow") + + runtime = json.loads(result.stdout) + runtime["metadata"] = {"name": runtime_name} + container = runtime["spec"]["template"]["spec"]["replicatedJobs"][0]["template"]["spec"][ + "template" + ]["spec"]["containers"][0] + container["image"] = "python:3.12-slim" + + kubectl(["apply", "-f", "-"], input_text=json.dumps(runtime)) + try: + yield runtime_name + finally: + kubectl( + ["delete", "clustertrainingruntime", runtime_name, "--ignore-not-found"], + check=False, + ) + + +def test_trainjob_live_runs_through_jobset_pipeline( + k8s_cluster: None, + glaas_health: str, + wheel_server: dict[str, str], # noqa: F811 + slim_training_runtime: str, + tmp_path_factory: pytest.TempPathFactory, +) -> None: + if not _crd_available("trainjobs.trainer.kubeflow.org"): + pytest.skip("trainer v2 not installed; run bootstrap_k8s.sh --with-kubeflow") + + project_dir = tmp_path_factory.mktemp("k8s-trainjob") + _write_project(project_dir, wheel_server["url"]) + + name = f"roar-tj-{uuid.uuid4().hex[:6]}" + (project_dir / "job.yaml").write_text( + TRAINJOB_TEMPLATE.format( + name=name, + namespace=NAMESPACE, + runtime=slim_training_runtime, + ), + encoding="utf-8", + ) + + try: + completed = _submit("job.yaml", cwd=project_dir) + logs = _pod_logs(f"jobset.sigs.k8s.io/jobset-name={name}") + run = { + "exit_code": completed.returncode, + "stdout": completed.stdout, + "stderr": completed.stderr, + "pod_logs": logs, + } + assert completed.returncode == 0, _describe(run) + + tasks = _task_rows(project_dir) + assert len(tasks) == 2, _describe(run) + + task_ids = {str(task["metadata"].get("k8s_task_id") or "") for task in tasks} + indices = {task_id.split(":")[2] for task_id in task_ids if task_id.count(":") >= 3} + assert indices == {"0", "1"}, ( + f"expected completion indices 0/1 across TrainJob nodes, got {task_ids}" + ) + assert all(":node:" in task_id for task_id in task_ids), task_ids + finally: + _cleanup(name, resource="trainjob") diff --git a/tests/backends/k8s/e2e/test_k8s_rayjob.py b/tests/backends/k8s/e2e/test_k8s_rayjob.py new file mode 100644 index 00000000..95909876 --- /dev/null +++ b/tests/backends/k8s/e2e/test_k8s_rayjob.py @@ -0,0 +1,229 @@ +"""Live RayJob delegation smoke on KubeRay. + +Requires ``bootstrap_k8s.sh --with-kuberay``; skips when the CRD is +absent. The k8s backend rewrites the RayJob (entrypoint through the Ray +driver entrypoint, runtime env with the roar wheel + Ray env contract, +Secret refs into the cluster pods) and delegates reconstitution to the +Ray backend — merged lineage lands as ``ray_task`` jobs under the k8s +submit node. + +This is the heaviest e2e in the harness (a multi-GB Ray image plus a +per-job pip env); everything pins to one node so the image pulls once. +""" + +from __future__ import annotations + +import sqlite3 +import subprocess +import uuid +from pathlib import Path + +import pytest + +from .conftest import NAMESPACE, kubectl +from .test_k8s_distributed import ( + PINNED_NODE, + _cleanup, + _describe, + _indent_script, + _pod_logs, + _submit, + _write_project, +) +from .test_k8s_product_path import wheel_server # noqa: F401 (module fixture reuse) + +pytestmark = [ + pytest.mark.e2e, + pytest.mark.k8s_e2e, + pytest.mark.timeout(1500), +] + +RAY_IMAGE = "rayproject/ray:2.46.0-py312-cpu" + +RAY_TRAIN_SCRIPT = """\ +from pathlib import Path + +import ray + +ray.init() + + +@ray.remote +def train_shard() -> int: + Path("/home/ray/task-out.bin").write_bytes(b"trained" * 4) + return 1 + + +assert ray.get(train_shard.remote()) == 1 +Path("/home/ray/driver-out.txt").write_text("done") +print("ray train done") +""" + +RAYJOB_TEMPLATE = """\ +apiVersion: v1 +kind: ConfigMap +metadata: + name: {name}-worker + namespace: {namespace} +data: + ray_train.py: | +{worker_script} +--- +apiVersion: ray.io/v1 +kind: RayJob +metadata: + name: {name} + namespace: {namespace} + labels: + app.kubernetes.io/part-of: roar-k8s-e2e +spec: + entrypoint: python /workload/ray_train.py + shutdownAfterJobFinishes: true + ttlSecondsAfterFinished: 1800 + submitterPodTemplate: + spec: + restartPolicy: Never + nodeSelector: + kubernetes.io/hostname: {pinned_node} + containers: + - name: ray-job-submitter + image: {ray_image} + rayClusterSpec: + rayVersion: "2.46.0" + headGroupSpec: + rayStartParams: + num-cpus: "1" + template: + spec: + nodeSelector: + kubernetes.io/hostname: {pinned_node} + volumes: + - name: workload + configMap: + name: {name}-worker + containers: + - name: ray-head + image: {ray_image} + volumeMounts: + - name: workload + mountPath: /workload + readOnly: true + resources: + requests: + cpu: 500m + memory: 1Gi + workerGroupSpecs: + - groupName: workers + replicas: 1 + rayStartParams: + num-cpus: "1" + template: + spec: + nodeSelector: + kubernetes.io/hostname: {pinned_node} + volumes: + - name: workload + configMap: + name: {name}-worker + containers: + - name: ray-worker + image: {ray_image} + volumeMounts: + - name: workload + mountPath: /workload + readOnly: true + resources: + requests: + cpu: 500m + memory: 1Gi +""" + + +def _query(project_dir: Path, sql: str, params: tuple = ()) -> list[sqlite3.Row]: + conn = sqlite3.connect(project_dir / ".roar" / "roar.db") + conn.row_factory = sqlite3.Row + try: + return conn.execute(sql, params).fetchall() + finally: + conn.close() + + +def test_rayjob_live_delegates_to_ray_backend( + k8s_cluster: None, + glaas_health: str, + wheel_server: dict[str, str], # noqa: F811 + tmp_path_factory: pytest.TempPathFactory, +) -> None: + if kubectl(["get", "crd", "rayjobs.ray.io"], check=False).returncode != 0: + pytest.skip("KubeRay not installed; run bootstrap_k8s.sh --with-kuberay") + + # One node, one pull: the Ray image is multi-GB. + pull = subprocess.run( + ["docker", "exec", PINNED_NODE, "crictl", "pull", f"docker.io/{RAY_IMAGE}"], + capture_output=True, + text=True, + timeout=900, + check=False, + ) + assert pull.returncode == 0, f"ray image pre-pull failed: {pull.stderr}" + + project_dir = tmp_path_factory.mktemp("k8s-rayjob") + _write_project(project_dir, wheel_server["url"]) + config_path = project_dir / ".roar" / "config.toml" + config_path.write_text( + config_path.read_text(encoding="utf-8").replace( + "wait_timeout_seconds = 420", + "wait_timeout_seconds = 900", + ), + encoding="utf-8", + ) + + name = f"roar-ray-{uuid.uuid4().hex[:6]}" + (project_dir / "rayjob.yaml").write_text( + RAYJOB_TEMPLATE.format( + name=name, + namespace=NAMESPACE, + pinned_node=PINNED_NODE, + ray_image=RAY_IMAGE, + worker_script=_indent_script(RAY_TRAIN_SCRIPT), + ), + encoding="utf-8", + ) + + try: + completed = _submit("rayjob.yaml", cwd=project_dir, timeout=1200) + logs = _pod_logs(f"ray.io/originated-from-cr-name={name}") + run = { + "exit_code": completed.returncode, + "stdout": completed.stdout, + "stderr": completed.stderr, + "pod_logs": logs, + } + assert completed.returncode == 0, _describe(run) + + submit_rows = _query( + project_dir, + "SELECT job_uid FROM jobs WHERE execution_backend = 'k8s' " + "AND execution_role = 'submit'", + ) + assert len(submit_rows) == 1, _describe(run) + + ray_tasks = _query( + project_dir, + "SELECT id FROM jobs WHERE job_type = 'ray_task'", + ) + assert ray_tasks, f"expected ray_task jobs from delegated reconstitution\n{_describe(run)}" + + # Delegation scope proven here: rewrite, Secret hygiene, in-pod pip + # bootstrap, fragment streaming from KubeRay pods, jobStatus wait, + # and Ray-reconstituter merge under the submit node. + # + # KNOWN GAP (Ray-backend follow-up, not k8s plumbing): per-task I/O + # fidelity inside KubeRay workers. Under setup-hook-only bootstrap + # (ROAR_WRAP intentionally off — see rayjob.py), fragments arrive as + # ray_task shells without file refs; the roar_worker task-boundary + # capture expects the py_executable/packaged-working_dir machinery + # of the native `roar run ray job submit` path. Once that lands, + # tighten this to assert task-out.bin in the ray_task outputs. + finally: + _cleanup(name, resource="rayjob") diff --git a/tests/backends/k8s/scripts/bootstrap_k8s.sh b/tests/backends/k8s/scripts/bootstrap_k8s.sh index 68044435..38f4b474 100755 --- a/tests/backends/k8s/scripts/bootstrap_k8s.sh +++ b/tests/backends/k8s/scripts/bootstrap_k8s.sh @@ -37,14 +37,21 @@ HOST_GLAAS_URL="${HOST_GLAAS_URL:-http://localhost:3001}" CLUSTER_GLAAS_PORT=3001 JOBSET_VERSION="v0.12.0" +TRAINING_OPERATOR_V1_REF="v1.9.2" +TRAINER_V2_REF="v2.2.1" +KUBERAY_VERSION="v1.4.0" WITH_MINIO=0 WITH_JOBSET=0 +WITH_KUBEFLOW=0 +WITH_KUBERAY=0 SKIP_GLAAS=0 for arg in "$@"; do case "$arg" in --with-minio) WITH_MINIO=1 ;; --with-jobset) WITH_JOBSET=1 ;; + --with-kubeflow) WITH_KUBEFLOW=1 ;; + --with-kuberay) WITH_KUBERAY=1 ;; --skip-glaas) SKIP_GLAAS=1 ;; *) echo "error: unknown flag: $arg" >&2 @@ -185,6 +192,30 @@ if ((WITH_JOBSET == 1)); then kubectl_ctx -n jobset-system rollout status deployment/jobset-controller-manager --timeout=300s fi +if ((WITH_KUBEFLOW == 1)); then + echo "▶ Installing Kubeflow training-operator v1 (${TRAINING_OPERATOR_V1_REF})" + kubectl_ctx apply --server-side \ + -k "github.com/kubeflow/training-operator/manifests/overlays/standalone?ref=${TRAINING_OPERATOR_V1_REF}" + kubectl_ctx -n kubeflow rollout status deployment/training-operator --timeout=300s + + echo "▶ Installing Kubeflow trainer v2 (${TRAINER_V2_REF})" + kubectl_ctx apply --server-side \ + -k "github.com/kubeflow/trainer.git/manifests/overlays/manager?ref=${TRAINER_V2_REF}" + kubectl_ctx -n kubeflow-system rollout status deployment/kubeflow-trainer-controller-manager --timeout=300s + echo "▶ Installing Kubeflow trainer v2 runtimes" + kubectl_ctx apply --server-side \ + -k "github.com/kubeflow/trainer.git/manifests/overlays/runtimes?ref=${TRAINER_V2_REF}" || \ + echo "warning: trainer runtimes overlay failed; TrainJob e2e will create its own runtime" >&2 +fi + +if ((WITH_KUBERAY == 1)); then + echo "▶ Installing KubeRay operator ${KUBERAY_VERSION}" + kubectl_ctx create -k "github.com/ray-project/kuberay/ray-operator/config/default?ref=${KUBERAY_VERSION}" 2>/dev/null || \ + kubectl_ctx apply --server-side -k "github.com/ray-project/kuberay/ray-operator/config/default?ref=${KUBERAY_VERSION}" + kubectl_ctx -n ray-system rollout status deployment/kuberay-operator --timeout=300s || \ + kubectl_ctx -n default rollout status deployment/kuberay-operator --timeout=300s +fi + echo echo "✓ Harness ready" echo " cluster: $CLUSTER_NAME (context $KUBE_CONTEXT)" diff --git a/tests/backends/k8s/unit/test_rayjob.py b/tests/backends/k8s/unit/test_rayjob.py new file mode 100644 index 00000000..a37a1f9b --- /dev/null +++ b/tests/backends/k8s/unit/test_rayjob.py @@ -0,0 +1,165 @@ +from __future__ import annotations + +import copy +from typing import Any + +import pytest +import yaml # type: ignore[import-untyped] + +from roar.backends.k8s.manifest import ( + K8sManifestError, + rewrite_manifest_for_lineage, + workload_kind_for_document, +) +from roar.backends.k8s.rayjob import rayjob_terminal_status + +RAYJOB_MANIFEST = { + "apiVersion": "ray.io/v1", + "kind": "RayJob", + "metadata": {"name": "ray-train", "namespace": "ml"}, + "spec": { + "entrypoint": "python train.py --epochs 3", + "runtimeEnvYAML": "pip:\n - pandas\nenv_vars:\n USER_VAR: keep\n", + "rayClusterSpec": { + "headGroupSpec": { + "rayStartParams": {}, + "template": { + "spec": {"containers": [{"name": "ray-head", "image": "rayproject/ray:2.46.0"}]} + }, + }, + "workerGroupSpecs": [ + { + "groupName": "workers", + "replicas": 1, + "template": { + "spec": { + "containers": [{"name": "ray-worker", "image": "rayproject/ray:2.46.0"}] + } + }, + } + ], + }, + }, +} + + +def _rewrite(documents: list[dict[str, Any]], **overrides: Any): + kwargs: dict[str, Any] = { + "secret_name": "roar-fragment-abc", + "session_id": "session-1", + "fragment_token": "ff" * 32, + "requirement": "roar-cli==0.3.7", + "cluster_glaas_url": "http://glaas:3001", + "tracer": "preload", + "parent_job_uid": "cafe0123", + } + kwargs.update(overrides) + return rewrite_manifest_for_lineage(documents, **kwargs) + + +def test_rayjob_recognized_as_workload() -> None: + workload = workload_kind_for_document(RAYJOB_MANIFEST) + assert workload is not None + assert workload.kind == "RayJob" + assert workload.kubectl_resource == "rayjobs.ray.io" + assert workload.rewrite_style == "rayjob" + + +def test_rayjob_entrypoint_wrapped_through_ray_driver() -> None: + rewrite = _rewrite([copy.deepcopy(RAYJOB_MANIFEST)]) + doc = next(d for d in rewrite.documents if d.get("kind") == "RayJob") + + assert doc["spec"]["entrypoint"] == ( + "python -m roar.execution.runtime.driver_entrypoint -- python train.py --epochs 3" + ) + assert rewrite.workload_kind == "RayJob" + assert "entrypoint" in rewrite.wrapped_containers + + +def test_rayjob_runtime_env_carries_ray_contract_without_secrets() -> None: + rewrite = _rewrite([copy.deepcopy(RAYJOB_MANIFEST)]) + doc = next(d for d in rewrite.documents if d.get("kind") == "RayJob") + runtime_env = yaml.safe_load(doc["spec"]["runtimeEnvYAML"]) + + assert "roar-cli==0.3.7" in runtime_env["pip"] + assert "pandas" in runtime_env["pip"] + assert ( + runtime_env["worker_process_setup_hook"] + == "roar.execution.runtime.worker_bootstrap.startup" + ) + + env_vars = runtime_env["env_vars"] + assert env_vars["USER_VAR"] == "keep" + assert env_vars["ROAR_EXECUTION_BACKEND"] == "ray" + assert env_vars["ROAR_JOB_ID"] == "cafe0123" + assert env_vars["GLAAS_URL"] == "http://glaas:3001" + assert env_vars["ROAR_RAY_NODE_AGENTS"] == "0" + # Credentials must never appear inline in the CR. + assert "ROAR_SESSION_ID" not in env_vars + assert "ROAR_FRAGMENT_TOKEN" not in env_vars + # No proxy runs in the pods: the local-proxy redirect must be stripped + # or user S3 traffic would hit a dead localhost port. + assert "AWS_ENDPOINT_URL" not in env_vars + assert "ROAR_PROXY_PORT" not in env_vars + # ROAR_WRAP stays off for RayJob v1 (pip-virtualenv .pth timing blows + # the worker registration timeout); the setup hook captures instead. + assert "ROAR_WRAP" not in env_vars + + +def test_rayjob_pods_get_secret_refs_and_secret_doc() -> None: + rewrite = _rewrite([copy.deepcopy(RAYJOB_MANIFEST)]) + doc = next(d for d in rewrite.documents if d.get("kind") == "RayJob") + + cluster = doc["spec"]["rayClusterSpec"] + pod_containers = [ + cluster["headGroupSpec"]["template"]["spec"]["containers"][0], + cluster["workerGroupSpecs"][0]["template"]["spec"]["containers"][0], + ] + for container in pod_containers: + env = {entry["name"]: entry for entry in container["env"]} + assert env["ROAR_SESSION_ID"]["valueFrom"]["secretKeyRef"]["name"] == "roar-fragment-abc" + assert env["ROAR_FRAGMENT_TOKEN"]["valueFrom"]["secretKeyRef"]["key"] == "token" + + secret = next(d for d in rewrite.documents if d.get("kind") == "Secret") + assert secret["metadata"]["namespace"] == "ml" + + +def test_rayjob_without_entrypoint_fails_actionably() -> None: + manifest = copy.deepcopy(RAYJOB_MANIFEST) + del manifest["spec"]["entrypoint"] + + with pytest.raises(K8sManifestError, match="entrypoint"): + _rewrite([manifest]) + + +def test_rayjob_terminal_status_mapping() -> None: + assert rayjob_terminal_status({"status": {"jobStatus": "SUCCEEDED"}}) == (True, "SUCCEEDED") + succeeded, message = rayjob_terminal_status( + {"status": {"jobStatus": "FAILED", "message": "boom"}} + ) + assert succeeded is False and message == "boom" + assert rayjob_terminal_status({"status": {"jobStatus": "RUNNING"}})[0] is None + assert rayjob_terminal_status({})[0] is None + + +def test_rayjob_plan_delegates_finalizer_to_ray( + job_manifest_path: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + import roar.backends.k8s.submit as submit_module + + monkeypatch.setattr( + submit_module, + "_register_fragment_session", + lambda *args, **kwargs: None, + ) + rayjob_path = job_manifest_path.parent / "rayjob.yaml" + rayjob_path.write_text(yaml.safe_dump(RAYJOB_MANIFEST, sort_keys=False), encoding="utf-8") + + plan = submit_module.plan_kubectl_job_submit_command( + ["kubectl", "apply", "-f", str(rayjob_path)] + ) + assert plan.backend_name == "k8s" + assert plan.session_id is not None + # RayJob reconstitution is delegated to the Ray backend explicitly, so + # the framework must not attach the k8s finalizer. + assert plan.finalize_run is not None diff --git a/tests/backends/k8s/unit/test_registry_retry.py b/tests/backends/k8s/unit/test_registry_retry.py new file mode 100644 index 00000000..32d9ec42 --- /dev/null +++ b/tests/backends/k8s/unit/test_registry_retry.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from roar.execution.framework import registry + + +def test_lookup_retries_transiently_skipped_builtin_imports() -> None: + """A backend whose plugin import failed early must recover on lookup. + + Sitecustomize (ROAR_WRAP=1) can trigger discovery at interpreter startup + before the runtime env's site-packages are importable (observed in Ray + pip virtualenvs); the failure must not poison the registry for the + lifetime of the process. + """ + registry._ensure_execution_backends_discovered() + saved_backends = list(registry._registered_execution_backends) + saved_skipped = dict(registry._skipped_builtin_backend_imports) + try: + # Simulate the poisoned state: the k8s plugin import "failed" early. + registry._registered_execution_backends[:] = [ + backend for backend in saved_backends if backend.name != "k8s" + ] + registry._skipped_builtin_backend_imports.clear() + registry._skipped_builtin_backend_imports["roar.backends.k8s.plugin"] = ( + "No module named 'cryptography'" + ) + + backend = registry.get_execution_backend("k8s") + assert backend.name == "k8s" + assert "roar.backends.k8s.plugin" not in registry._skipped_builtin_backend_imports + finally: + registry._registered_execution_backends[:] = saved_backends + registry._skipped_builtin_backend_imports.clear() + registry._skipped_builtin_backend_imports.update(saved_skipped) From 44a11a72453cd79a2afff388d633b0fca877e1fc Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Tue, 14 Jul 2026 15:11:49 +0000 Subject: [PATCH 08/31] fix(ray): capture pathlib I/O in workers and pin KubeRay e2e to ray 2.54 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes the per-task capture-fidelity gap on KubeRay found by the delegation live smoke, with two root causes isolated by probing a real worker (startup, open patch, executor patch, and task identity were all verified engaged — yet fragments arrived empty): - pathlib (Path.open/read_bytes/write_bytes) resolves io.open by module attribute, not the builtin, so the builtins-only open patch missed it; without the preload tracer (KubeRay pods) that I/O was invisible. roar_worker._startup now patches io.open alongside builtins.open. On preload-equipped native workers the duplicate python-level captures are absorbed by the existing native > python reconstitution dedup - ray 2.46 does not route task execution through the patched FunctionActorManager.get_execution_info, so the per-task flush never engaged; the KubeRay e2e image is pinned to the native harness's ray 2.54 (rayproject/ray:2.54.0-py312-cpu) where the roar_worker internals are supported The RayJob e2e's strict assertion (ray task file write present in ray_task outputs) is restored and green. Native-path regression check: identical compose-harness results with and without this change (5F/1P parity runs; the failures are a pre-existing runtime-env merge conflict on this dev box — job and driver ray.init envs both carrying ROAR_NO_TELEMETRY under a host ray 2.54 CLI — tracked as a separate follow-up). Ray + k8s unit suites and the full default gate are green. Verified live on KIND: 19/19 e2e including the strict RayJob smoke; default gate 2037 passed; ruff + mypy clean. Co-Authored-By: Claude Fable 5 --- docs/developer/k8s-integration.md | 11 +++++++++ roar/backends/ray/roar_worker.py | 5 ++++ tests/backends/k8s/e2e/test_k8s_rayjob.py | 30 +++++++++++++---------- 3 files changed, 33 insertions(+), 13 deletions(-) diff --git a/docs/developer/k8s-integration.md b/docs/developer/k8s-integration.md index f1f17ab9..d2ebbad3 100644 --- a/docs/developer/k8s-integration.md +++ b/docs/developer/k8s-integration.md @@ -36,6 +36,17 @@ reconstitution/attach delegate to the **Ray** backend's reconstituter, since the streamed fragments are Ray `TaskFragment` payloads that merge as `ray_task` jobs. +Per-task capture fidelity in KubeRay workers requires the Ray version the +roar_worker internals target (the native harness's pin, currently 2.54): +on 2.46, task execution was observed not to route through the patched +`FunctionActorManager.get_execution_info`, so task fragments arrived as +empty shells. Additionally, `roar_worker._startup` now patches `io.open` +alongside `builtins.open` — pathlib (`Path.read_bytes`/`write_bytes`) +resolves `io.open` by module attribute, and without the preload tracer +(KubeRay pods) those opens were invisible; on preload-equipped native +workers the duplicate python-level captures are absorbed by the existing +`native > python` dedup. + **Mounted storage** (`mount_map.py`): FUSE CSI mounts surface object I/O as local file syscalls under a mount path. The rewriter derives a per-container mount map (inline CSI volumes it can see — GCS FUSE, Mountpoint-for-S3 — diff --git a/roar/backends/ray/roar_worker.py b/roar/backends/ray/roar_worker.py index ef361db9..8379bfe3 100644 --- a/roar/backends/ray/roar_worker.py +++ b/roar/backends/ray/roar_worker.py @@ -6,6 +6,7 @@ import contextlib import functools import hashlib +import io import os import queue import re @@ -1630,6 +1631,10 @@ def _startup() -> None: if "libroar_tracer_preload" in os.environ.get("LD_PRELOAD", ""): _start_native_tracer_socket() builtins.open = _tracking_open + # pathlib (Path.open/read_bytes/write_bytes) and other stdlib callers + # resolve `io.open` by module attribute, not the builtin — without the + # preload tracer (e.g. KubeRay pods) those opens were invisible. + io.open = _tracking_open # type: ignore[assignment] _patch_subprocess_for_native_task_attribution() _patch_boto3() _patch_pandas_parquet() diff --git a/tests/backends/k8s/e2e/test_k8s_rayjob.py b/tests/backends/k8s/e2e/test_k8s_rayjob.py index 95909876..c5a62f04 100644 --- a/tests/backends/k8s/e2e/test_k8s_rayjob.py +++ b/tests/backends/k8s/e2e/test_k8s_rayjob.py @@ -38,7 +38,11 @@ pytest.mark.timeout(1500), ] -RAY_IMAGE = "rayproject/ray:2.46.0-py312-cpu" +# Pinned to the same Ray version as the native compose harness +# (tests/backends/ray/e2e/Dockerfile): the roar_worker per-task capture +# targets this version's internals, and 2.46 was observed not to route +# task execution through the patched FunctionActorManager path. +RAY_IMAGE = "rayproject/ray:2.54.0-py312-cpu" RAY_TRAIN_SCRIPT = """\ from pathlib import Path @@ -89,7 +93,7 @@ def train_shard() -> int: - name: ray-job-submitter image: {ray_image} rayClusterSpec: - rayVersion: "2.46.0" + rayVersion: "2.54.0" headGroupSpec: rayStartParams: num-cpus: "1" @@ -214,16 +218,16 @@ def test_rayjob_live_delegates_to_ray_backend( ) assert ray_tasks, f"expected ray_task jobs from delegated reconstitution\n{_describe(run)}" - # Delegation scope proven here: rewrite, Secret hygiene, in-pod pip - # bootstrap, fragment streaming from KubeRay pods, jobStatus wait, - # and Ray-reconstituter merge under the submit node. - # - # KNOWN GAP (Ray-backend follow-up, not k8s plumbing): per-task I/O - # fidelity inside KubeRay workers. Under setup-hook-only bootstrap - # (ROAR_WRAP intentionally off — see rayjob.py), fragments arrive as - # ray_task shells without file refs; the roar_worker task-boundary - # capture expects the py_executable/packaged-working_dir machinery - # of the native `roar run ray job submit` path. Once that lands, - # tighten this to assert task-out.bin in the ray_task outputs. + output_paths = { + str(row["path"]) + for row in _query( + project_dir, + "SELECT o.path FROM job_outputs o JOIN jobs j ON j.id = o.job_id " + "WHERE j.job_type = 'ray_task'", + ) + } + assert any(path.endswith("task-out.bin") for path in output_paths), ( + f"ray task file write missing from lineage: {output_paths}\n{_describe(run)}" + ) finally: _cleanup(name, resource="rayjob") From 21704334b91668f057b9f3912ccb1bb9625c9142 Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Tue, 14 Jul 2026 15:37:28 +0000 Subject: [PATCH 09/31] feat(k8s): add roar-runtime image and image-staged runtime mode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 3 slice 1: hermetic runtime staging (closes the Phase-0 "wheel is not hermetic" finding — pods no longer need network installs). - deploy/roar-runtime/Dockerfile: per-ABI dependency trees (cp310-cp313, resolved by uv --python-version without needing each interpreter), each with a generated top-level sitecustomize.py mirroring roar_inject.pth (PYTHONPATH staging doesn't process .pth files); the image doubles as the staging init container (default CMD copies /opt/roar-runtime into the shared mount) and the webhook server base - k8s.runtime_source=image + k8s.runtime_image config: the rewriter injects a roar-runtime-staging init container + emptyDir volume (idempotent by name for webhook reinvocation) and the wrapper activates the ABI-matched tree via PYTHONPATH instead of pip; TrainJob keeps the install path (no inline pod template for an init container); RayJob keeps its runtime-env pip mechanism - scripts/build_runtime_image.sh; .dockerignore wheel exception Verified live on KIND: image-mode e2e green in ~11s (vs ~16s+ for pip mode) — staging init container present, no pip in the wrapper, lineage captured; 4 new unit tests. Co-Authored-By: Claude Fable 5 --- .dockerignore | 2 + deploy/roar-runtime/Dockerfile | 50 ++++++ roar/backends/k8s/config.py | 19 +++ roar/backends/k8s/manifest.py | 105 ++++++++++++- roar/backends/k8s/submit.py | 2 + scripts/build_runtime_image.sh | 25 +++ tests/backends/k8s/e2e/test_k8s_phase3.py | 143 ++++++++++++++++++ tests/backends/k8s/unit/test_runtime_image.py | 88 +++++++++++ 8 files changed, 428 insertions(+), 6 deletions(-) create mode 100644 deploy/roar-runtime/Dockerfile create mode 100644 scripts/build_runtime_image.sh create mode 100644 tests/backends/k8s/e2e/test_k8s_phase3.py create mode 100644 tests/backends/k8s/unit/test_runtime_image.py diff --git a/.dockerignore b/.dockerignore index 017ced4d..780ad573 100644 --- a/.dockerignore +++ b/.dockerignore @@ -6,6 +6,8 @@ .roar .ralph dist +# roar-runtime image build needs the packaged wheel +!dist/roar_cli-*.whl __pycache__ **/__pycache__ *.pyc diff --git a/deploy/roar-runtime/Dockerfile b/deploy/roar-runtime/Dockerfile new file mode 100644 index 00000000..e18b8cc3 --- /dev/null +++ b/deploy/roar-runtime/Dockerfile @@ -0,0 +1,50 @@ +# roar-runtime: hermetic runtime staging image for Kubernetes lineage capture. +# +# Two roles: +# 1. Init-container staging (default CMD): copies self-contained per-ABI +# roar trees into a shared emptyDir so training pods need no network +# installs at startup (the OpenTelemetry auto-instrumentation pattern). +# 2. Webhook server base: `python -m roar.backends.k8s.webhook` runs the +# mutating admission injector (roar importable via the cp312 tree). +# +# Build context is the repo root and requires a packaged wheel in dist/ +# (scripts/build_wheel_with_bins.sh). See scripts/build_runtime_image.sh. + +FROM python:3.12-slim AS build + +RUN pip install --no-cache-dir uv + +COPY dist/ /tmp/dist/ + +# One self-contained dependency tree per supported CPython ABI. uv resolves +# each Python's wheels without needing that interpreter installed. The +# top-level sitecustomize.py mirrors roar_inject.pth: PYTHONPATH staging +# doesn't process .pth files, but Python imports `sitecustomize` from any +# sys.path entry — gated on ROAR_WRAP so it is inert outside instrumented +# process trees. +RUN set -e; \ + wheel="$(ls /tmp/dist/roar_cli-*.whl | head -1)"; \ + for version in 3.10 3.11 3.12 3.13; do \ + tag="cp$(echo "$version" | tr -d .)"; \ + uv pip install \ + --target "/opt/roar-runtime/$tag" \ + --python-version "$version" \ + --python-platform x86_64-unknown-linux-gnu \ + --link-mode=copy \ + "$wheel"; \ + printf '%s\n' \ + 'import os' \ + 'if os.environ.get("ROAR_WRAP") == "1":' \ + ' __import__("roar.execution.runtime.inject.sitecustomize")' \ + > "/opt/roar-runtime/$tag/sitecustomize.py"; \ + done + +FROM python:3.12-slim + +COPY --from=build /opt/roar-runtime /opt/roar-runtime + +# Webhook-server role: make roar importable for this image's own python. +ENV PYTHONPATH=/opt/roar-runtime/cp312 + +# Staging role (default): copy the runtime trees into the shared mount. +CMD ["/bin/sh", "-c", "cp -a /opt/roar-runtime/. ${ROAR_RUNTIME_TARGET:-/roar-runtime}/"] diff --git a/roar/backends/k8s/config.py b/roar/backends/k8s/config.py index 3d979c56..79bce43c 100644 --- a/roar/backends/k8s/config.py +++ b/roar/backends/k8s/config.py @@ -20,6 +20,12 @@ class K8sBackendConfig(BaseModel): enabled: bool = False tracer: str = "preload" + # How pods obtain the roar runtime: "install" pip-installs + # runtime_install_requirement at container start; "image" stages + # per-ABI trees from runtime_image via an init container (hermetic, + # no network at pod start). + runtime_source: str = "install" + runtime_image: str = "" runtime_install_requirement: str = "" cluster_glaas_url: str = "" bundle_dir: str = "" @@ -44,6 +50,19 @@ class K8sBackendConfig(BaseModel): default="preload", description="Tracer backend used inside instrumented pods (preload|ptrace|auto)", ), + "k8s.runtime_source": ConfigurableKeySpec( + value_type=str, + default="install", + description=( + "Runtime staging mode for pods: 'install' (pip at container start) or " + "'image' (init container copies per-ABI trees from k8s.runtime_image)" + ), + ), + "k8s.runtime_image": ConfigurableKeySpec( + value_type=str, + default="", + description="roar-runtime OCI image used when k8s.runtime_source = 'image'", + ), "k8s.runtime_install_requirement": ConfigurableKeySpec( value_type=str, default="", diff --git a/roar/backends/k8s/manifest.py b/roar/backends/k8s/manifest.py index 7f2a1687..3d985c19 100644 --- a/roar/backends/k8s/manifest.py +++ b/roar/backends/k8s/manifest.py @@ -16,7 +16,7 @@ import shlex from collections.abc import Callable -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from pathlib import Path from typing import Any @@ -24,7 +24,7 @@ from roar.backends.k8s.mount_map import build_container_mount_map, dump_mount_map -# sh -c script: "$0" is the synthetic argv0, "$@" is the original +# sh -c scripts: "$0" is the synthetic argv0, "$@" is the original # command+args. Lineage is best-effort by design: any failure to stage # the roar runtime falls back to running the original command # uninstrumented rather than failing the training job. @@ -35,6 +35,22 @@ exec python3 -m roar.backends.k8s.pod_entry "$@" """ +# Image-staged variant: the runtime tree was copied into the shared mount +# by the roar-runtime init container; pick the tree matching this +# container's Python ABI and activate it via PYTHONPATH. +_POD_WRAPPER_IMAGE_TEMPLATE = """\ +run_fallback() {{ echo "[roar-k8s] lineage runtime unavailable; running uninstrumented" >&2; exec "$@"; }} +command -v python3 >/dev/null 2>&1 || run_fallback "$@" +RT="{staging_mount}/cp$(python3 -c 'import sys; print("%d%d" % sys.version_info[:2])')" +[ -d "$RT" ] || run_fallback "$@" +export PYTHONPATH="$RT${{PYTHONPATH:+:$PYTHONPATH}}" +exec python3 -m roar.backends.k8s.pod_entry "$@" +""" + +_RUNTIME_STAGING_VOLUME = "roar-runtime" +_RUNTIME_STAGING_MOUNT = "/roar-runtime" +_RUNTIME_STAGING_INIT_CONTAINER = "roar-runtime-staging" + # Terminal workload conditions, unioned across kinds: Job uses # Complete/SuccessCriteriaMet + Failed/FailureTarget, JobSet uses # Completed/Failed, PyTorchJob v1 uses Succeeded/Failed, TrainJob uses @@ -206,6 +222,8 @@ def rewrite_manifest_for_lineage( parent_job_uid: str, bundle_dir: str = "", mount_map: dict[str, str] | None = None, + runtime_source: str = "install", + runtime_image: str = "", namespace_override: str | None = None, ) -> K8sManifestRewrite: """Return a rewritten copy of ``documents`` with lineage instrumentation. @@ -254,6 +272,8 @@ def rewrite_manifest_for_lineage( workload_name=workload_name, bundle_dir=bundle_dir, config_mount_map=dict(mount_map or {}), + runtime_source=runtime_source, + runtime_image=runtime_image, ) if workload.rewrite_style == "trainer_override": @@ -324,6 +344,12 @@ class _EnvContract: workload_name: str bundle_dir: str = "" config_mount_map: dict[str, str] = field(default_factory=dict) + runtime_source: str = "install" + runtime_image: str = "" + + @property + def image_staging(self) -> bool: + return self.runtime_source == "image" and bool(self.runtime_image) def _rewrite_pod_specs( @@ -347,6 +373,7 @@ def _rewrite_pod_specs( f"{workload.kind} {workload_name} pod template " f"{ref.role or ''} has no containers" ) + pod_wrapped = False for container in containers: if not isinstance(container, dict): continue @@ -358,8 +385,10 @@ def _rewrite_pod_specs( continue original = [str(part) for part in command] original.extend(str(part) for part in container.get("args", []) or []) - container["command"] = _wrapped_command(original, requirement=contract.requirement) + container["command"] = _wrapped_command(original, contract=contract) container.pop("args", None) + if contract.image_staging: + _add_staging_volume_mount(container) _inject_env_contract( container.setdefault("env", []), contract=contract, @@ -370,9 +399,65 @@ def _rewrite_pod_specs( ), ) wrapped.append(label) + pod_wrapped = True + if pod_wrapped and contract.image_staging: + _add_runtime_staging(ref.spec, runtime_image=contract.runtime_image) return wrapped, skipped +def _add_runtime_staging(pod_spec: dict[str, Any], *, runtime_image: str) -> None: + """Add the shared volume + init container that stage the roar runtime. + + Idempotent by name so webhook reinvocation and repeated rewrites are + safe. + """ + volumes = pod_spec.setdefault("volumes", []) + if isinstance(volumes, list) and not any( + isinstance(volume, dict) and volume.get("name") == _RUNTIME_STAGING_VOLUME + for volume in volumes + ): + volumes.append({"name": _RUNTIME_STAGING_VOLUME, "emptyDir": {}}) + + init_containers = pod_spec.setdefault("initContainers", []) + if isinstance(init_containers, list) and not any( + isinstance(container, dict) + and container.get("name") == _RUNTIME_STAGING_INIT_CONTAINER + for container in init_containers + ): + init_containers.append( + { + "name": _RUNTIME_STAGING_INIT_CONTAINER, + "image": runtime_image, + "command": [ + "/bin/sh", + "-c", + f"cp -a /opt/roar-runtime/. {_RUNTIME_STAGING_MOUNT}/", + ], + "volumeMounts": [ + { + "name": _RUNTIME_STAGING_VOLUME, + "mountPath": _RUNTIME_STAGING_MOUNT, + } + ], + } + ) + + +def _add_staging_volume_mount(container: dict[str, Any]) -> None: + mounts = container.setdefault("volumeMounts", []) + if isinstance(mounts, list) and not any( + isinstance(mount, dict) and mount.get("name") == _RUNTIME_STAGING_VOLUME + for mount in mounts + ): + mounts.append( + { + "name": _RUNTIME_STAGING_VOLUME, + "mountPath": _RUNTIME_STAGING_MOUNT, + "readOnly": True, + } + ) + + def _rewrite_trainjob( doc: dict[str, Any], *, @@ -403,7 +488,12 @@ def _rewrite_trainjob( original = [str(part) for part in command] original.extend(str(part) for part in trainer.get("args", []) or []) - trainer["command"] = _wrapped_command(original, requirement=contract.requirement) + # TrainJob has no inline pod template to attach an init container to, + # so runtime staging always uses the install path here. + install_contract = ( + replace(contract, runtime_source="install") if contract.image_staging else contract + ) + trainer["command"] = _wrapped_command(original, contract=install_contract) trainer.pop("args", None) env = trainer.setdefault("env", []) @@ -420,8 +510,11 @@ def _rewrite_trainjob( return ["node"], [] -def _wrapped_command(original: list[str], *, requirement: str) -> list[str]: - script = _POD_WRAPPER_TEMPLATE.format(requirement=shlex.quote(requirement)) +def _wrapped_command(original: list[str], *, contract: _EnvContract) -> list[str]: + if contract.image_staging: + script = _POD_WRAPPER_IMAGE_TEMPLATE.format(staging_mount=_RUNTIME_STAGING_MOUNT) + else: + script = _POD_WRAPPER_TEMPLATE.format(requirement=shlex.quote(contract.requirement)) return ["/bin/sh", "-c", script, "roar-k8s", *original] diff --git a/roar/backends/k8s/submit.py b/roar/backends/k8s/submit.py index 699d09b9..a2821d85 100644 --- a/roar/backends/k8s/submit.py +++ b/roar/backends/k8s/submit.py @@ -121,6 +121,8 @@ def plan_kubectl_job_submit_command(command: list[str]) -> ExecutionCommandPlan: parent_job_uid=parent_job_uid, bundle_dir=str(config.get("bundle_dir") or ""), mount_map=_config_mount_map(config), + runtime_source=str(config.get("runtime_source") or "install"), + runtime_image=str(config.get("runtime_image") or ""), namespace_override=_find_namespace_argument(command), ) if rewrite.skipped_containers: diff --git a/scripts/build_runtime_image.sh b/scripts/build_runtime_image.sh new file mode 100644 index 00000000..cb75aae0 --- /dev/null +++ b/scripts/build_runtime_image.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Builds the roar-runtime OCI image (init-container staging + webhook base). +# +# Usage: +# bash scripts/build_runtime_image.sh [TAG] # default: roar-runtime:dev +# +# Requires a packaged wheel in dist/ (scripts/build_wheel_with_bins.sh). + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ROOT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)" +TAG="${1:-roar-runtime:dev}" + +shopt -s nullglob +wheels=("$ROOT_DIR"/dist/roar_cli-*.whl) +shopt -u nullglob +if ((${#wheels[@]} == 0)); then + echo "error: no roar_cli wheel in $ROOT_DIR/dist" >&2 + echo "hint: build one first: bash scripts/build_wheel_with_bins.sh" >&2 + exit 1 +fi + +docker build -f "$ROOT_DIR/deploy/roar-runtime/Dockerfile" -t "$TAG" "$ROOT_DIR" +echo "✓ Built $TAG" diff --git a/tests/backends/k8s/e2e/test_k8s_phase3.py b/tests/backends/k8s/e2e/test_k8s_phase3.py new file mode 100644 index 00000000..220de7ba --- /dev/null +++ b/tests/backends/k8s/e2e/test_k8s_phase3.py @@ -0,0 +1,143 @@ +"""Phase-3 e2e: image-staged runtime (and, below, the webhook injector). + +Requires the roar-runtime image loaded into the cluster: + + bash scripts/build_runtime_image.sh + .tools/bin/kind load docker-image roar-runtime:dev --name roar-k8s-e2e +""" + +from __future__ import annotations + +import json +import sqlite3 +import uuid +from pathlib import Path + +import pytest + +from .conftest import KUBE_CONTEXT, NAMESPACE, kubectl +from .test_k8s_distributed import ( + SINGLE_JOB_TEMPLATE, + _cleanup, + _describe, + _pod_logs, + _submit, +) +from .test_k8s_product_path import wheel_server # noqa: F401 (module fixture reuse) + +pytestmark = [ + pytest.mark.e2e, + pytest.mark.k8s_e2e, + pytest.mark.timeout(900), +] + +RUNTIME_IMAGE = "roar-runtime:dev" + + +def _query(project_dir: Path, sql: str, params: tuple = ()) -> list[sqlite3.Row]: + conn = sqlite3.connect(project_dir / ".roar" / "roar.db") + conn.row_factory = sqlite3.Row + try: + return conn.execute(sql, params).fetchall() + finally: + conn.close() + + +def _runtime_image_loaded() -> bool: + import subprocess + + result = subprocess.run( + ["docker", "exec", "roar-k8s-e2e-worker", "crictl", "inspecti", "-q", + f"docker.io/library/{RUNTIME_IMAGE}"], + capture_output=True, + text=True, + check=False, + ) + return result.returncode == 0 + + +def _write_image_mode_project(project_dir: Path) -> None: + roar_dir = project_dir / ".roar" + roar_dir.mkdir(parents=True, exist_ok=True) + (roar_dir / "config.toml").write_text( + "\n".join( + [ + "[glaas]", + 'url = "http://localhost:3001"', + "", + "[k8s]", + "enabled = true", + 'runtime_source = "image"', + f'runtime_image = "{RUNTIME_IMAGE}"', + 'cluster_glaas_url = "http://glaas:3001"', + "wait_for_completion = true", + "wait_timeout_seconds = 420", + "poll_interval_seconds = 2.0", + ] + ) + + "\n", + encoding="utf-8", + ) + + +def test_image_staged_runtime_captures_without_network_install( + k8s_cluster: None, + glaas_health: str, + tmp_path_factory: pytest.TempPathFactory, +) -> None: + if not _runtime_image_loaded(): + pytest.skip( + "roar-runtime:dev not loaded; run scripts/build_runtime_image.sh " + "and kind load docker-image roar-runtime:dev --name roar-k8s-e2e" + ) + + project_dir = tmp_path_factory.mktemp("k8s-image-mode") + _write_image_mode_project(project_dir) + + job_name = f"roar-img-{uuid.uuid4().hex[:6]}" + (project_dir / "job.yaml").write_text( + SINGLE_JOB_TEMPLATE.format(name=job_name, namespace=NAMESPACE), + encoding="utf-8", + ) + + try: + completed = _submit("job.yaml", cwd=project_dir) + logs = _pod_logs(f"job-name={job_name}") + run = { + "exit_code": completed.returncode, + "stdout": completed.stdout, + "stderr": completed.stderr, + "pod_logs": logs, + } + assert completed.returncode == 0, _describe(run) + + # The submitted Job carries the staging init container. + job_doc = json.loads( + kubectl( + ["get", f"job/{job_name}", "-n", NAMESPACE, "-o", "json"] + ).stdout + ) + pod_spec = job_doc["spec"]["template"]["spec"] + init_names = [c["name"] for c in pod_spec.get("initContainers", [])] + assert "roar-runtime-staging" in init_names, init_names + assert not any( + "pip install" in part + for container in pod_spec["containers"] + for part in container.get("command", []) + ), "image mode must not pip install in the wrapper" + + tasks = _query(project_dir, "SELECT id FROM jobs WHERE job_type = 'k8s_task'") + assert len(tasks) == 1, _describe(run) + outputs = { + str(row["path"]) + for row in _query( + project_dir, + "SELECT path FROM job_outputs WHERE job_id = ?", + (tasks[0]["id"],), + ) + } + assert any(path.endswith("summary.json") for path in outputs), ( + f"{outputs}\n{_describe(run)}" + ) + finally: + _cleanup(job_name) diff --git a/tests/backends/k8s/unit/test_runtime_image.py b/tests/backends/k8s/unit/test_runtime_image.py new file mode 100644 index 00000000..9a69f74c --- /dev/null +++ b/tests/backends/k8s/unit/test_runtime_image.py @@ -0,0 +1,88 @@ +from __future__ import annotations + +import copy +from typing import Any + +from roar.backends.k8s.manifest import rewrite_manifest_for_lineage + +from .conftest import SINGLE_JOB_MANIFEST + + +def _rewrite_image_mode(documents: list[dict[str, Any]], **overrides: Any): + kwargs: dict[str, Any] = { + "secret_name": "roar-fragment-abc", + "session_id": "session-1", + "fragment_token": "ff" * 32, + "requirement": "roar-cli==0.3.7", + "cluster_glaas_url": "http://glaas:3001", + "tracer": "preload", + "parent_job_uid": "cafe0123", + "runtime_source": "image", + "runtime_image": "roar-runtime:dev", + } + kwargs.update(overrides) + return rewrite_manifest_for_lineage(documents, **kwargs) + + +def _job_pod_spec(rewrite) -> dict[str, Any]: + job = next(doc for doc in rewrite.documents if doc.get("kind") == "Job") + return job["spec"]["template"]["spec"] + + +def test_image_mode_adds_staging_init_container_and_volume() -> None: + rewrite = _rewrite_image_mode([copy.deepcopy(SINGLE_JOB_MANIFEST)]) + pod_spec = _job_pod_spec(rewrite) + + volumes = {volume["name"] for volume in pod_spec["volumes"]} + assert "roar-runtime" in volumes + + init = pod_spec["initContainers"][0] + assert init["name"] == "roar-runtime-staging" + assert init["image"] == "roar-runtime:dev" + assert "/opt/roar-runtime/." in init["command"][2] + + container = pod_spec["containers"][0] + mounts = {mount["name"]: mount for mount in container["volumeMounts"]} + assert mounts["roar-runtime"]["mountPath"] == "/roar-runtime" + assert mounts["roar-runtime"]["readOnly"] is True + + +def test_image_mode_wrapper_uses_pythonpath_not_pip() -> None: + rewrite = _rewrite_image_mode([copy.deepcopy(SINGLE_JOB_MANIFEST)]) + script = _job_pod_spec(rewrite)["containers"][0]["command"][2] + + assert "pip install" not in script + assert 'RT="/roar-runtime/cp' in script + assert "PYTHONPATH" in script + assert "roar.backends.k8s.pod_entry" in script + assert "run_fallback" in script + + +def test_image_mode_is_idempotent_for_webhook_reinvocation() -> None: + rewrite = _rewrite_image_mode([copy.deepcopy(SINGLE_JOB_MANIFEST)]) + job = next(doc for doc in rewrite.documents if doc.get("kind") == "Job") + + second = _rewrite_image_mode([copy.deepcopy(job)]) + pod_spec = _job_pod_spec(second) + staging_inits = [ + container + for container in pod_spec["initContainers"] + if container["name"] == "roar-runtime-staging" + ] + staging_volumes = [ + volume for volume in pod_spec["volumes"] if volume["name"] == "roar-runtime" + ] + assert len(staging_inits) == 1 + assert len(staging_volumes) == 1 + + +def test_install_mode_untouched_by_new_fields() -> None: + rewrite = _rewrite_image_mode( + [copy.deepcopy(SINGLE_JOB_MANIFEST)], + runtime_source="install", + runtime_image="", + ) + pod_spec = _job_pod_spec(rewrite) + assert "initContainers" not in pod_spec + script = pod_spec["containers"][0]["command"][2] + assert "pip install --quiet roar-cli==0.3.7" in script From 80ddc0a04bc9c06085c0b1db364150dc68d89ed4 Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Tue, 14 Jul 2026 16:24:36 +0000 Subject: [PATCH 10/31] feat(k8s): add mutating webhook injector for zero-touch lineage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 3 slice 2: the platform-install path. A stdlib HTTPS AdmissionReview server (roar/backends/k8s/webhook.py, served by the roar-runtime image) intercepts CREATE of all five workload kinds in namespaces labeled roar.glaas.ai/lineage=enabled and applies the same manifest rewrite as the CLI path — clients need no roar at all. - per admission: mints the fragment session against GLaaS, creates the credentials Secret through the k8s API (never embedded in the object), rewrites via the shared rewriter (image-staged runtime by default), annotates the workload (roar.glaas.ai/parent-uid, session-id, fragment-secret), returns a JSONPatch - never blocks admission: every internal failure returns allowed with a warning, paired with failurePolicy Ignore; dry-run admissions are side-effect-free (sideEffects NoneOnDryRun); idempotent under reinvocationPolicy IfNeeded via the parent-uid annotation - reconstitution is client-driven via the existing roar k8s attach (reads the webhook-created Secret) - harness: bootstrap --with-webhook builds/loads roar-runtime:dev and deploys via deploy_webhook.sh (self-signed CA, RBAC for secret creation, Deployment/Service/MutatingWebhookConfiguration); the image is always rebuilt (layer cache) — a stale image silently breaks the webhook - fix found live: the API server posts /mutate?timeout=10s — path matching must strip the query string Verified live on KIND: 22/22 e2e — zero-touch test does a plain kubectl apply in a labeled namespace (no roar client), asserts injection (annotations, staging init container, wrapped command), job completion, and attach-recovered lineage keyed to the webhook-minted parent uid; unlabeled-namespace control untouched. 7 new unit tests; default gate 2040 passed; ruff + mypy clean. Co-Authored-By: Claude Fable 5 --- docs/developer/k8s-integration.md | 31 +- roar/backends/k8s/manifest.py | 6 +- roar/backends/k8s/webhook.py | 301 ++++++++++++++++++ tests/backends/k8s/README.md | 10 +- tests/backends/k8s/e2e/test_k8s_phase3.py | 124 +++++++- tests/backends/k8s/scripts/bootstrap_k8s.sh | 13 + tests/backends/k8s/scripts/deploy_webhook.sh | 191 +++++++++++ tests/backends/k8s/unit/test_runtime_image.py | 4 +- tests/backends/k8s/unit/test_webhook.py | 170 ++++++++++ 9 files changed, 833 insertions(+), 17 deletions(-) create mode 100644 roar/backends/k8s/webhook.py create mode 100644 tests/backends/k8s/scripts/deploy_webhook.sh create mode 100644 tests/backends/k8s/unit/test_webhook.py diff --git a/docs/developer/k8s-integration.md b/docs/developer/k8s-integration.md index d2ebbad3..804723cb 100644 --- a/docs/developer/k8s-integration.md +++ b/docs/developer/k8s-integration.md @@ -15,12 +15,37 @@ identity comes from the downward API plus the completion-index/node-rank env chain (`JOB_COMPLETION_INDEX` → `PET_NODE_RANK` → pod-level `RANK`). Current phase status (see `design-docs/k8s-training-lineage-integration.md` -in the dev meta-repo): Phases 1–2 fully implemented and live-validated +in the dev meta-repo): Phases 1–3 implemented and live-validated (submit wrapping, operator adapters incl. real training-operator v1 and trainer v2 controllers, multi-pod capture, `roar k8s attach`, bundle-mode fallback, object-store I/O hooks, mount-map rewriting, retry-chaos -coverage, RayJob delegation). Remaining: the Phase-3 admission-webhook -injector (incl. the opt-in proxy sidecar). +coverage, RayJob delegation, the roar-runtime image, and the mutating +webhook injector). Remaining from Phase 3: the opt-in proxy sidecar and +a packaged Helm chart (the harness deploys via +`tests/backends/k8s/scripts/deploy_webhook.sh`). + +**Runtime staging modes**: `k8s.runtime_source = "install"` pip-installs +`k8s.runtime_install_requirement` at container start; `"image"` stages +hermetic per-ABI trees (cp310–cp313) from `k8s.runtime_image` via a +`roar-runtime-staging` init container + emptyDir — no network at pod +start. The image (`deploy/roar-runtime/Dockerfile`, +`scripts/build_runtime_image.sh`) ships a generated top-level +`sitecustomize.py` per tree because PYTHONPATH staging does not process +`.pth` files. TrainJob keeps the install path (no inline pod template); +RayJob keeps its runtime-env pip mechanism. + +**Webhook injector** (`webhook.py`): a stdlib HTTPS AdmissionReview +server (served by the roar-runtime image) intercepting CREATE of all +five workload kinds in namespaces labeled `roar.glaas.ai/lineage=enabled`. +It reuses the same manifest rewriter as the CLI path, mints the fragment +session against GLaaS, creates the credentials Secret through the k8s +API (never embedded in the object), annotates the workload +(`roar.glaas.ai/parent-uid`, `session-id`, `fragment-secret`), and +returns a JSONPatch. Idempotent under `reinvocationPolicy: IfNeeded` via +the parent-uid annotation; dry-run admissions are side-effect-free +(`sideEffects: NoneOnDryRun`); every internal failure returns allowed +with a warning and pairs with `failurePolicy: Ignore` — lineage never +blocks admission. Reconstitution is client-driven via `roar k8s attach`. **RayJob delegation** (`rayjob.py`): KubeRay overwrites container commands with `ray start --block`, so command-wrapping can't see user code. Instead diff --git a/roar/backends/k8s/manifest.py b/roar/backends/k8s/manifest.py index 3d985c19..b92d8b3a 100644 --- a/roar/backends/k8s/manifest.py +++ b/roar/backends/k8s/manifest.py @@ -420,8 +420,7 @@ def _add_runtime_staging(pod_spec: dict[str, Any], *, runtime_image: str) -> Non init_containers = pod_spec.setdefault("initContainers", []) if isinstance(init_containers, list) and not any( - isinstance(container, dict) - and container.get("name") == _RUNTIME_STAGING_INIT_CONTAINER + isinstance(container, dict) and container.get("name") == _RUNTIME_STAGING_INIT_CONTAINER for container in init_containers ): init_containers.append( @@ -446,8 +445,7 @@ def _add_runtime_staging(pod_spec: dict[str, Any], *, runtime_image: str) -> Non def _add_staging_volume_mount(container: dict[str, Any]) -> None: mounts = container.setdefault("volumeMounts", []) if isinstance(mounts, list) and not any( - isinstance(mount, dict) and mount.get("name") == _RUNTIME_STAGING_VOLUME - for mount in mounts + isinstance(mount, dict) and mount.get("name") == _RUNTIME_STAGING_VOLUME for mount in mounts ): mounts.append( { diff --git a/roar/backends/k8s/webhook.py b/roar/backends/k8s/webhook.py new file mode 100644 index 00000000..8b78b6d2 --- /dev/null +++ b/roar/backends/k8s/webhook.py @@ -0,0 +1,301 @@ +"""Mutating admission webhook: zero-touch lineage injection. + +Intercepts CREATE of supported training workloads (Job, JobSet, +PyTorchJob, TrainJob, RayJob) in opted-in namespaces and applies the +same manifest rewrite as the CLI path — the platform team installs one +thing and every workload in a labeled namespace gets lineage, with no +`roar run` on the submitting side. Reconstitution happens later via +``roar k8s attach`` (credentials live in the per-workload Secret the +webhook creates). + +Design constraints honored: +- never blocks admission: any internal failure returns allowed with a + warning (pair with ``failurePolicy: Ignore`` in the webhook config); +- idempotent under ``reinvocationPolicy: IfNeeded`` via the + ``roar.glaas.ai/parent-uid`` annotation; +- side effects (GLaaS session registration, Secret creation) are skipped + for dry-run admission (declare ``sideEffects: NoneOnDryRun``); +- stdlib only (http.server + ssl + urllib) so the roar-runtime image can + serve it without extra dependencies. +""" + +from __future__ import annotations + +import base64 +import json +import os +import secrets as secrets_module +import ssl +import sys +import urllib.request +from collections.abc import Callable +from dataclasses import dataclass, field +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Any + +OPT_OUT_LABEL = "roar.glaas.ai/lineage" +ANNOTATION_PARENT_UID = "roar.glaas.ai/parent-uid" +ANNOTATION_SESSION_ID = "roar.glaas.ai/session-id" +ANNOTATION_SECRET = "roar.glaas.ai/fragment-secret" + +_SERVICE_ACCOUNT_DIR = "/var/run/secrets/kubernetes.io/serviceaccount" + +SecretCreator = Callable[[str, str, dict[str, str]], None] +SessionRegistrar = Callable[[str, str], None] + + +@dataclass(frozen=True) +class WebhookSettings: + """Injection settings, sourced from the webhook Deployment's env.""" + + glaas_url: str + cluster_glaas_url: str + tracer: str = "preload" + runtime_source: str = "install" + runtime_image: str = "" + runtime_install_requirement: str = "roar-cli" + fragment_session_ttl_seconds: int = 86400 + bundle_dir: str = "" + mount_map: dict[str, str] = field(default_factory=dict) + + @classmethod + def from_environ(cls, environ: dict[str, str] | None = None) -> WebhookSettings: + env = os.environ if environ is None else environ + mount_map: dict[str, str] = {} + raw_mount_map = env.get("ROAR_WEBHOOK_MOUNT_MAP", "") + if raw_mount_map: + try: + parsed = json.loads(raw_mount_map) + if isinstance(parsed, dict): + mount_map = {str(k): str(v) for k, v in parsed.items()} + except json.JSONDecodeError: + pass + return cls( + glaas_url=env.get("ROAR_WEBHOOK_GLAAS_URL", "").strip(), + cluster_glaas_url=env.get("ROAR_WEBHOOK_CLUSTER_GLAAS_URL", "").strip(), + tracer=env.get("ROAR_WEBHOOK_TRACER", "preload").strip() or "preload", + runtime_source=env.get("ROAR_WEBHOOK_RUNTIME_SOURCE", "install").strip() or "install", + runtime_image=env.get("ROAR_WEBHOOK_RUNTIME_IMAGE", "").strip(), + runtime_install_requirement=env.get( + "ROAR_WEBHOOK_RUNTIME_REQUIREMENT", "roar-cli" + ).strip() + or "roar-cli", + fragment_session_ttl_seconds=int(env.get("ROAR_WEBHOOK_SESSION_TTL", "86400") or 86400), + bundle_dir=env.get("ROAR_WEBHOOK_BUNDLE_DIR", "").strip(), + mount_map=mount_map, + ) + + +def mutate_admission_review( + review: dict[str, Any], + *, + settings: WebhookSettings, + create_secret: SecretCreator, + register_session: SessionRegistrar, +) -> dict[str, Any]: + """Return the AdmissionReview response for one admission request.""" + request = review.get("request") or {} + uid = str(request.get("uid") or "") + + def _allow(warnings: list[str] | None = None) -> dict[str, Any]: + response: dict[str, Any] = {"uid": uid, "allowed": True} + if warnings: + response["warnings"] = warnings + return { + "apiVersion": "admission.k8s.io/v1", + "kind": "AdmissionReview", + "response": response, + } + + try: + from roar.backends.k8s.manifest import ( + rewrite_manifest_for_lineage, + workload_kind_for_document, + ) + from roar.execution.fragments.sessions import generate_fragment_session + + obj = request.get("object") + if not isinstance(obj, dict) or workload_kind_for_document(obj) is None: + return _allow() + + metadata = obj.get("metadata") or {} + labels = metadata.get("labels") or {} + annotations = metadata.get("annotations") or {} + if str(labels.get(OPT_OUT_LABEL, "")).lower() == "disabled": + return _allow() + if annotations.get(ANNOTATION_PARENT_UID): + return _allow() # already instrumented (reinvocation/replay) + if bool(request.get("dryRun")): + return _allow(["roar: dry-run admission is not instrumented"]) + + namespace = str(request.get("namespace") or metadata.get("namespace") or "default") + + session = generate_fragment_session() + register_session(session["session_id"], session["token_hash"]) + + parent_job_uid = secrets_module.token_hex(4) + secret_name = f"roar-fragment-{session['session_id'][:8]}" + create_secret( + namespace, + secret_name, + {"session_id": session["session_id"], "token": session["token"]}, + ) + + rewrite = rewrite_manifest_for_lineage( + [obj], + secret_name=secret_name, + session_id=None, # Secret created via the API, never embedded + fragment_token=None, + requirement=settings.runtime_install_requirement, + cluster_glaas_url=settings.cluster_glaas_url, + tracer=settings.tracer, + parent_job_uid=parent_job_uid, + bundle_dir=settings.bundle_dir, + mount_map=settings.mount_map, + runtime_source=settings.runtime_source, + runtime_image=settings.runtime_image, + namespace_override=namespace, + ) + rewritten = rewrite.documents[0] + + merged_annotations = dict(annotations) + merged_annotations[ANNOTATION_PARENT_UID] = parent_job_uid + merged_annotations[ANNOTATION_SESSION_ID] = session["session_id"] + merged_annotations[ANNOTATION_SECRET] = secret_name + + patch = [ + {"op": "replace", "path": "/spec", "value": rewritten["spec"]}, + {"op": "add", "path": "/metadata/annotations", "value": merged_annotations}, + ] + response: dict[str, Any] = { + "uid": uid, + "allowed": True, + "patchType": "JSONPatch", + "patch": base64.b64encode(json.dumps(patch).encode("utf-8")).decode("ascii"), + } + return { + "apiVersion": "admission.k8s.io/v1", + "kind": "AdmissionReview", + "response": response, + } + except Exception as exc: # never block admission on lineage failures + return _allow([f"roar: lineage injection skipped: {exc}"]) + + +def create_namespaced_secret(namespace: str, name: str, string_data: dict[str, str]) -> None: + """Create a Secret through the in-cluster API (service-account auth).""" + with open(f"{_SERVICE_ACCOUNT_DIR}/token", encoding="utf-8") as handle: + token = handle.read().strip() + + body = json.dumps( + { + "apiVersion": "v1", + "kind": "Secret", + "metadata": { + "name": name, + "namespace": namespace, + "labels": {"app.kubernetes.io/managed-by": "roar"}, + }, + "type": "Opaque", + "stringData": string_data, + } + ).encode("utf-8") + + request = urllib.request.Request( + url=f"https://kubernetes.default.svc/api/v1/namespaces/{namespace}/secrets", + data=body, + headers={ + "authorization": f"Bearer {token}", + "content-type": "application/json", + }, + method="POST", + ) + context = ssl.create_default_context(cafile=f"{_SERVICE_ACCOUNT_DIR}/ca.crt") + with urllib.request.urlopen(request, timeout=10, context=context) as response: + if response.status not in (200, 201): + raise RuntimeError(f"secret creation returned HTTP {response.status}") + + +def register_glaas_session(settings: WebhookSettings) -> SessionRegistrar: + def _register(session_id: str, token_hash: str) -> None: + from roar.integrations.glaas import GlaasClient + + client = GlaasClient(base_url=settings.glaas_url) + _result, error = client.register_fragment_session( + session_id=session_id, + token_hash=token_hash, + ttl_seconds=settings.fragment_session_ttl_seconds, + ) + if error: + raise RuntimeError(f"fragment session registration failed: {error}") + + return _register + + +class _WebhookHandler(BaseHTTPRequestHandler): + settings: WebhookSettings + + def do_GET(self) -> None: + if self.path == "/healthz": + self._respond(200, b"ok", content_type="text/plain") + else: + self._respond(404, b"not found", content_type="text/plain") + + def do_POST(self) -> None: + path = self.path.split("?", 1)[0] + if path != "/mutate": + self._respond(404, b"not found", content_type="text/plain") + return + length = int(self.headers.get("content-length") or 0) + try: + review = json.loads(self.rfile.read(length).decode("utf-8")) + except (json.JSONDecodeError, UnicodeDecodeError): + self._respond(400, b"invalid AdmissionReview", content_type="text/plain") + return + + result = mutate_admission_review( + review, + settings=self.settings, + create_secret=create_namespaced_secret, + register_session=register_glaas_session(self.settings), + ) + self._respond(200, json.dumps(result).encode("utf-8")) + + def _respond(self, status: int, body: bytes, content_type: str = "application/json") -> None: + self.send_response(status) + self.send_header("content-type", content_type) + self.send_header("content-length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, format: str, *args: Any) -> None: + print(f"[roar-webhook] {self.address_string()} {format % args}", file=sys.stderr) + + +def main(argv: list[str] | None = None) -> int: + import argparse + + parser = argparse.ArgumentParser(description="roar k8s lineage mutating webhook") + parser.add_argument("--port", type=int, default=8443) + parser.add_argument("--cert", required=True) + parser.add_argument("--key", required=True) + args = parser.parse_args(argv) + + settings = WebhookSettings.from_environ() + if not settings.glaas_url: + print("[roar-webhook] ROAR_WEBHOOK_GLAAS_URL is required", file=sys.stderr) + return 2 + + handler = type("BoundHandler", (_WebhookHandler,), {"settings": settings}) + server = ThreadingHTTPServer(("0.0.0.0", args.port), handler) + context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + context.load_cert_chain(certfile=args.cert, keyfile=args.key) + server.socket = context.wrap_socket(server.socket, server_side=True) + + print(f"[roar-webhook] serving on :{args.port}", file=sys.stderr) + server.serve_forever() + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/backends/k8s/README.md b/tests/backends/k8s/README.md index 7d041c44..6883421b 100644 --- a/tests/backends/k8s/README.md +++ b/tests/backends/k8s/README.md @@ -32,6 +32,11 @@ Three test layers share this harness: - `e2e/test_k8s_rayjob.py` — live RayJob delegation smoke on KubeRay (the heaviest test: multi-GB Ray image + per-job pip env). Skips without `--with-kuberay`. +- `e2e/test_k8s_phase3.py` — image-staged runtime (hermetic pod start via + the roar-runtime init container, no pip) and the zero-touch webhook + story: plain `kubectl apply` in a labeled namespace gets injected, and + `roar k8s attach` recovers the lineage; unlabeled namespaces stay + untouched. Skips without the image / `--with-webhook`. - `e2e/test_k8s_smoke.py` — the Phase-0 runtime diagnostic: fixtures hand-wrap the manifest (no backend involved) to isolate the runtime pieces (in-pod tracing, fragment streaming, identity contract) when the product @@ -55,8 +60,9 @@ bash scripts/build_wheel_with_bins.sh # create cluster + wire glaas + preflight # (--with-minio: S3 scenarios; --with-jobset: JobSet e2e; --with-kubeflow: -# PyTorchJob/TrainJob e2e; --with-kuberay: RayJob delegation e2e) -bash tests/backends/k8s/scripts/bootstrap_k8s.sh --with-jobset --with-minio --with-kubeflow --with-kuberay +# PyTorchJob/TrainJob e2e; --with-kuberay: RayJob delegation e2e; +# --with-webhook: builds/loads roar-runtime:dev + deploys the injector) +bash tests/backends/k8s/scripts/bootstrap_k8s.sh --with-jobset --with-minio --with-kubeflow --with-kuberay --with-webhook # run the smoke tests (addopts override needed: e2e dirs are ignored by default) pytest tests/backends/k8s/e2e -o addopts='' -m k8s_e2e -v diff --git a/tests/backends/k8s/e2e/test_k8s_phase3.py b/tests/backends/k8s/e2e/test_k8s_phase3.py index 220de7ba..0e75a70a 100644 --- a/tests/backends/k8s/e2e/test_k8s_phase3.py +++ b/tests/backends/k8s/e2e/test_k8s_phase3.py @@ -47,8 +47,15 @@ def _runtime_image_loaded() -> bool: import subprocess result = subprocess.run( - ["docker", "exec", "roar-k8s-e2e-worker", "crictl", "inspecti", "-q", - f"docker.io/library/{RUNTIME_IMAGE}"], + [ + "docker", + "exec", + "roar-k8s-e2e-worker", + "crictl", + "inspecti", + "-q", + f"docker.io/library/{RUNTIME_IMAGE}", + ], capture_output=True, text=True, check=False, @@ -113,9 +120,7 @@ def test_image_staged_runtime_captures_without_network_install( # The submitted Job carries the staging init container. job_doc = json.loads( - kubectl( - ["get", f"job/{job_name}", "-n", NAMESPACE, "-o", "json"] - ).stdout + kubectl(["get", f"job/{job_name}", "-n", NAMESPACE, "-o", "json"]).stdout ) pod_spec = job_doc["spec"]["template"]["spec"] init_names = [c["name"] for c in pod_spec.get("initContainers", [])] @@ -141,3 +146,112 @@ def test_image_staged_runtime_captures_without_network_install( ) finally: _cleanup(job_name) + + +def _webhook_deployed() -> bool: + result = kubectl(["get", "mutatingwebhookconfiguration", "roar-lineage-injector"], check=False) + return result.returncode == 0 + + +def _wait_job_terminal(job_name: str, namespace: str, timeout: int = 300) -> bool: + import time + + deadline = time.time() + timeout + while time.time() < deadline: + status = kubectl(["get", f"job/{job_name}", "-n", namespace, "-o", "json"], check=False) + if status.returncode == 0: + conditions = json.loads(status.stdout).get("status", {}).get("conditions") or [] + if any( + c.get("status") == "True" and c.get("type") in ("Complete", "Failed") + for c in conditions + ): + return any( + c.get("status") == "True" and c.get("type") == "Complete" for c in conditions + ) + import time as _time + + _time.sleep(3) + return False + + +def test_webhook_injects_lineage_zero_touch( + k8s_cluster: None, + glaas_health: str, + tmp_path_factory: pytest.TempPathFactory, +) -> None: + """The full zero-touch story: plain kubectl apply, no roar on the client. + + A labeled namespace gets automatic injection; lineage is recovered + afterwards with `roar k8s attach` using the webhook-created Secret. + """ + if not _webhook_deployed(): + pytest.skip("webhook not deployed; run bootstrap_k8s.sh --with-webhook") + if not _runtime_image_loaded(): + pytest.skip("roar-runtime:dev not loaded") + + from .test_k8s_distributed import _roar + + suffix = uuid.uuid4().hex[:6] + auto_ns = f"roar-e2e-auto-{suffix}" + job_name = f"roar-auto-{suffix}" + + kubectl(["create", "namespace", auto_ns]) + kubectl(["label", "namespace", auto_ns, "roar.glaas.ai/lineage=enabled"]) + try: + # Plain kubectl apply — the client knows nothing about roar. + manifest = SINGLE_JOB_TEMPLATE.format(name=job_name, namespace=auto_ns) + kubectl(["apply", "-f", "-"], input_text=manifest) + + job_doc = json.loads( + kubectl(["get", f"job/{job_name}", "-n", auto_ns, "-o", "json"]).stdout + ) + annotations = job_doc["metadata"].get("annotations") or {} + parent_uid = annotations.get("roar.glaas.ai/parent-uid") + assert parent_uid, f"webhook did not annotate the Job: {annotations}" + pod_spec = job_doc["spec"]["template"]["spec"] + assert "roar-runtime-staging" in [c["name"] for c in pod_spec.get("initContainers", [])] + assert "roar.backends.k8s.pod_entry" in pod_spec["containers"][0]["command"][2] + + assert _wait_job_terminal(job_name, auto_ns), ( + f"job did not complete:\n{_pod_logs(f'job-name={job_name}')}" + ) + + attach_dir = tmp_path_factory.mktemp("k8s-webhook-attach") + _write_image_mode_project(attach_dir) + attached = _roar( + ["k8s", "attach", f"job/{job_name}", "-n", auto_ns, "--context", KUBE_CONTEXT], + cwd=attach_dir, + ) + assert attached.returncode == 0, attached.stdout + attached.stderr + assert "lineage reconstituted" in attached.stdout, attached.stdout + + tasks = _query(attach_dir, "SELECT id FROM jobs WHERE job_type = 'k8s_task'") + assert len(tasks) == 1 + attach_rows = _query( + attach_dir, + "SELECT job_uid FROM jobs WHERE execution_role = 'attach'", + ) + assert attach_rows and attach_rows[0]["job_uid"] == parent_uid + finally: + kubectl(["delete", "namespace", auto_ns, "--ignore-not-found"], check=False) + + +def test_webhook_leaves_unlabeled_namespaces_untouched( + k8s_cluster: None, +) -> None: + if not _webhook_deployed(): + pytest.skip("webhook not deployed; run bootstrap_k8s.sh --with-webhook") + + job_name = f"roar-plain-{uuid.uuid4().hex[:6]}" + try: + manifest = SINGLE_JOB_TEMPLATE.format(name=job_name, namespace=NAMESPACE) + kubectl(["apply", "-f", "-"], input_text=manifest) + job_doc = json.loads( + kubectl(["get", f"job/{job_name}", "-n", NAMESPACE, "-o", "json"]).stdout + ) + annotations = job_doc["metadata"].get("annotations") or {} + assert "roar.glaas.ai/parent-uid" not in annotations + command = job_doc["spec"]["template"]["spec"]["containers"][0]["command"] + assert "roar.backends.k8s.pod_entry" not in " ".join(command) + finally: + kubectl(["delete", f"job/{job_name}", "-n", NAMESPACE, "--ignore-not-found"], check=False) diff --git a/tests/backends/k8s/scripts/bootstrap_k8s.sh b/tests/backends/k8s/scripts/bootstrap_k8s.sh index 38f4b474..e3832eb0 100755 --- a/tests/backends/k8s/scripts/bootstrap_k8s.sh +++ b/tests/backends/k8s/scripts/bootstrap_k8s.sh @@ -45,6 +45,7 @@ WITH_MINIO=0 WITH_JOBSET=0 WITH_KUBEFLOW=0 WITH_KUBERAY=0 +WITH_WEBHOOK=0 SKIP_GLAAS=0 for arg in "$@"; do case "$arg" in @@ -52,6 +53,7 @@ for arg in "$@"; do --with-jobset) WITH_JOBSET=1 ;; --with-kubeflow) WITH_KUBEFLOW=1 ;; --with-kuberay) WITH_KUBERAY=1 ;; + --with-webhook) WITH_WEBHOOK=1 ;; --skip-glaas) SKIP_GLAAS=1 ;; *) echo "error: unknown flag: $arg" >&2 @@ -208,6 +210,17 @@ if ((WITH_KUBEFLOW == 1)); then echo "warning: trainer runtimes overlay failed; TrainJob e2e will create its own runtime" >&2 fi +if ((WITH_WEBHOOK == 1)); then + echo "▶ Building and loading roar-runtime image" + # Always build: the docker layer cache makes this a no-op unless the + # wheel changed, and a stale image silently breaks the webhook. + bash "$REPO_ROOT/scripts/build_runtime_image.sh" + kind load docker-image roar-runtime:dev --name "$CLUSTER_NAME" + bash "$HARNESS_DIR/scripts/deploy_webhook.sh" + kubectl_ctx -n roar-system rollout restart deployment/roar-webhook + kubectl_ctx -n roar-system rollout status deployment/roar-webhook --timeout=180s +fi + if ((WITH_KUBERAY == 1)); then echo "▶ Installing KubeRay operator ${KUBERAY_VERSION}" kubectl_ctx create -k "github.com/ray-project/kuberay/ray-operator/config/default?ref=${KUBERAY_VERSION}" 2>/dev/null || \ diff --git a/tests/backends/k8s/scripts/deploy_webhook.sh b/tests/backends/k8s/scripts/deploy_webhook.sh new file mode 100644 index 00000000..b3a4951c --- /dev/null +++ b/tests/backends/k8s/scripts/deploy_webhook.sh @@ -0,0 +1,191 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Deploys the roar lineage mutating webhook into the KIND harness. +# +# - self-signed CA + server cert (no cert-manager dependency in the harness) +# - roar-runtime:dev image serves the webhook (built/loaded by caller) +# - namespaceSelector opt-in: roar.glaas.ai/lineage=enabled +# - failurePolicy Ignore: webhook outages never block workloads +# +# Env overrides: WEBHOOK_GLAAS_URL (server-visible from the webhook pod), +# WEBHOOK_CLUSTER_GLAAS_URL (visible from injected workload pods). + +HARNESS_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +TOOLS_BIN="$HARNESS_DIR/.tools/bin" +export PATH="$TOOLS_BIN:$PATH" + +CLUSTER_NAME="roar-k8s-e2e" +KUBE_CONTEXT="kind-${CLUSTER_NAME}" +WEBHOOK_NS="roar-system" +SERVICE="roar-webhook" +RUNTIME_IMAGE="${RUNTIME_IMAGE:-roar-runtime:dev}" +WEBHOOK_GLAAS_URL="${WEBHOOK_GLAAS_URL:-http://glaas.roar-e2e.svc.cluster.local:3001}" +WEBHOOK_CLUSTER_GLAAS_URL="${WEBHOOK_CLUSTER_GLAAS_URL:-http://glaas.roar-e2e.svc.cluster.local:3001}" + +kubectl_ctx() { kubectl --context "$KUBE_CONTEXT" "$@"; } + +echo "▶ Generating webhook TLS material" +CERT_DIR="$(mktemp -d)" +trap 'rm -rf "$CERT_DIR"' EXIT +openssl req -x509 -newkey rsa:2048 -nodes -days 30 \ + -keyout "$CERT_DIR/ca.key" -out "$CERT_DIR/ca.crt" \ + -subj "/CN=roar-webhook-ca" >/dev/null 2>&1 +openssl req -newkey rsa:2048 -nodes \ + -keyout "$CERT_DIR/tls.key" -out "$CERT_DIR/tls.csr" \ + -subj "/CN=${SERVICE}.${WEBHOOK_NS}.svc" >/dev/null 2>&1 +openssl x509 -req -in "$CERT_DIR/tls.csr" -days 30 \ + -CA "$CERT_DIR/ca.crt" -CAkey "$CERT_DIR/ca.key" -CAcreateserial \ + -out "$CERT_DIR/tls.crt" \ + -extfile <(printf "subjectAltName=DNS:%s.%s.svc,DNS:%s.%s.svc.cluster.local" \ + "$SERVICE" "$WEBHOOK_NS" "$SERVICE" "$WEBHOOK_NS") >/dev/null 2>&1 +CA_BUNDLE="$(base64 -w0 <"$CERT_DIR/ca.crt")" + +echo "▶ Deploying webhook to namespace $WEBHOOK_NS" +kubectl_ctx create namespace "$WEBHOOK_NS" --dry-run=client -o yaml | kubectl_ctx apply -f - +kubectl_ctx -n "$WEBHOOK_NS" create secret tls roar-webhook-tls \ + --cert="$CERT_DIR/tls.crt" --key="$CERT_DIR/tls.key" \ + --dry-run=client -o yaml | kubectl_ctx apply -f - + +kubectl_ctx apply -f - < None: for container in pod_spec["initContainers"] if container["name"] == "roar-runtime-staging" ] - staging_volumes = [ - volume for volume in pod_spec["volumes"] if volume["name"] == "roar-runtime" - ] + staging_volumes = [volume for volume in pod_spec["volumes"] if volume["name"] == "roar-runtime"] assert len(staging_inits) == 1 assert len(staging_volumes) == 1 diff --git a/tests/backends/k8s/unit/test_webhook.py b/tests/backends/k8s/unit/test_webhook.py new file mode 100644 index 00000000..d6fc86b2 --- /dev/null +++ b/tests/backends/k8s/unit/test_webhook.py @@ -0,0 +1,170 @@ +from __future__ import annotations + +import base64 +import copy +import json +from typing import Any + +from roar.backends.k8s.webhook import ( + ANNOTATION_PARENT_UID, + WebhookSettings, + mutate_admission_review, +) + +from .conftest import SINGLE_JOB_MANIFEST + +SETTINGS = WebhookSettings( + glaas_url="http://glaas.roar-e2e.svc.cluster.local:3001", + cluster_glaas_url="http://glaas.roar-e2e.svc.cluster.local:3001", + runtime_source="image", + runtime_image="roar-runtime:dev", +) + + +class _Spy: + def __init__(self, fail: bool = False) -> None: + self.calls: list[tuple] = [] + self.fail = fail + + def create_secret(self, namespace: str, name: str, data: dict[str, str]) -> None: + if self.fail: + raise RuntimeError("secret boom") + self.calls.append(("secret", namespace, name, sorted(data))) + + def register_session(self, session_id: str, token_hash: str) -> None: + if self.fail: + raise RuntimeError("glaas boom") + self.calls.append(("session", session_id)) + + +def _review(obj: dict[str, Any], *, dry_run: bool = False) -> dict[str, Any]: + return { + "apiVersion": "admission.k8s.io/v1", + "kind": "AdmissionReview", + "request": { + "uid": "req-1", + "namespace": "ml-auto", + "dryRun": dry_run, + "object": obj, + }, + } + + +def _decode_patch(result: dict[str, Any]) -> list[dict[str, Any]]: + return json.loads(base64.b64decode(result["response"]["patch"])) + + +def test_job_admission_is_instrumented() -> None: + spy = _Spy() + result = mutate_admission_review( + _review(copy.deepcopy(SINGLE_JOB_MANIFEST)), + settings=SETTINGS, + create_secret=spy.create_secret, + register_session=spy.register_session, + ) + + response = result["response"] + assert response["allowed"] is True + assert response["patchType"] == "JSONPatch" + + kinds = [call[0] for call in spy.calls] + assert kinds == ["session", "secret"] + secret_call = next(call for call in spy.calls if call[0] == "secret") + assert secret_call[1] == "ml-auto" # request namespace, not manifest's + assert secret_call[3] == ["session_id", "token"] + + patch = _decode_patch(result) + spec_patch = next(op for op in patch if op["path"] == "/spec") + command = spec_patch["value"]["template"]["spec"]["containers"][0]["command"] + assert command[:2] == ["/bin/sh", "-c"] + assert "roar.backends.k8s.pod_entry" in command[2] + init_names = [ + c["name"] for c in spec_patch["value"]["template"]["spec"].get("initContainers", []) + ] + assert "roar-runtime-staging" in init_names + + annotations_patch = next(op for op in patch if op["path"] == "/metadata/annotations") + assert ANNOTATION_PARENT_UID in annotations_patch["value"] + + +def test_non_workload_objects_pass_through() -> None: + spy = _Spy() + result = mutate_admission_review( + _review({"apiVersion": "v1", "kind": "ConfigMap", "metadata": {"name": "x"}}), + settings=SETTINGS, + create_secret=spy.create_secret, + register_session=spy.register_session, + ) + assert result["response"]["allowed"] is True + assert "patch" not in result["response"] + assert spy.calls == [] + + +def test_opt_out_label_skips_injection() -> None: + manifest = copy.deepcopy(SINGLE_JOB_MANIFEST) + manifest["metadata"].setdefault("labels", {})["roar.glaas.ai/lineage"] = "disabled" + spy = _Spy() + result = mutate_admission_review( + _review(manifest), + settings=SETTINGS, + create_secret=spy.create_secret, + register_session=spy.register_session, + ) + assert "patch" not in result["response"] + assert spy.calls == [] + + +def test_already_instrumented_is_idempotent() -> None: + manifest = copy.deepcopy(SINGLE_JOB_MANIFEST) + manifest["metadata"].setdefault("annotations", {})[ANNOTATION_PARENT_UID] = "cafe0123" + spy = _Spy() + result = mutate_admission_review( + _review(manifest), + settings=SETTINGS, + create_secret=spy.create_secret, + register_session=spy.register_session, + ) + assert "patch" not in result["response"] + assert spy.calls == [] + + +def test_dry_run_has_no_side_effects() -> None: + spy = _Spy() + result = mutate_admission_review( + _review(copy.deepcopy(SINGLE_JOB_MANIFEST), dry_run=True), + settings=SETTINGS, + create_secret=spy.create_secret, + register_session=spy.register_session, + ) + assert result["response"]["allowed"] is True + assert "patch" not in result["response"] + assert spy.calls == [] + + +def test_failures_never_block_admission() -> None: + spy = _Spy(fail=True) + result = mutate_admission_review( + _review(copy.deepcopy(SINGLE_JOB_MANIFEST)), + settings=SETTINGS, + create_secret=spy.create_secret, + register_session=spy.register_session, + ) + response = result["response"] + assert response["allowed"] is True + assert "patch" not in response + assert any("lineage injection skipped" in w for w in response.get("warnings", [])) + + +def test_settings_from_environ() -> None: + settings = WebhookSettings.from_environ( + { + "ROAR_WEBHOOK_GLAAS_URL": "http://glaas:3001", + "ROAR_WEBHOOK_CLUSTER_GLAAS_URL": "http://glaas.ns.svc:3001", + "ROAR_WEBHOOK_RUNTIME_SOURCE": "image", + "ROAR_WEBHOOK_RUNTIME_IMAGE": "roar-runtime:v1", + "ROAR_WEBHOOK_MOUNT_MAP": '{"/data": "gs://bucket"}', + } + ) + assert settings.glaas_url == "http://glaas:3001" + assert settings.runtime_image == "roar-runtime:v1" + assert settings.mount_map == {"/data": "gs://bucket"} From a32f2a8244de56bfb6efdad6f60797fc462c1885 Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Tue, 14 Jul 2026 16:47:41 +0000 Subject: [PATCH 11/31] fix(ray): correct the 2.46 capture diagnosis and warn on unpatchable executors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes the ray 2.46 circle-back — by disproving it. A local probe (worker_process_setup_hook wrapping FunctionActorManager methods) showed ray 2.46 DOES route task execution through get_execution_info; the empty task fragments previously blamed on 2.46 were entirely the pathlib/io.open capture gap, confounded because that fix landed together with the 2.54 image pin, which took the credit. - strict KubeRay e2e verified green on BOTH 2.46.0 and 2.54.0 with the final code; the e2e image is now overridable via ROAR_E2E_RAY_IMAGE (default stays the native harness's 2.54 pin) and rayVersion in the RayJob template tracks the image tag - detect-and-warn insurance: _patch_ray_task_execution_for_native_flush now prints a loud stderr warning (with the ray version) when the function manager cannot be imported or exposes neither get_execution_info nor _make_actor_method_executor, instead of silently degrading to task fragments without identity or file refs - developer doc and design doc corrected Verified: strict RayJob e2e 1/1 on 2.46 and 1/1 on 2.54; ray + k8s unit suites green; default gate 2040 passed; ruff + mypy clean. Co-Authored-By: Claude Fable 5 --- docs/developer/k8s-integration.md | 23 ++++++++++-------- roar/backends/ray/roar_worker.py | 29 ++++++++++++++++++++++- tests/backends/k8s/e2e/test_k8s_rayjob.py | 13 +++++----- 3 files changed, 48 insertions(+), 17 deletions(-) diff --git a/docs/developer/k8s-integration.md b/docs/developer/k8s-integration.md index 804723cb..02236d25 100644 --- a/docs/developer/k8s-integration.md +++ b/docs/developer/k8s-integration.md @@ -61,16 +61,19 @@ reconstitution/attach delegate to the **Ray** backend's reconstituter, since the streamed fragments are Ray `TaskFragment` payloads that merge as `ray_task` jobs. -Per-task capture fidelity in KubeRay workers requires the Ray version the -roar_worker internals target (the native harness's pin, currently 2.54): -on 2.46, task execution was observed not to route through the patched -`FunctionActorManager.get_execution_info`, so task fragments arrived as -empty shells. Additionally, `roar_worker._startup` now patches `io.open` -alongside `builtins.open` — pathlib (`Path.read_bytes`/`write_bytes`) -resolves `io.open` by module attribute, and without the preload tracer -(KubeRay pods) those opens were invisible; on preload-equipped native -workers the duplicate python-level captures are absorbed by the existing -`native > python` dedup. +Per-task capture fidelity in KubeRay workers is verified on ray 2.46 and +2.54 (the strict e2e runs against the native harness's 2.54 pin by +default; override with `ROAR_E2E_RAY_IMAGE` to smoke other versions). An +earlier diagnosis blamed 2.46's executor internals; empirically both +versions route task execution through the patched +`FunctionActorManager.get_execution_info` — the real gap was pathlib: +`roar_worker._startup` now patches `io.open` alongside `builtins.open` +(`Path.read_bytes`/`write_bytes` resolve `io.open` by module attribute, +and without the preload tracer those opens were invisible; on +preload-equipped native workers the duplicate python-level captures are +absorbed by the existing `native > python` dedup). If a future Ray moves +the executor internals, `roar_worker` now warns loudly instead of +silently emitting task fragments without identity or file refs. **Mounted storage** (`mount_map.py`): FUSE CSI mounts surface object I/O as local file syscalls under a mount path. The rewriter derives a per-container diff --git a/roar/backends/ray/roar_worker.py b/roar/backends/ray/roar_worker.py index 8379bfe3..e3d6f192 100644 --- a/roar/backends/ray/roar_worker.py +++ b/roar/backends/ray/roar_worker.py @@ -602,6 +602,20 @@ def _deactivate_threading_patch_for_native_task_attribution() -> None: threading.Thread.start = _real_thread_start # type: ignore[method-assign] +def _warn_task_capture_unavailable(reason: str) -> None: + try: + import ray + + version = getattr(ray, "__version__", "unknown") + except Exception: + version = "unknown" + print( + f"[roar] warning: per-task lineage capture is unavailable on ray {version} " + f"({reason}); task fragments may lack identity and file refs", + file=sys.stderr, + ) + + def _wrap_task_executor_for_native_flush( function: Callable[..., Any], *, @@ -654,7 +668,20 @@ def _wrapped(*args, **kwargs): def _patch_ray_task_execution_for_native_flush() -> None: try: from ray._private.function_manager import FunctionActorManager, FunctionExecutionInfo - except Exception: + except Exception as exc: + _warn_task_capture_unavailable(f"cannot import ray function manager: {exc}") + return + + if not callable(getattr(FunctionActorManager, "get_execution_info", None)) and not callable( + getattr(FunctionActorManager, "_make_actor_method_executor", None) + ): + # A future Ray moved the executor internals: task-boundary capture + # cannot engage, and task fragments would silently arrive without + # per-task identity or file refs. Fail loudly instead + # (verified engaging on 2.46 and 2.54). + _warn_task_capture_unavailable( + "FunctionActorManager has neither get_execution_info nor _make_actor_method_executor" + ) return current_get_execution_info = getattr(FunctionActorManager, "get_execution_info", None) diff --git a/tests/backends/k8s/e2e/test_k8s_rayjob.py b/tests/backends/k8s/e2e/test_k8s_rayjob.py index c5a62f04..7bdca011 100644 --- a/tests/backends/k8s/e2e/test_k8s_rayjob.py +++ b/tests/backends/k8s/e2e/test_k8s_rayjob.py @@ -13,6 +13,7 @@ from __future__ import annotations +import os import sqlite3 import subprocess import uuid @@ -38,11 +39,10 @@ pytest.mark.timeout(1500), ] -# Pinned to the same Ray version as the native compose harness -# (tests/backends/ray/e2e/Dockerfile): the roar_worker per-task capture -# targets this version's internals, and 2.46 was observed not to route -# task execution through the patched FunctionActorManager path. -RAY_IMAGE = "rayproject/ray:2.54.0-py312-cpu" +# Defaults to the native compose harness's pin +# (tests/backends/ray/e2e/Dockerfile); override with ROAR_E2E_RAY_IMAGE to +# smoke other Ray versions (verified: 2.46.0 and 2.54.0). +RAY_IMAGE = os.environ.get("ROAR_E2E_RAY_IMAGE", "rayproject/ray:2.54.0-py312-cpu") RAY_TRAIN_SCRIPT = """\ from pathlib import Path @@ -93,7 +93,7 @@ def train_shard() -> int: - name: ray-job-submitter image: {ray_image} rayClusterSpec: - rayVersion: "2.54.0" + rayVersion: "{ray_version}" headGroupSpec: rayStartParams: num-cpus: "1" @@ -190,6 +190,7 @@ def test_rayjob_live_delegates_to_ray_backend( pinned_node=PINNED_NODE, ray_image=RAY_IMAGE, worker_script=_indent_script(RAY_TRAIN_SCRIPT), + ray_version=RAY_IMAGE.split(":")[1].split("-")[0], ), encoding="utf-8", ) From e7342fc02895dfd1e9234b6ff0c42e8e7c628d70 Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Tue, 14 Jul 2026 17:41:18 +0000 Subject: [PATCH 12/31] feat(k8s): opt-in roar-proxy S3 sidecar for hook-invisible clients Add k8s.proxy_sidecar / k8s.proxy_upstream (and webhook ROAR_WEBHOOK_PROXY_SIDECAR / ROAR_WEBHOOK_PROXY_UPSTREAM): the rewriter appends a roar-s3-proxy native sidecar (init container with restartPolicy Always) running the roar-proxy binary from the staged runtime tree, so S3 traffic from clients the botocore hooks cannot see (non-Python binaries, raw HTTP) still lands as s3:// lineage refs. - Requires image staging; install mode warns and skips. - roar-proxy re-signs upstream requests, so the sidecar inherits the workload container's AWS_* credential env entries. - The proxy binds loopback only; the startup probe is an exec probe (kubelet tcpSocket probes dial the pod IP and never succeed). - Wrapped containers get AWS_ENDPOINT_URL=http://127.0.0.1:19191 only when the user has not set their own (explicit endpoints win). - pod_entry parses the proxy log (shared emptyDir) and folds refs into the fragment with capture_method="proxy", deduping objects the hooks already captured. - Image-mode fragments no longer report the staged /roar-runtime/** tree as inputs (filtered at reconstitution; capture stays raw). Live-validated on the KIND harness with a hook-invisible raw-urllib S3 client against MinIO (23/23 k8s e2e). Co-Authored-By: Claude Fable 5 --- docs/developer/k8s-integration.md | 43 ++++-- roar/backends/k8s/config.py | 23 +++ roar/backends/k8s/lineage.py | 22 +++ roar/backends/k8s/manifest.py | 145 ++++++++++++++++++ roar/backends/k8s/pod_entry.py | 65 ++++++++ roar/backends/k8s/submit.py | 10 ++ roar/backends/k8s/webhook.py | 7 + tests/backends/k8s/e2e/test_k8s_phase3.py | 123 +++++++++++++++ tests/backends/k8s/unit/test_proxy_sidecar.py | 145 ++++++++++++++++++ 9 files changed, 572 insertions(+), 11 deletions(-) create mode 100644 tests/backends/k8s/unit/test_proxy_sidecar.py diff --git a/docs/developer/k8s-integration.md b/docs/developer/k8s-integration.md index 02236d25..c25aa55d 100644 --- a/docs/developer/k8s-integration.md +++ b/docs/developer/k8s-integration.md @@ -19,8 +19,8 @@ in the dev meta-repo): Phases 1–3 implemented and live-validated (submit wrapping, operator adapters incl. real training-operator v1 and trainer v2 controllers, multi-pod capture, `roar k8s attach`, bundle-mode fallback, object-store I/O hooks, mount-map rewriting, retry-chaos -coverage, RayJob delegation, the roar-runtime image, and the mutating -webhook injector). Remaining from Phase 3: the opt-in proxy sidecar and +coverage, RayJob delegation, the roar-runtime image, the mutating +webhook injector, and the opt-in proxy sidecar). Remaining from Phase 3: a packaged Helm chart (the harness deploys via `tests/backends/k8s/scripts/deploy_webhook.sh`). @@ -46,6 +46,9 @@ the parent-uid annotation; dry-run admissions are side-effect-free (`sideEffects: NoneOnDryRun`); every internal failure returns allowed with a warning and pairs with `failurePolicy: Ignore` — lineage never blocks admission. Reconstitution is client-driven via `roar k8s attach`. +`ROAR_WEBHOOK_PROXY_SIDECAR=true` (+ optional +`ROAR_WEBHOOK_PROXY_UPSTREAM`) makes the injector add the proxy sidecar +to every workload it instruments. **RayJob delegation** (`rayjob.py`): KubeRay overwrites container commands with `ray start --block`, so command-wrapping can't see user code. Instead @@ -106,11 +109,26 @@ Two capture channels feed each pod's fragment: `job_inputs/job_outputs.byte_ranges`). Hooks no-op outside pods (env unset) and never raise into user code. - The `roar-proxy` S3 sidecar is deliberately not part of the CLI-side - backend: the hooks win on attribution and avoid `AWS_ENDPOINT_URL` - rewiring (which explicit-`endpoint_url` clients bypass). The proxy joins - in the Phase-3 webhook injector as an opt-in sidecar for non-Python S3 - clients (see the design doc's Phase 3). +- **Proxy sidecar** (opt-in, `k8s.proxy_sidecar = true`): for S3 clients + the botocore hooks can't see (Go/Rust/Java binaries, plain HTTP). The + rewriter appends a `roar-s3-proxy` **native sidecar** (an init container + with `restartPolicy: Always`, k8s ≥ 1.29/GA 1.33) running the + `roar-proxy` binary from the staged runtime tree — it therefore + **requires image staging** (`k8s.runtime_source = "image"`); with + install mode it warns and skips. Wrapped containers get + `AWS_ENDPOINT_URL=http://127.0.0.1:19191` (only when the user hasn't + set their own — explicit user endpoints win, and such clients simply + bypass the proxy). The proxy logs each request to a shared emptyDir + (`ROAR_K8S_PROXY_LOG`); `pod_entry` parses it via + `roar.execution.cluster.proxy.parse_log_line` and folds refs into the + fragment with `capture_method="proxy"`, skipping objects the hooks + already captured (hooks win on attribution). Two gotchas verified live: + `roar-proxy` is a **re-signing** reverse proxy, so the sidecar needs + AWS credentials — the rewriter copies the workload container's + `AWS_ACCESS_KEY_ID`/`AWS_SECRET_ACCESS_KEY`/session/role env entries + onto the sidecar (without them every forward 502s); and it binds + loopback only, so its startup probe must be an `exec` probe (kubelet + `tcpSocket` probes dial the pod IP, which never succeeds). Transport is streaming-first with a bundle fallback: when `k8s.bundle_dir` names a mounted shared volume and GLaaS is unreachable from the pod (probe @@ -181,11 +199,14 @@ recovered parent uid, and reconstitutes the streamed fragments. ## 4. Config Section `k8s` (registered `BackendConfigAdapter`): `enabled`, `tracer`, +`runtime_source` (`install`|`image`), `runtime_image`, `runtime_install_requirement`, `cluster_glaas_url`, `bundle_dir`, -`wait_for_completion`, `wait_timeout_seconds`, `poll_interval_seconds`, -`fragment_session_ttl_seconds`, plus the `[k8s.mount_map]` table -(config-file-only; maps mount paths to object-store URIs). Env overrides -beat config: `ROAR_CLUSTER_PIP_REQ`, `ROAR_CLUSTER_GLAAS_URL`. +`proxy_sidecar`, `proxy_upstream` (upstream S3 endpoint the proxy +forwards to; empty = AWS), `wait_for_completion`, `wait_timeout_seconds`, +`poll_interval_seconds`, `fragment_session_ttl_seconds`, plus the +`[k8s.mount_map]` table (config-file-only; maps mount paths to +object-store URIs). Env overrides beat config: `ROAR_CLUSTER_PIP_REQ`, +`ROAR_CLUSTER_GLAAS_URL`. ## 5. Tests diff --git a/roar/backends/k8s/config.py b/roar/backends/k8s/config.py index 79bce43c..a3e6b3b4 100644 --- a/roar/backends/k8s/config.py +++ b/roar/backends/k8s/config.py @@ -27,6 +27,11 @@ class K8sBackendConfig(BaseModel): runtime_source: str = "install" runtime_image: str = "" runtime_install_requirement: str = "" + # Opt-in roar-proxy S3 sidecar for clients the in-process hooks can't + # see (aws CLI, s5cmd, Go/Java SDKs). Requires runtime_source="image" + # (the sidecar runs the proxy binary from runtime_image). + proxy_sidecar: bool = False + proxy_upstream: str = "" cluster_glaas_url: str = "" bundle_dir: str = "" # Explicit mount-path -> object-store URI mapping for mounted storage @@ -80,6 +85,24 @@ class K8sBackendConfig(BaseModel): "host-visible glaas.url (ROAR_CLUSTER_GLAAS_URL env always wins)" ), ), + "k8s.proxy_sidecar": ConfigurableKeySpec( + value_type=bool, + default=False, + description=( + "Inject the roar-proxy S3 sidecar (native init container) to capture " + "S3 traffic from non-Python clients; requires k8s.runtime_source='image'. " + "In-process hooks remain the primary capture; user-set AWS_ENDPOINT_URL " + "always wins over the injected proxy redirect" + ), + ), + "k8s.proxy_upstream": ConfigurableKeySpec( + value_type=str, + default="", + description=( + "Upstream S3 endpoint the proxy sidecar forwards to (empty = real AWS S3); " + "set for MinIO/LocalStack-style deployments" + ), + ), "k8s.bundle_dir": ConfigurableKeySpec( value_type=str, default="", diff --git a/roar/backends/k8s/lineage.py b/roar/backends/k8s/lineage.py index 5f1658d6..cf95a52f 100644 --- a/roar/backends/k8s/lineage.py +++ b/roar/backends/k8s/lineage.py @@ -69,6 +69,7 @@ def collect_k8s_fragments( continue try: rewrite_fragment_paths(payload) + _drop_runtime_staging_noise(payload) parsed.append(ExecutionFragment.from_dict(payload)) except Exception: continue @@ -87,6 +88,27 @@ def collect_k8s_fragments( return len(parsed) +def _drop_runtime_staging_noise(fragment: dict[str, Any]) -> None: + """Filter roar's own staged runtime out of the captured signal. + + Image-staged pods import roar from the /roar-runtime emptyDir, and the + tracer faithfully records those reads. Like ignore_package_reads for + site-packages, they are runtime noise, not workload lineage — dropped + here at reconstitution (capture stays raw in the fragment stream). + """ + for list_key in ("reads", "writes"): + refs = fragment.get(list_key) + if not isinstance(refs, list): + continue + fragment[list_key] = [ + ref + for ref in refs + if not ( + isinstance(ref, dict) and str(ref.get("path") or "").startswith("/roar-runtime/") + ) + ] + + def resolve_active_session_context(db_path: str) -> tuple[int | None, int]: """Return ``(active_session_id, next_step_number)`` for fragment merges.""" if not db_path or not os.path.exists(db_path): diff --git a/roar/backends/k8s/manifest.py b/roar/backends/k8s/manifest.py index b92d8b3a..dd735adb 100644 --- a/roar/backends/k8s/manifest.py +++ b/roar/backends/k8s/manifest.py @@ -51,6 +51,14 @@ _RUNTIME_STAGING_MOUNT = "/roar-runtime" _RUNTIME_STAGING_INIT_CONTAINER = "roar-runtime-staging" +_PROXY_SIDECAR_NAME = "roar-s3-proxy" +_PROXY_LOG_VOLUME = "roar-proxy-log" +_PROXY_LOG_MOUNT = "/roar-proxy-log" +_PROXY_LOG_FILE = f"{_PROXY_LOG_MOUNT}/proxy.log" +_PROXY_PORT = 19191 +# Any ABI tree's copy works: the proxy is a standalone Rust binary. +_PROXY_BINARY = "/opt/roar-runtime/cp312/roar/bin/roar-proxy" + # Terminal workload conditions, unioned across kinds: Job uses # Complete/SuccessCriteriaMet + Failed/FailureTarget, JobSet uses # Completed/Failed, PyTorchJob v1 uses Succeeded/Failed, TrainJob uses @@ -224,6 +232,8 @@ def rewrite_manifest_for_lineage( mount_map: dict[str, str] | None = None, runtime_source: str = "install", runtime_image: str = "", + proxy_sidecar: bool = False, + proxy_upstream: str = "", namespace_override: str | None = None, ) -> K8sManifestRewrite: """Return a rewritten copy of ``documents`` with lineage instrumentation. @@ -274,6 +284,8 @@ def rewrite_manifest_for_lineage( config_mount_map=dict(mount_map or {}), runtime_source=runtime_source, runtime_image=runtime_image, + proxy_sidecar=proxy_sidecar, + proxy_upstream=proxy_upstream, ) if workload.rewrite_style == "trainer_override": @@ -346,11 +358,19 @@ class _EnvContract: config_mount_map: dict[str, str] = field(default_factory=dict) runtime_source: str = "install" runtime_image: str = "" + proxy_sidecar: bool = False + proxy_upstream: str = "" @property def image_staging(self) -> bool: return self.runtime_source == "image" and bool(self.runtime_image) + @property + def proxy_enabled(self) -> bool: + # The sidecar runs the proxy binary from the runtime image, so it + # is only available in image-staging mode. + return self.proxy_sidecar and self.image_staging + def _rewrite_pod_specs( pod_specs: list[PodSpecRef], @@ -389,6 +409,9 @@ def _rewrite_pod_specs( container.pop("args", None) if contract.image_staging: _add_staging_volume_mount(container) + if contract.proxy_enabled: + _add_proxy_log_mount(container) + _add_proxy_env(container) _inject_env_contract( container.setdefault("env", []), contract=contract, @@ -402,9 +425,131 @@ def _rewrite_pod_specs( pod_wrapped = True if pod_wrapped and contract.image_staging: _add_runtime_staging(ref.spec, runtime_image=contract.runtime_image) + if pod_wrapped and contract.proxy_enabled: + _add_proxy_sidecar( + ref.spec, + runtime_image=contract.runtime_image, + upstream=contract.proxy_upstream, + credential_env=_collect_aws_credential_env(containers), + ) return wrapped, skipped +# The proxy re-signs forwarded requests with its own credentials (clients +# hit localhost unauthenticated); it inherits the workload's AWS auth. +_AWS_CREDENTIAL_ENV_NAMES = ( + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "AWS_SESSION_TOKEN", + "AWS_REGION", + "AWS_DEFAULT_REGION", + "AWS_WEB_IDENTITY_TOKEN_FILE", + "AWS_ROLE_ARN", +) + + +def _collect_aws_credential_env(containers: list[Any]) -> list[dict[str, Any]]: + for container in containers: + if not isinstance(container, dict): + continue + entries = [ + entry + for entry in container.get("env") or [] + if isinstance(entry, dict) and entry.get("name") in _AWS_CREDENTIAL_ENV_NAMES + ] + if entries: + return [dict(entry) for entry in entries] + return [] + + +def _add_proxy_sidecar( + pod_spec: dict[str, Any], + *, + runtime_image: str, + upstream: str, + credential_env: list[dict[str, Any]] | None = None, +) -> None: + """Add the S3 proxy as a native sidecar (idempotent by name). + + A native sidecar (init container with restartPolicy Always, GA in + k8s 1.33) keeps serving while app containers run but never blocks + Job completion. The startupProbe orders it before the app containers + so early S3 calls don't race the listener. + """ + volumes = pod_spec.setdefault("volumes", []) + if isinstance(volumes, list) and not any( + isinstance(volume, dict) and volume.get("name") == _PROXY_LOG_VOLUME for volume in volumes + ): + volumes.append({"name": _PROXY_LOG_VOLUME, "emptyDir": {}}) + + init_containers = pod_spec.setdefault("initContainers", []) + if not isinstance(init_containers, list) or any( + isinstance(container, dict) and container.get("name") == _PROXY_SIDECAR_NAME + for container in init_containers + ): + return + + proxy_command = f"exec {_PROXY_BINARY} --port {_PROXY_PORT}" + if upstream: + proxy_command += f" --upstream {shlex.quote(upstream)}" + proxy_command += f" > {_PROXY_LOG_FILE} 2>&1" + + sidecar: dict[str, Any] = { + "name": _PROXY_SIDECAR_NAME, + "image": runtime_image, + "restartPolicy": "Always", + "command": ["/bin/sh", "-c", proxy_command], + "volumeMounts": [{"name": _PROXY_LOG_VOLUME, "mountPath": _PROXY_LOG_MOUNT}], + # The proxy binds loopback only, which the pod's shared network + # namespace makes reachable from app containers — but kubelet + # tcpSocket probes target the pod IP, so the probe must exec + # inside the netns instead. + "startupProbe": { + "exec": { + "command": [ + "python", + "-c", + f"import socket; socket.create_connection(('127.0.0.1', {_PROXY_PORT}), 1)", + ] + }, + "periodSeconds": 1, + "failureThreshold": 30, + }, + } + if credential_env: + sidecar["env"] = credential_env + init_containers.append(sidecar) + + +def _add_proxy_log_mount(container: dict[str, Any]) -> None: + mounts = container.setdefault("volumeMounts", []) + if isinstance(mounts, list) and not any( + isinstance(mount, dict) and mount.get("name") == _PROXY_LOG_VOLUME for mount in mounts + ): + mounts.append( + { + "name": _PROXY_LOG_VOLUME, + "mountPath": _PROXY_LOG_MOUNT, + "readOnly": True, + } + ) + + +def _add_proxy_env(container: dict[str, Any]) -> None: + env = container.setdefault("env", []) + if not isinstance(env, list): + return + existing = { + str(entry.get("name")) for entry in env if isinstance(entry, dict) and entry.get("name") + } + # User-set AWS_ENDPOINT_URL always wins (explicit endpoints bypass + # the proxy by design — the hooks still capture those clients). + if "AWS_ENDPOINT_URL" not in existing: + env.append({"name": "AWS_ENDPOINT_URL", "value": f"http://127.0.0.1:{_PROXY_PORT}"}) + if "ROAR_K8S_PROXY_LOG" not in existing: + env.append({"name": "ROAR_K8S_PROXY_LOG", "value": _PROXY_LOG_FILE}) + + def _add_runtime_staging(pod_spec: dict[str, Any], *, runtime_image: str) -> None: """Add the shared volume + init container that stage the roar runtime. diff --git a/roar/backends/k8s/pod_entry.py b/roar/backends/k8s/pod_entry.py index 3e1a12d2..8d9ef4ed 100644 --- a/roar/backends/k8s/pod_entry.py +++ b/roar/backends/k8s/pod_entry.py @@ -173,6 +173,15 @@ def _augment_with_object_io(fragments: list[dict]) -> None: events_path = Path(os.environ.get("ROAR_K8S_OBJECT_IO_FILE") or _object_io_events_path()) reads, writes = load_object_io_refs(events_path) + + proxy_reads, proxy_writes = _load_proxy_log_refs( + os.environ.get("ROAR_K8S_PROXY_LOG"), + seen_reads={ref["path"] for ref in reads}, + seen_writes={ref["path"] for ref in writes}, + ) + reads.extend(proxy_reads) + writes.extend(proxy_writes) + if not reads and not writes: return for fragment in fragments: @@ -180,6 +189,62 @@ def _augment_with_object_io(fragments: list[dict]) -> None: fragment.setdefault("writes", []).extend(writes) +def _load_proxy_log_refs( + log_path: str | None, + *, + seen_reads: set[str], + seen_writes: set[str], +) -> tuple[list[dict], list[dict]]: + """Parse the proxy sidecar's log into refs for hook-invisible clients. + + The in-process hooks are the primary S3 capture — proxy entries are + only added for paths the hooks did not already record. + """ + if not log_path: + return [], [] + path = Path(log_path) + if not path.is_file(): + return [], [] + + try: + from roar.execution.cluster.proxy import parse_log_line + except Exception: + return [], [] + + write_ops = {"PutObject", "CompleteMultipartUpload", "CopyObject"} + read_ops = {"GetObject"} + reads: dict[str, dict] = {} + writes: dict[str, dict] = {} + try: + lines = path.read_text(encoding="utf-8", errors="replace").splitlines() + except OSError: + return [], [] + for line in lines: + entry = parse_log_line(line) + if entry is None: + continue + if entry.operation in write_ops: + mode_seen, bucket_refs = seen_writes, writes + elif entry.operation in read_ops: + mode_seen, bucket_refs = seen_reads, reads + else: + continue + s3_path = f"s3://{entry.bucket}/{entry.key}" + if s3_path in mode_seen: + continue + ref: dict = { + "path": s3_path, + "hash": entry.etag or None, + "hash_algorithm": "etag" if entry.etag else "", + "size": int(entry.size_bytes or 0), + "capture_method": "proxy", + } + if entry.byte_ranges: + ref["byte_ranges"] = entry.byte_ranges + bucket_refs[s3_path] = ref + return list(reads.values()), list(writes.values()) + + def _emit_or_bundle(fragments: list[dict]) -> str: """Stream fragments; fall back to a bundle file when GLaaS is unreachable. diff --git a/roar/backends/k8s/submit.py b/roar/backends/k8s/submit.py index a2821d85..efb85e70 100644 --- a/roar/backends/k8s/submit.py +++ b/roar/backends/k8s/submit.py @@ -123,8 +123,18 @@ def plan_kubectl_job_submit_command(command: list[str]) -> ExecutionCommandPlan: mount_map=_config_mount_map(config), runtime_source=str(config.get("runtime_source") or "install"), runtime_image=str(config.get("runtime_image") or ""), + proxy_sidecar=bool(config.get("proxy_sidecar", False)), + proxy_upstream=str(config.get("proxy_upstream") or ""), namespace_override=_find_namespace_argument(command), ) + if bool(config.get("proxy_sidecar", False)) and ( + str(config.get("runtime_source") or "install") != "image" + or not str(config.get("runtime_image") or "") + ): + _warn( + "k8s.proxy_sidecar requires k8s.runtime_source='image' with " + "k8s.runtime_image set; sidecar not injected" + ) if rewrite.skipped_containers: _warn( "containers without an explicit command were left uninstrumented: " diff --git a/roar/backends/k8s/webhook.py b/roar/backends/k8s/webhook.py index 8b78b6d2..2ac42cf5 100644 --- a/roar/backends/k8s/webhook.py +++ b/roar/backends/k8s/webhook.py @@ -57,6 +57,8 @@ class WebhookSettings: fragment_session_ttl_seconds: int = 86400 bundle_dir: str = "" mount_map: dict[str, str] = field(default_factory=dict) + proxy_sidecar: bool = False + proxy_upstream: str = "" @classmethod def from_environ(cls, environ: dict[str, str] | None = None) -> WebhookSettings: @@ -83,6 +85,9 @@ def from_environ(cls, environ: dict[str, str] | None = None) -> WebhookSettings: fragment_session_ttl_seconds=int(env.get("ROAR_WEBHOOK_SESSION_TTL", "86400") or 86400), bundle_dir=env.get("ROAR_WEBHOOK_BUNDLE_DIR", "").strip(), mount_map=mount_map, + proxy_sidecar=env.get("ROAR_WEBHOOK_PROXY_SIDECAR", "").strip().lower() + in ("1", "true", "yes"), + proxy_upstream=env.get("ROAR_WEBHOOK_PROXY_UPSTREAM", "").strip(), ) @@ -154,6 +159,8 @@ def _allow(warnings: list[str] | None = None) -> dict[str, Any]: mount_map=settings.mount_map, runtime_source=settings.runtime_source, runtime_image=settings.runtime_image, + proxy_sidecar=settings.proxy_sidecar, + proxy_upstream=settings.proxy_upstream, namespace_override=namespace, ) rewritten = rewrite.documents[0] diff --git a/tests/backends/k8s/e2e/test_k8s_phase3.py b/tests/backends/k8s/e2e/test_k8s_phase3.py index 0e75a70a..7d77bd28 100644 --- a/tests/backends/k8s/e2e/test_k8s_phase3.py +++ b/tests/backends/k8s/e2e/test_k8s_phase3.py @@ -148,6 +148,129 @@ def test_image_staged_runtime_captures_without_network_install( _cleanup(job_name) +PROXY_JOB_TEMPLATE = """\ +apiVersion: batch/v1 +kind: Job +metadata: + name: {name} + namespace: {namespace} + labels: + app.kubernetes.io/part-of: roar-k8s-e2e +spec: + backoffLimit: 0 + ttlSecondsAfterFinished: 1800 + template: + spec: + restartPolicy: Never + volumes: + - name: work + emptyDir: {{}} + containers: + - name: trainer + image: python:3.12-slim + workingDir: /work + volumeMounts: + - name: work + mountPath: /work + env: + - name: AWS_ACCESS_KEY_ID + value: minioadmin + - name: AWS_SECRET_ACCESS_KEY + value: minioadmin + - name: AWS_DEFAULT_REGION + value: us-east-1 + command: + - python + - -c + - >- + import urllib.request; + data = urllib.request.urlopen('http://127.0.0.1:19191/{bucket}/datasets/train.csv', timeout=30).read(); + open('model.bin', 'wb').write(data * 2) +""" + + +def test_proxy_sidecar_captures_hook_invisible_s3_client( + k8s_cluster: None, + glaas_health: str, + tmp_path_factory: pytest.TempPathFactory, +) -> None: + """A raw-HTTP S3 client (no boto3 — the hooks are blind to it) reads + through the injected proxy sidecar, which re-signs the request with the + workload's inherited AWS credentials; the proxy log lands in lineage.""" + if not _runtime_image_loaded(): + pytest.skip("roar-runtime:dev not loaded") + if kubectl(["get", "deployment/minio", "-n", NAMESPACE], check=False).returncode != 0: + pytest.skip("MinIO not deployed; run bootstrap_k8s.sh --with-minio") + + import boto3 + from botocore.client import Config + + suffix = uuid.uuid4().hex[:6] + bucket = f"roar-proxy-{suffix}" + host_s3 = boto3.client( + "s3", + endpoint_url="http://localhost:39000", + aws_access_key_id="minioadmin", + aws_secret_access_key="minioadmin", + region_name="us-east-1", + config=Config(s3={"addressing_style": "path"}), + ) + host_s3.create_bucket(Bucket=bucket) + host_s3.put_object(Bucket=bucket, Key="datasets/train.csv", Body=b"x,y\n1.0,2.0\n2.0,3.9\n") + project_dir = tmp_path_factory.mktemp("k8s-proxy-mode") + _write_image_mode_project(project_dir) + config_path = project_dir / ".roar" / "config.toml" + config_path.write_text( + config_path.read_text(encoding="utf-8") + + 'proxy_sidecar = true\nproxy_upstream = "http://minio:9000"\n', + encoding="utf-8", + ) + + job_name = f"roar-proxy-{suffix}" + (project_dir / "job.yaml").write_text( + PROXY_JOB_TEMPLATE.format(name=job_name, namespace=NAMESPACE, bucket=bucket), + encoding="utf-8", + ) + + try: + completed = _submit("job.yaml", cwd=project_dir) + logs = _pod_logs(f"job-name={job_name}") + run = { + "exit_code": completed.returncode, + "stdout": completed.stdout, + "stderr": completed.stderr, + "pod_logs": logs, + } + assert completed.returncode == 0, _describe(run) + + tasks = _query(project_dir, "SELECT id FROM jobs WHERE job_type = 'k8s_task'") + assert len(tasks) == 1, _describe(run) + task_id = tasks[0]["id"] + + input_paths = { + str(row["path"]) + for row in _query( + project_dir, + "SELECT path FROM job_inputs WHERE job_id = ?", + (task_id,), + ) + } + assert f"s3://{bucket}/datasets/train.csv" in input_paths, ( + f"proxy-captured read missing: {input_paths}\n{_describe(run)}" + ) + output_paths = { + str(row["path"]) + for row in _query( + project_dir, + "SELECT path FROM job_outputs WHERE job_id = ?", + (task_id,), + ) + } + assert any(path.endswith("model.bin") for path in output_paths), output_paths + finally: + _cleanup(job_name) + + def _webhook_deployed() -> bool: result = kubectl(["get", "mutatingwebhookconfiguration", "roar-lineage-injector"], check=False) return result.returncode == 0 diff --git a/tests/backends/k8s/unit/test_proxy_sidecar.py b/tests/backends/k8s/unit/test_proxy_sidecar.py new file mode 100644 index 00000000..4b8d3b25 --- /dev/null +++ b/tests/backends/k8s/unit/test_proxy_sidecar.py @@ -0,0 +1,145 @@ +from __future__ import annotations + +import copy +from pathlib import Path +from typing import Any + +from roar.backends.k8s.manifest import rewrite_manifest_for_lineage +from roar.backends.k8s.pod_entry import _load_proxy_log_refs + +from .conftest import SINGLE_JOB_MANIFEST + + +def _rewrite_proxy_mode(documents: list[dict[str, Any]], **overrides: Any): + kwargs: dict[str, Any] = { + "secret_name": "roar-fragment-abc", + "session_id": "session-1", + "fragment_token": "ff" * 32, + "requirement": "roar-cli==0.3.7", + "cluster_glaas_url": "http://glaas:3001", + "tracer": "preload", + "parent_job_uid": "cafe0123", + "runtime_source": "image", + "runtime_image": "roar-runtime:dev", + "proxy_sidecar": True, + "proxy_upstream": "http://minio:9000", + } + kwargs.update(overrides) + return rewrite_manifest_for_lineage(documents, **kwargs) + + +def _job_pod_spec(rewrite) -> dict[str, Any]: + job = next(doc for doc in rewrite.documents if doc.get("kind") == "Job") + return job["spec"]["template"]["spec"] + + +def test_proxy_sidecar_injected_as_native_sidecar() -> None: + rewrite = _rewrite_proxy_mode([copy.deepcopy(SINGLE_JOB_MANIFEST)]) + pod_spec = _job_pod_spec(rewrite) + + sidecar = next(c for c in pod_spec["initContainers"] if c["name"] == "roar-s3-proxy") + assert sidecar["restartPolicy"] == "Always" + assert sidecar["image"] == "roar-runtime:dev" + assert "roar-proxy --port 19191" in sidecar["command"][2] + assert "--upstream http://minio:9000" in sidecar["command"][2] + probe_command = sidecar["startupProbe"]["exec"]["command"] + assert "socket.create_connection(('127.0.0.1', 19191)" in probe_command[2] + + volumes = {v["name"] for v in pod_spec["volumes"]} + assert "roar-proxy-log" in volumes + + container = pod_spec["containers"][0] + env = {entry["name"]: entry for entry in container["env"]} + assert env["AWS_ENDPOINT_URL"]["value"] == "http://127.0.0.1:19191" + assert env["ROAR_K8S_PROXY_LOG"]["value"] == "/roar-proxy-log/proxy.log" + mounts = {m["name"] for m in container["volumeMounts"]} + assert "roar-proxy-log" in mounts + + +def test_sidecar_inherits_workload_aws_credentials() -> None: + manifest = copy.deepcopy(SINGLE_JOB_MANIFEST) + manifest["spec"]["template"]["spec"]["containers"][0]["env"].extend( + [ + {"name": "AWS_ACCESS_KEY_ID", "value": "AKIA..."}, + { + "name": "AWS_SECRET_ACCESS_KEY", + "valueFrom": {"secretKeyRef": {"name": "aws", "key": "secret"}}, + }, + {"name": "AWS_DEFAULT_REGION", "value": "us-east-1"}, + {"name": "UNRELATED", "value": "no"}, + ] + ) + rewrite = _rewrite_proxy_mode([manifest]) + sidecar = next( + c for c in _job_pod_spec(rewrite)["initContainers"] if c["name"] == "roar-s3-proxy" + ) + names = [e["name"] for e in sidecar["env"]] + assert names == ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_DEFAULT_REGION"] + assert sidecar["env"][1]["valueFrom"]["secretKeyRef"]["name"] == "aws" + + +def test_user_endpoint_url_wins_over_proxy_redirect() -> None: + manifest = copy.deepcopy(SINGLE_JOB_MANIFEST) + manifest["spec"]["template"]["spec"]["containers"][0]["env"].append( + {"name": "AWS_ENDPOINT_URL", "value": "http://user-endpoint:9000"} + ) + rewrite = _rewrite_proxy_mode([manifest]) + container = _job_pod_spec(rewrite)["containers"][0] + endpoints = [e for e in container["env"] if e["name"] == "AWS_ENDPOINT_URL"] + assert endpoints == [{"name": "AWS_ENDPOINT_URL", "value": "http://user-endpoint:9000"}] + + +def test_proxy_requires_image_staging() -> None: + rewrite = _rewrite_proxy_mode( + [copy.deepcopy(SINGLE_JOB_MANIFEST)], + runtime_source="install", + runtime_image="", + ) + pod_spec = _job_pod_spec(rewrite) + assert "initContainers" not in pod_spec + env_names = {e["name"] for e in pod_spec["containers"][0]["env"]} + assert "AWS_ENDPOINT_URL" not in env_names + + +def test_proxy_sidecar_is_idempotent() -> None: + rewrite = _rewrite_proxy_mode([copy.deepcopy(SINGLE_JOB_MANIFEST)]) + job = next(doc for doc in rewrite.documents if doc.get("kind") == "Job") + second = _rewrite_proxy_mode([copy.deepcopy(job)]) + pod_spec = _job_pod_spec(second) + sidecars = [c for c in pod_spec["initContainers"] if c["name"] == "roar-s3-proxy"] + assert len(sidecars) == 1 + + +def test_proxy_log_refs_parse_and_defer_to_hooks(tmp_path: Path) -> None: + log = tmp_path / "proxy.log" + log.write_text( + "\n".join( + [ + "[S3:GetObject] s3://data/train.csv (100 bytes) etag=aaa session=s job=j", + "[S3:PutObject] s3://models/out.bin (64 bytes) etag=bbb", + "[S3:GetObject] s3://data/seen-by-hooks.csv etag=ccc", + "[S3:ListObjectsV2] s3://data/ ", + "not a log line", + ] + ), + encoding="utf-8", + ) + + reads, writes = _load_proxy_log_refs( + str(log), + seen_reads={"s3://data/seen-by-hooks.csv"}, + seen_writes=set(), + ) + assert [r["path"] for r in reads] == ["s3://data/train.csv"] + assert reads[0]["capture_method"] == "proxy" + assert reads[0]["hash"] == "aaa" + assert [w["path"] for w in writes] == ["s3://models/out.bin"] + assert writes[0]["size"] == 64 + + +def test_proxy_log_refs_missing_file() -> None: + assert _load_proxy_log_refs(None, seen_reads=set(), seen_writes=set()) == ([], []) + assert _load_proxy_log_refs("/nonexistent/proxy.log", seen_reads=set(), seen_writes=set()) == ( + [], + [], + ) From a43bb3e4cdebe2f3054ab416a87ca6ddd5f3228a Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Tue, 14 Jul 2026 17:41:47 +0000 Subject: [PATCH 13/31] feat(k8s): package the lineage webhook as a Helm chart Add deploy/charts/roar-lineage-webhook: webhook Deployment/Service, RBAC (Secret-create ClusterRole), and the MutatingWebhookConfiguration for all five workload kinds. TLS comes either from a supplied secret + tls.caBundle (required-value guarded) or from an optional cert-manager self-signed Issuer/Certificate with CA injection. injection.* values map 1:1 onto the ROAR_WEBHOOK_* env contract, including the proxy sidecar knobs. The harness deploy script now generates self-signed TLS material and installs this chart (helm downloaded into .tools/bin like kind/kubectl), so the zero-touch webhook e2e proves the packaged artifact end-to-end; the e2e webhook-presence check selects the configuration by chart labels instead of the old raw-manifest name. Validated: helm lint clean, cert-manager and manual-TLS renders, and 23/23 k8s e2e against the chart-deployed webhook. Co-Authored-By: Claude Fable 5 --- deploy/charts/roar-lineage-webhook/Chart.yaml | 14 ++ .../templates/_helpers.tpl | 30 +++ .../templates/certificate.yaml | 26 +++ .../templates/deployment.yaml | 95 +++++++++ .../roar-lineage-webhook/templates/rbac.yaml | 34 ++++ .../templates/webhook.yaml | 49 +++++ .../charts/roar-lineage-webhook/values.yaml | 60 ++++++ docs/developer/k8s-integration.md | 18 +- tests/backends/k8s/README.md | 12 +- tests/backends/k8s/e2e/test_k8s_phase3.py | 14 +- tests/backends/k8s/scripts/bootstrap_k8s.sh | 6 +- tests/backends/k8s/scripts/deploy_webhook.sh | 181 +++--------------- 12 files changed, 379 insertions(+), 160 deletions(-) create mode 100644 deploy/charts/roar-lineage-webhook/Chart.yaml create mode 100644 deploy/charts/roar-lineage-webhook/templates/_helpers.tpl create mode 100644 deploy/charts/roar-lineage-webhook/templates/certificate.yaml create mode 100644 deploy/charts/roar-lineage-webhook/templates/deployment.yaml create mode 100644 deploy/charts/roar-lineage-webhook/templates/rbac.yaml create mode 100644 deploy/charts/roar-lineage-webhook/templates/webhook.yaml create mode 100644 deploy/charts/roar-lineage-webhook/values.yaml mode change 100644 => 100755 tests/backends/k8s/scripts/deploy_webhook.sh diff --git a/deploy/charts/roar-lineage-webhook/Chart.yaml b/deploy/charts/roar-lineage-webhook/Chart.yaml new file mode 100644 index 00000000..958039ba --- /dev/null +++ b/deploy/charts/roar-lineage-webhook/Chart.yaml @@ -0,0 +1,14 @@ +apiVersion: v2 +name: roar-lineage-webhook +description: >- + Zero-touch ML training lineage for Kubernetes: a mutating admission + webhook that instruments Jobs, JobSets, PyTorchJobs, TrainJobs, and + RayJobs in opted-in namespaces with roar lineage capture. +type: application +version: 0.1.0 +appVersion: "0.3.7" +keywords: + - lineage + - provenance + - ml + - admission-webhook diff --git a/deploy/charts/roar-lineage-webhook/templates/_helpers.tpl b/deploy/charts/roar-lineage-webhook/templates/_helpers.tpl new file mode 100644 index 00000000..64e23afb --- /dev/null +++ b/deploy/charts/roar-lineage-webhook/templates/_helpers.tpl @@ -0,0 +1,30 @@ +{{- define "roar-webhook.fullname" -}} +{{- printf "%s-roar-lineage-webhook" .Release.Name | trunc 63 | trimSuffix "-" -}} +{{- end -}} + +{{- define "roar-webhook.labels" -}} +app.kubernetes.io/name: roar-lineage-webhook +app.kubernetes.io/instance: {{ .Release.Name }} +app.kubernetes.io/managed-by: {{ .Release.Service }} +{{- end -}} + +{{- define "roar-webhook.selectorLabels" -}} +app.kubernetes.io/name: roar-lineage-webhook +app.kubernetes.io/instance: {{ .Release.Name }} +{{- end -}} + +{{- define "roar-webhook.runtimeImage" -}} +{{- if .Values.injection.runtimeImage -}} +{{ .Values.injection.runtimeImage }} +{{- else -}} +{{ .Values.image.repository }}:{{ .Values.image.tag }} +{{- end -}} +{{- end -}} + +{{- define "roar-webhook.tlsSecretName" -}} +{{- if .Values.tls.secretName -}} +{{ .Values.tls.secretName }} +{{- else -}} +{{ include "roar-webhook.fullname" . }}-tls +{{- end -}} +{{- end -}} diff --git a/deploy/charts/roar-lineage-webhook/templates/certificate.yaml b/deploy/charts/roar-lineage-webhook/templates/certificate.yaml new file mode 100644 index 00000000..d5c63e25 --- /dev/null +++ b/deploy/charts/roar-lineage-webhook/templates/certificate.yaml @@ -0,0 +1,26 @@ +{{- if .Values.certManager.enabled }} +apiVersion: cert-manager.io/v1 +kind: Issuer +metadata: + name: {{ include "roar-webhook.fullname" . }}-selfsigned + namespace: {{ .Release.Namespace }} + labels: + {{- include "roar-webhook.labels" . | nindent 4 }} +spec: + selfSigned: {} +--- +apiVersion: cert-manager.io/v1 +kind: Certificate +metadata: + name: {{ include "roar-webhook.fullname" . }} + namespace: {{ .Release.Namespace }} + labels: + {{- include "roar-webhook.labels" . | nindent 4 }} +spec: + secretName: {{ include "roar-webhook.tlsSecretName" . }} + dnsNames: + - {{ include "roar-webhook.fullname" . }}.{{ .Release.Namespace }}.svc + - {{ include "roar-webhook.fullname" . }}.{{ .Release.Namespace }}.svc.cluster.local + issuerRef: + name: {{ include "roar-webhook.fullname" . }}-selfsigned +{{- end }} diff --git a/deploy/charts/roar-lineage-webhook/templates/deployment.yaml b/deploy/charts/roar-lineage-webhook/templates/deployment.yaml new file mode 100644 index 00000000..06f1d022 --- /dev/null +++ b/deploy/charts/roar-lineage-webhook/templates/deployment.yaml @@ -0,0 +1,95 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: {{ include "roar-webhook.fullname" . }} + namespace: {{ .Release.Namespace }} + labels: + {{- include "roar-webhook.labels" . | nindent 4 }} +spec: + replicas: {{ .Values.replicas }} + selector: + matchLabels: + {{- include "roar-webhook.selectorLabels" . | nindent 6 }} + template: + metadata: + labels: + {{- include "roar-webhook.selectorLabels" . | nindent 8 }} + spec: + serviceAccountName: {{ include "roar-webhook.fullname" . }} + containers: + - name: webhook + image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}" + imagePullPolicy: {{ .Values.image.pullPolicy }} + command: + - python + - -m + - roar.backends.k8s.webhook + - --port=8443 + - --cert=/tls/tls.crt + - --key=/tls/tls.key + env: + - name: ROAR_WEBHOOK_GLAAS_URL + value: {{ required "glaas.url is required" .Values.glaas.url | quote }} + - name: ROAR_WEBHOOK_CLUSTER_GLAAS_URL + value: {{ required "glaas.clusterUrl is required" .Values.glaas.clusterUrl | quote }} + - name: ROAR_WEBHOOK_TRACER + value: {{ .Values.injection.tracer | quote }} + - name: ROAR_WEBHOOK_RUNTIME_SOURCE + value: {{ .Values.injection.runtimeSource | quote }} + - name: ROAR_WEBHOOK_RUNTIME_IMAGE + value: {{ include "roar-webhook.runtimeImage" . | quote }} + - name: ROAR_WEBHOOK_RUNTIME_REQUIREMENT + value: {{ .Values.injection.runtimeRequirement | quote }} + - name: ROAR_WEBHOOK_SESSION_TTL + value: {{ .Values.injection.sessionTtlSeconds | quote }} + {{- if .Values.injection.proxySidecar }} + - name: ROAR_WEBHOOK_PROXY_SIDECAR + value: "true" + - name: ROAR_WEBHOOK_PROXY_UPSTREAM + value: {{ .Values.injection.proxyUpstream | quote }} + {{- end }} + - name: ROAR_NO_TELEMETRY + value: "1" + ports: + - containerPort: 8443 + volumeMounts: + - name: tls + mountPath: /tls + readOnly: true + readinessProbe: + httpGet: + path: /healthz + port: 8443 + scheme: HTTPS + initialDelaySeconds: 2 + periodSeconds: 3 + {{- with .Values.resources }} + resources: + {{- toYaml . | nindent 12 }} + {{- end }} + {{- with .Values.nodeSelector }} + nodeSelector: + {{- toYaml . | nindent 8 }} + {{- end }} + {{- with .Values.tolerations }} + tolerations: + {{- toYaml . | nindent 8 }} + {{- end }} + volumes: + - name: tls + secret: + secretName: {{ include "roar-webhook.tlsSecretName" . }} +--- +apiVersion: v1 +kind: Service +metadata: + name: {{ include "roar-webhook.fullname" . }} + namespace: {{ .Release.Namespace }} + labels: + {{- include "roar-webhook.labels" . | nindent 4 }} +spec: + selector: + {{- include "roar-webhook.selectorLabels" . | nindent 4 }} + ports: + - port: 443 + targetPort: 8443 diff --git a/deploy/charts/roar-lineage-webhook/templates/rbac.yaml b/deploy/charts/roar-lineage-webhook/templates/rbac.yaml new file mode 100644 index 00000000..7f2b89f9 --- /dev/null +++ b/deploy/charts/roar-lineage-webhook/templates/rbac.yaml @@ -0,0 +1,34 @@ +apiVersion: v1 +kind: ServiceAccount +metadata: + name: {{ include "roar-webhook.fullname" . }} + namespace: {{ .Release.Namespace }} + labels: + {{- include "roar-webhook.labels" . | nindent 4 }} +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + name: {{ include "roar-webhook.fullname" . }}-secrets + labels: + {{- include "roar-webhook.labels" . | nindent 4 }} +rules: + # Fragment-session credential Secrets created in opted-in namespaces. + - apiGroups: [""] + resources: ["secrets"] + verbs: ["create"] +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRoleBinding +metadata: + name: {{ include "roar-webhook.fullname" . }}-secrets + labels: + {{- include "roar-webhook.labels" . | nindent 4 }} +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: ClusterRole + name: {{ include "roar-webhook.fullname" . }}-secrets +subjects: + - kind: ServiceAccount + name: {{ include "roar-webhook.fullname" . }} + namespace: {{ .Release.Namespace }} diff --git a/deploy/charts/roar-lineage-webhook/templates/webhook.yaml b/deploy/charts/roar-lineage-webhook/templates/webhook.yaml new file mode 100644 index 00000000..2ce9ef2b --- /dev/null +++ b/deploy/charts/roar-lineage-webhook/templates/webhook.yaml @@ -0,0 +1,49 @@ +apiVersion: admissionregistration.k8s.io/v1 +kind: MutatingWebhookConfiguration +metadata: + name: {{ include "roar-webhook.fullname" . }} + labels: + {{- include "roar-webhook.labels" . | nindent 4 }} + {{- if .Values.certManager.enabled }} + annotations: + cert-manager.io/inject-ca-from: {{ .Release.Namespace }}/{{ include "roar-webhook.fullname" . }} + {{- end }} +webhooks: + - name: lineage.roar.glaas.ai + admissionReviewVersions: ["v1"] + sideEffects: NoneOnDryRun + failurePolicy: {{ .Values.webhook.failurePolicy }} + reinvocationPolicy: {{ .Values.webhook.reinvocationPolicy }} + timeoutSeconds: {{ .Values.webhook.timeoutSeconds }} + namespaceSelector: + matchLabels: + {{ .Values.namespaceSelector.key }}: {{ .Values.namespaceSelector.value }} + clientConfig: + service: + name: {{ include "roar-webhook.fullname" . }} + namespace: {{ .Release.Namespace }} + path: /mutate + {{- if not .Values.certManager.enabled }} + caBundle: {{ required "tls.caBundle is required unless certManager.enabled" .Values.tls.caBundle }} + {{- end }} + rules: + - apiGroups: ["batch"] + apiVersions: ["v1"] + operations: ["CREATE"] + resources: ["jobs"] + - apiGroups: ["jobset.x-k8s.io"] + apiVersions: ["*"] + operations: ["CREATE"] + resources: ["jobsets"] + - apiGroups: ["kubeflow.org"] + apiVersions: ["v1"] + operations: ["CREATE"] + resources: ["pytorchjobs"] + - apiGroups: ["trainer.kubeflow.org"] + apiVersions: ["*"] + operations: ["CREATE"] + resources: ["trainjobs"] + - apiGroups: ["ray.io"] + apiVersions: ["v1"] + operations: ["CREATE"] + resources: ["rayjobs"] diff --git a/deploy/charts/roar-lineage-webhook/values.yaml b/deploy/charts/roar-lineage-webhook/values.yaml new file mode 100644 index 00000000..0cfac8b0 --- /dev/null +++ b/deploy/charts/roar-lineage-webhook/values.yaml @@ -0,0 +1,60 @@ +# roar-lineage-webhook chart values. + +replicas: 1 + +image: + # The roar-runtime image (scripts/build_runtime_image.sh): serves the + # webhook and stages the runtime into instrumented pods. + repository: roar-runtime + tag: dev + pullPolicy: IfNotPresent + +glaas: + # GLaaS URL reachable from the webhook pod (fragment-session registration). + url: "" + # GLaaS URL reachable from instrumented workload pods (fragment streaming). + # Use a fully-qualified in-cluster or external URL — instrumented pods can + # live in any opted-in namespace. + clusterUrl: "" + +injection: + tracer: preload + # "image" stages the hermetic runtime via an init container (recommended); + # "install" pip-installs runtimeRequirement at container start. + runtimeSource: image + # Defaults to image.repository:image.tag when empty. + runtimeImage: "" + runtimeRequirement: roar-cli + sessionTtlSeconds: 86400 + # Opt-in roar-proxy S3 sidecar for clients the in-process hooks can't + # see; requires runtimeSource=image. + proxySidecar: false + proxyUpstream: "" + +namespaceSelector: + # Namespaces opt in by carrying this label. + key: roar.glaas.ai/lineage + value: enabled + +webhook: + failurePolicy: Ignore + timeoutSeconds: 10 + reinvocationPolicy: IfNeeded + +tls: + # Existing kubernetes.io/tls Secret with the serving cert for + # -roar-lineage-webhook..svc. Required unless + # certManager.enabled. + secretName: "" + # Base64 PEM CA bundle matching the serving cert. Required unless + # certManager.enabled (cert-manager then injects it). + caBundle: "" + +certManager: + # Manage the serving cert with cert-manager (self-signed Issuer + + # Certificate; the CA bundle is injected via cert-manager annotations). + enabled: false + +resources: {} +nodeSelector: {} +tolerations: [] diff --git a/docs/developer/k8s-integration.md b/docs/developer/k8s-integration.md index c25aa55d..aac2e34e 100644 --- a/docs/developer/k8s-integration.md +++ b/docs/developer/k8s-integration.md @@ -20,9 +20,10 @@ in the dev meta-repo): Phases 1–3 implemented and live-validated trainer v2 controllers, multi-pod capture, `roar k8s attach`, bundle-mode fallback, object-store I/O hooks, mount-map rewriting, retry-chaos coverage, RayJob delegation, the roar-runtime image, the mutating -webhook injector, and the opt-in proxy sidecar). Remaining from Phase 3: -a packaged Helm chart (the harness deploys via -`tests/backends/k8s/scripts/deploy_webhook.sh`). +webhook injector, the opt-in proxy sidecar, and the +`deploy/charts/roar-lineage-webhook` Helm chart — the harness's +`tests/backends/k8s/scripts/deploy_webhook.sh` installs the chart, so +the webhook e2e exercises the packaged artifact). **Runtime staging modes**: `k8s.runtime_source = "install"` pip-installs `k8s.runtime_install_requirement` at container start; `"image"` stages @@ -50,6 +51,17 @@ blocks admission. Reconstitution is client-driven via `roar k8s attach`. `ROAR_WEBHOOK_PROXY_UPSTREAM`) makes the injector add the proxy sidecar to every workload it instruments. +**Helm chart** (`deploy/charts/roar-lineage-webhook`): packages the +webhook Deployment/Service, RBAC (Secret-create ClusterRole), and the +`MutatingWebhookConfiguration`. Required values: `glaas.url`, +`glaas.clusterUrl`, and TLS — either `tls.caBundle` + a pre-created +`tls.secretName` (the harness path: openssl self-signed), or +`certManager.enabled=true` for a self-signed Issuer/Certificate with +CA injection. `injection.*` values map 1:1 onto the `ROAR_WEBHOOK_*` +env contract (tracer, runtimeSource/runtimeImage, runtimeRequirement, +sessionTtlSeconds, proxySidecar/proxyUpstream). The webhook service name +is `-roar-lineage-webhook`; certificate SANs must match it. + **RayJob delegation** (`rayjob.py`): KubeRay overwrites container commands with `ray start --block`, so command-wrapping can't see user code. Instead the RayJob rewrite reuses the Ray backend's runtime surface: the diff --git a/tests/backends/k8s/README.md b/tests/backends/k8s/README.md index 6883421b..e0cc54c5 100644 --- a/tests/backends/k8s/README.md +++ b/tests/backends/k8s/README.md @@ -33,10 +33,15 @@ Three test layers share this harness: heaviest test: multi-GB Ray image + per-job pip env). Skips without `--with-kuberay`. - `e2e/test_k8s_phase3.py` — image-staged runtime (hermetic pod start via - the roar-runtime init container, no pip) and the zero-touch webhook + the roar-runtime init container, no pip), the opt-in `roar-s3-proxy` + sidecar (a hook-invisible raw-HTTP S3 client's reads still land as + `s3://` lineage refs via proxy-log capture), and the zero-touch webhook story: plain `kubectl apply` in a labeled namespace gets injected, and `roar k8s attach` recovers the lineage; unlabeled namespaces stay - untouched. Skips without the image / `--with-webhook`. + untouched. The webhook is deployed through the + `deploy/charts/roar-lineage-webhook` Helm chart + (`scripts/deploy_webhook.sh`), so these tests exercise the packaged + chart. Skips without the image / `--with-webhook`. - `e2e/test_k8s_smoke.py` — the Phase-0 runtime diagnostic: fixtures hand-wrap the manifest (no backend involved) to isolate the runtime pieces (in-pod tracing, fragment streaming, identity contract) when the product @@ -50,7 +55,8 @@ live in `unit/` and run in the default gate — no cluster needed. - Docker - A packaged wheel: `bash scripts/build_wheel_with_bins.sh` (repo root) - Local glaas-api on `http://localhost:3001` (e.g. via pm2) -- `kind`/`kubectl` are downloaded automatically into `.tools/bin` if missing +- `kind`/`kubectl`/`helm` are downloaded automatically into `.tools/bin` if + missing ## Usage diff --git a/tests/backends/k8s/e2e/test_k8s_phase3.py b/tests/backends/k8s/e2e/test_k8s_phase3.py index 7d77bd28..6b86f562 100644 --- a/tests/backends/k8s/e2e/test_k8s_phase3.py +++ b/tests/backends/k8s/e2e/test_k8s_phase3.py @@ -272,8 +272,18 @@ def test_proxy_sidecar_captures_hook_invisible_s3_client( def _webhook_deployed() -> bool: - result = kubectl(["get", "mutatingwebhookconfiguration", "roar-lineage-injector"], check=False) - return result.returncode == 0 + result = kubectl( + [ + "get", + "mutatingwebhookconfiguration", + "-l", + "app.kubernetes.io/name=roar-lineage-webhook", + "-o", + "name", + ], + check=False, + ) + return result.returncode == 0 and bool(result.stdout.strip()) def _wait_job_terminal(job_name: str, namespace: str, timeout: int = 300) -> bool: diff --git a/tests/backends/k8s/scripts/bootstrap_k8s.sh b/tests/backends/k8s/scripts/bootstrap_k8s.sh index e3832eb0..15224e57 100755 --- a/tests/backends/k8s/scripts/bootstrap_k8s.sh +++ b/tests/backends/k8s/scripts/bootstrap_k8s.sh @@ -217,8 +217,10 @@ if ((WITH_WEBHOOK == 1)); then bash "$REPO_ROOT/scripts/build_runtime_image.sh" kind load docker-image roar-runtime:dev --name "$CLUSTER_NAME" bash "$HARNESS_DIR/scripts/deploy_webhook.sh" - kubectl_ctx -n roar-system rollout restart deployment/roar-webhook - kubectl_ctx -n roar-system rollout status deployment/roar-webhook --timeout=180s + # Restart so a freshly rebuilt roar-runtime:dev image is picked up even + # when the helm release itself is unchanged. + kubectl_ctx -n roar-system rollout restart deployment/roar-roar-lineage-webhook + kubectl_ctx -n roar-system rollout status deployment/roar-roar-lineage-webhook --timeout=180s fi if ((WITH_KUBERAY == 1)); then diff --git a/tests/backends/k8s/scripts/deploy_webhook.sh b/tests/backends/k8s/scripts/deploy_webhook.sh old mode 100644 new mode 100755 index b3a4951c..7dc3a667 --- a/tests/backends/k8s/scripts/deploy_webhook.sh +++ b/tests/backends/k8s/scripts/deploy_webhook.sh @@ -1,28 +1,39 @@ #!/usr/bin/env bash set -euo pipefail -# Deploys the roar lineage mutating webhook into the KIND harness. -# -# - self-signed CA + server cert (no cert-manager dependency in the harness) -# - roar-runtime:dev image serves the webhook (built/loaded by caller) -# - namespaceSelector opt-in: roar.glaas.ai/lineage=enabled -# - failurePolicy Ignore: webhook outages never block workloads +# Deploys the roar lineage mutating webhook into the KIND harness via the +# Helm chart (deploy/charts/roar-lineage-webhook) — the chart is the +# artifact under test; the harness only supplies self-signed TLS material +# (no cert-manager dependency here). # # Env overrides: WEBHOOK_GLAAS_URL (server-visible from the webhook pod), # WEBHOOK_CLUSTER_GLAAS_URL (visible from injected workload pods). HARNESS_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +REPO_ROOT="$(cd "$HARNESS_DIR/../../.." && pwd)" TOOLS_BIN="$HARNESS_DIR/.tools/bin" export PATH="$TOOLS_BIN:$PATH" +HELM_VERSION="v3.16.4" CLUSTER_NAME="roar-k8s-e2e" KUBE_CONTEXT="kind-${CLUSTER_NAME}" WEBHOOK_NS="roar-system" -SERVICE="roar-webhook" +RELEASE="roar" +# Must match the chart's fullname template: -roar-lineage-webhook +SERVICE="${RELEASE}-roar-lineage-webhook" +CHART_DIR="$REPO_ROOT/deploy/charts/roar-lineage-webhook" RUNTIME_IMAGE="${RUNTIME_IMAGE:-roar-runtime:dev}" WEBHOOK_GLAAS_URL="${WEBHOOK_GLAAS_URL:-http://glaas.roar-e2e.svc.cluster.local:3001}" WEBHOOK_CLUSTER_GLAAS_URL="${WEBHOOK_CLUSTER_GLAAS_URL:-http://glaas.roar-e2e.svc.cluster.local:3001}" +if ! command -v helm >/dev/null 2>&1; then + echo "▶ Downloading helm ${HELM_VERSION}" + tmp="$(mktemp -d)" + curl -fsSL "https://get.helm.sh/helm-${HELM_VERSION}-linux-amd64.tar.gz" | tar -xz -C "$tmp" + install -m 0755 "$tmp/linux-amd64/helm" "$TOOLS_BIN/helm" + rm -rf "$tmp" +fi + kubectl_ctx() { kubectl --context "$KUBE_CONTEXT" "$@"; } echo "▶ Generating webhook TLS material" @@ -41,151 +52,21 @@ openssl x509 -req -in "$CERT_DIR/tls.csr" -days 30 \ "$SERVICE" "$WEBHOOK_NS" "$SERVICE" "$WEBHOOK_NS") >/dev/null 2>&1 CA_BUNDLE="$(base64 -w0 <"$CERT_DIR/ca.crt")" -echo "▶ Deploying webhook to namespace $WEBHOOK_NS" +echo "▶ Installing chart ${CHART_DIR}" kubectl_ctx create namespace "$WEBHOOK_NS" --dry-run=client -o yaml | kubectl_ctx apply -f - -kubectl_ctx -n "$WEBHOOK_NS" create secret tls roar-webhook-tls \ +kubectl_ctx -n "$WEBHOOK_NS" create secret tls "${SERVICE}-tls" \ --cert="$CERT_DIR/tls.crt" --key="$CERT_DIR/tls.key" \ --dry-run=client -o yaml | kubectl_ctx apply -f - -kubectl_ctx apply -f - < Date: Tue, 14 Jul 2026 18:54:47 +0000 Subject: [PATCH 14/31] feat(fragments): renew expired fragment sessions on 403 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Long training jobs routinely outlive the fragment-session TTL, and an expired session 403s on both append and read — lineage was lost with no recovery path. Pair with the glaas-api renewal endpoint (POST /api/v1/fragments/sessions/:id/renew, token-authenticated and allowed even after expiry): - GlaasFragmentStreamer renews on a 403 batch POST and retries the chunk once; a second 403 gives up normally. Renewal is purely reactive — no client-side expiry clock. - Both fragment reconstituters (k8s, ray) renew-and-retry on a 403 fetch, so a late `roar k8s attach` recovers an expired session too. - GlaasClient.renew_fragment_session for host-side callers. - Streamer exposes delivered/failed/pending counters and close() reports flush success (used by the follow-up transport change). Degrades gracefully against a server without the endpoint (renew 404s -> previous behavior). Live-proven on the KIND harness: 60s-TTL session, 75s job, streamer renews mid-run and the finalizer reconstitutes (tests/backends/k8s/e2e/test_k8s_ttl_renewal.py; skips when the local glaas-api lacks the endpoint). Co-Authored-By: Claude Fable 5 --- docs/developer/k8s-integration.md | 10 + roar/backends/k8s/fragment_reconstituter.py | 19 ++ roar/backends/ray/fragment_reconstituter.py | 22 ++ roar/integrations/glaas/__init__.py | 1 + roar/integrations/glaas/client.py | 14 ++ roar/integrations/glaas/fragment_streamer.py | 69 +++++- tests/backends/k8s/README.md | 4 + .../backends/k8s/e2e/test_k8s_ttl_renewal.py | 234 ++++++++++++++++++ .../ray/unit/test_fragment_reconstituter.py | 52 ++++ .../glaas/test_fragment_streamer.py | 87 +++++++ 10 files changed, 509 insertions(+), 3 deletions(-) create mode 100644 tests/backends/k8s/e2e/test_k8s_ttl_renewal.py diff --git a/docs/developer/k8s-integration.md b/docs/developer/k8s-integration.md index aac2e34e..3bcfefc1 100644 --- a/docs/developer/k8s-integration.md +++ b/docs/developer/k8s-integration.md @@ -150,6 +150,16 @@ Note the fragment streamer swallows per-batch POST failures (reports "streamed" regardless), which is why the fallback needs its own probe — surfacing streamer failure counts is an open follow-up. +**Session TTL renewal**: fragment sessions are registered with +`k8s.fragment_session_ttl_seconds` (default 86400, server-capped at 7 +days) — training that outlives the TTL used to lose lineage because both +append and read 403 after expiry. The streamer now renews on 403 +(`POST /api/v1/fragments/sessions/{id}/renew`, token-authenticated, +allowed even after expiry) and retries the batch once; the k8s and ray +fragment reconstituters do the same on fetch, so a late `roar k8s attach` +can still recover an expired session. No clock tracking client-side — +renewal is purely reactive to the 403. + ## 2. Flow 1. **Match** (`roar/backends/k8s/submit.py`): binary `kubectl`, verb diff --git a/roar/backends/k8s/fragment_reconstituter.py b/roar/backends/k8s/fragment_reconstituter.py index a0271244..5c1ade8b 100644 --- a/roar/backends/k8s/fragment_reconstituter.py +++ b/roar/backends/k8s/fragment_reconstituter.py @@ -5,6 +5,7 @@ import base64 import json import sqlite3 +import urllib.error import urllib.request from dataclasses import dataclass from pathlib import Path @@ -13,6 +14,7 @@ from cryptography.hazmat.primitives.ciphers.aead import AESGCM from roar.backends.k8s.lineage import collect_k8s_fragments, resolve_active_session_context +from roar.integrations.glaas import renew_fragment_session def _get_logger(): @@ -104,6 +106,23 @@ def _fetch_batches(self) -> list[dict[str, Any]]: try: with urllib.request.urlopen(request, timeout=10) as response: payload = json.loads(response.read().decode("utf-8")) + except urllib.error.HTTPError as exc: + # An expired session 403s on read; renewal is token-authenticated + # and allowed after expiry, so attach can still recover lineage. + payload = None + if exc.code == 403 and renew_fragment_session( + self._glaas_url, self._session_id, self._token + ): + try: + with urllib.request.urlopen(request, timeout=10) as response: + payload = json.loads(response.read().decode("utf-8")) + except Exception as retry_exc: + exc = retry_exc # type: ignore[assignment] + if payload is None: + _get_logger().warning( + "Failed to fetch fragments for session %s: %s", self._session_id, exc + ) + return [] except Exception as exc: _get_logger().warning( "Failed to fetch fragments for session %s: %s", self._session_id, exc diff --git a/roar/backends/ray/fragment_reconstituter.py b/roar/backends/ray/fragment_reconstituter.py index a69fbd3d..c3658c1b 100644 --- a/roar/backends/ray/fragment_reconstituter.py +++ b/roar/backends/ray/fragment_reconstituter.py @@ -4,6 +4,7 @@ import json import os import sqlite3 +import urllib.error import urllib.request from dataclasses import dataclass from pathlib import Path @@ -11,6 +12,8 @@ from cryptography.hazmat.primitives.ciphers.aead import AESGCM +from roar.integrations.glaas import renew_fragment_session + from .collector import _resolve_active_session_context, collect_fragments from .fragment import derive_fragment_fallback_identity, derive_task_identity from .s3_key_paths import parse_s3_key_placeholder, s3_object_key @@ -112,6 +115,25 @@ def _fetch_batches(self) -> list[dict[str, Any]]: try: with urllib.request.urlopen(request, timeout=5) as response: payload = json.loads(response.read().decode("utf-8")) + except urllib.error.HTTPError as exc: + # An expired session 403s on read; renewal is token-authenticated + # and allowed after expiry, so attach can still recover lineage. + payload = None + if exc.code == 403 and renew_fragment_session( + self._glaas_url, self._session_id, self._token + ): + try: + with urllib.request.urlopen(request, timeout=5) as response: + payload = json.loads(response.read().decode("utf-8")) + except Exception as retry_exc: + exc = retry_exc # type: ignore[assignment] + if payload is None: + _get_logger().warning( + "Failed to fetch fragments for session %s: %s", + self._session_id, + exc, + ) + return [] except Exception as exc: _get_logger().warning( "Failed to fetch fragments for session %s: %s", diff --git a/roar/integrations/glaas/__init__.py b/roar/integrations/glaas/__init__.py index b332e4aa..b207c953 100644 --- a/roar/integrations/glaas/__init__.py +++ b/roar/integrations/glaas/__init__.py @@ -19,6 +19,7 @@ "get_glaas_url": ".auth", "make_auth_header": ".auth", "parse_json_response": ".transport", + "renew_fragment_session": ".fragment_streamer", "request_json": ".transport", "sign_payload": ".auth", } diff --git a/roar/integrations/glaas/client.py b/roar/integrations/glaas/client.py index 05f239dc..62adcbf6 100644 --- a/roar/integrations/glaas/client.py +++ b/roar/integrations/glaas/client.py @@ -595,6 +595,20 @@ def register_fragment_session( } return self._request("POST", "/api/v1/fragments/sessions", body) + def renew_fragment_session( + self, + session_id: str, + token: str, + ttl_seconds: int = 86400, + ) -> tuple[dict | None, str | None]: + """Extend a fragment session's expiry (allowed even after expiry).""" + encoded_token = urllib.parse.quote(token, safe="") + return self._request( + "POST", + f"/api/v1/fragments/sessions/{session_id}/renew?token={encoded_token}", + {"ttl_seconds": ttl_seconds}, + ) + def register_session( self, session_hash: str, diff --git a/roar/integrations/glaas/fragment_streamer.py b/roar/integrations/glaas/fragment_streamer.py index 0600bce4..c2237189 100644 --- a/roar/integrations/glaas/fragment_streamer.py +++ b/roar/integrations/glaas/fragment_streamer.py @@ -17,6 +17,41 @@ def _get_logger(): return get_logger() +DEFAULT_RENEW_TTL_SECONDS = 86400 + + +def renew_fragment_session( + glaas_url: str, + session_id: str, + token: str, + ttl_seconds: int = DEFAULT_RENEW_TTL_SECONDS, +) -> bool: + """Extend a fragment session's expiry to now + ttl_seconds. + + GLaaS allows renewal of an already-expired session (the token proves + ownership), so callers can retry once after a 403 instead of tracking + the expiry clock themselves. + """ + try: + request = urllib.request.Request( + url=f"{glaas_url.rstrip('/')}/api/v1/fragments/sessions/{session_id}/renew", + data=json.dumps({"ttl_seconds": int(ttl_seconds)}, separators=(",", ":")).encode( + "utf-8" + ), + headers={ + "content-type": "application/json", + "x-roar-fragment-token": token, + }, + method="POST", + ) + with urllib.request.urlopen(request, timeout=5): + pass + except Exception as exc: + _get_logger().warning("Failed to renew fragment session %s: %s", session_id, exc) + return False + return True + + class GlaasFragmentStreamer: def __init__( self, @@ -24,13 +59,22 @@ def __init__( token: str, glaas_url: str, flush_threshold: int = 50, + renew_ttl_seconds: int = DEFAULT_RENEW_TTL_SECONDS, ) -> None: self._session_id = session_id self._token = token self._glaas_url = glaas_url.rstrip("/") self._flush_threshold = max(1, int(flush_threshold)) + self._renew_ttl_seconds = renew_ttl_seconds self._buffer: list[dict[str, Any]] = [] self._next_sequence = 0 + self.delivered_batches = 0 + self.failed_batches = 0 + + @property + def pending_fragments(self) -> int: + """Fragments accepted via append_fragment but not yet delivered.""" + return len(self._buffer) def append_fragment(self, fragment_dict: dict[str, Any]) -> None: self._buffer.append(fragment_dict) @@ -50,6 +94,7 @@ def flush(self) -> bool: if ok: del self._buffer[:chunk_size] self._next_sequence += 1 + self.delivered_batches += 1 break if ( too_large @@ -60,6 +105,7 @@ def flush(self) -> bool: if too_large and chunk_size > 1: chunk_size = max(1, chunk_size // 2) continue + self.failed_batches += 1 return False return True @@ -87,7 +133,9 @@ def _split_oversized_fragment(self, fragment: dict[str, Any]) -> bool: self._buffer[:1] = replacement return True - def _post_chunk(self, chunk: list[dict[str, Any]]) -> tuple[bool, bool]: + def _post_chunk( + self, chunk: list[dict[str, Any]], after_renew: bool = False + ) -> tuple[bool, bool]: try: plaintext = json.dumps(chunk, separators=(",", ":")).encode("utf-8") key = bytes.fromhex(self._token) @@ -113,6 +161,20 @@ def _post_chunk(self, chunk: list[dict[str, Any]]) -> tuple[bool, bool]: with urllib.request.urlopen(request, timeout=5): pass except urllib.error.HTTPError as exc: + # 403 means the session expired mid-run (long training jobs + # routinely outlive the registration TTL). Renew once with the + # session token and retry; a second 403 gives up normally. + if ( + exc.code == 403 + and not after_renew + and renew_fragment_session( + self._glaas_url, + self._session_id, + self._token, + self._renew_ttl_seconds, + ) + ): + return self._post_chunk(chunk, after_renew=True) _get_logger().warning( "Failed to stream fragments for session %s sequence %d: %s", self._session_id, @@ -131,5 +193,6 @@ def _post_chunk(self, chunk: list[dict[str, Any]]) -> tuple[bool, bool]: return True, False - def close(self) -> None: - self.flush() + def close(self) -> bool: + """Flush remaining fragments; False when any batch stayed undelivered.""" + return self.flush() diff --git a/tests/backends/k8s/README.md b/tests/backends/k8s/README.md index e0cc54c5..f5370685 100644 --- a/tests/backends/k8s/README.md +++ b/tests/backends/k8s/README.md @@ -42,6 +42,10 @@ Three test layers share this harness: `deploy/charts/roar-lineage-webhook` Helm chart (`scripts/deploy_webhook.sh`), so these tests exercise the packaged chart. Skips without the image / `--with-webhook`. +- `e2e/test_k8s_ttl_renewal.py` — fragment-session TTL renewal: the + session is registered with the minimum 60s TTL and the Job outlives it; + the in-pod streamer renews on 403 and the finalizer still reconstitutes. + Skips when the local glaas-api lacks `POST .../sessions/:id/renew`. - `e2e/test_k8s_smoke.py` — the Phase-0 runtime diagnostic: fixtures hand-wrap the manifest (no backend involved) to isolate the runtime pieces (in-pod tracing, fragment streaming, identity contract) when the product diff --git a/tests/backends/k8s/e2e/test_k8s_ttl_renewal.py b/tests/backends/k8s/e2e/test_k8s_ttl_renewal.py new file mode 100644 index 00000000..959a24f0 --- /dev/null +++ b/tests/backends/k8s/e2e/test_k8s_ttl_renewal.py @@ -0,0 +1,234 @@ +"""Fragment-session TTL renewal e2e: training that outlives its session. + +Registers the fragment session with the minimum TTL (60s) and runs a Job +that finishes after the session has expired. Without renewal the pod's +fragment POSTs 403 and the finalizer's fetch 403s too — lineage is lost. +With renewal (streamer renew-on-403 + the glaas-api renew endpoint) the +run reconstitutes normally. + +Skips when the local glaas-api does not expose the renewal endpoint yet. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import secrets +import sqlite3 +import subprocess +import sys +import urllib.error +import urllib.request +import uuid +from pathlib import Path +from typing import Any + +import pytest + +from .conftest import HOST_GLAAS_URL, KUBE_CONTEXT, NAMESPACE, TOOLS_BIN, kubectl +from .test_k8s_product_path import wheel_server # noqa: F401 (module fixture reuse) + +pytestmark = [ + pytest.mark.e2e, + pytest.mark.k8s_e2e, + # The job deliberately sleeps past the 60s session TTL. + pytest.mark.timeout(900), +] + +SLEEP_SECONDS = 75 + +JOB_MANIFEST_TEMPLATE = """\ +apiVersion: batch/v1 +kind: Job +metadata: + name: {job_name} + namespace: {namespace} + labels: + app.kubernetes.io/part-of: roar-k8s-e2e +spec: + backoffLimit: 0 + ttlSecondsAfterFinished: 1800 + template: + spec: + restartPolicy: Never + volumes: + - name: work + emptyDir: {{}} + containers: + - name: trainer + image: python:3.12-slim + workingDir: /work + volumeMounts: + - name: work + mountPath: /work + command: + - python + - -c + - >- + import time; + time.sleep({sleep_seconds}); + open('model.bin', 'wb').write(b'slow-training' * 4) +""" + + +def _glaas_supports_renewal() -> bool: + token = secrets.token_bytes(32).hex() + session_id = str(uuid.uuid4()) + register = urllib.request.Request( + url=f"{HOST_GLAAS_URL}/api/v1/fragments/sessions", + data=json.dumps( + { + "session_id": session_id, + "token_hash": hashlib.sha256(token.encode("utf-8")).hexdigest(), + "ttl_seconds": 60, + } + ).encode("utf-8"), + headers={"content-type": "application/json"}, + method="POST", + ) + renew = urllib.request.Request( + url=f"{HOST_GLAAS_URL}/api/v1/fragments/sessions/{session_id}/renew", + data=json.dumps({"ttl_seconds": 60}).encode("utf-8"), + headers={ + "content-type": "application/json", + "x-roar-fragment-token": token, + }, + method="POST", + ) + try: + with urllib.request.urlopen(register, timeout=5): + pass + with urllib.request.urlopen(renew, timeout=5) as response: + return response.status == 200 + except urllib.error.HTTPError: + return False + except Exception: + return False + + +@pytest.fixture(scope="module") +def renewal_run( + k8s_cluster: None, + glaas_health: str, + wheel_server: dict[str, str], # noqa: F811 + tmp_path_factory: pytest.TempPathFactory, +) -> dict[str, Any]: + if not _glaas_supports_renewal(): + pytest.skip("local glaas-api lacks fragment-session renewal (POST .../renew)") + + project_dir = tmp_path_factory.mktemp("k8s-ttl-renewal") + roar_dir = project_dir / ".roar" + roar_dir.mkdir() + (roar_dir / "config.toml").write_text( + "\n".join( + [ + "[glaas]", + 'url = "http://localhost:3001"', + "", + "[k8s]", + "enabled = true", + f'runtime_install_requirement = "{wheel_server["url"]}"', + 'cluster_glaas_url = "http://glaas:3001"', + "fragment_session_ttl_seconds = 60", + "wait_for_completion = true", + "wait_timeout_seconds = 420", + "poll_interval_seconds = 2.0", + ] + ) + + "\n", + encoding="utf-8", + ) + + job_name = f"roar-ttl-renew-{uuid.uuid4().hex[:6]}" + manifest_path = project_dir / "job.yaml" + manifest_path.write_text( + JOB_MANIFEST_TEMPLATE.format( + job_name=job_name, + namespace=NAMESPACE, + sleep_seconds=SLEEP_SECONDS, + ), + encoding="utf-8", + ) + + env = dict(os.environ) + env["PATH"] = f"{TOOLS_BIN}{os.pathsep}{env.get('PATH', '')}" + env.pop("GLAAS_URL", None) + env.pop("ROAR_PROJECT_DIR", None) + + try: + completed = subprocess.run( + [ + sys.executable, + "-m", + "roar", + "run", + "kubectl", + "apply", + "--context", + KUBE_CONTEXT, + "-f", + "job.yaml", + ], + cwd=project_dir, + env=env, + capture_output=True, + text=True, + check=False, + timeout=800, + ) + pod_logs = kubectl( + ["logs", "-n", NAMESPACE, "-l", f"job-name={job_name}", "--tail=100"], + check=False, + ) + return { + "project_dir": project_dir, + "job_name": job_name, + "exit_code": completed.returncode, + "stdout": completed.stdout, + "stderr": completed.stderr, + "pod_logs": pod_logs.stdout + pod_logs.stderr, + } + finally: + kubectl(["delete", f"job/{job_name}", "-n", NAMESPACE, "--ignore-not-found"], check=False) + kubectl( + [ + "delete", + "secret", + "-n", + NAMESPACE, + "-l", + "app.kubernetes.io/managed-by=roar", + "--ignore-not-found", + ], + check=False, + ) + + +def _describe(run: dict[str, Any]) -> str: + return ( + f"exit={run['exit_code']}\nstdout:\n{run['stdout']}\nstderr:\n{run['stderr']}\n" + f"pod logs:\n{run.get('pod_logs', '')}" + ) + + +def test_lineage_survives_session_expiry(renewal_run: dict[str, Any]) -> None: + """The 60s session expires mid-run; renewal keeps stream + fetch working.""" + assert renewal_run["exit_code"] == 0, _describe(renewal_run) + combined = renewal_run["stdout"] + renewal_run["stderr"] + assert "lineage reconstituted" in combined, _describe(renewal_run) + + db_path = Path(renewal_run["project_dir"]) / ".roar" / "roar.db" + assert db_path.is_file(), _describe(renewal_run) + conn = sqlite3.connect(db_path) + conn.row_factory = sqlite3.Row + try: + tasks = conn.execute("SELECT id FROM jobs WHERE job_type = 'k8s_task'").fetchall() + assert len(tasks) == 1, _describe(renewal_run) + outputs = conn.execute( + "SELECT path FROM job_outputs WHERE job_id = ?", + (tasks[0]["id"],), + ).fetchall() + finally: + conn.close() + assert any(str(row["path"]).endswith("model.bin") for row in outputs), _describe(renewal_run) diff --git a/tests/backends/ray/unit/test_fragment_reconstituter.py b/tests/backends/ray/unit/test_fragment_reconstituter.py index 3b2f9676..92b0880f 100644 --- a/tests/backends/ray/unit/test_fragment_reconstituter.py +++ b/tests/backends/ray/unit/test_fragment_reconstituter.py @@ -3,6 +3,7 @@ import base64 import importlib import json +import urllib.error import urllib.request from pathlib import Path @@ -85,6 +86,57 @@ def _fake_urlopen(request: urllib.request.Request, timeout: int = 0): ] +def test_fetch_batches_renews_expired_session_and_retries( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + module = _module() + state = {"renewed": False} + response_body = _wrapped_fragments_payload([{"sequence": 0, "encrypted_batch": "batch-0"}]) + + def _fake_urlopen(request: urllib.request.Request, timeout: int = 0): + del timeout + if request.full_url.endswith("/renew"): + state["renewed"] = True + return _FakeHttpResponse(b"{}") + if not state["renewed"]: + raise urllib.error.HTTPError(request.full_url, 403, "Forbidden", hdrs=None, fp=None) + return _FakeHttpResponse(response_body) + + monkeypatch.setattr(module.urllib.request, "urlopen", _fake_urlopen) + reconstituter = module.FragmentReconstituter( + session_id="session-expired-read", + token="ab" * 32, + glaas_url="http://localhost:3001", + roar_db_path=tmp_path / ".roar" / "roar.db", + ) + + assert reconstituter._fetch_batches() == [{"sequence": 0, "encrypted_batch": "batch-0"}] + assert state["renewed"] is True + + +def test_fetch_batches_returns_empty_when_renew_fails( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + module = _module() + + def _fake_urlopen(request: urllib.request.Request, timeout: int = 0): + del timeout + code = 404 if request.full_url.endswith("/renew") else 403 + raise urllib.error.HTTPError(request.full_url, code, "nope", hdrs=None, fp=None) + + monkeypatch.setattr(module.urllib.request, "urlopen", _fake_urlopen) + reconstituter = module.FragmentReconstituter( + session_id="session-dead", + token="ab" * 32, + glaas_url="http://localhost:3001", + roar_db_path=tmp_path / ".roar" / "roar.db", + ) + + assert reconstituter._fetch_batches() == [] + + def test_fetch_batches_supports_flat_response_fallback( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, diff --git a/tests/integrations/glaas/test_fragment_streamer.py b/tests/integrations/glaas/test_fragment_streamer.py index 34d08012..64f86f4e 100644 --- a/tests/integrations/glaas/test_fragment_streamer.py +++ b/tests/integrations/glaas/test_fragment_streamer.py @@ -281,6 +281,93 @@ def _fake_urlopen(request: urllib.request.Request, timeout: int = 0): assert streamer._next_sequence == 4 +def test_post_403_renews_session_and_retries(monkeypatch: pytest.MonkeyPatch, token: str) -> None: + state = {"renewed": False} + renew_requests: list[urllib.request.Request] = [] + + def _fake_urlopen(request: urllib.request.Request, timeout: int = 0): + del timeout + if request.full_url.endswith("/renew"): + renew_requests.append(request) + state["renewed"] = True + return _FakeHttpResponse(status=200) + if not state["renewed"]: + raise urllib.error.HTTPError(request.full_url, 403, "Forbidden", hdrs=None, fp=None) + return _FakeHttpResponse(status=202) + + monkeypatch.setattr(urllib.request, "urlopen", _fake_urlopen) + + streamer = GlaasFragmentStreamer( + session_id="session-expired", + token=token, + glaas_url="http://localhost:3001", + renew_ttl_seconds=1234, + ) + streamer.append_fragment({"job_uid": "job-late"}) + + assert streamer.flush() is True + assert streamer._buffer == [] + assert streamer.delivered_batches == 1 + assert streamer.failed_batches == 0 + + assert len(renew_requests) == 1 + renew = renew_requests[0] + assert renew.full_url == ( + "http://localhost:3001/api/v1/fragments/sessions/session-expired/renew" + ) + headers = {name.lower(): value for name, value in renew.header_items()} + assert headers["x-roar-fragment-token"] == token + assert json.loads(renew.data.decode("utf-8")) == {"ttl_seconds": 1234} + + +def test_post_403_gives_up_when_renew_fails(monkeypatch: pytest.MonkeyPatch, token: str) -> None: + def _fake_urlopen(request: urllib.request.Request, timeout: int = 0): + del timeout + if request.full_url.endswith("/renew"): + raise urllib.error.HTTPError(request.full_url, 404, "Not Found", hdrs=None, fp=None) + raise urllib.error.HTTPError(request.full_url, 403, "Forbidden", hdrs=None, fp=None) + + monkeypatch.setattr(urllib.request, "urlopen", _fake_urlopen) + + streamer = GlaasFragmentStreamer( + session_id="session-gone", + token=token, + glaas_url="http://localhost:3001", + ) + fragment = {"job_uid": "job-lost"} + streamer.append_fragment(fragment) + + assert streamer.flush() is False + assert streamer._buffer == [fragment] + assert streamer.pending_fragments == 1 + assert streamer.failed_batches == 1 + + +def test_post_403_after_successful_renew_gives_up( + monkeypatch: pytest.MonkeyPatch, token: str +) -> None: + post_attempts = {"count": 0} + + def _fake_urlopen(request: urllib.request.Request, timeout: int = 0): + del timeout + if request.full_url.endswith("/renew"): + return _FakeHttpResponse(status=200) + post_attempts["count"] += 1 + raise urllib.error.HTTPError(request.full_url, 403, "Forbidden", hdrs=None, fp=None) + + monkeypatch.setattr(urllib.request, "urlopen", _fake_urlopen) + + streamer = GlaasFragmentStreamer( + session_id="session-still-forbidden", + token=token, + glaas_url="http://localhost:3001", + ) + streamer.append_fragment({"job_uid": "job-1"}) + + assert streamer.flush() is False + assert post_attempts["count"] == 2 + + def test_close_flushes_remaining_fragments(monkeypatch: pytest.MonkeyPatch, token: str) -> None: streamer = GlaasFragmentStreamer( session_id="session-close", From 951e1d442a363a5f15ce6bdfcf43a852b055e632 Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Tue, 14 Jul 2026 18:54:59 +0000 Subject: [PATCH 15/31] feat(fragments): report undelivered batches instead of claiming streamed MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit emit_fragment_dicts returned "streamed" whenever credentials were present, even when every batch POST failed — mid-run streaming failures were silently lost because callers (k8s pod_entry bundle fallback, ray local merge) only fall back on a non-"streamed" result. It now returns "streamed" only when the streamer delivered everything; partial or total failures log an undelivered-count warning and fall through to local_merge / local_fallback with all fragments (downstream merges dedupe by task identity, so overlap with delivered batches is absorbed). Co-Authored-By: Claude Fable 5 --- docs/developer/k8s-integration.md | 8 +- roar/backends/k8s/pod_entry.py | 7 +- roar/execution/fragments/transport.py | 21 +++- tests/execution/fragments/__init__.py | 0 tests/execution/fragments/test_transport.py | 106 ++++++++++++++++++++ 5 files changed, 134 insertions(+), 8 deletions(-) create mode 100644 tests/execution/fragments/__init__.py create mode 100644 tests/execution/fragments/test_transport.py diff --git a/docs/developer/k8s-integration.md b/docs/developer/k8s-integration.md index 3bcfefc1..3b82f78c 100644 --- a/docs/developer/k8s-integration.md +++ b/docs/developer/k8s-integration.md @@ -146,9 +146,11 @@ Transport is streaming-first with a bundle fallback: when `k8s.bundle_dir` names a mounted shared volume and GLaaS is unreachable from the pod (probe or non-streamed emit), `pod_entry` writes `roar-fragments-.json` there instead; `roar k8s ingest-bundles ` merges a host-visible copy later. -Note the fragment streamer swallows per-batch POST failures (reports -"streamed" regardless), which is why the fallback needs its own probe — -surfacing streamer failure counts is an open follow-up. +`emit_fragment_dicts` reports "streamed" only when every batch was +delivered (the streamer exposes delivered/failed/pending counts and logs +an undelivered summary), so mid-run streaming failures also reach the +fallback; the upfront reachability probe remains as a fast path that +skips per-batch POST timeouts when GLaaS is dark. **Session TTL renewal**: fragment sessions are registered with `k8s.fragment_session_ttl_seconds` (default 86400, server-capped at 7 diff --git a/roar/backends/k8s/pod_entry.py b/roar/backends/k8s/pod_entry.py index 8d9ef4ed..ed8595db 100644 --- a/roar/backends/k8s/pod_entry.py +++ b/roar/backends/k8s/pod_entry.py @@ -248,9 +248,10 @@ def _load_proxy_log_refs( def _emit_or_bundle(fragments: list[dict]) -> str: """Stream fragments; fall back to a bundle file when GLaaS is unreachable. - The streamer itself swallows per-batch POST failures, so a quick - reachability probe decides upfront; a non-"streamed" emit result also - falls back when a bundle directory is declared. + The reachability probe short-circuits the obviously-dark case without + paying per-batch POST timeouts; emit_fragment_dicts reports "streamed" + only when every batch was delivered, so mid-run streaming failures + (partial or total) also land in the bundle fallback. """ from roar.execution.fragments.transport import emit_fragment_dicts diff --git a/roar/execution/fragments/transport.py b/roar/execution/fragments/transport.py index cd0c0982..36be4535 100644 --- a/roar/execution/fragments/transport.py +++ b/roar/execution/fragments/transport.py @@ -11,6 +11,12 @@ FragmentEmitResult = Literal["streamed", "merged", "fallback", "skipped"] +def _get_logger(): + from roar.core.logging import get_logger + + return get_logger() + + def emit_fragment_dicts( fragments: Sequence[dict[str, Any]], *, @@ -28,6 +34,7 @@ def emit_fragment_dicts( glaas_url = str(resolved_env.get("GLAAS_URL") or "").strip() if session_id and token and glaas_url: + undelivered = len(normalized_fragments) try: streamer = GlaasFragmentStreamer( session_id=session_id, @@ -36,10 +43,20 @@ def emit_fragment_dicts( ) for fragment in normalized_fragments: streamer.append_fragment(fragment) - streamer.close() - return "streamed" + if streamer.close(): + return "streamed" + undelivered = streamer.pending_fragments except Exception: pass + # Partial deliveries fall through with everything: downstream merges + # dedupe by task identity, while dropping fragments here would lose + # lineage for good. + _get_logger().warning( + "Fragment streaming incomplete for session %s: %d of %d fragment(s) undelivered", + session_id, + undelivered, + len(normalized_fragments), + ) project_dir = str(resolved_env.get("ROAR_PROJECT_DIR", "")).strip() if local_merge is not None and project_dir: diff --git a/tests/execution/fragments/__init__.py b/tests/execution/fragments/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/execution/fragments/test_transport.py b/tests/execution/fragments/test_transport.py new file mode 100644 index 00000000..d26bd295 --- /dev/null +++ b/tests/execution/fragments/test_transport.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +from typing import Any + +import pytest + +from roar.execution.fragments import transport + + +class _FakeStreamer: + """Streamer double with a scripted close() outcome.""" + + def __init__( + self, + *, + close_result: bool = True, + pending_after_close: int = 0, + raise_on_append: bool = False, + ) -> None: + self._close_result = close_result + self._pending_after_close = pending_after_close + self._raise_on_append = raise_on_append + self.appended: list[dict[str, Any]] = [] + + def append_fragment(self, fragment: dict[str, Any]) -> None: + if self._raise_on_append: + raise RuntimeError("boom") + self.appended.append(fragment) + + def close(self) -> bool: + return self._close_result + + @property + def pending_fragments(self) -> int: + return self._pending_after_close + + +ENV = { + "ROAR_SESSION_ID": "session-1", + "ROAR_FRAGMENT_TOKEN": "ab" * 32, + "GLAAS_URL": "http://localhost:3001", + "ROAR_PROJECT_DIR": "/proj", +} + + +def _install_streamer(monkeypatch: pytest.MonkeyPatch, streamer: _FakeStreamer) -> None: + monkeypatch.setattr(transport, "GlaasFragmentStreamer", lambda **kwargs: streamer) + + +def test_emit_returns_streamed_only_when_fully_delivered( + monkeypatch: pytest.MonkeyPatch, +) -> None: + streamer = _FakeStreamer(close_result=True) + _install_streamer(monkeypatch, streamer) + + result = transport.emit_fragment_dicts([{"job_uid": "a"}], env=ENV) + + assert result == "streamed" + assert streamer.appended == [{"job_uid": "a"}] + + +def test_emit_falls_back_to_local_merge_on_undelivered_batches( + monkeypatch: pytest.MonkeyPatch, +) -> None: + streamer = _FakeStreamer(close_result=False, pending_after_close=2) + _install_streamer(monkeypatch, streamer) + merged: list[tuple[list[dict[str, Any]], str, str | None]] = [] + + result = transport.emit_fragment_dicts( + [{"job_uid": "a"}, {"job_uid": "b"}], + env=ENV, + local_merge=lambda fragments, project_dir, driver: merged.append( + (fragments, project_dir, driver) + ), + ) + + assert result == "merged" + assert merged == [([{"job_uid": "a"}, {"job_uid": "b"}], "/proj", None)] + + +def test_emit_falls_back_to_local_fallback_on_streamer_crash( + monkeypatch: pytest.MonkeyPatch, +) -> None: + streamer = _FakeStreamer(raise_on_append=True) + _install_streamer(monkeypatch, streamer) + fallback_calls: list[str] = [] + + result = transport.emit_fragment_dicts( + [{"job_uid": "a"}], + env={key: value for key, value in ENV.items() if key != "ROAR_PROJECT_DIR"}, + local_fallback=lambda: fallback_calls.append("fallback"), + ) + + assert result == "fallback" + assert fallback_calls == ["fallback"] + + +def test_emit_returns_skipped_when_no_fallback_available( + monkeypatch: pytest.MonkeyPatch, +) -> None: + streamer = _FakeStreamer(close_result=False, pending_after_close=1) + _install_streamer(monkeypatch, streamer) + + result = transport.emit_fragment_dicts([{"job_uid": "a"}], env=ENV) + + assert result == "skipped" From 487180b479ff88da636be361b9e8a4c9e4da584a Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Wed, 15 Jul 2026 16:08:33 +0000 Subject: [PATCH 16/31] fix(ray): stop duplicating roar env keys between Job and driver runtime envs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every `roar run ray job submit` whose driver called ray.init() died with "Failed to merge the Job's runtime env ... because of a conflict": the submit rewrite delivers ROAR_NO_TELEMETRY (plus the rest of the env contract) in the Job-level runtime env, and the patched driver-side ray.init re-added ROAR_NO_TELEMETRY — ray >= 2.4 rejects the Job/driver merge on ANY duplicated env-var key, even with identical values. The failure was roar conflicting with itself, not a ray version issue. The instrumented driver path now reads the job config ray injects (RAY_JOB_CONFIG_JSON_ENV_VAR) and only adds what the Job env lacks (same for py_executable — a job-level value wins with a warning). The non-instrumented path gets a generic guard that drops roar-added env keys the Job env already carries; user-supplied keys are left alone so genuine conflicts still surface through ray's own error. The merged runtime env that reaches workers is unchanged. Live-validated on the compose harness: the previously-failing runtime-env-conflict and native-tracing e2e pass, and the suite goes from dead-on-arrival instrumented submits to 57 passing (remaining failures are pre-existing capture-fidelity gaps newly reachable now that submits run, tracked separately). Co-Authored-By: Claude Fable 5 --- docs/developer/ray-integration.md | 10 ++ roar/backends/ray/runtime_hooks.py | 70 ++++++++- tests/backends/ray/unit/test_runtime_hooks.py | 143 ++++++++++++++++++ 3 files changed, 220 insertions(+), 3 deletions(-) diff --git a/docs/developer/ray-integration.md b/docs/developer/ray-integration.md index 7b8d32cc..b665efc2 100644 --- a/docs/developer/ray-integration.md +++ b/docs/developer/ray-integration.md @@ -117,6 +117,16 @@ This is the generalization seam for future local and distributed integrations. These two modules are the main backend-neutral extraction from the older inline Ray rewrite. +**Job/driver runtime-env merge safety**: Ray (≥ 2.4) refuses to merge the +submitted Job's runtime env with a driver's `ray.init` runtime env when any +field or env-var *key* appears in both — even with identical values +("Failed to merge the Job's runtime env"). Since the submit rewrite already +delivers the roar env contract at the Job level, the patched driver-side +`ray.init` reads the injected job config (`RAY_JOB_CONFIG_JSON_ENV_VAR`) +and only adds what the Job env lacks; roar-added duplicates are dropped +(`_drop_roar_env_keys_already_in_job_env`) while user-supplied keys are +left alone so genuine conflicts still surface through Ray's own error. + ### d. Driver bootstrap - `roar/execution/runtime/driver_entrypoint.py` diff --git a/roar/backends/ray/runtime_hooks.py b/roar/backends/ray/runtime_hooks.py index f7e5101f..3708a9ae 100644 --- a/roar/backends/ray/runtime_hooks.py +++ b/roar/backends/ray/runtime_hooks.py @@ -149,6 +149,9 @@ def _roar_ray_init(*args, **kwargs): ) runtime_env = prepare_worker_runtime_env(runtime_env, job_id) runtime_env = sanitize_worker_runtime_env_for_ray(ray_module, runtime_env) + runtime_env = _drop_roar_env_keys_already_in_job_env( + runtime_env, user_env_keys=set(env_vars) + ) kwargs["runtime_env"] = runtime_env result = real_ray_init(*args, **kwargs) register_pre_shutdown_ray_collection() @@ -609,16 +612,77 @@ def _client(service_name, *args, **kwargs): boto3.client = _client +# Ray's job manager exports the submitted Job's config to the driver under +# this literal name (the constant's value equals the constant's name). +_RAY_JOB_CONFIG_JSON_ENV_VAR = "RAY_JOB_CONFIG_JSON_ENV_VAR" + + +def _injected_job_runtime_env() -> dict[str, Any]: + """Runtime env the Job manager injected for this driver ({} outside a job).""" + raw = os.environ.get(_RAY_JOB_CONFIG_JSON_ENV_VAR, "") + if not raw: + return {} + try: + payload = json.loads(raw) + except (TypeError, ValueError): + return {} + if not isinstance(payload, dict): + return {} + runtime_env = payload.get("runtime_env") + return runtime_env if isinstance(runtime_env, dict) else {} + + +def _drop_roar_env_keys_already_in_job_env( + runtime_env: Mapping[str, Any], + *, + user_env_keys: set[str], +) -> dict[str, Any]: + """Keep the driver's ray.init runtime env mergeable inside a submitted Job. + + Roar-added env keys the Job's runtime env already carries are removed — + their values reach the driver and workers through the Job env anyway. + User-supplied keys stay untouched so a genuine conflict still surfaces + through Ray's own error. + """ + injected_env_keys = set((_injected_job_runtime_env().get("env_vars") or {}).keys()) + if not injected_env_keys: + return dict(runtime_env) + runtime_env_out = dict(runtime_env) + env_vars = dict(runtime_env_out.get("env_vars", {}) or {}) + for key in list(env_vars): + if key not in user_env_keys and key in injected_env_keys: + env_vars.pop(key) + if env_vars: + runtime_env_out["env_vars"] = env_vars + else: + runtime_env_out.pop("env_vars", None) + return runtime_env_out + + def _prepare_instrumented_job_worker_runtime_env( runtime_env: Mapping[str, Any] | None, job_id: str, ) -> dict[str, Any]: del job_id + # Ray refuses to merge the Job's runtime env with the driver's ray.init + # runtime env when any field or env-var key appears in both — even with + # identical values. The submit rewrite already delivers the roar env + # contract at the Job level, so only add what the Job env lacks. + injected = _injected_job_runtime_env() + injected_env_keys = set((injected.get("env_vars") or {}).keys()) runtime_env_out = dict(runtime_env or {}) env_vars = dict(runtime_env_out.get("env_vars", {}) or {}) - env_vars[ROAR_NO_TELEMETRY_ENV] = "1" - runtime_env_out["env_vars"] = env_vars - runtime_env_out["py_executable"] = WORKER_PY_EXECUTABLE + if ROAR_NO_TELEMETRY_ENV not in injected_env_keys: + env_vars[ROAR_NO_TELEMETRY_ENV] = "1" + if env_vars: + runtime_env_out["env_vars"] = env_vars + if "py_executable" in injected: + print( + "[roar] WARNING: the Job's runtime env already sets py_executable; " + "leaving it in place (per-task worker capture may be degraded)" + ) + else: + runtime_env_out["py_executable"] = WORKER_PY_EXECUTABLE return runtime_env_out diff --git a/tests/backends/ray/unit/test_runtime_hooks.py b/tests/backends/ray/unit/test_runtime_hooks.py index 94006092..ffb09934 100644 --- a/tests/backends/ray/unit/test_runtime_hooks.py +++ b/tests/backends/ray/unit/test_runtime_hooks.py @@ -419,6 +419,149 @@ def fake_ray_init(*_args, **kwargs): assert captured_runtime_env["env_vars"][ROAR_NO_TELEMETRY_ENV] == "1" +def _set_injected_job_config( + monkeypatch: pytest.MonkeyPatch, runtime_env: dict[str, object] +) -> None: + import json + + monkeypatch.setenv( + "RAY_JOB_CONFIG_JSON_ENV_VAR", + json.dumps({"runtime_env": runtime_env}), + ) + + +def test_instrumented_driver_env_skips_keys_the_job_env_already_carries( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _set_injected_job_config( + monkeypatch, + { + "env_vars": {ROAR_NO_TELEMETRY_ENV: "1", "ROAR_JOB_ID": "abc123"}, + "worker_process_setup_hook": "roar.execution.runtime.worker_bootstrap.startup", + }, + ) + + prepared = runtime_hooks._prepare_instrumented_job_worker_runtime_env({}, "abc123") + + # ray >= 2.4 rejects the Job/driver merge on ANY duplicated env key, so + # nothing the Job env already delivers may be re-added driver-side. + assert "env_vars" not in prepared + assert prepared["py_executable"] == "roar-worker" + + +def test_instrumented_driver_env_adds_contract_outside_a_submitted_job( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("RAY_JOB_CONFIG_JSON_ENV_VAR", raising=False) + + prepared = runtime_hooks._prepare_instrumented_job_worker_runtime_env({}, "abc123") + + assert prepared["env_vars"][ROAR_NO_TELEMETRY_ENV] == "1" + assert prepared["py_executable"] == "roar-worker" + + +def test_instrumented_driver_defers_to_job_level_py_executable( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _set_injected_job_config( + monkeypatch, + {"env_vars": {ROAR_NO_TELEMETRY_ENV: "1"}, "py_executable": "custom-exec"}, + ) + + prepared = runtime_hooks._prepare_instrumented_job_worker_runtime_env({}, "abc123") + + assert "py_executable" not in prepared + + +def test_instrumented_driver_preserves_user_runtime_env( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _set_injected_job_config(monkeypatch, {"env_vars": {ROAR_NO_TELEMETRY_ENV: "1"}}) + + prepared = runtime_hooks._prepare_instrumented_job_worker_runtime_env( + {"env_vars": {"USER_KEY": "value"}, "working_dir": "/tmp/user"}, + "abc123", + ) + + assert prepared["env_vars"] == {"USER_KEY": "value"} + assert prepared["working_dir"] == "/tmp/user" + + +def test_drop_roar_env_keys_keeps_user_keys_and_drops_roar_duplicates( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _set_injected_job_config( + monkeypatch, + {"env_vars": {"ROAR_JOB_ID": "job-level", "AWS_ENDPOINT_URL": "http://job:1"}}, + ) + + runtime_env = { + "env_vars": { + "ROAR_JOB_ID": "driver-added", + "AWS_ENDPOINT_URL": "http://user:2", + ROAR_NO_TELEMETRY_ENV: "1", + }, + "py_executable": "roar-worker", + } + + result = runtime_hooks._drop_roar_env_keys_already_in_job_env( + runtime_env, user_env_keys={"AWS_ENDPOINT_URL"} + ) + + # roar-added duplicate dropped; user-supplied duplicate kept (ray should + # surface that conflict); roar keys absent from the job env kept. + assert "ROAR_JOB_ID" not in result["env_vars"] + assert result["env_vars"]["AWS_ENDPOINT_URL"] == "http://user:2" + assert result["env_vars"][ROAR_NO_TELEMETRY_ENV] == "1" + assert result["py_executable"] == "roar-worker" + + +def test_patch_ray_init_merges_cleanly_inside_roar_submitted_job( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """End-to-end shape of the native-harness failure: strict ray merge.""" + import json + + job_runtime_env = { + "working_dir": "gcs://pkg.zip", + "env_vars": { + "AWS_ENDPOINT_URL": "http://127.0.0.1:34270", + "ROAR_EXECUTION_BACKEND": "ray", + "ROAR_JOB_ID": "a293c31c", + "ROAR_JOB_INSTRUMENTED": "1", + ROAR_NO_TELEMETRY_ENV: "1", + "ROAR_WRAP": "1", + }, + "worker_process_setup_hook": "roar.execution.runtime.worker_bootstrap.startup", + } + monkeypatch.setenv("RAY_JOB_CONFIG_JSON_ENV_VAR", json.dumps({"runtime_env": job_runtime_env})) + monkeypatch.setenv("ROAR_JOB_INSTRUMENTED", "1") + monkeypatch.setenv("ROAR_RAY_NODE_AGENTS", "0") + monkeypatch.setattr( + runtime_hooks, + "load_ray_config", + lambda: {"enabled": True, "pip_install": False}, + ) + monkeypatch.setattr(runtime_hooks, "register_pre_shutdown_ray_collection", lambda: None) + + def strict_merge_ray_init(*_args, **kwargs): + driver_env = dict(kwargs.get("runtime_env", {}) or {}) + parent = {key: value for key, value in job_runtime_env.items() if key != "env_vars"} + child = {key: value for key, value in driver_env.items() if key != "env_vars"} + if set(parent) & set(child): + raise ValueError("Failed to merge the Job's runtime env (field conflict)") + parent_env = dict(job_runtime_env.get("env_vars", {})) + child_env = dict(driver_env.get("env_vars", {}) or {}) + if set(parent_env) & set(child_env): + raise ValueError("Failed to merge the Job's runtime env (env key conflict)") + return "ok" + + fake_ray = SimpleNamespace(init=strict_merge_ray_init) + runtime_hooks.patch_ray_init(fake_ray) + + assert fake_ray.init() == "ok" + + def test_patch_ray_init_registers_pre_shutdown_collection_for_instrumented_jobs( monkeypatch: pytest.MonkeyPatch, ) -> None: From 0e1cdfbdfc2425b4d1590f2965192e8e10a2bc93 Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Wed, 15 Jul 2026 17:17:47 +0000 Subject: [PATCH 17/31] fix(k8s): merge 413-split fragment parts at reconstitution MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The streamer splits an oversized fragment into parts sharing one task identity after an HTTP 413; the keep-last-per-identity dedup then silently discarded every part but the last, losing most of a large task's reads/writes. Retries are already attempt-distinct in the identity contract (pod uid + restart attempt), so same-identity fragments can only be duplicate deliveries or split parts — dedup now unions reads/writes per path across them (last ref wins, last fragment's scalars win). No wire-format change. Co-Authored-By: Claude Fable 5 --- roar/backends/k8s/fragment_reconstituter.py | 42 +++++++-- .../k8s/unit/test_fragment_reconstituter.py | 94 +++++++++++++++++++ 2 files changed, 130 insertions(+), 6 deletions(-) create mode 100644 tests/backends/k8s/unit/test_fragment_reconstituter.py diff --git a/roar/backends/k8s/fragment_reconstituter.py b/roar/backends/k8s/fragment_reconstituter.py index 5c1ade8b..20da4cf2 100644 --- a/roar/backends/k8s/fragment_reconstituter.py +++ b/roar/backends/k8s/fragment_reconstituter.py @@ -167,18 +167,48 @@ def _decrypt_batch(self, batch: dict[str, Any], key: bytes) -> list[dict[str, An @staticmethod def _deduplicate_by_task_identity(fragments: list[dict[str, Any]]) -> list[dict[str, Any]]: - """Keep the last emission per task identity (batches are sequence-sorted). - - Re-emissions and retries replace earlier snapshots of the same task - rather than duplicating them; the merge engine additionally - reconciles job UIDs for anything that slips through. + """Coalesce fragments per task identity (batches are sequence-sorted). + + Retries are already attempt-distinct in the identity contract + (pod uid + restart attempt), so same-identity fragments are either + duplicate deliveries (stream + bundle) or parts of an oversized + fragment the streamer split after an HTTP 413. The last fragment's + scalar fields win; reads/writes are unioned across parts (per-path, + last ref wins) so splitting never discards lineage references. """ winners: dict[str, dict[str, Any]] = {} for index, fragment in enumerate(fragments): identity = str(fragment.get("task_identity") or "").strip() or f"fragment:{index}" - winners[identity] = fragment + previous = winners.get(identity) + if previous is None: + winners[identity] = dict(fragment) + continue + merged = dict(fragment) + for list_key in ("reads", "writes"): + merged[list_key] = K8sFragmentReconstituter._merge_refs_by_path( + previous.get(list_key), fragment.get(list_key) + ) + winners[identity] = merged return list(winners.values()) + @staticmethod + def _merge_refs_by_path(earlier: Any, later: Any) -> list[dict[str, Any]]: + merged: dict[str, dict[str, Any]] = {} + ordered: list[str] = [] + for refs in (earlier, later): + if not isinstance(refs, list): + continue + for ref in refs: + if not isinstance(ref, dict): + continue + path = str(ref.get("path") or "") + if not path: + continue + if path not in merged: + ordered.append(path) + merged[path] = ref + return [merged[path] for path in ordered] + @staticmethod def _sequence_key(batch: dict[str, Any]) -> int: sequence = batch.get("sequence") diff --git a/tests/backends/k8s/unit/test_fragment_reconstituter.py b/tests/backends/k8s/unit/test_fragment_reconstituter.py new file mode 100644 index 00000000..2dcfd2c5 --- /dev/null +++ b/tests/backends/k8s/unit/test_fragment_reconstituter.py @@ -0,0 +1,94 @@ +"""Unit tests for k8s fragment deduplication at reconstitution.""" + +from __future__ import annotations + +from roar.backends.k8s.fragment_reconstituter import K8sFragmentReconstituter + + +def _fragment(identity: str, reads: list[dict], writes: list[dict], **extra) -> dict: + return { + "task_identity": identity, + "reads": reads, + "writes": writes, + **extra, + } + + +def test_dedup_merges_oversized_split_parts() -> None: + """413-split parts share one identity; no reference may be discarded.""" + part_one = _fragment( + "pod-1:trainer:0:0", + reads=[{"path": "/data/a.csv", "hash": "aa"}], + writes=[{"path": "/out/model.bin", "hash": "mm"}], + ) + part_two = _fragment( + "pod-1:trainer:0:0", + reads=[{"path": "/data/b.csv", "hash": "bb"}], + writes=[{"path": "/out/metrics.json", "hash": "jj"}], + ) + + result = K8sFragmentReconstituter._deduplicate_by_task_identity([part_one, part_two]) + + assert len(result) == 1 + merged = result[0] + assert {ref["path"] for ref in merged["reads"]} == {"/data/a.csv", "/data/b.csv"} + assert {ref["path"] for ref in merged["writes"]} == { + "/out/model.bin", + "/out/metrics.json", + } + + +def test_dedup_collapses_duplicate_deliveries_without_duplicating_refs() -> None: + """Stream + bundle double-delivery of the same fragment stays one fragment.""" + fragment = _fragment( + "pod-1:trainer:0:0", + reads=[{"path": "/data/a.csv", "hash": "aa"}], + writes=[{"path": "/out/model.bin", "hash": "mm"}], + ) + + result = K8sFragmentReconstituter._deduplicate_by_task_identity( + [dict(fragment), dict(fragment)] + ) + + assert len(result) == 1 + assert len(result[0]["reads"]) == 1 + assert len(result[0]["writes"]) == 1 + + +def test_dedup_last_ref_wins_per_path_and_scalars_from_last_fragment() -> None: + earlier = _fragment( + "pod-1:trainer:0:0", + reads=[{"path": "/data/a.csv", "hash": "stale"}], + writes=[], + exit_code=1, + ) + later = _fragment( + "pod-1:trainer:0:0", + reads=[{"path": "/data/a.csv", "hash": "fresh"}], + writes=[], + exit_code=0, + ) + + result = K8sFragmentReconstituter._deduplicate_by_task_identity([earlier, later]) + + assert len(result) == 1 + assert result[0]["reads"] == [{"path": "/data/a.csv", "hash": "fresh"}] + assert result[0]["exit_code"] == 0 + + +def test_dedup_keeps_distinct_identities_separate() -> None: + rank0 = _fragment("pod-1:trainer:0:0", reads=[], writes=[{"path": "/out/w0"}]) + rank1 = _fragment("pod-2:trainer:1:0", reads=[], writes=[{"path": "/out/w1"}]) + + result = K8sFragmentReconstituter._deduplicate_by_task_identity([rank0, rank1]) + + assert len(result) == 2 + + +def test_dedup_fragments_without_identity_are_not_conflated() -> None: + anon_one = _fragment("", reads=[{"path": "/a"}], writes=[]) + anon_two = _fragment("", reads=[{"path": "/b"}], writes=[]) + + result = K8sFragmentReconstituter._deduplicate_by_task_identity([anon_one, anon_two]) + + assert len(result) == 2 From 25d7f96f57240b408dbad047287339b817037184 Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Wed, 15 Jul 2026 17:17:47 +0000 Subject: [PATCH 18/31] fix(k8s,ray): resolve plan-time .roar at the project root; match kubectl global flags MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both submit planners saved the fragment-session key (and k8s the prepared manifest) under a bare cwd/.roar, while the run finalizer loads the key from the context-resolved project root — submitting from a project subdirectory created a stray nested .roar and lineage silently never reconstituted locally. Planners now resolve the project .roar via the same upward walk the CLI context uses (ROAR_PROJECT_DIR override, bounded by the git root), shared as roar.execution.fragments.sessions.resolve_project_roar_dir. Covered by unit tests for both backends and a KIND e2e that submits from a nested directory. The kubectl matcher also accepted the verb only at argv[1], so common forms like `kubectl --context prod apply -f job.yaml` silently bypassed instrumentation. The verb now matches anywhere; the existing -f-manifest and single-supported-workload guards keep false positives out. Co-Authored-By: Claude Fable 5 --- roar/backends/k8s/submit.py | 13 +- roar/backends/ray/submit.py | 3 +- roar/execution/fragments/sessions.py | 30 ++++ .../k8s/e2e/test_k8s_subdir_submit.py | 150 ++++++++++++++++++ .../backends/k8s/unit/test_submit_planning.py | 57 +++++++ tests/execution/fragments/test_sessions.py | 63 ++++++++ 6 files changed, 312 insertions(+), 4 deletions(-) create mode 100644 tests/backends/k8s/e2e/test_k8s_subdir_submit.py create mode 100644 tests/execution/fragments/test_sessions.py diff --git a/roar/backends/k8s/submit.py b/roar/backends/k8s/submit.py index efb85e70..32a59792 100644 --- a/roar/backends/k8s/submit.py +++ b/roar/backends/k8s/submit.py @@ -17,7 +17,11 @@ load_manifest_documents, rewrite_manifest_for_lineage, ) -from roar.execution.fragments.sessions import generate_fragment_session, save_fragment_session +from roar.execution.fragments.sessions import ( + generate_fragment_session, + resolve_project_roar_dir, + save_fragment_session, +) from roar.execution.framework.contract import ExecutionCommandPlan _KUBECTL_VERBS = ("apply", "create") @@ -47,7 +51,10 @@ def matches_kubectl_job_submit_command(command: list[str]) -> bool: return False if Path(command[0]).name.lower() != "kubectl": return False - if command[1].lower() not in _KUBECTL_VERBS: + # Global flags may precede the verb (kubectl --context X apply -f ...), + # so accept the verb anywhere. The -f-points-at-a-manifest and + # single-supported-workload guards below keep false positives out. + if not any(arg.lower() in _KUBECTL_VERBS for arg in command[1:]): return False if not _k8s_backend_enabled(): return False @@ -77,7 +84,7 @@ def plan_kubectl_job_submit_command(command: list[str]) -> ExecutionCommandPlan: start_dir = os.environ.get("ROAR_PROJECT_DIR") or os.getcwd() config = load_k8s_backend_config(start_dir=start_dir) - roar_dir = Path(os.getcwd()) / ".roar" + roar_dir = resolve_project_roar_dir() glaas_url = _resolve_glaas_url() if glaas_url is None: diff --git a/roar/backends/ray/submit.py b/roar/backends/ray/submit.py index bc5b3773..e0042919 100644 --- a/roar/backends/ray/submit.py +++ b/roar/backends/ray/submit.py @@ -16,6 +16,7 @@ from roar.execution.fragments.sessions import ( generate_fragment_session as generate_fragment_key, ) +from roar.execution.fragments.sessions import resolve_project_roar_dir from roar.execution.fragments.sessions import save_fragment_session as save_key from roar.execution.framework.contract import ROAR_EXECUTION_BACKEND_ENV, ExecutionCommandPlan from roar.integrations.glaas import GlaasClient @@ -86,7 +87,7 @@ def plan_ray_job_submit_command(command: list[str]) -> ExecutionCommandPlan: except Exception: pass else: - save_key(Path(os.getcwd()) / ".roar", key) + save_key(resolve_project_roar_dir(), key) env_vars["ROAR_SESSION_ID"] = key["session_id"] env_vars["ROAR_FRAGMENT_TOKEN"] = key["token"] fragment_session_id = str(key["session_id"]) diff --git a/roar/execution/fragments/sessions.py b/roar/execution/fragments/sessions.py index 66c7586a..b10c4bef 100644 --- a/roar/execution/fragments/sessions.py +++ b/roar/execution/fragments/sessions.py @@ -2,13 +2,43 @@ import hashlib import json +import os import secrets import uuid +from collections.abc import Mapping from datetime import datetime, timezone from pathlib import Path from typing import Any +def resolve_project_roar_dir( + environ: Mapping[str, str] | None = None, + cwd: Path | None = None, +) -> Path: + """Locate the project's .roar directory the way the CLI context does. + + Honors ROAR_PROJECT_DIR, then walks upward from cwd looking for an + existing .roar directory (bounded by the enclosing git repository when + present); falls back to cwd/.roar. Submit planners must save plan-time + state (session keys, prepared manifests) here — the run finalizer loads + the session key from the context-resolved .roar, so saving under a bare + cwd/.roar strands the key when invoked from a project subdirectory. + """ + resolved_env = os.environ if environ is None else environ + override = str(resolved_env.get("ROAR_PROJECT_DIR") or "").strip() + if override: + return Path(override) / ".roar" + + base = Path.cwd() if cwd is None else Path(cwd) + for parent in [base, *base.parents]: + candidate = parent / ".roar" + if candidate.is_dir(): + return candidate + if (parent / ".git").exists(): + break + return base / ".roar" + + def generate_fragment_session() -> dict[str, str]: token = secrets.token_bytes(32).hex() return { diff --git a/tests/backends/k8s/e2e/test_k8s_subdir_submit.py b/tests/backends/k8s/e2e/test_k8s_subdir_submit.py new file mode 100644 index 00000000..059eda5c --- /dev/null +++ b/tests/backends/k8s/e2e/test_k8s_subdir_submit.py @@ -0,0 +1,150 @@ +"""Subdirectory-invocation e2e: plan-time state must reach the root .roar. + +Regression for the cwd/.roar bug: submitting from a nested project +directory used to save the fragment-session key (and prepared manifest) +under a freshly created `/.roar`, while the finalizer loads the +key from the context-resolved project root — lineage was silently never +reconstituted. The submit itself succeeded, which is what made it easy +to miss. +""" + +from __future__ import annotations + +import os +import sqlite3 +import subprocess +import sys +import uuid +from pathlib import Path +from typing import Any + +import pytest + +from .conftest import KUBE_CONTEXT, NAMESPACE, TOOLS_BIN, kubectl +from .test_k8s_product_path import JOB_MANIFEST_TEMPLATE, wheel_server # noqa: F401 + +pytestmark = [ + pytest.mark.e2e, + pytest.mark.k8s_e2e, + pytest.mark.timeout(900), +] + + +@pytest.fixture(scope="module") +def subdir_run( + k8s_cluster: None, + glaas_health: str, + wheel_server: dict[str, str], # noqa: F811 + tmp_path_factory: pytest.TempPathFactory, +) -> dict[str, Any]: + project_dir = tmp_path_factory.mktemp("k8s-subdir") + subprocess.run(["git", "init", "-q"], cwd=project_dir, check=True) + roar_dir = project_dir / ".roar" + roar_dir.mkdir() + (roar_dir / "config.toml").write_text( + "\n".join( + [ + "[glaas]", + 'url = "http://localhost:3001"', + "", + "[k8s]", + "enabled = true", + f'runtime_install_requirement = "{wheel_server["url"]}"', + 'cluster_glaas_url = "http://glaas:3001"', + "wait_for_completion = true", + "wait_timeout_seconds = 300", + "poll_interval_seconds = 2.0", + ] + ) + + "\n", + encoding="utf-8", + ) + + nested = project_dir / "experiments" / "run-1" + nested.mkdir(parents=True) + + job_name = f"roar-subdir-{uuid.uuid4().hex[:6]}" + manifest_path = nested / "job.yaml" + manifest_path.write_text( + JOB_MANIFEST_TEMPLATE.format(job_name=job_name, namespace=NAMESPACE), + encoding="utf-8", + ) + + (project_dir / ".gitignore").write_text(".roar/\n", encoding="utf-8") + subprocess.run(["git", "config", "user.email", "e2e@example.com"], cwd=project_dir, check=True) + subprocess.run(["git", "config", "user.name", "E2E"], cwd=project_dir, check=True) + subprocess.run(["git", "add", "-A"], cwd=project_dir, check=True) + subprocess.run(["git", "commit", "-q", "-m", "init"], cwd=project_dir, check=True) + + env = dict(os.environ) + env["PATH"] = f"{TOOLS_BIN}{os.pathsep}{env.get('PATH', '')}" + env.pop("GLAAS_URL", None) + env.pop("ROAR_PROJECT_DIR", None) + + try: + completed = subprocess.run( + [ + sys.executable, + "-m", + "roar", + "run", + "kubectl", + "apply", + "--context", + KUBE_CONTEXT, + "-f", + "job.yaml", + ], + cwd=nested, + env=env, + capture_output=True, + text=True, + check=False, + timeout=600, + ) + return { + "project_dir": project_dir, + "nested_dir": nested, + "job_name": job_name, + "exit_code": completed.returncode, + "stdout": completed.stdout, + "stderr": completed.stderr, + } + finally: + kubectl(["delete", f"job/{job_name}", "-n", NAMESPACE, "--ignore-not-found"], check=False) + kubectl( + [ + "delete", + "secret", + "-n", + NAMESPACE, + "-l", + "app.kubernetes.io/managed-by=roar", + "--ignore-not-found", + ], + check=False, + ) + + +def _describe(run: dict[str, Any]) -> str: + return f"exit={run['exit_code']}\nstdout:\n{run['stdout']}\nstderr:\n{run['stderr']}" + + +def test_subdir_submit_reconstitutes_into_project_root(subdir_run: dict[str, Any]) -> None: + assert subdir_run["exit_code"] == 0, _describe(subdir_run) + combined = subdir_run["stdout"] + subdir_run["stderr"] + assert "lineage reconstituted" in combined, _describe(subdir_run) + assert "failed to load fragment session" not in combined, _describe(subdir_run) + + # No stray .roar in the invocation directory. + assert not (Path(subdir_run["nested_dir"]) / ".roar").exists(), _describe(subdir_run) + + db_path = Path(subdir_run["project_dir"]) / ".roar" / "roar.db" + assert db_path.is_file(), _describe(subdir_run) + conn = sqlite3.connect(db_path) + conn.row_factory = sqlite3.Row + try: + tasks = conn.execute("SELECT id FROM jobs WHERE job_type = 'k8s_task'").fetchall() + finally: + conn.close() + assert len(tasks) == 1, _describe(subdir_run) diff --git a/tests/backends/k8s/unit/test_submit_planning.py b/tests/backends/k8s/unit/test_submit_planning.py index 6dec5384..af94695b 100644 --- a/tests/backends/k8s/unit/test_submit_planning.py +++ b/tests/backends/k8s/unit/test_submit_planning.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +import subprocess from pathlib import Path import pytest @@ -21,6 +22,35 @@ def test_matches_kubectl_apply_with_job_manifest(job_manifest_path: Path) -> Non ) +def test_matches_kubectl_with_global_flags_before_verb(job_manifest_path: Path) -> None: + """Global flags routinely precede the verb; they must not bypass lineage.""" + assert matches_kubectl_job_submit_command( + ["kubectl", "--context", "prod-cluster", "apply", "-f", str(job_manifest_path)] + ) + assert matches_kubectl_job_submit_command( + ["kubectl", "-n", "ml", "create", "-f", str(job_manifest_path)] + ) + assert matches_kubectl_job_submit_command( + [ + "kubectl", + "--kubeconfig", + "/tmp/kc", + "--context", + "c", + "apply", + "-f", + str(job_manifest_path), + ] + ) + + +def test_does_not_match_non_submit_verbs_with_manifest(job_manifest_path: Path) -> None: + assert not matches_kubectl_job_submit_command( + ["kubectl", "delete", "-f", str(job_manifest_path)] + ) + assert not matches_kubectl_job_submit_command(["kubectl", "diff", "-f", str(job_manifest_path)]) + + def test_does_not_match_when_disabled(job_manifest_path: Path) -> None: config_path = job_manifest_path.parent / ".roar" / "config.toml" config_path.write_text("[k8s]\nenabled = false\n", encoding="utf-8") @@ -115,6 +145,33 @@ def fake_register(glaas_url: str, session_id: str, token_hash: str, ttl: int = 8 assert key_payload["session_id"] == plan.session_id +def test_plan_from_subdirectory_saves_state_in_project_root( + job_manifest_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Plan-time state must land in the .roar the finalizer will read.""" + project_root = job_manifest_path.parent + # Real projects are git repos; config discovery bounds its upward walk + # by the repo root, so the fixture needs one to search past the subdir. + subprocess.run(["git", "init", "-q"], cwd=project_root, check=True) + nested = project_root / "jobs" / "nested" + nested.mkdir(parents=True) + monkeypatch.chdir(nested) + monkeypatch.delenv("ROAR_PROJECT_DIR", raising=False) + monkeypatch.setattr( + "roar.backends.k8s.submit._register_fragment_session", + lambda *args, **kwargs: None, + ) + + plan = plan_kubectl_job_submit_command(["kubectl", "apply", "-f", str(job_manifest_path)]) + + assert plan.session_id + key_path = project_root / ".roar" / "fragment-sessions" / f"{plan.session_id}.key" + assert key_path.is_file() + prepared_path = Path(plan.command[plan.command.index("-f") + 1]) + assert prepared_path.parent == project_root / ".roar" / "k8s" / "prepared" + assert not (nested / ".roar").exists() + + def test_plan_degrades_to_uninstrumented_when_registration_fails( job_manifest_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/tests/execution/fragments/test_sessions.py b/tests/execution/fragments/test_sessions.py new file mode 100644 index 00000000..ee99448c --- /dev/null +++ b/tests/execution/fragments/test_sessions.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +from roar.execution.fragments.sessions import resolve_project_roar_dir + + +def test_resolves_project_root_from_subdirectory(tmp_path: Path) -> None: + (tmp_path / ".roar").mkdir() + nested = tmp_path / "jobs" / "nested" + nested.mkdir(parents=True) + + assert resolve_project_roar_dir(environ={}, cwd=nested) == tmp_path / ".roar" + + +def test_project_dir_env_override_wins(tmp_path: Path) -> None: + (tmp_path / ".roar").mkdir() + other = tmp_path / "elsewhere" + other.mkdir() + + resolved = resolve_project_roar_dir(environ={"ROAR_PROJECT_DIR": str(other)}, cwd=tmp_path) + + assert resolved == other / ".roar" + + +def test_falls_back_to_cwd_when_no_project_found(tmp_path: Path) -> None: + nested = tmp_path / "nested" + nested.mkdir() + + assert resolve_project_roar_dir(environ={}, cwd=nested) == nested / ".roar" + + +def test_walk_stops_at_git_repository_boundary(tmp_path: Path) -> None: + # A .roar above an unrelated git repo must not be adopted. + (tmp_path / ".roar").mkdir() + repo = tmp_path / "unrelated-repo" + (repo / ".git").mkdir(parents=True) + inner = repo / "src" + inner.mkdir() + + assert resolve_project_roar_dir(environ={}, cwd=inner) == inner / ".roar" + + +def test_git_root_with_roar_dir_is_used(tmp_path: Path) -> None: + (tmp_path / ".git").mkdir() + (tmp_path / ".roar").mkdir() + nested = tmp_path / "src" + nested.mkdir() + + assert resolve_project_roar_dir(environ={}, cwd=nested) == tmp_path / ".roar" + + +@pytest.mark.parametrize("value", ["", " "]) +def test_blank_project_dir_env_is_ignored(tmp_path: Path, value: str) -> None: + (tmp_path / ".roar").mkdir() + nested = tmp_path / "sub" + nested.mkdir() + + resolved = resolve_project_roar_dir(environ={"ROAR_PROJECT_DIR": value}, cwd=nested) + + assert resolved == tmp_path / ".roar" From c843de461a2e4bb82ce2347a9270930dceecc316 Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Wed, 15 Jul 2026 17:17:47 +0000 Subject: [PATCH 19/31] fix(k8s): parent RayJob ray tasks to the recorded submit job MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The RayJob rewrite set ROAR_JOB_ID to the k8s parent uid but never passed driver_job_uid to merge_worker_bootstrap_env, so workers had no ROAR_DRIVER_JOB_UID and their fragments merged without a parent — ray_task rows carried no DAG edge back to the k8s submit job (the merge prefers the fragment-level parent, which was empty, and the host-side fallback env is unset in this flow). The live RayJob e2e now asserts every ray_task row parents to the recorded submit job, which is exactly the assertion whose absence let this slip. Co-Authored-By: Claude Fable 5 --- roar/backends/k8s/rayjob.py | 4 ++++ tests/backends/k8s/e2e/test_k8s_rayjob.py | 13 ++++++++++++- tests/backends/k8s/unit/test_rayjob.py | 3 +++ 3 files changed, 19 insertions(+), 1 deletion(-) diff --git a/roar/backends/k8s/rayjob.py b/roar/backends/k8s/rayjob.py index 8b626736..a570dedc 100644 --- a/roar/backends/k8s/rayjob.py +++ b/roar/backends/k8s/rayjob.py @@ -141,6 +141,10 @@ def _merged_runtime_env_yaml( dict(runtime_env.get("env_vars") or {}), build_submit_source_environ(context), job_id=context.job_id, + # Ray workers stamp ROAR_DRIVER_JOB_UID into each fragment's + # parent_job_uid; without it the ray_task rows merge with no DAG + # edge back to the recorded k8s submit job. + driver_job_uid=contract.parent_job_uid, overwrite_existing=True, ) env_vars[ROAR_EXECUTION_BACKEND_ENV] = "ray" diff --git a/tests/backends/k8s/e2e/test_k8s_rayjob.py b/tests/backends/k8s/e2e/test_k8s_rayjob.py index 7bdca011..154baae5 100644 --- a/tests/backends/k8s/e2e/test_k8s_rayjob.py +++ b/tests/backends/k8s/e2e/test_k8s_rayjob.py @@ -215,10 +215,21 @@ def test_rayjob_live_delegates_to_ray_backend( ray_tasks = _query( project_dir, - "SELECT id FROM jobs WHERE job_type = 'ray_task'", + "SELECT id, parent_job_uid FROM jobs WHERE job_type = 'ray_task'", ) assert ray_tasks, f"expected ray_task jobs from delegated reconstitution\n{_describe(run)}" + # Every ray task must carry the DAG edge back to the recorded k8s + # submit job (regression: driver_job_uid was not threaded into the + # worker env, so tasks merged parentless and no test noticed). + submit_uid = str(submit_rows[0]["job_uid"]) + orphaned = [ + row["id"] for row in ray_tasks if str(row["parent_job_uid"] or "") != submit_uid + ] + assert not orphaned, ( + f"ray_task jobs not parented to submit job {submit_uid}: {orphaned}\n{_describe(run)}" + ) + output_paths = { str(row["path"]) for row in _query( diff --git a/tests/backends/k8s/unit/test_rayjob.py b/tests/backends/k8s/unit/test_rayjob.py index a37a1f9b..bf0f9062 100644 --- a/tests/backends/k8s/unit/test_rayjob.py +++ b/tests/backends/k8s/unit/test_rayjob.py @@ -92,6 +92,9 @@ def test_rayjob_runtime_env_carries_ray_contract_without_secrets() -> None: assert env_vars["USER_VAR"] == "keep" assert env_vars["ROAR_EXECUTION_BACKEND"] == "ray" assert env_vars["ROAR_JOB_ID"] == "cafe0123" + # Workers stamp this into each fragment's parent_job_uid; it is the DAG + # edge from ray_task rows back to the recorded k8s submit job. + assert env_vars["ROAR_DRIVER_JOB_UID"] == "cafe0123" assert env_vars["GLAAS_URL"] == "http://glaas:3001" assert env_vars["ROAR_RAY_NODE_AGENTS"] == "0" # Credentials must never appear inline in the CR. From e631e50994afe8d8ee49837ba7537f6a8bf296ae Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Wed, 15 Jul 2026 17:18:03 +0000 Subject: [PATCH 20/31] feat(k8s): webhook creates no resources before the rewrite; aged-Secret cleanup MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The admission handler registered the GLaaS session and created the credentials Secret before attempting the manifest rewrite, so a rewrite failure (the step most likely to fail) orphaned both. The rewrite — a pure function of the admitted object — now runs first; external resources are only created once it succeeds. Credential Secrets can never carry an ownerReference (the workload has no UID at CREATE admission), so the chart gains a secretCleanup CronJob (roar.backends.k8s.secret_cleanup, runs on the runtime image) deleting roar-fragment-* Secrets older than the maximum session TTL. It uses its own ServiceAccount with list/delete only; the webhook's stays create-only. Running pods are unaffected (env resolves at pod start) — only late attach credential recovery ages out with the session itself. Co-Authored-By: Claude Fable 5 --- .../templates/secret-cleanup.yaml | 70 ++++++++++++ .../charts/roar-lineage-webhook/values.yaml | 13 +++ roar/backends/k8s/secret_cleanup.py | 102 ++++++++++++++++++ roar/backends/k8s/webhook.py | 19 ++-- .../backends/k8s/unit/test_secret_cleanup.py | 49 +++++++++ tests/backends/k8s/unit/test_webhook.py | 24 +++++ 6 files changed, 270 insertions(+), 7 deletions(-) create mode 100644 deploy/charts/roar-lineage-webhook/templates/secret-cleanup.yaml create mode 100644 roar/backends/k8s/secret_cleanup.py create mode 100644 tests/backends/k8s/unit/test_secret_cleanup.py diff --git a/deploy/charts/roar-lineage-webhook/templates/secret-cleanup.yaml b/deploy/charts/roar-lineage-webhook/templates/secret-cleanup.yaml new file mode 100644 index 00000000..b35b257f --- /dev/null +++ b/deploy/charts/roar-lineage-webhook/templates/secret-cleanup.yaml @@ -0,0 +1,70 @@ +{{- if .Values.secretCleanup.enabled }} +apiVersion: v1 +kind: ServiceAccount +metadata: + name: {{ include "roar-webhook.fullname" . }}-secret-cleanup + namespace: {{ .Release.Namespace }} + labels: + {{- include "roar-webhook.labels" . | nindent 4 }} +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + name: {{ include "roar-webhook.fullname" . }}-secret-cleanup + labels: + {{- include "roar-webhook.labels" . | nindent 4 }} +rules: + - apiGroups: [""] + resources: ["secrets"] + verbs: ["list", "delete"] +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRoleBinding +metadata: + name: {{ include "roar-webhook.fullname" . }}-secret-cleanup + labels: + {{- include "roar-webhook.labels" . | nindent 4 }} +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: ClusterRole + name: {{ include "roar-webhook.fullname" . }}-secret-cleanup +subjects: + - kind: ServiceAccount + name: {{ include "roar-webhook.fullname" . }}-secret-cleanup + namespace: {{ .Release.Namespace }} +--- +apiVersion: batch/v1 +kind: CronJob +metadata: + name: {{ include "roar-webhook.fullname" . }}-secret-cleanup + namespace: {{ .Release.Namespace }} + labels: + {{- include "roar-webhook.labels" . | nindent 4 }} +spec: + schedule: {{ .Values.secretCleanup.schedule | quote }} + concurrencyPolicy: Forbid + successfulJobsHistoryLimit: 1 + failedJobsHistoryLimit: 3 + jobTemplate: + spec: + backoffLimit: 1 + ttlSecondsAfterFinished: 86400 + template: + metadata: + labels: + {{- include "roar-webhook.selectorLabels" . | nindent 12 }} + spec: + serviceAccountName: {{ include "roar-webhook.fullname" . }}-secret-cleanup + restartPolicy: Never + containers: + - name: cleanup + image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}" + imagePullPolicy: {{ .Values.image.pullPolicy }} + env: + - name: ROAR_SECRET_MAX_AGE_SECONDS + value: {{ .Values.secretCleanup.maxAgeSeconds | quote }} + command: + - python + - -m + - roar.backends.k8s.secret_cleanup +{{- end }} diff --git a/deploy/charts/roar-lineage-webhook/values.yaml b/deploy/charts/roar-lineage-webhook/values.yaml index 0cfac8b0..4a365a2b 100644 --- a/deploy/charts/roar-lineage-webhook/values.yaml +++ b/deploy/charts/roar-lineage-webhook/values.yaml @@ -55,6 +55,19 @@ certManager: # Certificate; the CA bundle is injected via cert-manager annotations). enabled: false +secretCleanup: + # Credential Secrets cannot carry an ownerReference (the workload has no + # UID at CREATE admission), so a CronJob deletes roar-managed fragment + # Secrets older than maxAgeSeconds. Runs under its own ServiceAccount; + # the webhook's stays create-only. + enabled: true + schedule: "0 3 * * *" + # Default matches the maximum fragment-session TTL (7 days): a Secret + # older than this belongs to a session the server no longer honors. + # Running pods are unaffected (env is resolved at pod start); only + # late `roar k8s attach` credential recovery ages out with it. + maxAgeSeconds: 604800 + resources: {} nodeSelector: {} tolerations: [] diff --git a/roar/backends/k8s/secret_cleanup.py b/roar/backends/k8s/secret_cleanup.py new file mode 100644 index 00000000..f2d40ab2 --- /dev/null +++ b/roar/backends/k8s/secret_cleanup.py @@ -0,0 +1,102 @@ +"""Delete aged roar fragment-credential Secrets (webhook CronJob entrypoint). + +Credential Secrets cannot carry an ownerReference: at CREATE admission the +workload has no UID yet, so Kubernetes garbage collection can never adopt +them. This module lists roar-managed Secrets cluster-wide and deletes those +older than ROAR_SECRET_MAX_AGE_SECONDS (default: the 7-day maximum +fragment-session TTL — a Secret older than that belongs to a session the +server no longer honors). Running pods are unaffected (env is resolved at +pod start); only late `roar k8s attach` credential recovery ages out. + +Runs in-cluster under a dedicated ServiceAccount with list/delete on +Secrets (see the chart's secret-cleanup template). +""" + +from __future__ import annotations + +import json +import os +import ssl +import sys +import urllib.parse +import urllib.request +from datetime import datetime, timezone +from typing import Any + +_SERVICE_ACCOUNT_DIR = "/var/run/secrets/kubernetes.io/serviceaccount" +_API_BASE = "https://kubernetes.default.svc" +_MANAGED_BY_SELECTOR = "app.kubernetes.io/managed-by=roar" +_SECRET_NAME_PREFIX = "roar-fragment-" +DEFAULT_MAX_AGE_SECONDS = 604800 + + +def _api_request(path: str, method: str = "GET") -> dict[str, Any]: + with open(f"{_SERVICE_ACCOUNT_DIR}/token", encoding="utf-8") as handle: + token = handle.read().strip() + request = urllib.request.Request( + url=f"{_API_BASE}{path}", + headers={"authorization": f"Bearer {token}"}, + method=method, + ) + context = ssl.create_default_context(cafile=f"{_SERVICE_ACCOUNT_DIR}/ca.crt") + with urllib.request.urlopen(request, timeout=30, context=context) as response: + payload = json.loads(response.read().decode("utf-8")) + return payload if isinstance(payload, dict) else {} + + +def _parse_timestamp(raw: Any) -> datetime | None: + if not isinstance(raw, str) or not raw: + return None + try: + return datetime.fromisoformat(raw.replace("Z", "+00:00")) + except ValueError: + return None + + +def select_expired_secrets( + items: list[dict[str, Any]], + *, + max_age_seconds: int, + now: datetime, +) -> list[tuple[str, str]]: + """Return (namespace, name) for roar fragment Secrets older than max age.""" + expired: list[tuple[str, str]] = [] + for item in items: + metadata = item.get("metadata") or {} + name = str(metadata.get("name") or "") + namespace = str(metadata.get("namespace") or "") + if not name.startswith(_SECRET_NAME_PREFIX) or not namespace: + continue + created = _parse_timestamp(metadata.get("creationTimestamp")) + if created is None: + continue + if (now - created).total_seconds() > max_age_seconds: + expired.append((namespace, name)) + return expired + + +def main() -> int: + max_age = int(os.environ.get("ROAR_SECRET_MAX_AGE_SECONDS", DEFAULT_MAX_AGE_SECONDS)) + selector = urllib.parse.quote(_MANAGED_BY_SELECTOR, safe="=") + listing = _api_request(f"/api/v1/secrets?labelSelector={selector}") + items = [item for item in listing.get("items") or [] if isinstance(item, dict)] + + expired = select_expired_secrets(items, max_age_seconds=max_age, now=datetime.now(timezone.utc)) + failures = 0 + for namespace, name in expired: + try: + _api_request(f"/api/v1/namespaces/{namespace}/secrets/{name}", method="DELETE") + print(f"[roar-secret-cleanup] deleted {namespace}/{name}") + except Exception as exc: + failures += 1 + print(f"[roar-secret-cleanup] failed to delete {namespace}/{name}: {exc}") + + print( + f"[roar-secret-cleanup] scanned {len(items)} secret(s), " + f"deleted {len(expired) - failures}, failed {failures}" + ) + return 1 if failures else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/roar/backends/k8s/webhook.py b/roar/backends/k8s/webhook.py index 2ac42cf5..fd2f844e 100644 --- a/roar/backends/k8s/webhook.py +++ b/roar/backends/k8s/webhook.py @@ -136,16 +136,14 @@ def _allow(warnings: list[str] | None = None) -> dict[str, Any]: namespace = str(request.get("namespace") or metadata.get("namespace") or "default") session = generate_fragment_session() - register_session(session["session_id"], session["token_hash"]) - parent_job_uid = secrets_module.token_hex(4) secret_name = f"roar-fragment-{session['session_id'][:8]}" - create_secret( - namespace, - secret_name, - {"session_id": session["session_id"], "token": session["token"]}, - ) + # Rewrite first: it is a pure function of the admitted object, and + # it is the step most likely to fail (unsupported shape, no + # wrappable containers). Only create external resources — the GLaaS + # session and the credentials Secret — once the rewrite succeeded, + # so a failed admission leaves nothing behind. rewrite = rewrite_manifest_for_lineage( [obj], secret_name=secret_name, @@ -165,6 +163,13 @@ def _allow(warnings: list[str] | None = None) -> dict[str, Any]: ) rewritten = rewrite.documents[0] + register_session(session["session_id"], session["token_hash"]) + create_secret( + namespace, + secret_name, + {"session_id": session["session_id"], "token": session["token"]}, + ) + merged_annotations = dict(annotations) merged_annotations[ANNOTATION_PARENT_UID] = parent_job_uid merged_annotations[ANNOTATION_SESSION_ID] = session["session_id"] diff --git a/tests/backends/k8s/unit/test_secret_cleanup.py b/tests/backends/k8s/unit/test_secret_cleanup.py new file mode 100644 index 00000000..b7093c29 --- /dev/null +++ b/tests/backends/k8s/unit/test_secret_cleanup.py @@ -0,0 +1,49 @@ +"""Unit tests for the fragment-Secret cleanup selection logic.""" + +from __future__ import annotations + +from datetime import datetime, timedelta, timezone + +from roar.backends.k8s.secret_cleanup import DEFAULT_MAX_AGE_SECONDS, select_expired_secrets + +NOW = datetime(2026, 7, 15, 12, 0, 0, tzinfo=timezone.utc) + + +def _secret(name: str, namespace: str, age_seconds: int) -> dict: + created = NOW - timedelta(seconds=age_seconds) + return { + "metadata": { + "name": name, + "namespace": namespace, + "creationTimestamp": created.strftime("%Y-%m-%dT%H:%M:%SZ"), + } + } + + +def test_selects_only_aged_fragment_secrets() -> None: + items = [ + _secret("roar-fragment-old00000", "ml", DEFAULT_MAX_AGE_SECONDS + 3600), + _secret("roar-fragment-fresh000", "ml", 3600), + ] + + expired = select_expired_secrets(items, max_age_seconds=DEFAULT_MAX_AGE_SECONDS, now=NOW) + + assert expired == [("ml", "roar-fragment-old00000")] + + +def test_never_touches_non_fragment_secrets_even_if_labeled() -> None: + items = [ + _secret("some-other-secret", "ml", DEFAULT_MAX_AGE_SECONDS * 10), + ] + + assert select_expired_secrets(items, max_age_seconds=1, now=NOW) == [] + + +def test_skips_malformed_entries() -> None: + items = [ + {"metadata": {"name": "roar-fragment-x", "namespace": ""}}, + {"metadata": {"name": "roar-fragment-y", "namespace": "ml"}}, # no timestamp + {"metadata": {"name": "roar-fragment-z", "namespace": "ml", "creationTimestamp": "bogus"}}, + ] + + assert select_expired_secrets(items, max_age_seconds=1, now=NOW) == [] diff --git a/tests/backends/k8s/unit/test_webhook.py b/tests/backends/k8s/unit/test_webhook.py index d6fc86b2..dc92330a 100644 --- a/tests/backends/k8s/unit/test_webhook.py +++ b/tests/backends/k8s/unit/test_webhook.py @@ -141,6 +141,30 @@ def test_dry_run_has_no_side_effects() -> None: assert spy.calls == [] +def test_rewrite_failure_creates_no_session_or_secret() -> None: + """The rewrite runs before any external resource is created, so a + workload the rewriter rejects must not leave an orphaned GLaaS session + or credentials Secret behind (regression: side effects used to come + first).""" + spy = _Spy() + no_command = copy.deepcopy(SINGLE_JOB_MANIFEST) + for container in no_command["spec"]["template"]["spec"]["containers"]: + container.pop("command", None) + + result = mutate_admission_review( + _review(no_command), + settings=SETTINGS, + create_secret=spy.create_secret, + register_session=spy.register_session, + ) + + response = result["response"] + assert response["allowed"] is True + assert "patch" not in response + assert response.get("warnings") + assert spy.calls == [] + + def test_failures_never_block_admission() -> None: spy = _Spy(fail=True) result = mutate_admission_review( From 80c29c2cc4266c0328db85655198e9f7fcc0fc05 Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Wed, 15 Jul 2026 17:18:03 +0000 Subject: [PATCH 21/31] feat(k8s): pod-wrapper import probe; rayjob attach aliases MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The pod wrappers guarded python3/pip/staged-tree presence but exec'd pod_entry unconditionally — an import failure past exec (broken install, ABI mismatch such as the x86-64 staged tree on an arm64 node, unsupported interpreter) killed the container with no way back. An import probe now runs as the last step before exec and falls back to the original command uninstrumented, which also makes non-x86-64 platforms degrade gracefully instead of failing training. The remaining unisolated window (a crash inside pod_entry after the probe) is documented instead of overclaimed. Also adds the missing rayjob/rayjobs attach aliases (a WORKLOAD_KINDS-completeness test keeps future kinds honest) and mentions rayjob in the attach help text. Co-Authored-By: Claude Fable 5 --- docs/developer/k8s-integration.md | 22 +++++++++++++++++++--- roar/backends/k8s/attach.py | 2 ++ roar/backends/k8s/manifest.py | 8 +++++++- roar/cli/commands/k8s.py | 7 ++++--- tests/backends/k8s/unit/test_attach.py | 13 +++++++++++++ 5 files changed, 45 insertions(+), 7 deletions(-) diff --git a/docs/developer/k8s-integration.md b/docs/developer/k8s-integration.md index 3b82f78c..5bb47906 100644 --- a/docs/developer/k8s-integration.md +++ b/docs/developer/k8s-integration.md @@ -51,6 +51,13 @@ blocks admission. Reconstitution is client-driven via `roar k8s attach`. `ROAR_WEBHOOK_PROXY_UPSTREAM`) makes the injector add the proxy sidecar to every workload it instruments. +The webhook rewrites the manifest *before* creating the GLaaS session or +the credentials Secret, so a rejected/unrewritable workload leaves +nothing behind. Secrets cannot carry an ownerReference (no workload UID +exists at CREATE admission); the chart ships a `secretCleanup` CronJob +(own ServiceAccount, list/delete only, name-prefix guarded) that removes +roar fragment Secrets older than the maximum session TTL. + **Helm chart** (`deploy/charts/roar-lineage-webhook`): packages the webhook Deployment/Service, RBAC (Secret-create ClusterRole), and the `MutatingWebhookConfiguration`. Required values: `glaas.url`, @@ -165,8 +172,11 @@ renewal is purely reactive to the 403. ## 2. Flow 1. **Match** (`roar/backends/k8s/submit.py`): binary `kubectl`, verb - `apply|create`, `-f` pointing at a manifest containing exactly one - supported workload; gated by `k8s.enabled` (default off). + `apply|create` anywhere in the argument list (global flags like + `--context` may precede it), `-f` pointing at a manifest containing + exactly one supported workload; gated by `k8s.enabled` (default off). + Stdin (`-f -`), URLs, directories, and kustomize sources are outside + the matching contract and pass through uninstrumented. 2. **Plan** (same module): pre-registers a GLaaS fragment session (saves the `.key` under `.roar/fragment-sessions/`), rewrites the manifest (`manifest.py`), writes it 0600 under `.roar/k8s/prepared/` with a @@ -182,7 +192,13 @@ renewal is purely reactive to the 403. Each container with an explicit `command` gets a `/bin/sh -c` script that pip-installs the roar runtime (`k8s.runtime_install_requirement`, default pinned `roar-cli`) and execs `python3 -m roar.backends.k8s.pod_entry "$@"`; - on any bootstrap failure it falls back to exec'ing the original command + bootstrap failures — missing python3, failed install, absent staged + tree, and (via an import probe just before the exec) broken installs, + ABI mismatches, or unsupported interpreters — fall back to exec'ing + the original command uninstrumented. A crash *inside* `pod_entry` + after the probe passes is the remaining unisolated window; past that + point the training exit code always wins. The runtime image stages + x86-64 (glibc) trees only — other platforms fail the probe and run uninstrumented. Injects the env contract (`GLAAS_URL` = cluster-visible URL, Secret-backed `ROAR_SESSION_ID`/`ROAR_FRAGMENT_TOKEN`, downward-API identity fields) and appends the Secret document. Containers without an diff --git a/roar/backends/k8s/attach.py b/roar/backends/k8s/attach.py index 80fcf710..78a7a9a0 100644 --- a/roar/backends/k8s/attach.py +++ b/roar/backends/k8s/attach.py @@ -60,6 +60,8 @@ class K8sAttachResult: "pytorchjobs": "pytorchjobs.kubeflow.org", "trainjob": "trainjobs.trainer.kubeflow.org", "trainjobs": "trainjobs.trainer.kubeflow.org", + "rayjob": "rayjobs.ray.io", + "rayjobs": "rayjobs.ray.io", } diff --git a/roar/backends/k8s/manifest.py b/roar/backends/k8s/manifest.py index dd735adb..966f9fe2 100644 --- a/roar/backends/k8s/manifest.py +++ b/roar/backends/k8s/manifest.py @@ -27,11 +27,16 @@ # sh -c scripts: "$0" is the synthetic argv0, "$@" is the original # command+args. Lineage is best-effort by design: any failure to stage # the roar runtime falls back to running the original command -# uninstrumented rather than failing the training job. +# uninstrumented rather than failing the training job. The import probe +# is the last line of defense before exec — it catches broken installs, +# ABI mismatches (e.g. an x86-64 staged tree on an arm64 node), and +# unsupported interpreters, all of which would otherwise kill the +# container after exec with no way back. _POD_WRAPPER_TEMPLATE = """\ run_fallback() {{ echo "[roar-k8s] lineage runtime unavailable; running uninstrumented" >&2; exec "$@"; }} command -v python3 >/dev/null 2>&1 || run_fallback "$@" python3 -m pip install --quiet {requirement} || run_fallback "$@" +python3 -c 'import roar.backends.k8s.pod_entry' || run_fallback "$@" exec python3 -m roar.backends.k8s.pod_entry "$@" """ @@ -44,6 +49,7 @@ RT="{staging_mount}/cp$(python3 -c 'import sys; print("%d%d" % sys.version_info[:2])')" [ -d "$RT" ] || run_fallback "$@" export PYTHONPATH="$RT${{PYTHONPATH:+:$PYTHONPATH}}" +python3 -c 'import roar.backends.k8s.pod_entry' || run_fallback "$@" exec python3 -m roar.backends.k8s.pod_entry "$@" """ diff --git a/roar/cli/commands/k8s.py b/roar/cli/commands/k8s.py index bde8d564..aa9201a3 100644 --- a/roar/cli/commands/k8s.py +++ b/roar/cli/commands/k8s.py @@ -65,9 +65,10 @@ def k8s_attach( ) -> None: """Reconstitute lineage from an already-submitted workload. - WORKLOAD is a name or KIND/NAME (job, jobset, pytorchjob, trainjob) that - was instrumented at submit time. Credentials come from a locally saved - fragment-session key, the cluster Secret, or --session-file. + WORKLOAD is a name or KIND/NAME (job, jobset, pytorchjob, trainjob, + rayjob) that was instrumented at submit time. Credentials come from a + locally saved fragment-session key, the cluster Secret, or + --session-file. """ global_flags = ["--context", kube_context] if kube_context else [] try: diff --git a/tests/backends/k8s/unit/test_attach.py b/tests/backends/k8s/unit/test_attach.py index e008683d..572b7da0 100644 --- a/tests/backends/k8s/unit/test_attach.py +++ b/tests/backends/k8s/unit/test_attach.py @@ -51,6 +51,19 @@ def test_attach_recovers_identity_from_jobset_and_trainjob() -> None: assert _session_secret_name(env) == "roar-fragment-deadbeef" +def test_kind_aliases_cover_every_supported_workload() -> None: + """Every WORKLOAD_KINDS entry must be reachable via KIND/NAME attach.""" + from roar.backends.k8s.attach import _KIND_ALIASES + from roar.backends.k8s.manifest import WORKLOAD_KINDS + + aliased_resources = set(_KIND_ALIASES.values()) + for kind in WORKLOAD_KINDS: + assert kind.kubectl_resource in aliased_resources, ( + f"{kind.kind} ({kind.kubectl_resource}) has no attach alias" + ) + assert kind.kind.lower() in _KIND_ALIASES + + def test_attach_reports_uninstrumented_workload() -> None: doc = copy.deepcopy(SINGLE_JOB_MANIFEST) workload_kind = workload_kind_for_document(doc) From 1378e2ffc8385b7cbd35d0c54bc1a7cf40d1f868 Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Wed, 15 Jul 2026 18:26:07 +0000 Subject: [PATCH 22/31] fix(k8s): preserve user RayJob runtime env during instrumentation The RayJob rewrite destroyed three kinds of user runtime configuration: pip dicts were flattened to package lists (losing pip_check/pip_version), a user worker_process_setup_hook was silently replaced with roar's, and a user-supplied AWS_ENDPOINT_URL (MinIO, LocalStack, private S3) was stripped along with roar's dead local-proxy redirect. - keep the pip dict shape and append the roar requirement to packages; reject pip requirements-file references with an actionable error - carry a displaced user setup hook as ROAR_USER_SETUP_HOOK; roar's worker_bootstrap.startup chains it after capture installs, and user hook failures propagate exactly as they would uninstrumented - drop the local-proxy redirect from the merge source instead of popping the merged result, so user endpoints survive the rewrite Co-Authored-By: Claude Fable 5 --- roar/backends/k8s/rayjob.py | 45 ++++++++--- roar/execution/runtime/worker_bootstrap.py | 22 ++++++ tests/backends/k8s/unit/test_rayjob.py | 75 +++++++++++++++++++ .../runtime/test_worker_bootstrap.py | 48 ++++++++++++ 4 files changed, 178 insertions(+), 12 deletions(-) diff --git a/roar/backends/k8s/rayjob.py b/roar/backends/k8s/rayjob.py index a570dedc..54e10787 100644 --- a/roar/backends/k8s/rayjob.py +++ b/roar/backends/k8s/rayjob.py @@ -35,6 +35,7 @@ _DRIVER_ENTRYPOINT_PREFIX = "python -m roar.execution.runtime.driver_entrypoint -- " _WORKER_SETUP_HOOK = "roar.execution.runtime.worker_bootstrap.startup" +_USER_SETUP_HOOK_ENV = "ROAR_USER_SETUP_HOOK" RAYJOB_SUCCESS_STATUSES = frozenset({"SUCCEEDED"}) RAYJOB_FAILURE_STATUSES = frozenset({"FAILED", "STOPPED"}) @@ -118,15 +119,28 @@ def _merged_runtime_env_yaml( runtime_env = loaded pip = runtime_env.get("pip") - packages: list[str] = [] - if isinstance(pip, dict): - packages = [str(item) for item in pip.get("packages") or []] - elif isinstance(pip, list): - packages = [str(item) for item in pip] + if isinstance(pip, str): + raise K8sManifestError( + f"RayJob {workload_name} references a pip requirements file " + f"({pip!r}) in runtimeEnvYAML; roar cannot append its runtime " + "requirement to a file reference. Use the list or dict form of pip." + ) + pip_options = dict(pip) if isinstance(pip, dict) else None + if pip_options is not None: + packages = [str(item) for item in pip_options.get("packages") or []] + else: + packages = [str(item) for item in pip] if isinstance(pip, list) else [] if contract.requirement and contract.requirement not in packages: packages.append(contract.requirement) - runtime_env["pip"] = packages - + if pip_options is not None: + # Preserve the dict form: pip_check/pip_version and any future + # options must survive instrumentation. + pip_options["packages"] = packages + runtime_env["pip"] = pip_options + else: + runtime_env["pip"] = packages + + user_setup_hook = str(runtime_env.get("worker_process_setup_hook") or "").strip() runtime_env["worker_process_setup_hook"] = _WORKER_SETUP_HOOK # The Ray env contract, keyed to the k8s parent uid so Ray fragments @@ -137,9 +151,16 @@ def _merged_runtime_env_yaml( host_glaas_url=None, job_id=contract.parent_job_uid, ) + source_environ = build_submit_source_environ(context) + # No proxy runs in the pods, so the Ray submit contract's local-proxy + # redirect would point S3 traffic at a dead localhost port. Dropping it + # from the merge source (rather than popping the merged result) + # preserves a user-supplied AWS_ENDPOINT_URL such as MinIO. + source_environ.pop("AWS_ENDPOINT_URL", None) + source_environ.pop("ROAR_PROXY_PORT", None) env_vars = merge_worker_bootstrap_env( dict(runtime_env.get("env_vars") or {}), - build_submit_source_environ(context), + source_environ, job_id=context.job_id, # Ray workers stamp ROAR_DRIVER_JOB_UID into each fragment's # parent_job_uid; without it the ray_task rows merge with no DAG @@ -152,10 +173,10 @@ def _merged_runtime_env_yaml( # Per the proxy decision, node agents/proxy sidecars stay off for # RayJob delegation v1; in-process hooks are the capture surface. env_vars["ROAR_RAY_NODE_AGENTS"] = "0" - # No proxy runs in the pods, so the merged local-proxy redirect would - # point user S3 traffic at a dead localhost port — strip it. - env_vars.pop("AWS_ENDPOINT_URL", None) - env_vars.pop("ROAR_PROXY_PORT", None) + if user_setup_hook and user_setup_hook != _WORKER_SETUP_HOOK: + # roar's startup hook runs the displaced user hook after capture is + # installed (worker_bootstrap._run_user_setup_hook). + env_vars[_USER_SETUP_HOOK_ENV] = user_setup_hook # ROAR_WRAP is deliberately NOT set: in Ray pip virtualenvs the # roar_inject.pth fires at worker interpreter startup before the # virtualenv's site-packages are importable, and the sitecustomize diff --git a/roar/execution/runtime/worker_bootstrap.py b/roar/execution/runtime/worker_bootstrap.py index 6a5cab4e..adc4168e 100644 --- a/roar/execution/runtime/worker_bootstrap.py +++ b/roar/execution/runtime/worker_bootstrap.py @@ -17,6 +17,7 @@ WORKER_SETUP_HOOK = "roar.execution.runtime.worker_bootstrap.startup" WORKER_PY_EXECUTABLE = "roar-worker" +USER_SETUP_HOOK_ENV = "ROAR_USER_SETUP_HOOK" def _resolve_execution_backend_name(environ: Mapping[str, str]) -> str: @@ -120,6 +121,27 @@ def startup() -> None: if distributed is None: raise RuntimeError(f"execution backend {backend.name!r} does not support worker bootstrap") distributed.worker_bootstrap.startup() + _run_user_setup_hook(os.environ) + + +def _run_user_setup_hook(environ: Mapping[str, str]) -> None: + """Chain a user worker_process_setup_hook displaced by roar's. + + Instrumented submits (RayJob rewrite) replace the user's setup hook + with roar's and carry the original as ROAR_USER_SETUP_HOOK. User-hook + failures propagate: without roar that failure would have failed the + worker, and instrumentation must not change that contract. + """ + hook_path = str(environ.get(USER_SETUP_HOOK_ENV) or "").strip() + if not hook_path or hook_path == WORKER_SETUP_HOOK: + return + module_name, separator, attribute = hook_path.rpartition(".") + if not separator: + raise RuntimeError(f"{USER_SETUP_HOOK_ENV}={hook_path!r} is not a module.attribute path") + import importlib + + hook = getattr(importlib.import_module(module_name), attribute) + hook() def run_worker_entrypoint(argv: list[str] | None = None) -> None: diff --git a/tests/backends/k8s/unit/test_rayjob.py b/tests/backends/k8s/unit/test_rayjob.py index bf0f9062..aee5e787 100644 --- a/tests/backends/k8s/unit/test_rayjob.py +++ b/tests/backends/k8s/unit/test_rayjob.py @@ -127,6 +127,81 @@ def test_rayjob_pods_get_secret_refs_and_secret_doc() -> None: assert secret["metadata"]["namespace"] == "ml" +def test_rayjob_preserves_pip_dict_options() -> None: + manifest = copy.deepcopy(RAYJOB_MANIFEST) + manifest["spec"]["runtimeEnvYAML"] = yaml.safe_dump( + { + "pip": {"packages": ["pandas"], "pip_check": False, "pip_version": "==23.3.1"}, + "env_vars": {"USER_VAR": "keep"}, + }, + sort_keys=False, + ) + rewrite = _rewrite([manifest]) + doc = next(d for d in rewrite.documents if d.get("kind") == "RayJob") + runtime_env = yaml.safe_load(doc["spec"]["runtimeEnvYAML"]) + + pip = runtime_env["pip"] + assert pip["pip_check"] is False + assert pip["pip_version"] == "==23.3.1" + assert pip["packages"] == ["pandas", "roar-cli==0.3.7"] + + +def test_rayjob_rejects_pip_requirements_file_reference() -> None: + manifest = copy.deepcopy(RAYJOB_MANIFEST) + manifest["spec"]["runtimeEnvYAML"] = "pip: requirements.txt\n" + + with pytest.raises(K8sManifestError, match="requirements"): + _rewrite([manifest]) + + +def test_rayjob_chains_user_worker_setup_hook() -> None: + manifest = copy.deepcopy(RAYJOB_MANIFEST) + manifest["spec"]["runtimeEnvYAML"] = yaml.safe_dump( + {"worker_process_setup_hook": "my_pkg.hooks.setup"}, + sort_keys=False, + ) + rewrite = _rewrite([manifest]) + doc = next(d for d in rewrite.documents if d.get("kind") == "RayJob") + runtime_env = yaml.safe_load(doc["spec"]["runtimeEnvYAML"]) + + assert ( + runtime_env["worker_process_setup_hook"] + == "roar.execution.runtime.worker_bootstrap.startup" + ) + # roar's startup hook invokes the displaced user hook after capture is + # installed (worker_bootstrap.startup reads ROAR_USER_SETUP_HOOK). + assert runtime_env["env_vars"]["ROAR_USER_SETUP_HOOK"] == "my_pkg.hooks.setup" + + +def test_rayjob_does_not_chain_roar_hook_to_itself() -> None: + manifest = copy.deepcopy(RAYJOB_MANIFEST) + manifest["spec"]["runtimeEnvYAML"] = yaml.safe_dump( + {"worker_process_setup_hook": "roar.execution.runtime.worker_bootstrap.startup"}, + sort_keys=False, + ) + rewrite = _rewrite([manifest]) + doc = next(d for d in rewrite.documents if d.get("kind") == "RayJob") + runtime_env = yaml.safe_load(doc["spec"]["runtimeEnvYAML"]) + + assert "ROAR_USER_SETUP_HOOK" not in runtime_env["env_vars"] + + +def test_rayjob_preserves_user_aws_endpoint_url() -> None: + manifest = copy.deepcopy(RAYJOB_MANIFEST) + manifest["spec"]["runtimeEnvYAML"] = yaml.safe_dump( + {"env_vars": {"AWS_ENDPOINT_URL": "http://minio:9000"}}, + sort_keys=False, + ) + rewrite = _rewrite([manifest]) + doc = next(d for d in rewrite.documents if d.get("kind") == "RayJob") + env_vars = yaml.safe_load(doc["spec"]["runtimeEnvYAML"])["env_vars"] + + # Only roar's own local-proxy redirect is dropped (no proxy runs in the + # pods); a user-supplied object-store endpoint must survive the rewrite. + assert env_vars["AWS_ENDPOINT_URL"] == "http://minio:9000" + assert "ROAR_PROXY_PORT" not in env_vars + + def test_rayjob_without_entrypoint_fails_actionably() -> None: manifest = copy.deepcopy(RAYJOB_MANIFEST) del manifest["spec"]["entrypoint"] diff --git a/tests/execution/runtime/test_worker_bootstrap.py b/tests/execution/runtime/test_worker_bootstrap.py index 968734d3..4cfd594e 100644 --- a/tests/execution/runtime/test_worker_bootstrap.py +++ b/tests/execution/runtime/test_worker_bootstrap.py @@ -109,3 +109,51 @@ def test_build_packaged_worker_runtime_env_preserves_non_local_working_dir( ) assert prepared["working_dir"] == "s3://bucket/cloud-demo" + + +def test_run_user_setup_hook_invokes_displaced_hook(monkeypatch) -> None: + import sys + import types + + from roar.execution.runtime.worker_bootstrap import _run_user_setup_hook + + calls: list[str] = [] + module = types.ModuleType("fake_user_hooks") + module.setup = lambda: calls.append("ran") # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "fake_user_hooks", module) + + _run_user_setup_hook({"ROAR_USER_SETUP_HOOK": "fake_user_hooks.setup"}) + + assert calls == ["ran"] + + +def test_run_user_setup_hook_noop_without_env() -> None: + from roar.execution.runtime.worker_bootstrap import _run_user_setup_hook + + _run_user_setup_hook({}) + _run_user_setup_hook({"ROAR_USER_SETUP_HOOK": " "}) + # roar's own hook must never chain to itself. + _run_user_setup_hook( + {"ROAR_USER_SETUP_HOOK": "roar.execution.runtime.worker_bootstrap.startup"} + ) + + +def test_run_user_setup_hook_propagates_user_failure(monkeypatch) -> None: + import sys + import types + + from roar.execution.runtime.worker_bootstrap import _run_user_setup_hook + + def boom() -> None: + raise RuntimeError("user hook failed") + + module = types.ModuleType("fake_failing_hooks") + module.setup = boom # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "fake_failing_hooks", module) + + # Without roar the user's hook failure would have failed the worker; + # instrumentation must not swallow it. + import pytest + + with pytest.raises(RuntimeError, match="user hook failed"): + _run_user_setup_hook({"ROAR_USER_SETUP_HOOK": "fake_failing_hooks.setup"}) From f01145b0379606de56433ea12d8e023e3d4791ba Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Wed, 15 Jul 2026 18:35:24 +0000 Subject: [PATCH 23/31] fix(k8s): fall back to uninstrumented run when roar fails preflight MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit pod_entry returned roar run's exit code directly, so a tracer/runtime failure before the workload launched failed the training job — breaking the documented lineage-never-blocks-training contract (ARM nodes and unsupported native boundaries were the likely triggers). roar run already distinguishes pre-launch setup failures from workload exits (ExecutionReport.setup_error: nonzero exit with no job created). Expose it across the subprocess boundary: when ROAR_RUN_REPORT_FILE is set, the run service writes {exit_code, setup_error, tracer_backend} before finalizers run. pod_entry reruns the original command uninstrumented only on a positively reported setup failure — a missing or unreadable report is ambiguous (roar may have died after the workload ran) and must propagate rather than risk double-running non-idempotent training. New KIND chaos e2e forces the failure with k8s.tracer=ebpf (valid mode, cannot pass preflight in an unprivileged pod) and proves the Job still completes with the workload's own outcome and no fabricated lineage. Co-Authored-By: Claude Fable 5 --- roar/application/run/execution.py | 26 +++ roar/application/run/service.py | 5 + roar/backends/k8s/pod_entry.py | 34 ++++ tests/application/run/test_execution.py | 47 +++++ .../backends/k8s/e2e/test_k8s_tracer_chaos.py | 177 ++++++++++++++++++ tests/backends/k8s/unit/test_pod_entry.py | 108 +++++++++++ 6 files changed, 397 insertions(+) create mode 100644 tests/backends/k8s/e2e/test_k8s_tracer_chaos.py create mode 100644 tests/backends/k8s/unit/test_pod_entry.py diff --git a/roar/application/run/execution.py b/roar/application/run/execution.py index ae835ddd..b6028496 100644 --- a/roar/application/run/execution.py +++ b/roar/application/run/execution.py @@ -28,6 +28,32 @@ class ExecutionReport: setup_error: bool = False +RUN_REPORT_FILE_ENV = "ROAR_RUN_REPORT_FILE" + + +def write_run_report_file(report: ExecutionReport) -> None: + """Write a machine-readable run outcome when ROAR_RUN_REPORT_FILE is set. + + Wrappers that cannot see past the exit code (the k8s pod entrypoint) use + ``setup_error`` to distinguish "roar failed before launching the + workload" from "the workload ran and failed" — only the former is safe + to rerun uninstrumented. Advisory only: never fails the run. + """ + target = str(os.environ.get(RUN_REPORT_FILE_ENV) or "").strip() + if not target: + return + import contextlib + import json + + payload = { + "exit_code": report.exit_code, + "setup_error": report.setup_error, + "tracer_backend": report.tracer_backend, + } + with contextlib.suppress(OSError): + Path(target).write_text(json.dumps(payload), encoding="utf-8") + + def validate_git_clean( *, verb: str = "run", diff --git a/roar/application/run/service.py b/roar/application/run/service.py index b9efe929..331647a6 100644 --- a/roar/application/run/service.py +++ b/roar/application/run/service.py @@ -17,6 +17,7 @@ execute_and_report, get_hash_algorithms, validate_git_clean, + write_run_report_file, ) from .requests import BuildRequest, RunRequest from .verbosity import resolve_verbosity @@ -150,6 +151,10 @@ def _execute_tracked_command( ) raise + # Written before finalizers so a late finalizer crash cannot leave a + # wrapper believing the workload never launched. + write_run_report_file(report) + try: if planned.finalize_run: planned.finalize_run(cast(Any, _FinalizerContext(roar_dir=roar_dir))) diff --git a/roar/backends/k8s/pod_entry.py b/roar/backends/k8s/pod_entry.py index ed8595db..f4f27604 100644 --- a/roar/backends/k8s/pod_entry.py +++ b/roar/backends/k8s/pod_entry.py @@ -64,11 +64,20 @@ def _run_traced(command: list[str]) -> int: # I/O capture) install themselves; events land next to the local db. child_env.setdefault("ROAR_WRAP", "1") child_env.setdefault("ROAR_K8S_OBJECT_IO_FILE", str(_object_io_events_path())) + report_path = _run_report_path() + _remove_stale_report(report_path) + child_env["ROAR_RUN_REPORT_FILE"] = str(report_path) run = subprocess.run( [sys.executable, "-m", "roar", "run", "--tracer", tracer, *command], env=child_env, check=False, ) + if run.returncode != 0 and _reported_setup_error(report_path): + print( + "[roar-k8s] roar run failed before launching the workload; running uninstrumented", + file=sys.stderr, + ) + return _run_uninstrumented(command) return run.returncode @@ -76,6 +85,31 @@ def _object_io_events_path() -> Path: return Path.cwd() / ".roar" / "k8s-object-io.jsonl" +def _run_report_path() -> Path: + return Path.cwd() / ".roar" / "k8s-run-report.json" + + +def _remove_stale_report(report_path: Path) -> None: + import contextlib + + with contextlib.suppress(OSError): + report_path.unlink() + + +def _reported_setup_error(report_path: Path) -> bool: + """True only when roar positively reported a pre-launch setup failure. + + A missing or unreadable report is ambiguous — roar may have crashed + after the workload ran — so it must NOT trigger a rerun: double-running + non-idempotent training is worse than a lost-lineage failure. + """ + try: + payload = json.loads(report_path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return False + return isinstance(payload, dict) and bool(payload.get("setup_error")) + + def _run_uninstrumented(command: list[str]) -> int: completed = subprocess.run(command, check=False) return completed.returncode diff --git a/tests/application/run/test_execution.py b/tests/application/run/test_execution.py index 370934a2..eda68a74 100644 --- a/tests/application/run/test_execution.py +++ b/tests/application/run/test_execution.py @@ -41,3 +41,50 @@ def test_raises_on_dirty_tree_inside_repo(tmp_path, monkeypatch) -> None: with pytest.raises(ValueError) as excinfo: validate_git_clean(verb="run", args=["python", "train.py"]) assert str(excinfo.value).startswith("Run blocked:") + + +def test_write_run_report_file_writes_payload(tmp_path, monkeypatch) -> None: + from roar.application.run.execution import ExecutionReport, write_run_report_file + + target = tmp_path / "report.json" + monkeypatch.setenv("ROAR_RUN_REPORT_FILE", str(target)) + + write_run_report_file(ExecutionReport(exit_code=3, tracer_backend="preload")) + + import json + + payload = json.loads(target.read_text(encoding="utf-8")) + assert payload == {"exit_code": 3, "setup_error": False, "tracer_backend": "preload"} + + +def test_write_run_report_file_marks_setup_error(tmp_path, monkeypatch) -> None: + from roar.application.run.execution import ExecutionReport, write_run_report_file + + target = tmp_path / "report.json" + monkeypatch.setenv("ROAR_RUN_REPORT_FILE", str(target)) + + write_run_report_file(ExecutionReport(exit_code=1, setup_error=True)) + + import json + + assert json.loads(target.read_text(encoding="utf-8"))["setup_error"] is True + + +def test_write_run_report_file_noop_without_env(tmp_path, monkeypatch) -> None: + from roar.application.run.execution import ExecutionReport, write_run_report_file + + monkeypatch.delenv("ROAR_RUN_REPORT_FILE", raising=False) + monkeypatch.chdir(tmp_path) + + write_run_report_file(ExecutionReport(exit_code=0)) + + assert list(tmp_path.iterdir()) == [] + + +def test_write_run_report_file_never_raises_on_unwritable_path(monkeypatch) -> None: + from roar.application.run.execution import ExecutionReport, write_run_report_file + + monkeypatch.setenv("ROAR_RUN_REPORT_FILE", "/nonexistent-dir/report.json") + + # Advisory only: an unwritable report path must not fail the run. + write_run_report_file(ExecutionReport(exit_code=0)) diff --git a/tests/backends/k8s/e2e/test_k8s_tracer_chaos.py b/tests/backends/k8s/e2e/test_k8s_tracer_chaos.py new file mode 100644 index 00000000..0583da24 --- /dev/null +++ b/tests/backends/k8s/e2e/test_k8s_tracer_chaos.py @@ -0,0 +1,177 @@ +"""Tracer-chaos e2e: a roar setup failure must never block training. + +Regression for the workload-safety contract: `k8s.tracer = "ebpf"` is a +valid tracer mode, but inside an unprivileged pod the ebpf backend +cannot pass preflight (and the staged runtime ships no ebpf tracer), so +in-pod `roar run --tracer ebpf` fails before launching the workload. +pod_entry must detect the pre-launch failure through the run report +(ROAR_RUN_REPORT_FILE) and rerun the original command uninstrumented — +the Job completes with the workload's own exit code and no lineage, +rather than failing training over a lineage problem. +""" + +from __future__ import annotations + +import os +import sqlite3 +import subprocess +import sys +import uuid +from pathlib import Path +from typing import Any + +import pytest + +from .conftest import KUBE_CONTEXT, NAMESPACE, TOOLS_BIN, kubectl +from .test_k8s_product_path import wheel_server # noqa: F401 + +pytestmark = [ + pytest.mark.e2e, + pytest.mark.k8s_e2e, + pytest.mark.timeout(900), +] + +CHAOS_MANIFEST_TEMPLATE = """\ +apiVersion: batch/v1 +kind: Job +metadata: + name: {job_name} + namespace: {namespace} + labels: + app.kubernetes.io/part-of: roar-k8s-e2e +spec: + backoffLimit: 0 + ttlSecondsAfterFinished: 1800 + template: + spec: + restartPolicy: Never + containers: + - name: trainer + image: python:3.12-slim + workingDir: /tmp + command: + - python + - -c + - "print('CHAOS_TRAIN_OK')" +""" + + +@pytest.fixture(scope="module") +def chaos_run( + k8s_cluster: None, + glaas_health: str, + wheel_server: dict[str, str], # noqa: F811 + tmp_path_factory: pytest.TempPathFactory, +) -> dict[str, Any]: + project_dir = tmp_path_factory.mktemp("k8s-tracer-chaos") + roar_dir = project_dir / ".roar" + roar_dir.mkdir() + (roar_dir / "config.toml").write_text( + "\n".join( + [ + "[glaas]", + 'url = "http://localhost:3001"', + "", + "[k8s]", + "enabled = true", + f'runtime_install_requirement = "{wheel_server["url"]}"', + 'cluster_glaas_url = "http://glaas:3001"', + # Valid tracer mode that cannot work in an unprivileged + # pod: forces an in-pod pre-launch setup failure. + 'tracer = "ebpf"', + "wait_for_completion = true", + "wait_timeout_seconds = 300", + "poll_interval_seconds = 2.0", + ] + ) + + "\n", + encoding="utf-8", + ) + + job_name = f"roar-chaos-{uuid.uuid4().hex[:6]}" + manifest_path = project_dir / "job.yaml" + manifest_path.write_text( + CHAOS_MANIFEST_TEMPLATE.format(job_name=job_name, namespace=NAMESPACE), + encoding="utf-8", + ) + + env = dict(os.environ) + env["PATH"] = f"{TOOLS_BIN}{os.pathsep}{env.get('PATH', '')}" + env.pop("GLAAS_URL", None) + env.pop("ROAR_PROJECT_DIR", None) + + try: + completed = subprocess.run( + [ + sys.executable, + "-m", + "roar", + "run", + "kubectl", + "apply", + "--context", + KUBE_CONTEXT, + "-f", + "job.yaml", + ], + cwd=project_dir, + env=env, + capture_output=True, + text=True, + check=False, + timeout=600, + ) + pod_logs = kubectl( + ["logs", "-n", NAMESPACE, "-l", f"job-name={job_name}", "--tail=100"], + check=False, + ) + return { + "project_dir": project_dir, + "job_name": job_name, + "exit_code": completed.returncode, + "stdout": completed.stdout, + "stderr": completed.stderr, + "pod_logs": pod_logs.stdout + pod_logs.stderr, + } + finally: + kubectl(["delete", f"job/{job_name}", "-n", NAMESPACE, "--ignore-not-found"], check=False) + kubectl( + [ + "delete", + "secret", + "-n", + NAMESPACE, + "-l", + "app.kubernetes.io/managed-by=roar", + "--ignore-not-found", + ], + check=False, + ) + + +def _describe(run: dict[str, Any]) -> str: + return ( + f"exit={run['exit_code']}\nstdout:\n{run['stdout']}\nstderr:\n{run['stderr']}\n" + f"pod logs:\n{run.get('pod_logs', '')}" + ) + + +def test_workload_survives_tracer_setup_failure(chaos_run: dict[str, Any]) -> None: + # The Job must complete even though roar's tracer cannot start: the + # submit (which waits for completion) exits 0 on the workload's own + # success. + assert chaos_run["exit_code"] == 0, _describe(chaos_run) + assert "CHAOS_TRAIN_OK" in chaos_run["pod_logs"], _describe(chaos_run) + assert "running uninstrumented" in chaos_run["pod_logs"], _describe(chaos_run) + + +def test_no_lineage_is_fabricated_for_uninstrumented_run(chaos_run: dict[str, Any]) -> None: + db_path = Path(chaos_run["project_dir"]) / ".roar" / "roar.db" + assert db_path.is_file(), _describe(chaos_run) + conn = sqlite3.connect(db_path) + conn.row_factory = sqlite3.Row + try: + tasks = conn.execute("SELECT id FROM jobs WHERE job_type = 'k8s_task'").fetchall() + finally: + conn.close() + assert tasks == [], _describe(chaos_run) diff --git a/tests/backends/k8s/unit/test_pod_entry.py b/tests/backends/k8s/unit/test_pod_entry.py new file mode 100644 index 00000000..f527327d --- /dev/null +++ b/tests/backends/k8s/unit/test_pod_entry.py @@ -0,0 +1,108 @@ +"""Pod-entry workload-safety contract: lineage never blocks training. + +``_run_traced`` shells out to ``roar run`` and cannot see past its exit +code, so it asks for a run report (ROAR_RUN_REPORT_FILE) and reruns the +original command uninstrumented only when roar positively reported a +pre-launch setup failure. A missing report is ambiguous — roar may have +crashed after the workload ran — and must never trigger a rerun. +""" + +from __future__ import annotations + +import json +import subprocess +from pathlib import Path +from typing import Any + +import pytest + +from roar.backends.k8s import pod_entry + + +class _FakeRun: + """Stub subprocess.run capturing invocations of roar run + fallback.""" + + def __init__(self, *, roar_exit: int, report: dict[str, Any] | None) -> None: + self.roar_exit = roar_exit + self.report = report + self.calls: list[list[str]] = [] + + def __call__(self, args: list[str], **kwargs: Any) -> subprocess.CompletedProcess: + self.calls.append(list(args)) + if "roar" in args and "run" in args: + env = kwargs.get("env") or {} + target = env.get("ROAR_RUN_REPORT_FILE") + if target and self.report is not None: + Path(target).write_text(json.dumps(self.report), encoding="utf-8") + return subprocess.CompletedProcess(args, self.roar_exit) + return subprocess.CompletedProcess(args, 7) + + +@pytest.fixture +def pod_workdir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("ROAR_K8S_WORKDIR", str(tmp_path)) + (tmp_path / ".roar").mkdir() + return tmp_path + + +def test_setup_error_report_reruns_uninstrumented( + pod_workdir: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + fake = _FakeRun(roar_exit=1, report={"exit_code": 1, "setup_error": True}) + monkeypatch.setattr(pod_entry.subprocess, "run", fake) + + exit_code = pod_entry._run_traced(["python", "train.py"]) + + # Second call is the raw command; its exit code (7) is the pod's. + assert len(fake.calls) == 2 + assert fake.calls[1] == ["python", "train.py"] + assert exit_code == 7 + + +def test_workload_failure_is_propagated_without_rerun( + pod_workdir: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + fake = _FakeRun(roar_exit=3, report={"exit_code": 3, "setup_error": False}) + monkeypatch.setattr(pod_entry.subprocess, "run", fake) + + exit_code = pod_entry._run_traced(["python", "train.py"]) + + assert len(fake.calls) == 1 + assert exit_code == 3 + + +def test_missing_report_never_triggers_rerun( + pod_workdir: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + # roar died without writing a report: ambiguous, the workload may have + # already run — propagate the failure rather than risk a double-run. + fake = _FakeRun(roar_exit=1, report=None) + monkeypatch.setattr(pod_entry.subprocess, "run", fake) + + exit_code = pod_entry._run_traced(["python", "train.py"]) + + assert len(fake.calls) == 1 + assert exit_code == 1 + + +def test_stale_report_from_prior_attempt_is_ignored( + pod_workdir: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + stale = pod_workdir / ".roar" / "k8s-run-report.json" + stale.write_text(json.dumps({"exit_code": 1, "setup_error": True}), encoding="utf-8") + fake = _FakeRun(roar_exit=1, report=None) + monkeypatch.setattr(pod_entry.subprocess, "run", fake) + + exit_code = pod_entry._run_traced(["python", "train.py"]) + + assert len(fake.calls) == 1 + assert exit_code == 1 + + +def test_successful_run_ignores_report(pod_workdir: Path, monkeypatch: pytest.MonkeyPatch) -> None: + fake = _FakeRun(roar_exit=0, report={"exit_code": 0, "setup_error": False}) + monkeypatch.setattr(pod_entry.subprocess, "run", fake) + + assert pod_entry._run_traced(["python", "train.py"]) == 0 + assert len(fake.calls) == 1 From 01dd9ff960a9d8cbd5cc8617858896fc40654359 Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Wed, 15 Jul 2026 18:39:37 +0000 Subject: [PATCH 24/31] fix(k8s): per-container bundle identity and multi-manifest submit guard Bundle fallback derived its filename from the pod name alone, so two wrapped containers in one pod (or an in-place restart) overwrote each other's fragments last-writer-wins. Bundle names now carry pod-container-attempt identity (attempt resolution mirrors the task_id contract) and are written atomically via os.replace so ingest never reads a torn bundle. The submit matcher also accepted commands with several -f/--filename arguments but rewrote only the first: one workload got lineage and completion polling while the rest were applied unrewritten and unwaited. Multiple manifest arguments now decline k8s instrumentation with an actionable warning and leave the kubectl command untouched. Co-Authored-By: Claude Fable 5 --- docs/developer/k8s-integration.md | 2 +- roar/backends/k8s/bundles.py | 45 ++++++++++++++----- roar/backends/k8s/config.py | 3 +- roar/backends/k8s/pod_entry.py | 12 ++++- roar/backends/k8s/submit.py | 43 +++++++++++++----- tests/backends/k8s/unit/test_bundles.py | 41 +++++++++++++++-- .../backends/k8s/unit/test_submit_planning.py | 25 +++++++++++ 7 files changed, 142 insertions(+), 29 deletions(-) diff --git a/docs/developer/k8s-integration.md b/docs/developer/k8s-integration.md index 5bb47906..7bd0236c 100644 --- a/docs/developer/k8s-integration.md +++ b/docs/developer/k8s-integration.md @@ -151,7 +151,7 @@ Two capture channels feed each pod's fragment: Transport is streaming-first with a bundle fallback: when `k8s.bundle_dir` names a mounted shared volume and GLaaS is unreachable from the pod (probe -or non-streamed emit), `pod_entry` writes `roar-fragments-.json` there +or non-streamed emit), `pod_entry` writes `roar-fragments---.json` there instead; `roar k8s ingest-bundles ` merges a host-visible copy later. `emit_fragment_dicts` reports "streamed" only when every batch was delivered (the streamer exposes delivered/failed/pending counts and logs diff --git a/roar/backends/k8s/bundles.py b/roar/backends/k8s/bundles.py index e6d1373e..d3bbd5c0 100644 --- a/roar/backends/k8s/bundles.py +++ b/roar/backends/k8s/bundles.py @@ -1,16 +1,17 @@ """Bundle-mode fragment fallback for GLaaS-less pods. When pods cannot reach GLaaS, the pod entrypoint writes its execution -fragment as ``roar-fragments-.json`` into ``ROAR_K8S_BUNDLE_DIR`` -(a mounted shared volume declared via ``k8s.bundle_dir``). Someone with -access to that volume later runs ``roar k8s ingest-bundles `` to -merge the bundles into the local lineage DB — the OSMO bundle pattern, -k8s-shaped. +fragment as ``roar-fragments---.json`` into +``ROAR_K8S_BUNDLE_DIR`` (a mounted shared volume declared via +``k8s.bundle_dir``). Someone with access to that volume later runs +``roar k8s ingest-bundles `` to merge the bundles into the local +lineage DB — the OSMO bundle pattern, k8s-shaped. """ from __future__ import annotations import json +import os import re from dataclasses import dataclass from pathlib import Path @@ -32,18 +33,40 @@ class K8sBundleIngestResult: bundle_paths: list[str] -def bundle_filename_for_pod(pod_name: str) -> str: - safe = re.sub(r"[^A-Za-z0-9_.-]", "-", pod_name.strip() or "pod") +def bundle_filename(pod_name: str, *, container: str = "main", attempt: str = "0") -> str: + """Bundle name carrying pod + container + attempt identity. + + Every wrapped container in a pod writes its own bundle; pod name alone + made co-located containers (and in-place restarts) overwrite each other. + """ + parts = ( + pod_name.strip() or "pod", + container.strip() or "main", + str(attempt).strip() or "0", + ) + safe = re.sub(r"[^A-Za-z0-9_.-]", "-", "-".join(parts)) return f"{BUNDLE_FILENAME_PREFIX}{safe}.json" -def write_fragment_bundle(bundle_dir: Path, pod_name: str, fragments: list[dict[str, Any]]) -> Path: +def write_fragment_bundle( + bundle_dir: Path, + pod_name: str, + fragments: list[dict[str, Any]], + *, + container: str = "main", + attempt: str = "0", +) -> Path: bundle_dir.mkdir(parents=True, exist_ok=True) - target = bundle_dir / bundle_filename_for_pod(pod_name) - target.write_text( + target = bundle_dir / bundle_filename(pod_name, container=container, attempt=attempt) + # Atomic replace: a reader (or a concurrent ingest) must never see a + # torn bundle, and the .tmp suffix keeps partial writes out of the + # discovery glob. + temp = target.with_name(f"{target.name}.tmp-{os.getpid()}") + temp.write_text( json.dumps({"fragments": fragments}, indent=2, sort_keys=True) + "\n", encoding="utf-8", ) + os.replace(temp, target) return target @@ -105,7 +128,7 @@ def ingest_fragment_bundles( "BUNDLE_FILENAME_PREFIX", "K8sBundleError", "K8sBundleIngestResult", - "bundle_filename_for_pod", + "bundle_filename", "discover_fragment_bundles", "ingest_fragment_bundles", "write_fragment_bundle", diff --git a/roar/backends/k8s/config.py b/roar/backends/k8s/config.py index a3e6b3b4..7e5e3888 100644 --- a/roar/backends/k8s/config.py +++ b/roar/backends/k8s/config.py @@ -108,7 +108,8 @@ class K8sBackendConfig(BaseModel): default="", description=( "In-pod directory (a mounted shared volume) where pods write " - "roar-fragments-.json bundles when GLaaS streaming is unavailable; " + "roar-fragments---.json bundles when GLaaS " + "streaming is unavailable; " "ingest later with `roar k8s ingest-bundles`" ), ), diff --git a/roar/backends/k8s/pod_entry.py b/roar/backends/k8s/pod_entry.py index f4f27604..85213c76 100644 --- a/roar/backends/k8s/pod_entry.py +++ b/roar/backends/k8s/pod_entry.py @@ -318,7 +318,17 @@ def _write_bundle(bundle_dir: str, fragments: list[dict]) -> str: from roar.backends.k8s.bundles import write_fragment_bundle pod_name = str(os.environ.get("ROAR_K8S_POD_NAME") or "pod").strip() or "pod" - target = write_fragment_bundle(Path(bundle_dir), pod_name, fragments) + container = str(os.environ.get("ROAR_K8S_CONTAINER") or "main").strip() or "main" + # Attempt resolution mirrors task_identity_from_environment so the + # bundle name and the fragment task_id agree on which attempt this is. + attempt = ( + str(os.environ.get("ROAR_K8S_RESTART_ATTEMPT") or "").strip() + or str(os.environ.get("TORCHELASTIC_RESTART_COUNT") or "").strip() + or "0" + ) + target = write_fragment_bundle( + Path(bundle_dir), pod_name, fragments, container=container, attempt=attempt + ) print(f"[roar-k8s] GLaaS unavailable; wrote fragment bundle to {target}") return "bundled" diff --git a/roar/backends/k8s/submit.py b/roar/backends/k8s/submit.py index 32a59792..1672fd24 100644 --- a/roar/backends/k8s/submit.py +++ b/roar/backends/k8s/submit.py @@ -59,10 +59,20 @@ def matches_kubectl_job_submit_command(command: list[str]) -> bool: if not _k8s_backend_enabled(): return False - filename = _find_filename_argument(command) - if filename is None: + manifests = _find_filename_arguments(command) + if not manifests: return False - manifest_path = Path(filename[1]) + if len(manifests) > 1: + # Rewriting only the first manifest would instrument one workload + # and silently pass the rest through unwaited; decline the whole + # submit so kubectl behaves exactly as the user wrote it. + _warn( + "multiple -f/--filename arguments are not supported by the k8s " + "lineage backend; running without k8s instrumentation (apply " + "each manifest in its own command to capture lineage)" + ) + return False + manifest_path = Path(manifests[0][1]) if not manifest_path.is_file(): return False @@ -252,17 +262,26 @@ def _k8s_backend_enabled() -> bool: return bool(load_k8s_backend_config(start_dir=start_dir).get("enabled", False)) -def _find_filename_argument(command: list[str]) -> tuple[int, str] | None: - for index, arg in enumerate(command[2:], start=2): +def _find_filename_arguments(command: list[str]) -> list[tuple[int, str]]: + """All (value_index, value) pairs for -f/--filename in the command.""" + found: list[tuple[int, str]] = [] + index = 2 + while index < len(command): + arg = command[index] if arg in ("-f", "--filename"): if index + 1 < len(command): - return index + 1, command[index + 1] - return None - if arg.startswith("--filename="): - return index, arg.split("=", 1)[1] - if arg.startswith("-f="): - return index, arg.split("=", 1)[1] - return None + found.append((index + 1, command[index + 1])) + index += 2 + continue + elif arg.startswith(("--filename=", "-f=")): + found.append((index, arg.split("=", 1)[1])) + index += 1 + return found + + +def _find_filename_argument(command: list[str]) -> tuple[int, str] | None: + found = _find_filename_arguments(command) + return found[0] if found else None def _replace_filename_argument(command: list[str], prepared_path: str) -> list[str]: diff --git a/tests/backends/k8s/unit/test_bundles.py b/tests/backends/k8s/unit/test_bundles.py index 2e04adab..3729010a 100644 --- a/tests/backends/k8s/unit/test_bundles.py +++ b/tests/backends/k8s/unit/test_bundles.py @@ -1,5 +1,6 @@ from __future__ import annotations +import json import sqlite3 from pathlib import Path @@ -7,7 +8,7 @@ from roar.backends.k8s.bundles import ( K8sBundleError, - bundle_filename_for_pod, + bundle_filename, discover_fragment_bundles, ingest_fragment_bundles, write_fragment_bundle, @@ -42,8 +43,8 @@ def _fragment(pod: str, index: str) -> dict: def test_bundle_filename_sanitizes_pod_names() -> None: - assert bundle_filename_for_pod("train-0/pod x") == "roar-fragments-train-0-pod-x.json" - assert bundle_filename_for_pod("") == "roar-fragments-pod.json" + assert bundle_filename("train-0/pod x") == "roar-fragments-train-0-pod-x-main-0.json" + assert bundle_filename("") == "roar-fragments-pod-main-0.json" def test_write_and_discover_bundles(tmp_path: Path) -> None: @@ -100,3 +101,37 @@ def test_ingest_empty_directory_fails_actionably(tmp_path: Path) -> None: empty.mkdir() with pytest.raises(K8sBundleError, match="no roar-fragments-"): ingest_fragment_bundles(roar_dir=tmp_path / ".roar", directory=empty) + + +def test_bundle_filenames_are_distinct_per_container_and_attempt() -> None: + from roar.backends.k8s.bundles import bundle_filename + + trainer = bundle_filename("pod-0", container="trainer", attempt="0") + sidecar = bundle_filename("pod-0", container="sidecar", attempt="0") + retry = bundle_filename("pod-0", container="trainer", attempt="1") + + assert len({trainer, sidecar, retry}) == 3 + assert trainer == "roar-fragments-pod-0-trainer-0.json" + + +def test_multi_container_bundles_do_not_overwrite(tmp_path: Path) -> None: + bundle_dir = tmp_path / "bundles" + first = write_fragment_bundle( + bundle_dir, "pod-0", [_fragment("pod-0", "0")], container="trainer" + ) + second = write_fragment_bundle( + bundle_dir, "pod-0", [_fragment("pod-0", "1")], container="sidecar" + ) + + assert first != second + assert sorted(discover_fragment_bundles(bundle_dir)) == sorted([first, second]) + # Both containers' fragments survive: last-writer-wins is the bug. + assert json.loads(first.read_text(encoding="utf-8"))["fragments"][0]["task_id"].endswith(":0:0") + + +def test_write_fragment_bundle_leaves_no_temp_files(tmp_path: Path) -> None: + bundle_dir = tmp_path / "bundles" + write_fragment_bundle(bundle_dir, "pod-0", [_fragment("pod-0", "0")]) + + leftovers = [path.name for path in bundle_dir.iterdir()] + assert leftovers == ["roar-fragments-pod-0-main-0.json"] diff --git a/tests/backends/k8s/unit/test_submit_planning.py b/tests/backends/k8s/unit/test_submit_planning.py index af94695b..1fdf66fc 100644 --- a/tests/backends/k8s/unit/test_submit_planning.py +++ b/tests/backends/k8s/unit/test_submit_planning.py @@ -203,3 +203,28 @@ def test_resolve_runtime_requirement_precedence(monkeypatch: pytest.MonkeyPatch) == "roar-cli==1.0" ) assert resolve_runtime_requirement({}).startswith("roar-cli") + + +def test_declines_multiple_manifest_arguments_with_warning( + job_manifest_path: Path, capsys: pytest.CaptureFixture[str] +) -> None: + second = job_manifest_path.parent / "job-b.yaml" + manifest = yaml.safe_load(job_manifest_path.read_text(encoding="utf-8")) + manifest["metadata"]["name"] = "train-demo-b" + second.write_text(yaml.safe_dump(manifest, sort_keys=False), encoding="utf-8") + + # Rewriting only the first manifest would instrument one workload and + # silently pass the other through unwaited — decline the whole submit + # instead so kubectl behavior is untouched. + assert not matches_kubectl_job_submit_command( + ["kubectl", "apply", "-f", str(job_manifest_path), "-f", str(second)] + ) + assert "multiple -f/--filename" in capsys.readouterr().err + + assert not matches_kubectl_job_submit_command( + ["kubectl", "apply", "-f", str(job_manifest_path), f"--filename={second}"] + ) + + +def test_single_manifest_still_matches_after_multi_f_guard(job_manifest_path: Path) -> None: + assert matches_kubectl_job_submit_command(["kubectl", "apply", "-f", str(job_manifest_path)]) From 72f6df973ff71b5f4928a448f4219351bd057e45 Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Wed, 15 Jul 2026 18:41:17 +0000 Subject: [PATCH 25/31] fix(fragments): owner-only session keys; drop query-string renewal MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Session .key files carry the raw fragment token but were written with write_text, inheriting the umask (0664 observed) — group/other could read the credential. Keys are now created 0600 from the first byte and swapped in with os.replace so a rewrite of an existing over-permissive key also ends owner-only. Also removes the unused GlaasClient.renew_fragment_session, which put the raw token in the URL query string where access logs and proxies capture it. The live renewal path (fragment_streamer) authenticates with the x-roar-fragment-token header and stays. Co-Authored-By: Claude Fable 5 --- roar/execution/fragments/sessions.py | 10 +++++- roar/integrations/glaas/client.py | 17 +++------ tests/backends/ray/unit/test_fragment_key.py | 38 ++++++++++++++++++++ 3 files changed, 51 insertions(+), 14 deletions(-) diff --git a/roar/execution/fragments/sessions.py b/roar/execution/fragments/sessions.py index b10c4bef..bbed6943 100644 --- a/roar/execution/fragments/sessions.py +++ b/roar/execution/fragments/sessions.py @@ -56,7 +56,15 @@ def fragment_session_path(roar_dir: Path, session_id: str) -> Path: def save_fragment_session(roar_dir: Path, payload: dict[str, Any]) -> Path: path = fragment_session_path(roar_dir, str(payload["session_id"])) path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(json.dumps(payload, separators=(",", ":")), encoding="utf-8") + # The key file carries the raw fragment token: owner-only from the + # first byte (write_text would inherit the umask, leaving group/other + # read), and swapped in atomically so a crash never leaves a partial + # or over-permissive key behind. + temp = path.with_name(f"{path.name}.tmp-{os.getpid()}") + fd = os.open(temp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + with os.fdopen(fd, "w", encoding="utf-8") as handle: + handle.write(json.dumps(payload, separators=(",", ":"))) + os.replace(temp, path) return path diff --git a/roar/integrations/glaas/client.py b/roar/integrations/glaas/client.py index 62adcbf6..470167b1 100644 --- a/roar/integrations/glaas/client.py +++ b/roar/integrations/glaas/client.py @@ -595,19 +595,10 @@ def register_fragment_session( } return self._request("POST", "/api/v1/fragments/sessions", body) - def renew_fragment_session( - self, - session_id: str, - token: str, - ttl_seconds: int = 86400, - ) -> tuple[dict | None, str | None]: - """Extend a fragment session's expiry (allowed even after expiry).""" - encoded_token = urllib.parse.quote(token, safe="") - return self._request( - "POST", - f"/api/v1/fragments/sessions/{session_id}/renew?token={encoded_token}", - {"ttl_seconds": ttl_seconds}, - ) + # Fragment-session renewal lives in fragment_streamer.renew_fragment_session: + # it authenticates with the x-roar-fragment-token header. A client-method + # variant here once put the raw token in the query string, where access + # logs and proxies capture it — do not reintroduce it. def register_session( self, diff --git a/tests/backends/ray/unit/test_fragment_key.py b/tests/backends/ray/unit/test_fragment_key.py index eb6db921..cade870e 100644 --- a/tests/backends/ray/unit/test_fragment_key.py +++ b/tests/backends/ray/unit/test_fragment_key.py @@ -58,3 +58,41 @@ def test_load_key_round_trip_restores_session_id_and_token(tmp_path: Path) -> No def test_missing_key_raises_file_not_found(tmp_path: Path) -> None: with pytest.raises(FileNotFoundError): load_fragment_session(tmp_path / ".roar", "00000000-0000-4000-8000-000000000000") + + +def test_key_file_is_owner_only_even_under_permissive_umask(tmp_path: Path) -> None: + import os + + key = generate_fragment_session() + previous_umask = os.umask(0o000) + try: + path = save_fragment_session(tmp_path / ".roar", key) + finally: + os.umask(previous_umask) + + # The key file carries the raw fragment token; write_text would have + # inherited the umask and left it group/other readable. + assert (path.stat().st_mode & 0o777) == 0o600 + + +def test_key_file_rewrite_stays_owner_only_and_leaves_no_temp(tmp_path: Path) -> None: + roar_dir = tmp_path / ".roar" + key = generate_fragment_session() + first = save_fragment_session(roar_dir, key) + first.chmod(0o664) + + second = save_fragment_session(roar_dir, key) + + assert second == first + assert (second.stat().st_mode & 0o777) == 0o600 + assert [path.name for path in second.parent.iterdir()] == [second.name] + + +def test_glaas_client_has_no_query_string_renewal() -> None: + # Regression: a GlaasClient.renew_fragment_session variant once put the + # raw fragment token in the URL query string, where access logs and + # proxies capture it. Renewal lives in fragment_streamer with the + # x-roar-fragment-token header. + from roar.integrations.glaas import GlaasClient + + assert not hasattr(GlaasClient, "renew_fragment_session") From 41cae186640bbc4c854e503846862662ae970d56 Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Wed, 15 Jul 2026 18:42:51 +0000 Subject: [PATCH 26/31] fix(charts): ship the Secret-cleanup CronJob disabled by default MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The cleanup CronJob defaulted on, granting its ServiceAccount cluster-wide list/delete on Secrets — list responses include Secret data, so a default install was one compromised cleanup image away from cluster-wide credential read. Its fixed-age policy also conflicts with session renewal: a renewed workload outliving maxAgeSeconds loses the Secret its retry/scale-up pods reference, and attach recovery with it. The CronJob is now opt-in, with both trade-offs documented in values.yaml, the module docstring, and the developer doc. Liveness-aware (controller-grade) cleanup remains deliberately deferred. Co-Authored-By: Claude Fable 5 --- .../charts/roar-lineage-webhook/values.yaml | 28 +++++++++++++------ docs/developer/k8s-integration.md | 13 +++++++-- roar/backends/k8s/secret_cleanup.py | 17 ++++++----- 3 files changed, 40 insertions(+), 18 deletions(-) diff --git a/deploy/charts/roar-lineage-webhook/values.yaml b/deploy/charts/roar-lineage-webhook/values.yaml index 4a365a2b..280f2160 100644 --- a/deploy/charts/roar-lineage-webhook/values.yaml +++ b/deploy/charts/roar-lineage-webhook/values.yaml @@ -57,15 +57,27 @@ certManager: secretCleanup: # Credential Secrets cannot carry an ownerReference (the workload has no - # UID at CREATE admission), so a CronJob deletes roar-managed fragment - # Secrets older than maxAgeSeconds. Runs under its own ServiceAccount; - # the webhook's stays create-only. - enabled: true + # UID at CREATE admission), so an optional CronJob deletes roar-managed + # fragment Secrets older than maxAgeSeconds. Runs under its own + # ServiceAccount; the webhook's stays create-only. + # + # Off by default — enabling it is a real privilege and reliability + # trade-off; read before flipping: + # - RBAC: the CronJob's ServiceAccount gets cluster-wide list/delete on + # Secrets. Kubernetes list responses include Secret data, so this is + # equivalent to cluster-wide Secret READ; the roar label selector does + # not narrow what the account is authorized to request. Treat the + # cleanup image and ServiceAccount as privileged security components. + # - Renewal conflict: fragment sessions are renewed on 403 and can + # outlive any fixed maxAgeSeconds. Deleting an aged Secret does not + # affect running containers (env resolves at pod start), but a retry + # pod, scale-up, or new Ray worker that references the deleted Secret + # cannot start, and cluster-based `roar k8s attach` recovery is lost. + # Keep it disabled unless orphaned credential Secrets are a bigger + # problem in your cluster than the above; prefer a maxAgeSeconds well + # past your longest training run when you do enable it. + enabled: false schedule: "0 3 * * *" - # Default matches the maximum fragment-session TTL (7 days): a Secret - # older than this belongs to a session the server no longer honors. - # Running pods are unaffected (env is resolved at pod start); only - # late `roar k8s attach` credential recovery ages out with it. maxAgeSeconds: 604800 resources: {} diff --git a/docs/developer/k8s-integration.md b/docs/developer/k8s-integration.md index 7bd0236c..18b9624f 100644 --- a/docs/developer/k8s-integration.md +++ b/docs/developer/k8s-integration.md @@ -54,9 +54,16 @@ to every workload it instruments. The webhook rewrites the manifest *before* creating the GLaaS session or the credentials Secret, so a rejected/unrewritable workload leaves nothing behind. Secrets cannot carry an ownerReference (no workload UID -exists at CREATE admission); the chart ships a `secretCleanup` CronJob -(own ServiceAccount, list/delete only, name-prefix guarded) that removes -roar fragment Secrets older than the maximum session TTL. +exists at CREATE admission); the chart ships an **opt-in** `secretCleanup` +CronJob (own ServiceAccount, name-prefix guarded) that removes roar +fragment Secrets older than `maxAgeSeconds`. It is disabled by default +for two reasons documented in `values.yaml`: its ServiceAccount needs +cluster-wide list/delete on Secrets (list responses carry Secret data, +so this is cluster-wide Secret read), and age-based deletion conflicts +with session renewal — a renewed long-running workload whose Secret ages +out cannot start retry/scale-up pods and loses `roar k8s attach` +recovery. Enable it only after accepting both, with `maxAgeSeconds` +comfortably past the longest expected training run. **Helm chart** (`deploy/charts/roar-lineage-webhook`): packages the webhook Deployment/Service, RBAC (Secret-create ClusterRole), and the diff --git a/roar/backends/k8s/secret_cleanup.py b/roar/backends/k8s/secret_cleanup.py index f2d40ab2..707e100a 100644 --- a/roar/backends/k8s/secret_cleanup.py +++ b/roar/backends/k8s/secret_cleanup.py @@ -3,13 +3,16 @@ Credential Secrets cannot carry an ownerReference: at CREATE admission the workload has no UID yet, so Kubernetes garbage collection can never adopt them. This module lists roar-managed Secrets cluster-wide and deletes those -older than ROAR_SECRET_MAX_AGE_SECONDS (default: the 7-day maximum -fragment-session TTL — a Secret older than that belongs to a session the -server no longer honors). Running pods are unaffected (env is resolved at -pod start); only late `roar k8s attach` credential recovery ages out. - -Runs in-cluster under a dedicated ServiceAccount with list/delete on -Secrets (see the chart's secret-cleanup template). +older than ROAR_SECRET_MAX_AGE_SECONDS (default 7 days). + +Age is a heuristic, not a session-liveness signal: sessions renew on 403 +and can outlive any fixed TTL. Running containers are unaffected (env +resolves at pod start), but retry/scale-up pods referencing a deleted +Secret cannot start, and cluster-based `roar k8s attach` recovery is +lost. The chart therefore ships this CronJob disabled by default; its +ServiceAccount needs cluster-wide list/delete on Secrets (list responses +include Secret data), so treat it as a privileged component. See the +chart values for the full trade-off. """ from __future__ import annotations From bd5bb075c5f88b1e54eb89b1e38c3ef02cfc7314 Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Thu, 16 Jul 2026 17:09:47 +0000 Subject: [PATCH 27/31] fix(k8s): parse the kubectl subcommand structurally The verb-anywhere matcher accepted global-flag VALUES as the verb: 'kubectl --context apply delete -f job.yaml' and 'kubectl delete -f job.yaml --namespace create' both matched the submit backend, which would register a session, rewrite a manifest being deleted, and wait on a workload that no longer exists. The matcher now finds the first non-flag token after skipping known value-taking kubectl global flags (space and = forms); both review commands are regression tests. Co-Authored-By: Claude Fable 5 --- roar/backends/k8s/submit.py | 58 +++++++++++++++++-- .../backends/k8s/unit/test_submit_planning.py | 27 +++++++++ 2 files changed, 81 insertions(+), 4 deletions(-) diff --git a/roar/backends/k8s/submit.py b/roar/backends/k8s/submit.py index 1672fd24..a5639430 100644 --- a/roar/backends/k8s/submit.py +++ b/roar/backends/k8s/submit.py @@ -25,6 +25,39 @@ from roar.execution.framework.contract import ExecutionCommandPlan _KUBECTL_VERBS = ("apply", "create") +# kubectl global flags that consume the next argument when written in the +# space-separated form (kubectl options). Needed to find the real +# subcommand: in `kubectl --context apply delete -f x`, "apply" is a flag +# value, not the verb. +_KUBECTL_GLOBAL_VALUE_FLAGS = frozenset( + { + "--as", + "--as-group", + "--as-uid", + "--cache-dir", + "--certificate-authority", + "--client-certificate", + "--client-key", + "--cluster", + "--context", + "--kubeconfig", + "-n", + "--namespace", + "--password", + "--profile", + "--profile-output", + "--request-timeout", + "-s", + "--server", + "--tls-server-name", + "--token", + "--user", + "--username", + "-v", + "--v", + "--vmodule", + } +) SUBMIT_CONTEXT_SUFFIX = ".context.json" @@ -51,10 +84,7 @@ def matches_kubectl_job_submit_command(command: list[str]) -> bool: return False if Path(command[0]).name.lower() != "kubectl": return False - # Global flags may precede the verb (kubectl --context X apply -f ...), - # so accept the verb anywhere. The -f-points-at-a-manifest and - # single-supported-workload guards below keep false positives out. - if not any(arg.lower() in _KUBECTL_VERBS for arg in command[1:]): + if _kubectl_subcommand(command) not in _KUBECTL_VERBS: return False if not _k8s_backend_enabled(): return False @@ -262,6 +292,26 @@ def _k8s_backend_enabled() -> bool: return bool(load_k8s_backend_config(start_dir=start_dir).get("enabled", False)) +def _kubectl_subcommand(command: list[str]) -> str | None: + """The first non-flag token after skipping global flag/value pairs. + + Unknown bare flags are treated as boolean, so a value-taking global + flag missing from _KUBECTL_GLOBAL_VALUE_FLAGS would have its value + misread as the subcommand. Keep the set in sync with `kubectl options` + when new globals appear. + """ + index = 1 + while index < len(command): + arg = command[index] + if not arg.startswith("-"): + return arg.lower() + if "=" not in arg and arg in _KUBECTL_GLOBAL_VALUE_FLAGS: + index += 2 + continue + index += 1 + return None + + def _find_filename_arguments(command: list[str]) -> list[tuple[int, str]]: """All (value_index, value) pairs for -f/--filename in the command.""" found: list[tuple[int, str]] = [] diff --git a/tests/backends/k8s/unit/test_submit_planning.py b/tests/backends/k8s/unit/test_submit_planning.py index 1fdf66fc..5e88b33f 100644 --- a/tests/backends/k8s/unit/test_submit_planning.py +++ b/tests/backends/k8s/unit/test_submit_planning.py @@ -228,3 +228,30 @@ def test_declines_multiple_manifest_arguments_with_warning( def test_single_manifest_still_matches_after_multi_f_guard(job_manifest_path: Path) -> None: assert matches_kubectl_job_submit_command(["kubectl", "apply", "-f", str(job_manifest_path)]) + + +def test_verb_as_global_flag_value_does_not_match(job_manifest_path: Path) -> None: + # "apply" here is the VALUE of --context; the real subcommand is delete. + # Matching would rewrite and wait on a workload that is being deleted. + assert not matches_kubectl_job_submit_command( + ["kubectl", "--context", "apply", "delete", "-f", str(job_manifest_path)] + ) + # "create" is the value of --namespace, after the real delete verb. + assert not matches_kubectl_job_submit_command( + ["kubectl", "delete", "-f", str(job_manifest_path), "--namespace", "create"] + ) + + +def test_verb_after_boolean_global_flag_matches(job_manifest_path: Path) -> None: + assert matches_kubectl_job_submit_command( + ["kubectl", "--insecure-skip-tls-verify", "apply", "-f", str(job_manifest_path)] + ) + + +def test_equals_form_global_flag_values_are_skipped(job_manifest_path: Path) -> None: + assert matches_kubectl_job_submit_command( + ["kubectl", "--context=prod", "apply", "-f", str(job_manifest_path)] + ) + assert not matches_kubectl_job_submit_command( + ["kubectl", "--context=prod", "delete", "-f", str(job_manifest_path)] + ) From 274871c1173fc4498047bb66a6281b812e45168f Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Thu, 16 Jul 2026 17:13:39 +0000 Subject: [PATCH 28/31] fix(k8s): make roar k8s prepare mirror the managed rewrite config prepare passed only requirement, GLaaS URL, tracer, Secret name, and parent UID to the rewriter, so image-staged, bundle-enabled, mount-mapped, proxy, or namespace-overridden configurations previewed differently from what the managed submit would actually apply. Both paths now share manifest_rewrite_config_kwargs (single source of truth for config-derived rewrite kwargs) and the proxy-sidecar cannot-inject warning; prepare gains -n/--namespace mirroring kubectl apply -n semantics. Co-Authored-By: Claude Fable 5 --- roar/backends/k8s/submit.py | 52 +++++++--- roar/cli/commands/k8s.py | 22 +++- tests/backends/k8s/unit/test_prepare_cli.py | 106 ++++++++++++++++++++ 3 files changed, 163 insertions(+), 17 deletions(-) create mode 100644 tests/backends/k8s/unit/test_prepare_cli.py diff --git a/roar/backends/k8s/submit.py b/roar/backends/k8s/submit.py index a5639430..a9e9ccbf 100644 --- a/roar/backends/k8s/submit.py +++ b/roar/backends/k8s/submit.py @@ -8,6 +8,7 @@ import sys from dataclasses import asdict, dataclass from pathlib import Path +from typing import Any from roar.backends.k8s.config import load_k8s_backend_config from roar.backends.k8s.manifest import ( @@ -164,24 +165,13 @@ def plan_kubectl_job_submit_command(command: list[str]) -> ExecutionCommandPlan: fragment_token=session["token"], requirement=resolve_runtime_requirement(config), cluster_glaas_url=_resolve_cluster_glaas_url(config, glaas_url), - tracer=str(config.get("tracer") or "preload"), parent_job_uid=parent_job_uid, - bundle_dir=str(config.get("bundle_dir") or ""), - mount_map=_config_mount_map(config), - runtime_source=str(config.get("runtime_source") or "install"), - runtime_image=str(config.get("runtime_image") or ""), - proxy_sidecar=bool(config.get("proxy_sidecar", False)), - proxy_upstream=str(config.get("proxy_upstream") or ""), namespace_override=_find_namespace_argument(command), + **manifest_rewrite_config_kwargs(config), ) - if bool(config.get("proxy_sidecar", False)) and ( - str(config.get("runtime_source") or "install") != "image" - or not str(config.get("runtime_image") or "") - ): - _warn( - "k8s.proxy_sidecar requires k8s.runtime_source='image' with " - "k8s.runtime_image set; sidecar not injected" - ) + proxy_warning = proxy_sidecar_config_warning(config) + if proxy_warning: + _warn(proxy_warning) if rewrite.skipped_containers: _warn( "containers without an explicit command were left uninstrumented: " @@ -358,6 +348,36 @@ def _find_namespace_argument(command: list[str]) -> str | None: return None +def manifest_rewrite_config_kwargs(config: dict) -> dict[str, Any]: + """Config-derived rewrite kwargs shared by managed submit and `roar k8s prepare`. + + One source of truth: prepare must preview exactly the rewrite the + managed path would submit (bundle, mount map, runtime staging, proxy). + """ + return { + "tracer": str(config.get("tracer") or "preload"), + "bundle_dir": str(config.get("bundle_dir") or ""), + "mount_map": _config_mount_map(config), + "runtime_source": str(config.get("runtime_source") or "install"), + "runtime_image": str(config.get("runtime_image") or ""), + "proxy_sidecar": bool(config.get("proxy_sidecar", False)), + "proxy_upstream": str(config.get("proxy_upstream") or ""), + } + + +def proxy_sidecar_config_warning(config: dict) -> str | None: + """Warning when proxy_sidecar is configured but cannot be injected.""" + if bool(config.get("proxy_sidecar", False)) and ( + str(config.get("runtime_source") or "install") != "image" + or not str(config.get("runtime_image") or "") + ): + return ( + "k8s.proxy_sidecar requires k8s.runtime_source='image' with " + "k8s.runtime_image set; sidecar not injected" + ) + return None + + def _config_mount_map(config: dict) -> dict[str, str]: raw = config.get("mount_map") if not isinstance(raw, dict): @@ -412,7 +432,9 @@ def _warn(message: str) -> None: "K8sSubmitContext", "discard_submit_context", "load_submit_context", + "manifest_rewrite_config_kwargs", "matches_kubectl_job_submit_command", "plan_kubectl_job_submit_command", + "proxy_sidecar_config_warning", "resolve_runtime_requirement", ] diff --git a/roar/cli/commands/k8s.py b/roar/cli/commands/k8s.py index aa9201a3..1c7ea32f 100644 --- a/roar/cli/commands/k8s.py +++ b/roar/cli/commands/k8s.py @@ -14,7 +14,11 @@ load_manifest_documents, rewrite_manifest_for_lineage, ) -from ...backends.k8s.submit import resolve_runtime_requirement +from ...backends.k8s.submit import ( + manifest_rewrite_config_kwargs, + proxy_sidecar_config_warning, + resolve_runtime_requirement, +) from ..context import RoarContext from ..decorators import require_init @@ -150,18 +154,27 @@ def k8s_ingest_bundles(ctx: RoarContext, directory: Path) -> None: default=None, help="Cluster-visible GLaaS URL injected into pods (defaults to k8s.cluster_glaas_url)", ) +@click.option( + "-n", + "--namespace", + default=None, + help="Namespace override, mirroring `kubectl apply -n` on the managed path", +) def k8s_prepare( manifest_path: Path, output_path: Path, secret_name: str, requirement: str | None, cluster_glaas_url: str | None, + namespace: str | None, ) -> None: """Rewrite a Job manifest with lineage instrumentation for inspection. Unlike the managed `roar run kubectl apply -f ...` path, no fragment session is registered and no Secret is embedded: create the Secret named by --secret-name (keys: session_id, token) before applying the output. + Everything else mirrors the managed rewrite: bundle, mount-map, runtime + staging, and proxy-sidecar settings are read from the same k8s config. """ config = load_k8s_backend_config() resolved_requirement = requirement or resolve_runtime_requirement(config) @@ -178,12 +191,17 @@ def k8s_prepare( fragment_token=None, requirement=resolved_requirement, cluster_glaas_url=resolved_cluster_glaas, - tracer=str(config.get("tracer") or "preload"), parent_job_uid=secrets.token_hex(4), + namespace_override=namespace, + **manifest_rewrite_config_kwargs(config), ) except K8sManifestError as exc: raise click.ClickException(str(exc)) from exc + proxy_warning = proxy_sidecar_config_warning(config) + if proxy_warning: + click.echo(f"warning: {proxy_warning}", err=True) + output_path.parent.mkdir(parents=True, exist_ok=True) output_path.write_text(dump_manifest_documents(rewrite.documents), encoding="utf-8") diff --git a/tests/backends/k8s/unit/test_prepare_cli.py b/tests/backends/k8s/unit/test_prepare_cli.py new file mode 100644 index 00000000..0cd54e16 --- /dev/null +++ b/tests/backends/k8s/unit/test_prepare_cli.py @@ -0,0 +1,106 @@ +"""`roar k8s prepare` must preview exactly what the managed path submits. + +Regression for config drift: prepare used to pass only requirement, +GLaaS URL, tracer, Secret name, and parent UID to the rewriter, so an +image-staged, bundle-enabled, mount-mapped, or namespace-overridden +configuration previewed differently from what `roar run kubectl apply` +would actually apply. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +import yaml # type: ignore[import-untyped] +from click.testing import CliRunner + +from roar.cli.commands.k8s import k8s + +from .conftest import SINGLE_JOB_MANIFEST + + +@pytest.fixture() +def configured_project(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + roar_dir = tmp_path / ".roar" + roar_dir.mkdir() + (roar_dir / "config.toml").write_text( + "\n".join( + [ + "[k8s]", + "enabled = true", + 'bundle_dir = "/mnt/bundles"', + 'runtime_source = "image"', + 'runtime_image = "roar-runtime:pinned"', + "", + "[k8s.mount_map]", + '"/data" = "pvc://training-data"', + "", + "[glaas]", + 'url = "http://localhost:3001"', + ] + ) + + "\n", + encoding="utf-8", + ) + manifest = tmp_path / "job.yaml" + manifest.write_text(yaml.safe_dump(SINGLE_JOB_MANIFEST, sort_keys=False), encoding="utf-8") + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("ROAR_PROJECT_DIR", str(tmp_path)) + return tmp_path + + +def test_prepare_mirrors_managed_rewrite_configuration(configured_project: Path) -> None: + output_path = configured_project / "prepared.yaml" + result = CliRunner().invoke( + k8s, + [ + "prepare", + "-f", + "job.yaml", + "-o", + str(output_path), + "-n", + "override-ns", + ], + ) + assert result.exit_code == 0, result.output + + documents = list(yaml.safe_load_all(output_path.read_text(encoding="utf-8"))) + job = next(doc for doc in documents if doc and doc.get("kind") == "Job") + pod_spec = job["spec"]["template"]["spec"] + + # namespace_override feeds the resolved namespace (Secret placement and + # the recorded context), exactly as `kubectl apply -n` does on the + # managed path; the workload's own metadata is never mutated. + assert "(namespace override-ns)" in result.output + + # Image staging: the managed path would add the runtime-staging init + # container with the pinned image; prepare must show the same. + init_names = [c.get("name") for c in pod_spec.get("initContainers", [])] + assert "roar-runtime-staging" in init_names, result.output + + env = { + entry.get("name"): entry.get("value") + for entry in pod_spec["containers"][0].get("env", []) + if isinstance(entry, dict) + } + assert env.get("ROAR_K8S_BUNDLE_DIR") == "/mnt/bundles" + assert "pvc://training-data" in str(env.get("ROAR_K8S_MOUNT_MAP")) + + +def test_prepare_warns_when_proxy_sidecar_cannot_inject( + configured_project: Path, +) -> None: + config_path = configured_project / ".roar" / "config.toml" + config_path.write_text( + '[k8s]\nenabled = true\nproxy_sidecar = true\n\n[glaas]\nurl = "http://localhost:3001"\n', + encoding="utf-8", + ) + output_path = configured_project / "prepared.yaml" + result = CliRunner().invoke( + k8s, + ["prepare", "-f", "job.yaml", "-o", str(output_path)], + ) + assert result.exit_code == 0, result.output + assert "proxy_sidecar requires" in result.output From f4ced1846fe0dfa7441c789a7f13fb48993a408c Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Thu, 16 Jul 2026 17:14:14 +0000 Subject: [PATCH 29/31] fix(k8s): require exactly one wheel when building the runtime image MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Dockerfile globs dist/, so multiple wheels silently produced an image containing whichever sorted first — typically a stale build. The build script now fails with the wheel list and prints which wheel a successful build used. Co-Authored-By: Claude Fable 5 --- scripts/build_runtime_image.sh | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/scripts/build_runtime_image.sh b/scripts/build_runtime_image.sh index cb75aae0..797a1f74 100644 --- a/scripts/build_runtime_image.sh +++ b/scripts/build_runtime_image.sh @@ -20,6 +20,15 @@ if ((${#wheels[@]} == 0)); then echo "hint: build one first: bash scripts/build_wheel_with_bins.sh" >&2 exit 1 fi +if ((${#wheels[@]} > 1)); then + # The Dockerfile globs dist/ too; with several wheels present the image + # would silently contain whichever sorts first — likely a stale build. + echo "error: ${#wheels[@]} roar_cli wheels in $ROOT_DIR/dist; exactly one is required:" >&2 + printf ' %s\n' "${wheels[@]}" >&2 + echo "hint: remove stale wheels, then rebuild: bash scripts/build_wheel_with_bins.sh" >&2 + exit 1 +fi +echo "▶ Building $TAG from ${wheels[0]##*/}" docker build -f "$ROOT_DIR/deploy/roar-runtime/Dockerfile" -t "$TAG" "$ROOT_DIR" echo "✓ Built $TAG" From 0e7ad68ee9f21941a893ed653f53cc7559e14c7f Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Thu, 16 Jul 2026 17:14:44 +0000 Subject: [PATCH 30/31] docs(k8s): state the local-first boundary and structural matcher Managed capture requires host-side GLaaS at submit time; bundle_dir covers pod-side outages only. Documented explicitly (with the prepare escape hatch) until the deferred local-first design lands, and the match description updated for structural verb parsing. Co-Authored-By: Claude Fable 5 --- docs/developer/k8s-integration.md | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/docs/developer/k8s-integration.md b/docs/developer/k8s-integration.md index 18b9624f..af37d742 100644 --- a/docs/developer/k8s-integration.md +++ b/docs/developer/k8s-integration.md @@ -166,6 +166,17 @@ an undelivered summary), so mid-run streaming failures also reach the fallback; the upfront reachability probe remains as a fast path that skips per-batch POST timeouts when GLaaS is dark. +**Local-first boundary (current phase):** managed Kubernetes capture +requires GLaaS to be reachable *from the submitting host* at submit +time — the fragment session is registered before the manifest is +rewritten, and a failed registration submits the workload +uninstrumented (with a warning). `k8s.bundle_dir` therefore covers +*pod-side* GLaaS outages only; it is not an offline mode. When no GLaaS +is available anywhere, `roar k8s prepare` plus an out-of-band Secret is +the manual path. A true bundle-only/local-first plan (instrument without +pre-registering a shared session, register at ingest/publish time) is +deliberately deferred to the local-first design phase. + **Session TTL renewal**: fragment sessions are registered with `k8s.fragment_session_ttl_seconds` (default 86400, server-capped at 7 days) — training that outlives the TTL used to lose lineage because both @@ -178,10 +189,12 @@ renewal is purely reactive to the 403. ## 2. Flow -1. **Match** (`roar/backends/k8s/submit.py`): binary `kubectl`, verb - `apply|create` anywhere in the argument list (global flags like - `--context` may precede it), `-f` pointing at a manifest containing - exactly one supported workload; gated by `k8s.enabled` (default off). +1. **Match** (`roar/backends/k8s/submit.py`): binary `kubectl`, subcommand + `apply|create` found structurally (known value-taking global flags like + `--context X` are skipped with their values, so a flag value is never + mistaken for the verb), exactly one `-f` pointing at a manifest + containing exactly one supported workload; gated by `k8s.enabled` + (default off). Stdin (`-f -`), URLs, directories, and kustomize sources are outside the matching contract and pass through uninstrumented. 2. **Plan** (same module): pre-registers a GLaaS fragment session (saves the From 86ad3812427bfecef26e3343cae25caf7325462b Mon Sep 17 00:00:00 2001 From: Trevor Basinger Date: Thu, 16 Jul 2026 17:23:53 +0000 Subject: [PATCH 31/31] fix(k8s): isolate in-pod roar state per pod/container/attempt MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit pod_entry kept .roar/roar.db, object-io events, and the run report in the workload's workdir. On a shared or persistent volume (emptyDir shared by co-located wrapped containers, PVC surviving retries) every participant shared one sqlite db and one events file — racing writes and cross-attributing lineage between containers and attempts. roar state now lives in pod-local storage keyed by the identity contract (pod_uid:container:completion_index:restart_attempt); the workload keeps its own cwd. RoarContext.create now honors an existing ROAR_PROJECT_DIR — the contract config loaders and fragment planners already follow — gated on the directory existing so Ray-propagated host paths inside pods never resolve to a phantom project. An unusable state dir falls back to an uninstrumented run rather than failing training. New shared-workdir KIND e2e: two wrapped containers on one emptyDir volume reconstitute as two distinct k8s_tasks, each owning only its own outputs. Co-Authored-By: Claude Fable 5 --- roar/backends/k8s/pod_entry.py | 37 ++- roar/cli/context.py | 15 ++ .../k8s/e2e/test_k8s_shared_workdir.py | 230 ++++++++++++++++++ tests/backends/k8s/unit/test_pod_entry.py | 93 +++++-- tests/unit/test_cli_context.py | 28 +++ 5 files changed, 385 insertions(+), 18 deletions(-) create mode 100644 tests/backends/k8s/e2e/test_k8s_shared_workdir.py diff --git a/roar/backends/k8s/pod_entry.py b/roar/backends/k8s/pod_entry.py index 85213c76..8cd0be61 100644 --- a/roar/backends/k8s/pod_entry.py +++ b/roar/backends/k8s/pod_entry.py @@ -15,6 +15,7 @@ import json import os +import re import subprocess import sys import tempfile @@ -43,9 +44,22 @@ def _run_traced(command: list[str]) -> int: print(f"[roar-k8s] cannot use workdir {workdir}: {exc}", file=sys.stderr) return _run_uninstrumented(command) - if not (workdir / ".roar").is_dir(): + # roar state (db, object-io events, run report) is isolated per + # pod/container/attempt in pod-local storage — never in the workdir. + # A workdir on a shared or persistent volume would otherwise share + # .roar/roar.db across containers, pods, and retries, cross-attributing + # lineage (and racing sqlite). + state_dir = _state_root() + try: + state_dir.mkdir(parents=True, exist_ok=True) + except OSError as exc: + print(f"[roar-k8s] cannot create state dir {state_dir}: {exc}", file=sys.stderr) + return _run_uninstrumented(command) + + if not (state_dir / ".roar").is_dir(): init = subprocess.run( [sys.executable, "-m", "roar", "init", "-n"], + cwd=state_dir, capture_output=True, text=True, check=False, @@ -59,6 +73,9 @@ def _run_traced(command: list[str]) -> int: tracer = str(os.environ.get("ROAR_K8S_TRACER") or "preload").strip() or "preload" child_env = dict(os.environ) + # The CLI context honors ROAR_PROJECT_DIR, so roar run records into the + # isolated state dir while the workload keeps the user's workdir as cwd. + child_env["ROAR_PROJECT_DIR"] = str(state_dir) # Activate the roar_inject.pth sitecustomize in every Python child so # the k8s backend's runtime-import hooks (botocore/aiobotocore object # I/O capture) install themselves; events land next to the local db. @@ -81,12 +98,24 @@ def _run_traced(command: list[str]) -> int: return run.returncode +def _state_root() -> Path: + """Pod-local roar state directory keyed by the task identity contract.""" + override = str(os.environ.get("ROAR_K8S_STATE_DIR") or "").strip() + base = Path(override) if override else Path(tempfile.gettempdir()) / "roar-k8s" + task_id, _task_name = task_identity_from_environment() + return base / re.sub(r"[^A-Za-z0-9_.-]", "-", task_id) + + +def _state_roar_dir() -> Path: + return _state_root() / ".roar" + + def _object_io_events_path() -> Path: - return Path.cwd() / ".roar" / "k8s-object-io.jsonl" + return _state_roar_dir() / "k8s-object-io.jsonl" def _run_report_path() -> Path: - return Path.cwd() / ".roar" / "k8s-run-report.json" + return _state_roar_dir() / "k8s-run-report.json" def _remove_stale_report(report_path: Path) -> None: @@ -159,7 +188,7 @@ def _emit_lineage_best_effort() -> None: with tempfile.TemporaryDirectory(prefix="roar-k8s-") as tmp: bundle_path = Path(tmp) / "roar-fragments.json" export_local_job_fragment_bundle( - roar_dir=Path.cwd() / ".roar", + roar_dir=_state_roar_dir(), output_path=bundle_path, backend_name="k8s", task_id=task_id, diff --git a/roar/cli/context.py b/roar/cli/context.py index ac7ec757..348eb447 100644 --- a/roar/cli/context.py +++ b/roar/cli/context.py @@ -70,6 +70,21 @@ def create(cls, cwd: Path | None = None) -> RoarContext: if cwd is None: cwd = Path.cwd() + # ROAR_PROJECT_DIR pins the project explicitly — the same contract + # the config loaders and fragment planners already honor + # (resolve_project_roar_dir). Gated on the directory existing so + # environments that propagate a host path into a pod (Ray worker + # env) fall back to the cwd walk instead of a phantom project. + override = str(os.environ.get("ROAR_PROJECT_DIR") or "").strip() + if override and Path(override).is_dir(): + project_dir = Path(override) + return cls( + roar_dir=project_dir / ".roar", + repo_root=cls._get_repo_root(project_dir), + cwd=cwd, + is_interactive=sys.stdin.isatty(), + ) + # Get VCS provider and find repo root. repo_root = cls._get_repo_root(cwd) diff --git a/tests/backends/k8s/e2e/test_k8s_shared_workdir.py b/tests/backends/k8s/e2e/test_k8s_shared_workdir.py new file mode 100644 index 00000000..f1c13d5a --- /dev/null +++ b/tests/backends/k8s/e2e/test_k8s_shared_workdir.py @@ -0,0 +1,230 @@ +"""Shared-workdir e2e: co-located wrapped containers must not share state. + +Regression for in-pod state isolation: two wrapped containers sharing a +workdir volume used to share `/.roar/roar.db` (and the fixed +object-io events path), racing sqlite and cross-attributing lineage. +roar state now lives in a pod-local directory keyed by +pod_uid:container:completion_index:restart_attempt, so each container's +fragment carries only its own writes while the shared volume stays a +plain data surface. +""" + +from __future__ import annotations + +import os +import sqlite3 +import subprocess +import sys +import uuid +from pathlib import Path +from typing import Any + +import pytest + +from .conftest import KUBE_CONTEXT, NAMESPACE, TOOLS_BIN, kubectl +from .test_k8s_product_path import wheel_server # noqa: F401 + +pytestmark = [ + pytest.mark.e2e, + pytest.mark.k8s_e2e, + pytest.mark.timeout(900), +] + +SHARED_WORKDIR_MANIFEST = """\ +apiVersion: batch/v1 +kind: Job +metadata: + name: {job_name} + namespace: {namespace} + labels: + app.kubernetes.io/part-of: roar-k8s-e2e +spec: + backoffLimit: 0 + ttlSecondsAfterFinished: 1800 + template: + spec: + restartPolicy: Never + volumes: + - name: work + emptyDir: {{}} + containers: + - name: trainer-a + image: python:3.12-slim + workingDir: /work + volumeMounts: + - name: work + mountPath: /work + command: + - python + - -c + - "open('out-a.bin', 'wb').write(b'alpha' * 8)" + - name: trainer-b + image: python:3.12-slim + workingDir: /work + volumeMounts: + - name: work + mountPath: /work + command: + - python + - -c + - "open('out-b.bin', 'wb').write(b'bravo' * 8)" +""" + + +@pytest.fixture(scope="module") +def shared_workdir_run( + k8s_cluster: None, + glaas_health: str, + wheel_server: dict[str, str], # noqa: F811 + tmp_path_factory: pytest.TempPathFactory, +) -> dict[str, Any]: + project_dir = tmp_path_factory.mktemp("k8s-shared-workdir") + roar_dir = project_dir / ".roar" + roar_dir.mkdir() + (roar_dir / "config.toml").write_text( + "\n".join( + [ + "[glaas]", + 'url = "http://localhost:3001"', + "", + "[k8s]", + "enabled = true", + f'runtime_install_requirement = "{wheel_server["url"]}"', + 'cluster_glaas_url = "http://glaas:3001"', + "wait_for_completion = true", + "wait_timeout_seconds = 300", + "poll_interval_seconds = 2.0", + ] + ) + + "\n", + encoding="utf-8", + ) + + job_name = f"roar-shared-{uuid.uuid4().hex[:6]}" + (project_dir / "job.yaml").write_text( + SHARED_WORKDIR_MANIFEST.format(job_name=job_name, namespace=NAMESPACE), + encoding="utf-8", + ) + + env = dict(os.environ) + env["PATH"] = f"{TOOLS_BIN}{os.pathsep}{env.get('PATH', '')}" + env.pop("GLAAS_URL", None) + env.pop("ROAR_PROJECT_DIR", None) + + try: + completed = subprocess.run( + [ + sys.executable, + "-m", + "roar", + "run", + "kubectl", + "apply", + "--context", + KUBE_CONTEXT, + "-f", + "job.yaml", + ], + cwd=project_dir, + env=env, + capture_output=True, + text=True, + check=False, + timeout=600, + ) + pod_logs = kubectl( + [ + "logs", + "-n", + NAMESPACE, + "-l", + f"job-name={job_name}", + "--all-containers", + "--tail=100", + ], + check=False, + ) + return { + "project_dir": project_dir, + "job_name": job_name, + "exit_code": completed.returncode, + "stdout": completed.stdout, + "stderr": completed.stderr, + "pod_logs": pod_logs.stdout + pod_logs.stderr, + } + finally: + kubectl(["delete", f"job/{job_name}", "-n", NAMESPACE, "--ignore-not-found"], check=False) + kubectl( + [ + "delete", + "secret", + "-n", + NAMESPACE, + "-l", + "app.kubernetes.io/managed-by=roar", + "--ignore-not-found", + ], + check=False, + ) + + +def _describe(run: dict[str, Any]) -> str: + return ( + f"exit={run['exit_code']}\nstdout:\n{run['stdout']}\nstderr:\n{run['stderr']}\n" + f"pod logs:\n{run.get('pod_logs', '')}" + ) + + +def _query(run: dict[str, Any], sql: str, params: tuple = ()) -> list[sqlite3.Row]: + db_path = Path(run["project_dir"]) / ".roar" / "roar.db" + assert db_path.is_file(), f"no roar.db produced\n{_describe(run)}" + conn = sqlite3.connect(db_path) + conn.row_factory = sqlite3.Row + try: + return conn.execute(sql, params).fetchall() + finally: + conn.close() + + +def test_both_containers_reconstitute_as_distinct_tasks( + shared_workdir_run: dict[str, Any], +) -> None: + assert shared_workdir_run["exit_code"] == 0, _describe(shared_workdir_run) + + tasks = _query( + shared_workdir_run, + "SELECT id, command FROM jobs WHERE job_type = 'k8s_task' ORDER BY command", + ) + assert len(tasks) == 2, _describe(shared_workdir_run) + # The recorded command is k8s_task:/ — one per container. + assert tasks[0]["command"].endswith("/trainer-a"), _describe(shared_workdir_run) + assert tasks[1]["command"].endswith("/trainer-b"), _describe(shared_workdir_run) + + +def test_outputs_are_attributed_to_their_own_container( + shared_workdir_run: dict[str, Any], +) -> None: + tasks = _query( + shared_workdir_run, + "SELECT id, command FROM jobs WHERE job_type = 'k8s_task' ORDER BY command", + ) + assert len(tasks) == 2, _describe(shared_workdir_run) + + outputs_by_container: dict[str, set[str]] = {} + for task in tasks: + container = "trainer-a" if task["command"].endswith("/trainer-a") else "trainer-b" + rows = _query( + shared_workdir_run, + "SELECT path FROM job_outputs WHERE job_id = ?", + (task["id"],), + ) + outputs_by_container[container] = {str(row["path"]) for row in rows} + + a_paths = {p for p in outputs_by_container["trainer-a"] if p.endswith(".bin")} + b_paths = {p for p in outputs_by_container["trainer-b"] if p.endswith(".bin")} + assert any(p.endswith("out-a.bin") for p in a_paths), _describe(shared_workdir_run) + assert any(p.endswith("out-b.bin") for p in b_paths), _describe(shared_workdir_run) + # The cross-attribution bug: with a shared /.roar both tasks + # exported from one db and each claimed the other's writes. + assert not any(p.endswith("out-b.bin") for p in a_paths), _describe(shared_workdir_run) + assert not any(p.endswith("out-a.bin") for p in b_paths), _describe(shared_workdir_run) diff --git a/tests/backends/k8s/unit/test_pod_entry.py b/tests/backends/k8s/unit/test_pod_entry.py index f527327d..f5b9b046 100644 --- a/tests/backends/k8s/unit/test_pod_entry.py +++ b/tests/backends/k8s/unit/test_pod_entry.py @@ -1,10 +1,15 @@ -"""Pod-entry workload-safety contract: lineage never blocks training. +"""Pod-entry contracts: lineage never blocks training, state never shared. ``_run_traced`` shells out to ``roar run`` and cannot see past its exit code, so it asks for a run report (ROAR_RUN_REPORT_FILE) and reruns the original command uninstrumented only when roar positively reported a pre-launch setup failure. A missing report is ambiguous — roar may have crashed after the workload ran — and must never trigger a rerun. + +roar state (db, object-io events, run report) lives in a pod-local +directory keyed by the task identity contract, never in the workdir: a +shared or persistent workdir volume must not share .roar/roar.db across +containers, pods, or retry attempts. """ from __future__ import annotations @@ -20,30 +25,46 @@ class _FakeRun: - """Stub subprocess.run capturing invocations of roar run + fallback.""" + """Stub subprocess.run for roar init, roar run, and the fallback.""" def __init__(self, *, roar_exit: int, report: dict[str, Any] | None) -> None: self.roar_exit = roar_exit self.report = report self.calls: list[list[str]] = [] + self.run_envs: list[dict[str, str]] = [] + self.init_cwds: list[str] = [] def __call__(self, args: list[str], **kwargs: Any) -> subprocess.CompletedProcess: self.calls.append(list(args)) + if "roar" in args and "init" in args: + cwd = Path(str(kwargs.get("cwd") or Path.cwd())) + self.init_cwds.append(str(cwd)) + (cwd / ".roar").mkdir(parents=True, exist_ok=True) + return subprocess.CompletedProcess(args, 0) if "roar" in args and "run" in args: - env = kwargs.get("env") or {} + env = dict(kwargs.get("env") or {}) + self.run_envs.append(env) target = env.get("ROAR_RUN_REPORT_FILE") if target and self.report is not None: Path(target).write_text(json.dumps(self.report), encoding="utf-8") return subprocess.CompletedProcess(args, self.roar_exit) return subprocess.CompletedProcess(args, 7) + @property + def fallback_calls(self) -> list[list[str]]: + return [args for args in self.calls if "roar" not in args] + @pytest.fixture def pod_workdir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + workdir = tmp_path / "work" + workdir.mkdir() monkeypatch.chdir(tmp_path) - monkeypatch.setenv("ROAR_K8S_WORKDIR", str(tmp_path)) - (tmp_path / ".roar").mkdir() - return tmp_path + monkeypatch.setenv("ROAR_K8S_WORKDIR", str(workdir)) + monkeypatch.setenv("ROAR_K8S_STATE_DIR", str(tmp_path / "state")) + monkeypatch.setenv("ROAR_K8S_POD_UID", "pod-uid-1") + monkeypatch.setenv("ROAR_K8S_CONTAINER", "trainer") + return workdir def test_setup_error_report_reruns_uninstrumented( @@ -54,9 +75,8 @@ def test_setup_error_report_reruns_uninstrumented( exit_code = pod_entry._run_traced(["python", "train.py"]) - # Second call is the raw command; its exit code (7) is the pod's. - assert len(fake.calls) == 2 - assert fake.calls[1] == ["python", "train.py"] + # The fallback runs the raw command; its exit code (7) is the pod's. + assert fake.fallback_calls == [["python", "train.py"]] assert exit_code == 7 @@ -68,7 +88,7 @@ def test_workload_failure_is_propagated_without_rerun( exit_code = pod_entry._run_traced(["python", "train.py"]) - assert len(fake.calls) == 1 + assert fake.fallback_calls == [] assert exit_code == 3 @@ -82,21 +102,22 @@ def test_missing_report_never_triggers_rerun( exit_code = pod_entry._run_traced(["python", "train.py"]) - assert len(fake.calls) == 1 + assert fake.fallback_calls == [] assert exit_code == 1 def test_stale_report_from_prior_attempt_is_ignored( pod_workdir: Path, monkeypatch: pytest.MonkeyPatch ) -> None: - stale = pod_workdir / ".roar" / "k8s-run-report.json" + stale = pod_entry._run_report_path() + stale.parent.mkdir(parents=True, exist_ok=True) stale.write_text(json.dumps({"exit_code": 1, "setup_error": True}), encoding="utf-8") fake = _FakeRun(roar_exit=1, report=None) monkeypatch.setattr(pod_entry.subprocess, "run", fake) exit_code = pod_entry._run_traced(["python", "train.py"]) - assert len(fake.calls) == 1 + assert fake.fallback_calls == [] assert exit_code == 1 @@ -105,4 +126,48 @@ def test_successful_run_ignores_report(pod_workdir: Path, monkeypatch: pytest.Mo monkeypatch.setattr(pod_entry.subprocess, "run", fake) assert pod_entry._run_traced(["python", "train.py"]) == 0 - assert len(fake.calls) == 1 + assert fake.fallback_calls == [] + + +def test_roar_state_is_isolated_from_the_workdir( + pod_workdir: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + fake = _FakeRun(roar_exit=0, report={"exit_code": 0, "setup_error": False}) + monkeypatch.setattr(pod_entry.subprocess, "run", fake) + + pod_entry._run_traced(["python", "train.py"]) + + # init targeted the identity-scoped state dir, not the workdir; the + # run child was pointed at it via ROAR_PROJECT_DIR. + state_dir = Path(fake.init_cwds[0]) + assert state_dir == pod_entry._state_root() + assert "pod-uid-1" in state_dir.name and "trainer" in state_dir.name + assert fake.run_envs[0]["ROAR_PROJECT_DIR"] == str(state_dir) + assert not (pod_workdir / ".roar").exists() + + +def test_state_dirs_differ_per_container_and_attempt( + pod_workdir: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + trainer = pod_entry._state_root() + monkeypatch.setenv("ROAR_K8S_CONTAINER", "sidecar") + sidecar = pod_entry._state_root() + monkeypatch.setenv("ROAR_K8S_RESTART_ATTEMPT", "1") + retry = pod_entry._state_root() + + assert len({trainer, sidecar, retry}) == 3 + + +def test_unusable_state_dir_falls_back_uninstrumented( + pod_workdir: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + blocker = tmp_path / "blocked" + blocker.write_text("", encoding="utf-8") + monkeypatch.setenv("ROAR_K8S_STATE_DIR", str(blocker)) + fake = _FakeRun(roar_exit=0, report=None) + monkeypatch.setattr(pod_entry.subprocess, "run", fake) + + exit_code = pod_entry._run_traced(["python", "train.py"]) + + assert fake.fallback_calls == [["python", "train.py"]] + assert exit_code == 7 diff --git a/tests/unit/test_cli_context.py b/tests/unit/test_cli_context.py index de015080..51a71f01 100644 --- a/tests/unit/test_cli_context.py +++ b/tests/unit/test_cli_context.py @@ -36,3 +36,31 @@ def test_create_falls_back_to_cwd_roar_dir_when_uninitialized(tmp_path: Path) -> assert ctx.roar_dir == cwd / ".roar" assert ctx.is_initialized is False + + +def test_create_honors_existing_roar_project_dir(tmp_path, monkeypatch) -> None: + project = tmp_path / "project" + (project / ".roar").mkdir(parents=True) + elsewhere = tmp_path / "elsewhere" + elsewhere.mkdir() + monkeypatch.setenv("ROAR_PROJECT_DIR", str(project)) + monkeypatch.chdir(elsewhere) + + ctx = RoarContext.create() + + # Same contract as resolve_project_roar_dir: the env override pins the + # project regardless of cwd (the k8s pod entrypoint relies on this to + # keep roar state out of shared workdirs). + assert ctx.roar_dir == project / ".roar" + + +def test_create_ignores_nonexistent_roar_project_dir(tmp_path, monkeypatch) -> None: + # Ray worker env propagates a HOST project path into pods; when that + # path does not exist the cwd walk must win, not a phantom project. + (tmp_path / ".roar").mkdir() + monkeypatch.setenv("ROAR_PROJECT_DIR", "/nonexistent/host/project") + monkeypatch.chdir(tmp_path) + + ctx = RoarContext.create() + + assert ctx.roar_dir == tmp_path / ".roar"