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/.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/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/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/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..280f2160 --- /dev/null +++ b/deploy/charts/roar-lineage-webhook/values.yaml @@ -0,0 +1,85 @@ +# 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 + +secretCleanup: + # Credential Secrets cannot carry an ownerReference (the workload has no + # 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 * * *" + maxAgeSeconds: 604800 + +resources: {} +nodeSelector: {} +tolerations: [] 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/docs/developer/k8s-integration.md b/docs/developer/k8s-integration.md new file mode 100644 index 00000000..af37d742 --- /dev/null +++ b/docs/developer/k8s-integration.md @@ -0,0 +1,290 @@ +# Kubernetes Integration (Developer) + +## 1. High-level summary + +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. + +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–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, the mutating +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 +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`. +`ROAR_WEBHOOK_PROXY_SIDECAR=true` (+ optional +`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 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 +`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 +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. + +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 +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: + +- **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. 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. + +- **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 +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 +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 +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`, 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 + `.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`): 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 "$@"`; + 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 + 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` + 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. + +`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`, +`runtime_source` (`install`|`image`), `runtime_image`, +`runtime_install_requirement`, `cluster_glaas_url`, `bundle_dir`, +`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 + +- 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/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/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/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/__init__.py b/roar/backends/k8s/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/roar/backends/k8s/attach.py b/roar/backends/k8s/attach.py new file mode 100644 index 00000000..78a7a9a0 --- /dev/null +++ b/roar/backends/k8s/attach.py @@ -0,0 +1,402 @@ +"""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", + "rayjob": "rayjobs.ray.io", + "rayjobs": "rayjobs.ray.io", +} + + +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 + + 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 " + "(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 = _reconstituter_factory(workload_kind)( + 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 _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, +) -> 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)] + + 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): + 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(_is_instrumented(entry) 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/bundles.py b/roar/backends/k8s/bundles.py new file mode 100644 index 00000000..d3bbd5c0 --- /dev/null +++ b/roar/backends/k8s/bundles.py @@ -0,0 +1,135 @@ +"""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 os +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(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]], + *, + container: str = "main", + attempt: str = "0", +) -> Path: + bundle_dir.mkdir(parents=True, exist_ok=True) + 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 + + +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", + "discover_fragment_bundles", + "ingest_fragment_bundles", + "write_fragment_bundle", +] diff --git a/roar/backends/k8s/config.py b/roar/backends/k8s/config.py new file mode 100644 index 00000000..7e5e3888 --- /dev/null +++ b/roar/backends/k8s/config.py @@ -0,0 +1,194 @@ +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" + # 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 = "" + # 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 + # 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) + 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_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="", + 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.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="", + 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, + 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 + +# 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" +""" + + +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..20da4cf2 --- /dev/null +++ b/roar/backends/k8s/fragment_reconstituter.py @@ -0,0 +1,261 @@ +"""Fetch, decrypt, and merge k8s fragment sessions from GLaaS.""" + +from __future__ import annotations + +import base64 +import json +import sqlite3 +import urllib.error +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 +from roar.integrations.glaas import renew_fragment_session + + +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 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 + ) + 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]]: + """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}" + 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") + 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..2cbaca62 --- /dev/null +++ b/roar/backends/k8s/host_execution.py @@ -0,0 +1,196 @@ +"""Host-side execution for kubectl workload 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.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 +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 + + +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_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)), + ) + 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 _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( + { + "workload_kind": submit_context.workload_kind, + "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..cf95a52f --- /dev/null +++ b/roar/backends/k8s/lineage.py @@ -0,0 +1,144 @@ +"""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.""" + 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) + _drop_runtime_staging_noise(payload) + 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 _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): + 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..966f9fe2 --- /dev/null +++ b/roar/backends/k8s/manifest.py @@ -0,0 +1,756 @@ +"""Kubernetes workload manifest rewriting for lineage instrumentation. + +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. 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, replace +from pathlib import Path +from typing import Any + +import yaml # type: ignore[import-untyped] + +from roar.backends.k8s.mount_map import build_container_mount_map, dump_mount_map + +# 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. 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 "$@" +""" + +# 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}}" +python3 -c 'import roar.backends.k8s.pod_entry' || run_fallback "$@" +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" + +_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 +# 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 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]: + 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 _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") + 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, + 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", + ), +) + + +@dataclass(frozen=True) +class K8sManifestRewrite: + documents: list[dict[str, Any]] + workload_kind: str + kubectl_resource: str + 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 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( + 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, + bundle_dir: str = "", + 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. + + 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). + """ + 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( + f"manifest contains {len(workloads)} workloads ({', '.join(names)}); " + "roar instruments exactly one workload per submit" + ) + + 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] + doc = _deep_copy(source_doc) + rewritten_documents[doc_index] = doc + + 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 {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" + + 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, + bundle_dir=bundle_dir, + 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": + 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, + workload_name=workload_name, + contract=contract, + ) + + if not wrapped: + raise K8sManifestError( + 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)" + ) + + 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, + workload_kind=workload.kind, + kubectl_resource=workload.kubectl_resource, + job_name=workload_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) + + +@dataclass(frozen=True) +class _EnvContract: + secret_name: str + requirement: str + cluster_glaas_url: str + tracer: str + parent_job_uid: str + workload_name: str + bundle_dir: str = "" + 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], + *, + 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" + ) + pod_wrapped = False + 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, contract=contract) + 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, + container_name=container_name, + role=ref.role, + mount_map_entries=build_container_mount_map( + ref.spec, container, contract.config_mount_map + ), + ) + wrapped.append(label) + 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. + + 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], + *, + 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 []) + # 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", []) + 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="", + # TrainJob has no visible pod spec; only explicit config entries apply. + mount_map_entries=build_container_mount_map({}, {}, contract.config_mount_map), + ) + return ["node"], [] + + +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] + + +def _inject_env_contract( + env: list[Any], + *, + 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") + 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": 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", 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) + add_value("ROAR_K8S_CONTAINER", 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") + 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 _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 + for key in path: + node = node.get(key) if isinstance(node, dict) else None + return node if isinstance(node, dict) else None + + +def _deep_copy(document: dict[str, Any]) -> dict[str, Any]: + import copy + + return copy.deepcopy(document) + + +__all__ = [ + "WORKLOAD_FAILURE_CONDITIONS", + "WORKLOAD_KINDS", + "WORKLOAD_SUCCESS_CONDITIONS", + "K8sManifestError", + "K8sManifestRewrite", + "WorkloadKind", + "dump_manifest_documents", + "find_workload_documents", + "load_manifest_documents", + "rewrite_manifest_for_lineage", + "workload_kind_for_document", +] 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/object_io.py b/roar/backends/k8s/object_io.py new file mode 100644 index 00000000..80d8c2ce --- /dev/null +++ b/roar/backends/k8s/object_io.py @@ -0,0 +1,220 @@ +"""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: 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 + 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 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: + 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 + 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: 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 + + +__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 new file mode 100644 index 00000000..8eeb1522 --- /dev/null +++ b/roar/backends/k8s/plugin.py @@ -0,0 +1,114 @@ +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.object_io import patch_imported_module +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, + RuntimeImportAdapter, + 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, + ), + # 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",), + 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..8cd0be61 --- /dev/null +++ b/roar/backends/k8s/pod_entry.py @@ -0,0 +1,366 @@ +"""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 re +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) + + # 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, + ) + 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" + 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. + 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 + + +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 _state_roar_dir() / "k8s-object-io.jsonl" + + +def _run_report_path() -> Path: + return _state_roar_dir() / "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 + + +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" + # 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 = ( + 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 + + 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=_state_roar_dir(), + 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")) + + 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"), + "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_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) + + 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: + fragment.setdefault("reads", []).extend(reads) + 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. + + 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 + + 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" + 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" + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/roar/backends/k8s/rayjob.py b/roar/backends/k8s/rayjob.py new file mode 100644 index 00000000..54e10787 --- /dev/null +++ b/roar/backends/k8s/rayjob.py @@ -0,0 +1,230 @@ +"""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" +_USER_SETUP_HOOK_ENV = "ROAR_USER_SETUP_HOOK" + +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") + 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) + 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 + # 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, + ) + 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 {}), + 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 + # 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" + 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" + 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 + # 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/secret_cleanup.py b/roar/backends/k8s/secret_cleanup.py new file mode 100644 index 00000000..707e100a --- /dev/null +++ b/roar/backends/k8s/secret_cleanup.py @@ -0,0 +1,105 @@ +"""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 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 + +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/submit.py b/roar/backends/k8s/submit.py new file mode 100644 index 00000000..a9e9ccbf --- /dev/null +++ b/roar/backends/k8s/submit.py @@ -0,0 +1,440 @@ +"""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 typing import Any + +from roar.backends.k8s.config import load_k8s_backend_config +from roar.backends.k8s.manifest import ( + K8sManifestError, + dump_manifest_documents, + find_workload_documents, + load_manifest_documents, + rewrite_manifest_for_lineage, +) +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") +# 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" + + +@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 + workload_kind: str + kubectl_resource: 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 _kubectl_subcommand(command) not in _KUBECTL_VERBS: + return False + if not _k8s_backend_enabled(): + return False + + manifests = _find_filename_arguments(command) + if not manifests: + return False + 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 + + try: + documents = load_manifest_documents(manifest_path) + except K8sManifestError: + return False + return len(find_workload_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 = resolve_project_roar_dir() + + 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), + parent_job_uid=parent_job_uid, + namespace_override=_find_namespace_argument(command), + **manifest_rewrite_config_kwargs(config), + ) + 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: " + + ", ".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), + workload_kind=rewrite.workload_kind, + kubectl_resource=rewrite.kubectl_resource, + 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)) + + # 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, + ) + + +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 _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]] = [] + index = 2 + while index < len(command): + arg = command[index] + if arg in ("-f", "--filename"): + if index + 1 < len(command): + 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]: + 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 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): + 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 + + 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", + "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/backends/k8s/webhook.py b/roar/backends/k8s/webhook.py new file mode 100644 index 00000000..fd2f844e --- /dev/null +++ b/roar/backends/k8s/webhook.py @@ -0,0 +1,313 @@ +"""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) + proxy_sidecar: bool = False + proxy_upstream: str = "" + + @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, + proxy_sidecar=env.get("ROAR_WEBHOOK_PROXY_SIDECAR", "").strip().lower() + in ("1", "true", "yes"), + proxy_upstream=env.get("ROAR_WEBHOOK_PROXY_UPSTREAM", "").strip(), + ) + + +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() + parent_job_uid = secrets_module.token_hex(4) + secret_name = f"roar-fragment-{session['session_id'][:8]}" + + # 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, + 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, + proxy_sidecar=settings.proxy_sidecar, + proxy_upstream=settings.proxy_upstream, + namespace_override=namespace, + ) + 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"] + 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/roar/backends/k8s/workload_wait.py b/roar/backends/k8s/workload_wait.py new file mode 100644 index 00000000..e09c14d5 --- /dev/null +++ b/roar/backends/k8s/workload_wait.py @@ -0,0 +1,133 @@ +"""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; 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 []: + 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/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/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/backends/ray/roar_worker.py b/roar/backends/ray/roar_worker.py index ef361db9..e3d6f192 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 @@ -601,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], *, @@ -653,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) @@ -1630,6 +1658,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/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/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/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..1c7ea32f --- /dev/null +++ b/roar/cli/commands/k8s.py @@ -0,0 +1,226 @@ +from __future__ import annotations + +import secrets +from pathlib import Path + +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, + dump_manifest_documents, + load_manifest_documents, + rewrite_manifest_for_lineage, +) +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 + + +@click.group("k8s", invoke_without_command=True) +@click.pass_context +def k8s(ctx: click.Context) -> None: + """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, + 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: + 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("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", + "--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)", +) +@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) + 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, + 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") + + click.echo(f"Prepared manifest written to {output_path}") + 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)}") + 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/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/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/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/roar/execution/fragments/sessions.py b/roar/execution/fragments/sessions.py index 66c7586a..bbed6943 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 { @@ -26,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/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/roar/execution/framework/registry.py b/roar/execution/framework/registry.py index 5c0d584b..f083041e 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 @@ -71,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/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/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..470167b1 100644 --- a/roar/integrations/glaas/client.py +++ b/roar/integrations/glaas/client.py @@ -595,6 +595,11 @@ def register_fragment_session( } return self._request("POST", "/api/v1/fragments/sessions", body) + # 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, 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/scripts/build_runtime_image.sh b/scripts/build_runtime_image.sh new file mode 100644 index 00000000..797a1f74 --- /dev/null +++ b/scripts/build_runtime_image.sh @@ -0,0 +1,34 @@ +#!/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 +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" 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/README.md b/tests/backends/k8s/README.md new file mode 100644 index 00000000..f5370685 --- /dev/null +++ b/tests/backends/k8s/README.md @@ -0,0 +1,109 @@ +# k8s Lineage E2E Harness (Tier 1) + +KIND-based harness for pressure-testing roar lineage capture in Kubernetes +training pods (`design-docs/k8s-training-lineage-integration.md`). + +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_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_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_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_phase3.py` — image-staged runtime (hermetic pod start via + 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. 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_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 + 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 + +- 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`/`helm` 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 +# (--with-minio: S3 scenarios; --with-jobset: JobSet e2e; --with-kubeflow: +# 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 + +# 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_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/e2e/test_k8s_distributed.py b/tests/backends/k8s/e2e/test_k8s_distributed.py new file mode 100644 index 00000000..22660617 --- /dev/null +++ b/tests/backends/k8s/e2e/test_k8s_distributed.py @@ -0,0 +1,534 @@ +"""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, timeout: int = 700) -> subprocess.CompletedProcess[str]: + return subprocess.run( + [sys.executable, "-m", "roar", *args], + cwd=cwd, + env=_roar_env(), + capture_output=True, + text=True, + check=False, + timeout=timeout, + ) + + +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, + ) + + +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/e2e/test_k8s_fallback_s3.py b/tests/backends/k8s/e2e/test_k8s_fallback_s3.py new file mode 100644 index 00000000..3b32220a --- /dev/null +++ b/tests/backends/k8s/e2e/test_k8s_fallback_s3.py @@ -0,0 +1,358 @@ +"""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"}), +) + +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 +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" + + 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/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_phase3.py b/tests/backends/k8s/e2e/test_k8s_phase3.py new file mode 100644 index 00000000..6b86f562 --- /dev/null +++ b/tests/backends/k8s/e2e/test_k8s_phase3.py @@ -0,0 +1,390 @@ +"""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) + + +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", + "-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: + 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/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/e2e/test_k8s_rayjob.py b/tests/backends/k8s/e2e/test_k8s_rayjob.py new file mode 100644 index 00000000..154baae5 --- /dev/null +++ b/tests/backends/k8s/e2e/test_k8s_rayjob.py @@ -0,0 +1,245 @@ +"""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 os +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), +] + +# 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 + +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: "{ray_version}" + 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), + ray_version=RAY_IMAGE.split(":")[1].split("-")[0], + ), + 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, 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( + 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") 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/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/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/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/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/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..15224e57 --- /dev/null +++ b/tests/backends/k8s/scripts/bootstrap_k8s.sh @@ -0,0 +1,246 @@ +#!/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). +# 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] [--with-jobset] [--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 + +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 +WITH_WEBHOOK=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 ;; + --with-webhook) WITH_WEBHOOK=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 + +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 + +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_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" + # 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 + 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)" +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/deploy_webhook.sh b/tests/backends/k8s/scripts/deploy_webhook.sh new file mode 100755 index 00000000..7dc3a667 --- /dev/null +++ b/tests/backends/k8s/scripts/deploy_webhook.sh @@ -0,0 +1,72 @@ +#!/usr/bin/env bash +set -euo pipefail + +# 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" +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" +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 "▶ 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 "${SERVICE}-tls" \ + --cert="$CERT_DIR/tls.crt" --key="$CERT_DIR/tls.key" \ + --dry-run=client -o yaml | kubectl_ctx apply -f - + +helm upgrade --install "$RELEASE" "$CHART_DIR" \ + --kube-context "$KUBE_CONTEXT" \ + --namespace "$WEBHOOK_NS" \ + --set image.repository="${RUNTIME_IMAGE%%:*}" \ + --set image.tag="${RUNTIME_IMAGE##*:}" \ + --set glaas.url="$WEBHOOK_GLAAS_URL" \ + --set glaas.clusterUrl="$WEBHOOK_CLUSTER_GLAAS_URL" \ + --set tls.secretName="${SERVICE}-tls" \ + --set tls.caBundle="$CA_BUNDLE" \ + --wait --timeout 180s + +echo "✓ Webhook chart deployed (opt-in: label namespaces with roar.glaas.ai/lineage=enabled)" 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/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_attach.py b/tests/backends/k8s/unit/test_attach.py new file mode 100644 index 00000000..572b7da0 --- /dev/null +++ b/tests/backends/k8s/unit/test_attach.py @@ -0,0 +1,88 @@ +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_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) + 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_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_bundles.py b/tests/backends/k8s/unit/test_bundles.py new file mode 100644 index 00000000..3729010a --- /dev/null +++ b/tests/backends/k8s/unit/test_bundles.py @@ -0,0 +1,137 @@ +from __future__ import annotations + +import json +import sqlite3 +from pathlib import Path + +import pytest + +from roar.backends.k8s.bundles import ( + K8sBundleError, + bundle_filename, + 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("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: + 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" + 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) + assert result.bundles_ingested == 2 + 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() + + +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) + + +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_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 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..c70c034a --- /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 workload"): + _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_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" 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..16105e6e --- /dev/null +++ b/tests/backends/k8s/unit/test_object_io.py @@ -0,0 +1,160 @@ +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 == [] + + +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 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..f5b9b046 --- /dev/null +++ b/tests/backends/k8s/unit/test_pod_entry.py @@ -0,0 +1,173 @@ +"""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 + +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 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 = 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(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( + 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"]) + + # 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 + + +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 fake.fallback_calls == [] + 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 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_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 fake.fallback_calls == [] + 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 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/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 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()) == ( + [], + [], + ) diff --git a/tests/backends/k8s/unit/test_rayjob.py b/tests/backends/k8s/unit/test_rayjob.py new file mode 100644 index 00000000..aee5e787 --- /dev/null +++ b/tests/backends/k8s/unit/test_rayjob.py @@ -0,0 +1,243 @@ +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" + # 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. + 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_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"] + + 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) 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..e24a283d --- /dev/null +++ b/tests/backends/k8s/unit/test_runtime_image.py @@ -0,0 +1,86 @@ +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 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_submit_planning.py b/tests/backends/k8s/unit/test_submit_planning.py new file mode 100644 index 00000000..5e88b33f --- /dev/null +++ b/tests/backends/k8s/unit/test_submit_planning.py @@ -0,0 +1,257 @@ +from __future__ import annotations + +import json +import subprocess +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_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") + 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_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: + 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") + + +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)]) + + +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)] + ) diff --git a/tests/backends/k8s/unit/test_webhook.py b/tests/backends/k8s/unit/test_webhook.py new file mode 100644 index 00000000..dc92330a --- /dev/null +++ b/tests/backends/k8s/unit/test_webhook.py @@ -0,0 +1,194 @@ +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_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( + _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"} 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)]) 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() 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") 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/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: 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_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" 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" 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"}) 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", 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"