diff --git a/src/compute-plane-services/nvsnap/cmd/agent/main.go b/src/compute-plane-services/nvsnap/cmd/agent/main.go index 86317c0212..b7e27605d4 100644 --- a/src/compute-plane-services/nvsnap/cmd/agent/main.go +++ b/src/compute-plane-services/nvsnap/cmd/agent/main.go @@ -167,6 +167,13 @@ func main() { // Empty disables the inject (back-compat: PVC mount still works // once the rox PVC binds; kubelet may stall in // ContainerCreating in the meantime). + // One-downloader-per-hash admission election for chart-shaped model + // workloads (docs/proposals/helm-chart-cache-election.md). Needs L2. + flag.BoolVar(&config.Election.Enabled, "election", false, + "Elect one capture leader per cache hash at admission and gate the other model workers until the capture is promoted (needs L2)") + flag.DurationVar(&config.Election.Deadline, "election-deadline", 0, + "Bound on a leader's cold start plus capture; past it nvsnap-server evicts the gated followers for re-election (default 60m)") + flag.StringVar(&config.Webhook.L2WaitImage, "webhook-l2-wait-image", "", "Image ref for the nvsnap-l2-wait init container injected onto restore pods (nvsnap#147)") diff --git a/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/templates/agent-daemonset.yaml b/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/templates/agent-daemonset.yaml index a9fc51e89b..47c763558e 100644 --- a/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/templates/agent-daemonset.yaml +++ b/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/templates/agent-daemonset.yaml @@ -112,6 +112,11 @@ spec: # stamped in the manifest. - --cachedir-env-file=/etc/nvsnap/cachedir-env/env {{- end }} + {{- if and .Values.agent.election .Values.agent.election.enabled }} + # one-downloader election (values: agent.election) + - --election + - --election-deadline={{ .Values.agent.election.deadline | default "60m" }} + {{- end }} - --rootfs-cm-namespace={{ .Release.Namespace }} # nvsnap#194: OverlayFS scratch root. MUST match the # nvsnap-overlays volumeMount path below and the hostPath diff --git a/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/templates/agent-rbac.yaml b/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/templates/agent-rbac.yaml index 13416bb931..ece0b2b4fd 100644 --- a/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/templates/agent-rbac.yaml +++ b/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/templates/agent-rbac.yaml @@ -86,6 +86,13 @@ rules: - apiGroups: ["snapshot.storage.k8s.io"] resources: ["volumesnapshots"] verbs: ["get", "list", "watch", "create", "delete"] + # volumesnapshotcontents: a namespace-local restore claim on + # snapshot-clone storage is a clone of a pre-provisioned snapshot whose + # content points at the promote's CSI snapshot handle + # (SnapshotClonePromoter.EnsureClaim). + - apiGroups: ["snapshot.storage.k8s.io"] + resources: ["volumesnapshotcontents"] + verbs: ["get", "list", "watch", "create", "delete"] # storageclasses.get: the L2 backend probes the configured SC at # startup (internal/agent/l2_integration.go validateL2StorageClass) # to verify it exists and log its provisioner. Without this rule diff --git a/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/templates/server.yaml b/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/templates/server.yaml index bda5ecea18..a7b200f1b5 100644 --- a/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/templates/server.yaml +++ b/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/templates/server.yaml @@ -64,9 +64,11 @@ rules: # Lease ('nvsnap-promote-') to serialize concurrent promote # attempts on the same hash. Deleted on hash delete so a future # re-capture isn't blocked by a stale lease holder. + # get/list/watch: the election reconciler lists nvsnap-capture- + # Leases to check leader liveness and deletes them on release. - apiGroups: ["coordination.k8s.io"] resources: ["leases"] - verbs: ["delete"] + verbs: ["get", "list", "watch", "delete"] # services.get: the observability proxy (internal/server/observability_proxy.go) # probes whether the grafana / jaeger / prometheus Services exist # before rendering the matching UI nav tile. Without this rule the diff --git a/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/values.yaml b/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/values.yaml index bc5a67fcaf..a1508d9774 100644 --- a/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/values.yaml +++ b/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/values.yaml @@ -107,6 +107,20 @@ agent: # capture (no gemm capture, no prewarm). podCacheDir: "/opt/nvsnap" + # One downloader per cache hash for chart-shaped model workloads + # (docs/proposals/helm-chart-cache-election.md). At admission the + # webhook composes the cache hash of every GPU pod that names a model; + # if a promoted capture exists the pod restores from it, otherwise one + # pod is elected to capture and the rest are held by a scheduling gate + # (no node, no GPU) until nvsnap-server sees the promote and releases + # them. Needs agent.l2. Off until qualified on a cluster. + election: + enabled: false + # Bound on the leader's cold start plus capture. Past it the server + # evicts the gated followers so their controller recreates them and a + # new leader is elected. + deadline: 60m + # cacheEnvTemplate is the cachedir env set (nvsnap #244), rendered into # the nvsnap-cachedir-env ConfigMap when podCacheDir is set. The agent # webhook reads it on CAPTURE inject — edit it (helm value or diff --git a/src/compute-plane-services/nvsnap/deploy/k8s/charts/vllm-workers/Chart.yaml b/src/compute-plane-services/nvsnap/deploy/k8s/charts/vllm-workers/Chart.yaml new file mode 100644 index 0000000000..18b13661d3 --- /dev/null +++ b/src/compute-plane-services/nvsnap/deploy/k8s/charts/vllm-workers/Chart.yaml @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +apiVersion: v2 +name: vllm-workers +version: 0.1.0 +description: >- + Stock vLLM Deployment with N identical GPU workers and no nvsnap markers. + Test chart for the one-downloader election (docs/proposals/helm-chart-cache-election.md). diff --git a/src/compute-plane-services/nvsnap/deploy/k8s/charts/vllm-workers/templates/deployment.yaml b/src/compute-plane-services/nvsnap/deploy/k8s/charts/vllm-workers/templates/deployment.yaml new file mode 100644 index 0000000000..32b08b2928 --- /dev/null +++ b/src/compute-plane-services/nvsnap/deploy/k8s/charts/vllm-workers/templates/deployment.yaml @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +apiVersion: apps/v1 +kind: Deployment +metadata: + name: {{ .Release.Name }} + labels: { app: {{ .Release.Name }} } +spec: + replicas: {{ .Values.replicas }} + selector: + matchLabels: { app: {{ .Release.Name }} } + template: + metadata: + labels: + app: {{ .Release.Name }} + spec: + automountServiceAccountToken: false + {{- with .Values.excludeNodes }} + affinity: + nodeAffinity: + requiredDuringSchedulingIgnoredDuringExecution: + nodeSelectorTerms: + - matchExpressions: + - key: kubernetes.io/hostname + operator: NotIn + values: {{ toJson . }} + {{- end }} + tolerations: + - { key: "nvidia.com/gpu", operator: Exists, effect: NoSchedule } + {{- with .Values.imagePullSecret }} + imagePullSecrets: + - name: {{ . }} + {{- end }} + containers: + - name: vllm + image: {{ .Values.image }} + imagePullPolicy: IfNotPresent + command: ["/bin/bash", "-lc"] + args: + - | + set -e + nohup setsid vllm serve --model {{ .Values.model }} --host 0.0.0.0 --port 8000 --max-model-len {{ .Values.maxModelLen }} --tensor-parallel-size {{ .Values.tensorParallel }} --gpu-memory-utilization {{ .Values.gpuMemoryUtilization }} {{ .Values.extraArgs }} > /vllm.out 2>&1 < /dev/null & + tail -F /vllm.out & + while true; do sleep 30; done + env: + - { name: PYTHONUNBUFFERED, value: "1" } + - { name: HF_HOME, value: "/root/.cache/huggingface" } + - { name: VLLM_ENABLE_V1_MULTIPROCESSING, value: "1" } + - { name: HF_HUB_DISABLE_XET, value: "1" } + {{- with .Values.env }} + {{- toYaml . | nindent 12 }} + {{- end }} + ports: [{ containerPort: 8000, name: http }] + readinessProbe: + httpGet: { path: /v1/models, port: 8000 } + initialDelaySeconds: 20 + periodSeconds: 5 + failureThreshold: {{ .Values.readinessFailureThreshold }} + resources: + limits: { nvidia.com/gpu: {{ .Values.tensorParallel | quote }} } + volumeMounts: [{ name: shm, mountPath: /dev/shm }] + volumes: + - name: shm + emptyDir: { medium: Memory, sizeLimit: {{ .Values.shmSize }} } diff --git a/src/compute-plane-services/nvsnap/deploy/k8s/charts/vllm-workers/values.yaml b/src/compute-plane-services/nvsnap/deploy/k8s/charts/vllm-workers/values.yaml new file mode 100644 index 0000000000..840b4ea06b --- /dev/null +++ b/src/compute-plane-services/nvsnap/deploy/k8s/charts/vllm-workers/values.yaml @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# A plain vLLM chart, deliberately free of nvsnap labels and annotations: +# the election has to work on what customers deploy. Every worker is +# identical; the webhook decides per pod who downloads. +replicas: 2 +image: vllm/vllm-openai:v0.20.0 +imagePullSecret: "" # e.g. nvsnap-pull-secret +model: Qwen/Qwen2.5-32B-Instruct +tensorParallel: 4 # also the GPU request per worker +maxModelLen: 4096 +gpuMemoryUtilization: 0.85 +extraArgs: "" # appended to vllm serve, e.g. "--enforce-eager" +env: [] # extra env, e.g. [{name: HF_TOKEN, valueFrom: {secretKeyRef: {name: hf-token, key: token}}}] +shmSize: 32Gi +readinessFailureThreshold: 400 # x 5 s; cold starts of large models take minutes +excludeNodes: [] # hostnames to keep the workers off diff --git a/src/compute-plane-services/nvsnap/docs/BENCHMARK.md b/src/compute-plane-services/nvsnap/docs/BENCHMARK.md index a5fc566765..3cd1a29835 100644 --- a/src/compute-plane-services/nvsnap/docs/BENCHMARK.md +++ b/src/compute-plane-services/nvsnap/docs/BENCHMARK.md @@ -42,6 +42,9 @@ recompiling. | Model | Engine | Cold | Restore | Speedup | |---|---|---:|---:|---:| +| Qwen3-235B-A22B FP8 (chart, 2 replicas, 2 nodes) | vLLM TP=8 EP | 870 s | 367 s | 2.4x | +| Qwen2.5-32B-Instruct (chart, 2 replicas) | vLLM TP=4 | 325 s | 170 s | 1.9x | +| Llama-3.1-70B-Instruct | vLLM TP=4 | 632 s | 246 s | 2.6x | | DeepSeek-V4-Flash | SGLang TP=8 | 1162 s | 289 s | 4.0x | | gemma-4-31B-it | SGLang | 492 s | 109 s | 4.5x | | gpt-oss-120b | vLLM TP=4 | ~550 s | 171 s | 3.2x | @@ -51,6 +54,16 @@ recompiling. | e5-mistral-7b-instruct | vLLM | 73 s | 90 s | -- | | whisper-large-v3 | NIM (Riva) | 72 s | 74 s | -- | +The two rows above the June table are from dev1 (EKS, NVMesh, 2026-09-25); +the Qwen row is a two-replica Helm chart through the admission election +(`docs/proposals/helm-chart-cache-election.md`): first deploy pays one cold +leader plus capture, every later start of the chart is the restore number. +Phase split for Qwen3-235B, cold to warm: model download and load 641 s to +17 s, torch.compile 94 s to 15 s, init engine 292 s to 48 s; the warm 367 s is +mostly the 230 s prewarm read of 237 GB. Qwen2.5-32B: model download and load 247 s to 6 s, +torch.compile 57 s to 8.5 s, init engine 77 s to 22 s, CUDA graph capture +10 s either way. Llama-70B: model load 481 s to 14 s, compile 65 s to 3 s. + Wins scale with the cold JIT/compile cost; neutral on compile-light and framework-bound workloads. Two implementation details enable the cache reuse: diff --git a/src/compute-plane-services/nvsnap/docs/proposals/helm-chart-cache-election.md b/src/compute-plane-services/nvsnap/docs/proposals/helm-chart-cache-election.md new file mode 100644 index 0000000000..590e9013f7 --- /dev/null +++ b/src/compute-plane-services/nvsnap/docs/proposals/helm-chart-cache-election.md @@ -0,0 +1,298 @@ +# Helm chart cache reuse: one downloader per hash + +Goal: when a chart deploys N model workers, exactly one downloads and +initializes; the rest start from its cache. If the cache already exists, +every worker starts from it. No change to the customer's chart. + +Status: implemented on branch `nvsnap/helm-election`, verified on dev1 2026-09-25 (results below). +Supersedes the `nvsnap.io/restore-from: "auto"` annotation, which required +a template opt-in and could not be a capture source (issue #2099). + +```mermaid +sequenceDiagram + participant RS as ReplicaSet + participant WH as webhook (agent) + participant L as Lease nvsnap-capture- + participant A as agent watcher + participant S as nvsnap-server + RS->>WH: create pod 1 (GPU + model) + WH->>WH: hash = compose(spec); stamp nvsnap.io/hash + WH->>L: create (holder = pod 1 UID) + L-->>WH: created + WH-->>RS: leader: capture decoration, capture label + RS->>WH: create pod 2..N + WH->>L: create + L-->>WH: AlreadyExists + WH-->>RS: follower: rox mount + schedulingGate, label nvsnap.io/gated + A->>A: pod 1 Ready, capture, promote rox- + A->>S: pvc-state ready + S->>RS: remove gate on pods with label hash= + RS->>RS: pods 2..N schedule, mount rox, start warm +``` + +## Decision at admission + +For every pod created, the webhook runs this before any annotation logic. + +1. Classify. A pod is a model worker when it requests `nvidia.com/gpu` + and the main container has an inferable model identity: `--model` or + `--model-path` in args, `HF_MODEL_ID`/`MODEL_ID`/`NIM_MODEL_NAME` env, + or a NIM image. Anything else (routers, frontends, etcd, nats) is not + touched. An explicit `nvsnap.io/restore-from` still takes the old path. +2. Hash once. Compose the hash from the incoming spec (image tag, model + identity, command, args, cache-relevant env, driver major, format + version). Stamp it as annotation `nvsnap.io/hash` (full) and label + `nvsnap.io/hash` (short, selectable). The watcher reads the annotation + instead of recomposing from the live pod, which removed the mismatch + between admission-time and capture-time inputs (image digest, injected + env). +3. Ready: a usable manifest exists and the rox PVC is Bound. Restore + decoration, same as an explicit-hash restore today. Done. +4. Elect. Try to create Lease `nvsnap-capture-` in nvsnap-system + with the pod UID as holder and a deadline. Create is atomic, so N + simultaneous admissions produce one winner. + - Winner (leader): capture decoration (cache env, `/opt/nvsnap` + emptyDir), label `nvsnap.io/capture=true`, annotation + `nvsnap.io/role=leader`. Runs at once. + - Loser (follower): restore decoration against the deterministic claim + name `rox-` even though it is not Bound yet, plus + `schedulingGates: [nvsnap.io/wait-for-cache]`, label + `nvsnap.io/gated=true`, annotation `nvsnap.io/role=follower`. Holds no + node and no GPU. +5. Fail open. Any error in classify, compose, or the Lease call admits the + pod unchanged. The webhook stays an optimization. + +Followers need ReadOnlyMany storage (shared-volume or snapshot-clone with +`readOnlyMany: true`). On per-pod-clone storage there is no claim to name +ahead of time; followers are admitted unchanged and start cold. + +## Release and failure + +nvsnap-server already receives the promote state per hash. + +- `ready`: list pods with labels `nvsnap.io/hash=,nvsnap.io/gated=true` + across namespaces, patch `spec.schedulingGates: []` and + `nvsnap.io/gated=false`. The rox is Bound at this point, so the follower + schedules, mounts it, and starts warm. The `nvsnap-l2-wait` init container + stays as a second check and exits at once. +- `failed`, or the leader pod is gone or Failed before commit, or the Lease + deadline passes: the server reconciler deletes the Lease and deletes the + gated followers that have a controller owner. Their controller recreates + them, admission runs again with no Lease, and a new leader is elected. + Gated pods without an owner are left in place and logged. A follower's + volumes reference a PVC that will never exist, so recreation is the only + way to change its fate; pod volumes are immutable after creation. + +The leader Lease is not renewed. Its deadline is the admission time plus a +configurable bound (default 60 minutes) that covers a cold start plus a +capture; the reconciler treats an expired Lease like a failed leader. + +## What already exists + +| Piece | Where | +|---|---| +| Hash composition | `internal/rootfsonly/composer.go` | +| Capture decoration | `internal/webhook/cachedir.go` `cacheDirCapturePatches` | +| Restore decoration | `internal/webhook/cachedir.go` `tryL2CacheDir` | +| Per-hash Lease pattern | `internal/checkpointstore/percapture_pvc.go` `acquireLease` | +| Promote state endpoint | `internal/server/sources.go` `updatePVCPromoteStateByHash` | +| Server reconcile loop | `internal/server/reconciler.go` | +| Wait init container | `internal/webhook/l2_mount.go` `buildL2WaitContainer` | + +## What changes + +- `internal/rootfsonly/composer.go`: export `InferModelID`, add the env + sources. +- `internal/election` (new): classifier and `LeaseElector`. +- `internal/webhook`: the decision above, follower decoration built from a + claim name rather than a `Mount` call, pod identity logged as + `generateName+UID` because Deployment pods have no name at admission. +- `internal/rootfsonly/watcher.go`, `orchestrator.go`: honour the stamped + hash (`CaptureRequest.Hash`). +- `internal/server`: ungate on `ready`, leader liveness in the reconciler. +- Helm: `agent.election.enabled` (default false until qualified), + `agent.election.leaseTimeout`; server RBAC gains `leases` get/list/watch. + +## Disaggregated workers share one cache + +Dynamo prefill and decode workers run the same image and download the +same model; they differ only in a role flag (`--is-prefill-worker`, +`--is-decode-worker`, `--disaggregation-mode`) and in the KV-transfer and +discovery wiring (`--kv-transfer-config`, `DYN_*`, `ETCD_ENDPOINTS`, +`NATS_SERVER`). The hash leaves those out (`stripRoleFlags`, +`internal/rootfsonly/composer.go`), so both roles elect one leader and share +one rox. The model tree is identical across roles and mounts read-only; +compile caches live in the per-pod writable shadow, so a role that needs +different kernels recompiles into it. Flags that change the download +(`--model`, `--revision`, `--tokenizer`, quantization) stay in the hash, +because a pod restoring from a tree that lacks its files would fail on the +read-only mount. + +## Claims across namespaces + +NVCF runs every chart in its own namespace, so the capture and the +restores usually live apart, and a PVC is namespaced. The promoter grows +`EnsureClaim(hash, ns)`: make `rox-` exist in `ns` from the promoted +artifact. Two strategies, the same shape NVCA uses for one model volume +across tenant namespaces (`pkg/storage/modelcache.go`): + +- shared-volume (NVMesh, EFS, Filestore): one more secondary static PV, + `nvsnap-ro-pv--`, with the CSI handle rewritten for that + namespace, pre-bound to a `rox-` claim there. Zero copy. +- snapshot-clone with ReadOnlyMany (Hyperdisk ML): a VolumeSnapshot is + namespaced and a clone must name one in its own namespace, so the + promote's snapshot handle is re-exposed there through a pre-provisioned + VolumeSnapshotContent + VolumeSnapshot pair (Retain), then cloned. +- per-pod clone: no shared claim exists to reproduce; `ErrUnsupported`, + the pod starts cold. + +Two call sites. `Mount` on a claim miss mints the claim and retries, which +covers a restore admitted after the promote in any namespace. The promote +itself, before publishing `ready`, mints the claim in every namespace that +already holds pods stamped with the hash (gated followers admitted before +the promote), so they find it bound when the server releases them. All +minted objects carry `nvsnap.io/hash-short` and `nvsnap.io/namespace`; +`Delete` lists and reaps them. + +A restore namespace also needs the agent token Secret and the restore-pod +NetworkPolicy, both already fanned out by `agent.l2.restoreNamespaces`. +That policy selects every pod in the namespace and allows egress only to +nvsnap-server, so in a namespace with no other egress allows it becomes +default-deny for everything else, DNS included, and a restored vLLM cannot +resolve huggingface.co (it lists the repo even with a warm cache). NVCF +function namespaces carry NVCA's egress allows; a bare test namespace does +not. Seen on dev1 2026-09-25; the chart should allow DNS in that policy or +document the requirement. + +## Grove and the scheduling gate + +Under the Dynamo operator, Grove gang scheduling owns `schedulingGates` and +rewrote the follower's list at creation (managedFields: `grove-operator`, +same second), removing `nvsnap.io/wait-for-cache`. The follower then sits +`Unschedulable` on volume binding instead, because `rox-` does not +exist yet; it still holds no node and no GPU, and schedules when the promote +creates the claim. Same outcome, noisier events. Stock charts keep the gate. + +## Verification + +Unit: classifier matrix, election win/lose/error, follower patch shape +(gate, claim name, no GPU held), watcher uses the stamped hash, server +ungate and eviction against a fake clientset. + +E2E: `scripts/test-election-e2e.sh` drives `deploy/k8s/charts/vllm-workers` +(a stock vLLM Deployment with no nvsnap markers; model, TP and replicas are +values). Historically the TinyLlama TP=2 chart from the scale-up test, `replicas=1` then +`scale 2`, and `replicas=2` from the start. Expected: one capture, second +pod `SchedulingGated` until `ready`, then Ready without a download. Then +`helm uninstall` and reinstall: both pods restore. + +## Results, dev1 2026-09-25 + +Stock Deployment chart, TinyLlama-1.1B on vLLM v0.20.0 TP=2, no nvsnap +labels or annotations in the template, agent v0.2.75-election2, server +v0.0.32-election2, NVMesh shared-volume storage. + +``` +helm install replicas=2 14:57:37 webhook (same second): leader + follower, hash c8bbf555 + follower: SchedulingGated, no node, claim rox-c8bbf555... +leader Ready +70s +server: capture promoted; followers released +124s (released=1) +follower Ready +53s after release, node different from the leader + 0 download lines; inits nvsnap-l2-wait, nvsnap-seed-cache, nvsnap-prewarm + serves "The capital of France is" -> " Paris." +kubectl scale replicas=3 third pod admitted role=restore, no election; Ready +53s +helm uninstall + install both pods role=restore; both Ready +50s +``` + +Cold start of the same pod on the same node without nvsnap: 94 s. + +Two things the first runs taught: + +- A pod has no UID and no name when a mutating webhook sees its CREATE. + The Lease holder is an id the webhook mints and stamps on the leader as + `nvsnap.io/election-id`; the unit fixture had carried a UID and hid this. +- A privileged test pod without `CUDA_VISIBLE_DEVICES` uses whichever GPUs + it likes, so the device plugin hands "free" GPUs that are 70 GiB full to + the next pod. The election leader landed on such a node once and vLLM + refused to start. Node hygiene, not an election failure; the run was + repeated with that node excluded. + +### Cross-namespace, dev1 2026-09-25 + +Capture only in `nvsnap-system` (hash c8bbf555); the same stock chart +installed in `nvsnap-xns`, replicas=2, admitted in the same second, agent +v0.2.75-election6, NVMesh shared-volume: + +``` +webhook "L2 claim minted in restore namespace" namespace=nvsnap-xns, then both + "election: promoted capture exists; restoring" +claim nvsnap-xns/rox-c8bbf555... Bound to nvsnap-ro-pv-c8bbf555...-1cfa3cd7 +PV handle single-zone-cluster:csi-...:nvsnap-xns, ro, [ro norecovery nouuid] +mount /dev/nvmesh/csi-... on /opt/nvsnap type xfs (ro,nouuid,norecovery), 2.1G model tree +pods both role=restore, Ready +60s, 0 download lines, serve " Paris." +``` + +The first attempt exposed a create-then-label race between the two +admissions ("the object has been modified"); the per-namespace PV is now +created already labelled and the fresh concurrent mint passed. + +### Dynamo disaggregated, dev1 2026-09-25 + +NVCF Dynamo sample (vllm-runtime 1.1.1, Qwen3-0.6B, frontend + prefill + +decode). With the role-neutral hash both workers compose 18313f93: prefill +elected leader, decode follower, frontend ignored. The gang did not +schedule: `kai-scheduler` places the podgang `myllm-0` as a unit and the +follower is unschedulable until `rox-` exists, so the leader is held +with it. See "Grove and the scheduling gate"; the resolution for gang +scheduling is a design decision recorded in issue #2099. + +### Qwen2.5-32B-Instruct TP=4 chart, dev1 2026-09-25 + +Stock Deployment, vLLM v0.20.0, torch.compile on (no enforce-eager), +replicas=2 on one node, agent v0.2.75-election6, NVMesh. Capture 65.7 GB, +805 files. + +``` +first deploy, replicas=2 + t+0 leader + follower admitted (same second); follower gated, no GPU + t+325s leader Ready (cold) + t+439s capture copied + promoted, follower released + t+600s follower Ready (161 s after release), 0 downloads, serves + one download of 65.7 GB instead of two +uninstall + reinstall, replicas=2 + t+170s both Ready (prewarm init 83 s each), both role=restore + +engine phases cold (leader) warm (follower) + model load + download 247 s 6.4 s + torch.compile 57 s 8.5 s (graphs loaded from cache, 2.7 s each range) + init engine total 77 s 22 s + CUDA graph capture 10 s 10 s (not cacheable) +``` + +### Qwen3-235B-A22B-Instruct-2507-FP8 TP=8 chart, two nodes, dev1 2026-09-25 + +Stock chart (`deploy/k8s/charts/vllm-workers`, `--enable-expert-parallel`; +the FP8 block quantization does not split the MoE intermediate eight ways +without it), replicas=2, one worker per 8x H100 node, agent v0.2.75-election6, +NVMesh. Capture 237.0 GB, 2006 files. + +``` +first deploy, replicas=2 + t+0 leader + follower admitted; follower gated, no GPU + t+870s leader Ready (cold: 641 s download+load, 94 s compile, 292 s init, 42 s graphs) + t+~1165s capture committed + promoted, follower released + t+1524s follower Ready (359 s after release), 0 downloads, serves + one download of 237 GB instead of two +uninstall + reinstall, replicas=2, both nodes reading the same rox at once + t+367s both Ready; prewarm init 230 s each (237 GB), engine 17 s load, 15 s compile, 48 s init + +engine phases cold (leader) warm (follower / reinstall) + model load + download 641 s 16-17 s + torch.compile 94 s 15 s + init engine total 292 s 48 s + CUDA graph capture 42 s 28 s +``` + +At this size the warm start is the volume read: 230 s of the 367 s is the +prewarm sweep at ~1 GB/s, shared by two readers. The engine work that the +cache removes went from 641+94 s to 17+15 s. diff --git a/src/compute-plane-services/nvsnap/docs/proposals/helm-shared-model-volume.md b/src/compute-plane-services/nvsnap/docs/proposals/helm-shared-model-volume.md new file mode 100644 index 0000000000..f25b2d7763 --- /dev/null +++ b/src/compute-plane-services/nvsnap/docs/proposals/helm-shared-model-volume.md @@ -0,0 +1,160 @@ +# Helm functions: every expensive artifact produced once per cluster + +Goal: for any Helm chart of GPU model workers, on any registry, any +deployment shape and any scheduler, the model is downloaded once per cluster +and the compile artifacts (torch.compile, Triton, FlashInfer, DeepGEMM, CUDA +JIT) are built once per cluster, and every other pod consumes them. No chart +change. No pod is ever held back from scheduling. + +Status: design, 2026-09-26. Supersedes the gate-and-promote path of +`helm-chart-cache-election.md` for Helm functions. Issue #2099. + +```mermaid +sequenceDiagram + participant P0 as pod-0 (writer) + participant P1 as pod-1..N (readers, any node, any namespace) + participant WH as webhook + participant A as agent + participant V as model volume nvsnap-model- + WH->>WH: identity, landing path, group; elect writer (Lease) + WH-->>P0: init "download" -> V (rw), main mounts V ro + WH-->>P1: landing path -> V (DFS) or emptyDir + wait init (NVMesh) + Note over P0,P1: all pods schedule immediately + P0->>V: download; init exits 0 + A->>V: complete: marker (DFS) / ro attach + bind into P1 (NVMesh) + P1->>P1: wait init sees marker, engine starts + Note over P0,P1: engines start together; multi-node groups form as today +``` + +## Two artifacts, two lifecycles + +| Artifact | Producer | Written when | Consumers may run concurrently with producer? | +|---|---|---|---| +| Model tree | one download step | before any engine starts; immutable afterwards | no: it is complete before consumers need it | +| Compile caches | every rank of every pod | during engine init | yes: ranks of one multi-node instance compile at the same time | + +This split is the whole design. The model is a write-once file set, so it +can be produced by one pod and attached read-only by all others on any +storage, including block. Compile caches are produced concurrently by +ranks that cannot wait for each other, so sharing them within one first +start needs a shared writable filesystem; sharing them across starts does +not. + +## The five mechanisms + +1. Identity and landing path (webhook). Identity is a URI derived from the + pod: `hf://org/repo[@rev]` from `--model`, `--model-path`, `--model=`, + positional `vllm serve `, `HF_MODEL_ID`, `MODEL_ID`; `ngc://org/team/ + model:ver` from an init's `NGC_MODEL_NAME` or `ngc registry model + download-version`; `s3://` from `aws s3 sync`; KServe `storageUri`; + `nim://image@profile`. Members of a group with no identity of their own + (LWS `--headless` workers) inherit from the group's template via owner + references (`leaderworkerset.sigs.k8s.io/group-key`, StatefulSet, + `grove.io/podgang`). Landing path is where the download writes: + `HF_HOME`, `NIM_CACHE_PATH`, the init's dest, `/mnt/models`. If the + volume there is already a PVC, hostPath or OCI image, skip the pod. + Role flags and wiring env stay out of the hash (`stripRoleFlags`). + +2. One download step per identity per cluster (webhook + Lease). If the + chart downloads in an init container, that init is the download step. + If the engine downloads itself, the webhook injects an init + (`huggingface-cli download ` into the landing path) and starts the + engine offline. The Lease `nvsnap-model-` elects the writer among + concurrent admissions; the writer's init runs the download, its main + container mounts the model read-only. Every other pod is a reader and + its download init is replaced by a wait. + +3. Model volume per identity, immutable after download (agent + storage + profile). Distributed filesystem: one RWX volume; writer and readers + mount it at admission; completion is a marker file the writer's init + writes on exit 0. NVMesh: the writer's PVC is the artifact; on the + init's exit 0 the agent creates the read-only static PV and marks it + complete; readers attach it read-only. Later deployments, other + namespaces (`EnsureClaim`, done) and new versions of the function all + attach the same volume. There is no capture copy of the model anymore. + +4. Readers never block scheduling (webhook + agent). On a distributed + filesystem the RWX claim exists and is bound at admission, so the pod + schedules. On NVMesh, before the download is complete, a reader cannot + reference a bindable claim, so it gets an emptyDir at the landing path + plus `nvsnap-wait-model`; once complete, the agent on the reader's node + attaches the read-only volume (mount-holder, existing) and bind-mounts + it over the emptyDir, then drops the marker the wait init is polling. + After completion, NVMesh readers reference the read-only claim directly. + No pod ever needs network access to nvsnap; function namespaces block it. + +5. Compile caches (webhook env + agent). All caches are redirected to a + cache location keyed by image digest plus identity plus role-neutral + args. Distributed filesystem: `/cache//`, read-write for + every pod; the engines' `filelock` and atomic replace make identical + compiles converge; `cacheMode: shadow` in the profile keeps a per-pod + writable copy seeded from it where the operator does not trust the + filesystem's locks (Lustre needs `-o flock`; the agent checks). NVMesh: + each pod compiles into a local emptyDir on the first start; after the + writer is Ready the agent captures its cache dir into a read-only cache + volume `nvsnap-cache-` (the existing capture path, now caches + only, hundreds of MB); every later pod mounts it read-only with a + writable shadow (today's seed init). + +## Every scenario, same mechanisms + +Shapes: D = single-pod Deployment replicas; M = multi-pod instance (LWS, +StatefulSet, Dynamo podgang, prefill+decode). Storage: DFS, NVMesh. State: +first = nothing exists; concurrent = download in flight; later = complete. + +| Shape / storage / state | Download | Compile | Pods held? | +|---|---|---|---| +| D, DFS, first | 1 (writer init); readers wait on marker then start | once, shared cache, engines lock | no | +| M, DFS, first | 1; readers wait on marker; group forms when all engines start | ranks share cache; duplicates limited to races within one start | no | +| any, DFS, concurrent or later | 0 | 0 | no | +| D, NVMesh, first | 1; readers get bind-mounted ro volume on completion | writer compiles; readers compile locally once, then cache volume exists | no | +| M, NVMesh, first | 1; readers bind-mounted on completion; group forms | each pod compiles once (concurrent ranks, no shared fs) | no | +| any, NVMesh, later | 0 (ro claim at admission) | 0 (cache volume ro + shadow) | no | +| any, other namespace, later | 0 (`EnsureClaim`) | 0 | no | +| neither storage | nvsnap does nothing for Helm | | no | + +The one row that does not reach "once per cluster" is M on NVMesh on the +first start of a model, for compile only: each pod of that instance builds +its kernels once, because its ranks need the binaries while the other +pod's ranks are still building them and there is no shared filesystem +between them. Everything after that first start is a full hit. + +## Failure behaviour (decided: always fall back, never deadlock) + +| Failure | Effect | Recovery | +|---|---|---| +| writer dies before complete | readers' wait reaches the Lease deadline | readers download locally (NVMesh) or into the volume (DFS; per-file atomic); Lease expires; next admission elects a new writer | +| volume full | writer's download fails, init restarts | same as above; retention by last use with a size budget is part of this design's follow-up, since a full volume fails every writer | +| agent down on a reader node (NVMesh) | no bind arrives | wait deadline, local download | +| writer pod restarts after complete | volume immutable, unaffected | none needed | +| identity changes (revision, quantization, image) | different URI or cache key | separate volume; old one ages out | +| gang scheduler | readers always schedulable (RWX bound, or emptyDir); writer PVC binds in seconds on Immediate storage classes | none needed | + +## What changes in the code + +Stays: classifier (extended per mechanism 1), role-neutral hash, Lease +election (elects the writer), `EnsureClaim`, storage profiles, cache env +injection and seed init, mount-holder and bind injection (L1), the server +reconciler, `vllm-workers` chart and runner. + +New: identity from init containers and group inheritance; download-init +injection for engine-internal downloads; init wrapping for marker and +wait; per-identity model volume created at admission from the profile's +class (RWX on DFS, RWO writer PVC on NVMesh); agent completion handler +(init exit 0 -> marker / ro PV + bind); cache volume capture (caches only) +on NVMesh; `cacheMode`; last-use labels. + +Removed for Helm: `schedulingGates`, promote-to-ROX of the whole tree, +`nvsnap-l2-wait`, `restore-from`. + +## Build order and verification + +1. Identity, landing path, group inheritance; tests from real specs + (prd11 NGC function, Dynamo sample, LWS example, plain Deployment). +2. Model volume and writer path on both storage classes; readers on DFS + (marker wait). +3. NVMesh readers: agent completion, ro attach, bind into emptyDir. +4. Compile caches: shared on DFS, cache-volume capture on NVMesh. +5. Retire the gate; e2e on dev1 in all matrix rows that dev1 can host + (NVMesh; DFS stands in with an NFS class), each measured cold, first + deploy with two pods per instance, redeploy, second namespace. diff --git a/src/compute-plane-services/nvsnap/internal/agent/BUILD.bazel b/src/compute-plane-services/nvsnap/internal/agent/BUILD.bazel index c700911aa0..9f0c3148b0 100644 --- a/src/compute-plane-services/nvsnap/internal/agent/BUILD.bazel +++ b/src/compute-plane-services/nvsnap/internal/agent/BUILD.bazel @@ -49,6 +49,7 @@ go_library( "//src/compute-plane-services/nvsnap/internal/criu", "//src/compute-plane-services/nvsnap/internal/criu/mountinfo", "//src/compute-plane-services/nvsnap/internal/cuda", + "//src/compute-plane-services/nvsnap/internal/election", "//src/compute-plane-services/nvsnap/internal/metrics", "//src/compute-plane-services/nvsnap/internal/objectstore", "//src/compute-plane-services/nvsnap/internal/rootfsonly", @@ -90,6 +91,7 @@ go_test( "checkpoint_plan_a_test.go", "fsstore_test.go", "l2_integration_test.go", + "l2_profile_prewarm_test.go", "l2_promote_async_test.go", "l2_writer_test.go", "nim_backend_test.go", diff --git a/src/compute-plane-services/nvsnap/internal/agent/agent.go b/src/compute-plane-services/nvsnap/internal/agent/agent.go index 1c9f5ee23d..5983b54d20 100644 --- a/src/compute-plane-services/nvsnap/internal/agent/agent.go +++ b/src/compute-plane-services/nvsnap/internal/agent/agent.go @@ -43,6 +43,7 @@ import ( "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/containerd" "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/criu" "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/cuda" + "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/election" "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/metrics" "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/objectstore" "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/runtime" @@ -158,6 +159,11 @@ type Config struct { // empty. L2 L2BackendConfig + // Election turns on the one-downloader-per-hash admission election + // for chart-shaped model workloads. Needs L2 (the followers mount the + // promoted rox); ignored when L2 is off. + Election ElectionConfig + // Replication is the opt-in cross-cluster replication config (the L4 // tier). See docs/design/cross-cluster-replication.md. When // Replication.ObjectStore.Provider AND HomeBucket are both non-empty, @@ -205,6 +211,17 @@ type ObjectStoreConfig struct { PeerBuckets []string } +// ElectionConfig configures the admission election +// (docs/proposals/helm-chart-cache-election.md). +type ElectionConfig struct { + // Enabled turns the election on. Default off until qualified. + Enabled bool + // Deadline bounds a leader's cold start plus capture; past it the + // server evicts the gated followers for re-election. Zero means + // election.DefaultDeadline. + Deadline time.Duration +} + // L2BackendConfig is the per-capture PVC L2 backend (nvsnap#63). See // docs/L2-PVC-CRIU-DESIGN.md. type L2BackendConfig struct { @@ -274,6 +291,9 @@ type Agent struct { // its prewarm policy; the promoter strategy is already baked into // l2Backend. l2Profile *checkpointstore.StorageProfile + // elector is the admission election, built with the L2 backend when + // Election.Enabled; nil keeps the webhook on its explicit paths. + elector election.Elector // kubeClient is the shared K8s API client used by the rootfs-only // capture watcher AND the admission-webhook cascade-fetch path diff --git a/src/compute-plane-services/nvsnap/internal/agent/l2_integration.go b/src/compute-plane-services/nvsnap/internal/agent/l2_integration.go index 8f3ba19ef9..d73b2e3b39 100644 --- a/src/compute-plane-services/nvsnap/internal/agent/l2_integration.go +++ b/src/compute-plane-services/nvsnap/internal/agent/l2_integration.go @@ -38,6 +38,7 @@ import ( "k8s.io/client-go/rest" "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/checkpointstore" + "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/election" ) // storageProfilesConfigMap is the optional per-cluster overlay that @@ -176,6 +177,10 @@ func (a *Agent) startL2Backend(_ context.Context, cfg L2BackendConfig) (checkpoi // behavior) — back-compat for clusters with no profile match. promoter, profile := resolveL2Promoter(context.Background(), kc, dyn, cfg.StorageClass, cfg.Namespace, log) a.l2Profile = profile + if a.config.Election.Enabled { + a.elector = &election.LeaseElector{KubeClient: kc, Namespace: cfg.Namespace, Deadline: a.config.Election.Deadline} + log.WithField("deadline", a.config.Election.Deadline).Info("admission election enabled (one downloader per hash)") + } // SnapshotClass is only meaningful for the snapshot-clone strategy. // Shared-volume backends (NVMesh/EFS/Filestore) never snapshot — they diff --git a/src/compute-plane-services/nvsnap/internal/agent/webhook_integration.go b/src/compute-plane-services/nvsnap/internal/agent/webhook_integration.go index dbc0fabd01..ef7c51d914 100644 --- a/src/compute-plane-services/nvsnap/internal/agent/webhook_integration.go +++ b/src/compute-plane-services/nvsnap/internal/agent/webhook_integration.go @@ -152,6 +152,9 @@ func (a *Agent) startWebhook(ctx context.Context, cfg WebhookConfig, backend che // Storage profile of the L2 StorageClass: the cachedir restore // reads its prewarm policy (on/off, reader count). nil = defaults. StorageProfile: a.l2Profile, + // One-downloader election for chart pods; nil when off or L2 is + // off (docs/proposals/helm-chart-cache-election.md). + Elector: a.elector, Composer: &rootfsonly.HashInputComposer{ CUDADriverMajor: a.config.RootfsCapture.CUDADriverMajor, }, diff --git a/src/compute-plane-services/nvsnap/internal/checkpointstore/mounter.go b/src/compute-plane-services/nvsnap/internal/checkpointstore/mounter.go index 77956c2460..fd921b7fe8 100644 --- a/src/compute-plane-services/nvsnap/internal/checkpointstore/mounter.go +++ b/src/compute-plane-services/nvsnap/internal/checkpointstore/mounter.go @@ -19,6 +19,7 @@ package checkpointstore import ( "context" + "errors" corev1 "k8s.io/api/core/v1" ) @@ -84,3 +85,18 @@ type Backend interface { Store Mounter } + +// ErrUnsupported means the storage strategy cannot provide what was asked +// (for example a namespace-local shared claim on per-pod-clone storage). +var ErrUnsupported = errors.New("checkpointstore: unsupported by this storage strategy") + +// PendingMounter is implemented by backends that can name the restore +// claim for a hash before it exists. The election decorates follower pods +// against that name and gates their scheduling until the promote binds +// it, so the mount spec has to be known ahead of the artifact. +type PendingMounter interface { + // PendingMountSpec returns the mount a restore pod will use once the + // artifact for hash is promoted, and false when the storage cannot + // name one ahead of time (per-pod clone strategies). + PendingMountSpec(hash string, vol VolumeMeta) (PodMount, bool) +} diff --git a/src/compute-plane-services/nvsnap/internal/checkpointstore/percapture_pvc.go b/src/compute-plane-services/nvsnap/internal/checkpointstore/percapture_pvc.go index f5896db74c..315564d034 100644 --- a/src/compute-plane-services/nvsnap/internal/checkpointstore/percapture_pvc.go +++ b/src/compute-plane-services/nvsnap/internal/checkpointstore/percapture_pvc.go @@ -48,6 +48,7 @@ import ( "context" "errors" "fmt" + "io" "strings" "time" @@ -425,6 +426,9 @@ func (b *PerCapturePVCBackend) Put(ctx context.Context, hash string, sources []C if publishName == "" { publishName = roxName } + // Gated followers in other namespaces were admitted against + // rox- in their own namespace; make it exist before ready. + b.ensureClaimsForStampedPods(ctx, hash, ns) if err := b.setState(hash, pvcStateReady, publishName); err != nil { return Manifest{}, fmt.Errorf("publish ready: %w", err) } @@ -845,6 +849,32 @@ func (b *PerCapturePVCBackend) Get(ctx context.Context, hash, _ string) (Manifes return b.Stat(ctx, hash) } +// PendingMountSpec names the shared ReadOnlyMany claim for hash without +// requiring it to exist. Only ROX-capable strategies (shared-volume, or +// snapshot-clone with readOnlyMany) have a single deterministic claim; +// per-pod clones are minted at mount time and cannot be named early. +func (b *PerCapturePVCBackend) PendingMountSpec(hash string, vol VolumeMeta) (PodMount, bool) { + b.applyDefaults() + caps := b.Promoter.Caps() + if !caps.ReadOnlyMany && !caps.SharedVolume { + return PodMount{}, false + } + volName := "nvsnap-checkpoint" + if vol.Name != "" { + volName = vol.Name + } + claim := "rox-" + ShortHash(hash) + return PodMount{ + Volume: corev1.Volume{ + Name: volName, + VolumeSource: corev1.VolumeSource{ + PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{ClaimName: claim, ReadOnly: true}, + }, + }, + VolumeMount: corev1.VolumeMount{Name: volName, MountPath: vol.MountPath, ReadOnly: true}, + }, true +} + // Mount returns the pod-spec fragments needed to mount the rox- // PVC ReadOnly at vol.MountPath. Called by the admission webhook // when stamping a restore-from pod. @@ -871,9 +901,58 @@ func (b *PerCapturePVCBackend) Mount(ctx context.Context, hash string, vol Volum } // Storage-specific: shared-ROX returns the one rox- claim; // per-pod clones a fresh RWO PVC; shared-volume binds a static PV. + pm, err := b.Promoter.MountSpec(ctx, hash, vol) + if !errors.Is(err, ErrNotFound) { + return pm, err + } + // No claim in this namespace. NVCF runs each chart in its own + // namespace, so a restore usually lands where the capture never was: + // mint the namespace-local claim from the promoted artifact and try + // again. ErrNotFound from EnsureClaim means nothing is promoted yet; + // ErrUnsupported means per-pod-clone storage, which has no shared + // claim to bind; both leave the caller on its cold path. + if eerr := b.Promoter.EnsureClaim(ctx, hash, vol.Namespace); eerr != nil { + if errors.Is(eerr, ErrNotFound) || errors.Is(eerr, ErrUnsupported) { + return PodMount{}, ErrNotFound + } + return PodMount{}, fmt.Errorf("ensure claim in %s: %w", vol.Namespace, eerr) + } + b.log().WithFields(logrus.Fields{"hash": ShortHash(hash), "namespace": vol.Namespace}). + Info("L2 claim minted in restore namespace") return b.Promoter.MountSpec(ctx, hash, vol) } +// ensureClaimsForStampedPods mints the namespace-local claim in every +// namespace that already holds pods stamped with hash (election +// followers waiting on their gate, admitted before the promote existed), +// so they find rox- bound when nvsnap-server releases them. +// Best-effort: a failure is logged, the promote still publishes ready, +// and the affected pod's own Mount path retries on its next admission. +func (b *PerCapturePVCBackend) ensureClaimsForStampedPods(ctx context.Context, hash, captureNS string) { + pods, err := b.KubeClient.CoreV1().Pods("").List(ctx, metav1.ListOptions{LabelSelector: "nvsnap.io/hash=" + ShortHash(hash)}) + if err != nil { + b.log().WithError(err).WithField("hash", ShortHash(hash)).Warn("list stamped pods for cross-namespace claims failed") + return + } + seen := map[string]bool{captureNS: true} + for i := range pods.Items { + ns := pods.Items[i].Namespace + if seen[ns] { + continue + } + seen[ns] = true + if err := b.Promoter.EnsureClaim(ctx, hash, ns); err != nil { + if errors.Is(err, ErrUnsupported) { + continue + } + b.log().WithError(err).WithFields(logrus.Fields{"hash": ShortHash(hash), "namespace": ns}). + Warn("cross-namespace claim for stamped pods failed") + continue + } + b.log().WithFields(logrus.Fields{"hash": ShortHash(hash), "namespace": ns}).Info("L2 claim minted for stamped pods") + } +} + // Delete removes the rox PVC + any leftover rwx PVC + snapshot from // the backend's default namespace. Called by the retention controller // / cascade-delete (nvsnap#74); idempotent. @@ -912,3 +991,13 @@ func (b *PerCapturePVCBackend) Delete(ctx context.Context, hash string) error { // with any in-flight Jobs from the legacy code path; that subcommand // is no longer reachable from production callers and is scheduled // for removal in a follow-up cleanup commit. + +// log returns the backend logger, or a discard logger when none is set. +func (b *PerCapturePVCBackend) log() logrus.FieldLogger { + if b.Log != nil { + return b.Log + } + l := logrus.New() + l.SetOutput(io.Discard) + return l +} diff --git a/src/compute-plane-services/nvsnap/internal/checkpointstore/percapture_pvc_test.go b/src/compute-plane-services/nvsnap/internal/checkpointstore/percapture_pvc_test.go index 07c2d82e80..3beb86ee6d 100644 --- a/src/compute-plane-services/nvsnap/internal/checkpointstore/percapture_pvc_test.go +++ b/src/compute-plane-services/nvsnap/internal/checkpointstore/percapture_pvc_test.go @@ -42,6 +42,7 @@ import ( corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/runtime/schema" dynamicfake "k8s.io/client-go/dynamic/fake" kubefake "k8s.io/client-go/kubernetes/fake" ) @@ -102,8 +103,11 @@ func newTestBackend(t *testing.T, opts ...func(*PerCapturePVCBackend)) *PerCaptu t.Helper() scheme := runtime.NewScheme() b := &PerCapturePVCBackend{ - KubeClient: kubefake.NewSimpleClientset(), - DynClient: dynamicfake.NewSimpleDynamicClient(scheme), + KubeClient: kubefake.NewSimpleClientset(), + DynClient: dynamicfake.NewSimpleDynamicClientWithCustomListKinds(scheme, map[schema.GroupVersionResource]string{ + volumeSnapshotGVR: "VolumeSnapshotList", + volumeSnapshotContentGVR: "VolumeSnapshotContentList", + }), Catalog: &stubCatalog{}, Namespace: "nvsnap-system", StorageClass: "hyperdisk-ml", diff --git a/src/compute-plane-services/nvsnap/internal/checkpointstore/promoter.go b/src/compute-plane-services/nvsnap/internal/checkpointstore/promoter.go index 206882bfce..a3193afbe0 100644 --- a/src/compute-plane-services/nvsnap/internal/checkpointstore/promoter.go +++ b/src/compute-plane-services/nvsnap/internal/checkpointstore/promoter.go @@ -106,8 +106,18 @@ type Promoter interface { // when no promoted artifact exists for hash. MountSpec(ctx context.Context, hash string, vol VolumeMeta) (PodMount, error) + // EnsureClaim makes the shared restore claim for hash exist in ns, + // minted from the promoted artifact, so a pod in a namespace other + // than the capture's can mount it. NVCF runs every chart in its own + // namespace, so the capture and the restores usually live apart. + // Idempotent. Returns ErrNotFound when nothing is promoted yet and + // ErrUnsupported when the strategy has no shared artifact to bind + // (per-pod clone). + EnsureClaim(ctx context.Context, hash, ns string) error + // Delete reclaims every artifact this strategy created for hash in - // the given namespace. Idempotent (NotFound == success). + // the given namespace and every namespace-local claim minted by + // EnsureClaim. Idempotent (NotFound == success). Delete(ctx context.Context, hash, namespace string) error // Caps advertises capabilities for startup validation + logging. diff --git a/src/compute-plane-services/nvsnap/internal/checkpointstore/promoter_namespace_test.go b/src/compute-plane-services/nvsnap/internal/checkpointstore/promoter_namespace_test.go new file mode 100644 index 0000000000..55fd88196f --- /dev/null +++ b/src/compute-plane-services/nvsnap/internal/checkpointstore/promoter_namespace_test.go @@ -0,0 +1,236 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package checkpointstore + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/sirupsen/logrus" + corev1 "k8s.io/api/core/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + "k8s.io/apimachinery/pkg/api/resource" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/runtime/schema" + dynamicfake "k8s.io/client-go/dynamic/fake" + "k8s.io/client-go/kubernetes/fake" + k8stesting "k8s.io/client-go/testing" +) + +const xnsHash = "abcdefabcdefabcdefabcdefabcdefabcdefabcdefabcdefabcdefabcdefabcd" + +// promotedSharedFixture is the cluster after a shared-volume promote in +// "capture-ns": primary PV (Retain), secondary reader PV, rox claim. +func promotedSharedFixture() *fake.Clientset { + sc := "nvmesh-sc" + primary := &corev1.PersistentVolume{ + ObjectMeta: metav1.ObjectMeta{Name: "pv-primary", Labels: map[string]string{"nvsnap.io/role": "writer"}}, + Spec: corev1.PersistentVolumeSpec{ + Capacity: corev1.ResourceList{corev1.ResourceStorage: resource.MustParse("96Gi")}, + PersistentVolumeReclaimPolicy: corev1.PersistentVolumeReclaimRetain, + PersistentVolumeSource: corev1.PersistentVolumeSource{CSI: &corev1.CSIPersistentVolumeSource{Driver: "nvmesh-csi.excelero.com", VolumeHandle: "nvmesh-vol:proj:capture-ns"}}, + }, + } + sec := primary.DeepCopy() + sec.ObjectMeta = metav1.ObjectMeta{Name: secondaryPVName(xnsHash), Labels: map[string]string{"nvsnap.io/role": "reader-shared"}} + sec.Spec.ClaimRef = &corev1.ObjectReference{Kind: "PersistentVolumeClaim", Name: sharedROXName(xnsHash), Namespace: "capture-ns"} + rox := &corev1.PersistentVolumeClaim{ + ObjectMeta: metav1.ObjectMeta{Name: sharedROXName(xnsHash), Namespace: "capture-ns", Labels: map[string]string{labelHashShort: ShortHash(xnsHash)}}, + Spec: corev1.PersistentVolumeClaimSpec{VolumeName: sec.Name, StorageClassName: &sc}, + } + return fake.NewSimpleClientset(primary, sec, rox) +} + +func sharedPromoter(kc *fake.Clientset) *SharedVolumePromoter { + tx, _ := LookupVolumeHandleTransform("nvmesh") + return &SharedVolumePromoter{KubeClient: kc, StorageClass: "nvmesh-sc", Transform: tx, MountOptions: []string{"ro", "norecovery", "nouuid"}, Log: logrus.New()} +} + +func TestSharedVolume_EnsureClaimMintsNamespaceLocalClaim(t *testing.T) { + kc := promotedSharedFixture() + p := sharedPromoter(kc) + ctx := context.Background() + if err := p.EnsureClaim(ctx, xnsHash, "fn-ns"); err != nil { + t.Fatal(err) + } + pvc, err := kc.CoreV1().PersistentVolumeClaims("fn-ns").Get(ctx, sharedROXName(xnsHash), metav1.GetOptions{}) + if err != nil { + t.Fatalf("rox claim not minted in fn-ns: %v", err) + } + pvName := namespacedSecondaryPVName(xnsHash, "fn-ns") + if pvc.Spec.VolumeName != pvName || pvc.Spec.AccessModes[0] != corev1.ReadOnlyMany { + t.Errorf("claim must bind the per-namespace PV read-only: %+v", pvc.Spec) + } + pv, err := kc.CoreV1().PersistentVolumes().Get(ctx, pvName, metav1.GetOptions{}) + if err != nil { + t.Fatal(err) + } + if pv.Spec.CSI.VolumeHandle != "nvmesh-vol:proj:fn-ns" { + t.Errorf("handle must be rewritten for the consumer namespace, got %q", pv.Spec.CSI.VolumeHandle) + } + if pv.Spec.ClaimRef == nil || pv.Spec.ClaimRef.Namespace != "fn-ns" || !pv.Spec.CSI.ReadOnly || pv.Spec.PersistentVolumeReclaimPolicy != corev1.PersistentVolumeReclaimRetain { + t.Errorf("per-namespace PV must be pre-bound to fn-ns, read-only, retained: %+v", pv.Spec) + } + if pv.Labels[labelHashShort] != ShortHash(xnsHash) || pv.Labels[labelNamespace] != "fn-ns" { + t.Errorf("per-namespace PV must carry hash and namespace labels for Delete: %v", pv.Labels) + } + // Idempotent, and a second namespace gets its own PV. + if err := p.EnsureClaim(ctx, xnsHash, "fn-ns"); err != nil { + t.Errorf("second EnsureClaim must be a no-op, got %v", err) + } + if err := p.EnsureClaim(ctx, xnsHash, "other-ns"); err != nil { + t.Fatal(err) + } + pvs, _ := kc.CoreV1().PersistentVolumes().List(ctx, metav1.ListOptions{LabelSelector: labelHashShort + "=" + ShortHash(xnsHash)}) + if len(pvs.Items) != 2 { + t.Errorf("want one per-namespace PV per consumer namespace, got %d", len(pvs.Items)) + } + // The capture namespace already has the promote's own claim: no-op. + if err := p.EnsureClaim(ctx, xnsHash, "capture-ns"); err != nil { + t.Errorf("capture namespace must be a no-op, got %v", err) + } +} + +// Concurrent admissions of one chart all run EnsureClaim; the fake +// rejects every PV update so the test fails if any step depends on one. +func TestSharedVolume_EnsureClaimNeedsNoPVUpdate(t *testing.T) { + kc := promotedSharedFixture() + kc.PrependReactor("update", "persistentvolumes", func(k8stesting.Action) (bool, runtime.Object, error) { + return true, nil, apierrors.NewConflict(schema.GroupResource{Resource: "persistentvolumes"}, "pv", errors.New("the object has been modified")) + }) + p := sharedPromoter(kc) + ctx := context.Background() + for range 3 { + if err := p.EnsureClaim(ctx, xnsHash, "fn-ns"); err != nil { + t.Fatalf("EnsureClaim must not depend on a PV update: %v", err) + } + } + pv, err := kc.CoreV1().PersistentVolumes().Get(ctx, namespacedSecondaryPVName(xnsHash, "fn-ns"), metav1.GetOptions{}) + if err != nil { + t.Fatal(err) + } + if pv.Labels[labelHashShort] != ShortHash(xnsHash) || pv.Labels[labelNamespace] != "fn-ns" { + t.Errorf("per-namespace PV must be created already labelled: %v", pv.Labels) + } +} + +func TestSharedVolume_EnsureClaimBeforePromoteIsNotFound(t *testing.T) { + p := sharedPromoter(fake.NewSimpleClientset()) + if err := p.EnsureClaim(context.Background(), xnsHash, "fn-ns"); !errors.Is(err, ErrNotFound) { + t.Errorf("nothing promoted: want ErrNotFound, got %v", err) + } +} + +func TestSharedVolume_DeleteReapsNamespacedClaims(t *testing.T) { + kc := promotedSharedFixture() + p := sharedPromoter(kc) + ctx := context.Background() + for _, ns := range []string{"fn-a", "fn-b"} { + if err := p.EnsureClaim(ctx, xnsHash, ns); err != nil { + t.Fatal(err) + } + } + if err := p.Delete(ctx, xnsHash, "capture-ns"); err != nil { + t.Fatal(err) + } + for _, ns := range []string{"capture-ns", "fn-a", "fn-b"} { + if _, err := kc.CoreV1().PersistentVolumeClaims(ns).Get(ctx, sharedROXName(xnsHash), metav1.GetOptions{}); !apierrors.IsNotFound(err) { + t.Errorf("%s: rox claim must be gone, err=%v", ns, err) + } + } + pvs, _ := kc.CoreV1().PersistentVolumes().List(ctx, metav1.ListOptions{}) + if len(pvs.Items) != 0 { + names := []string{} + for _, pv := range pvs.Items { + names = append(names, pv.Name) + } + t.Errorf("all PVs must be gone after Delete, left %v", names) + } +} + +// Mount in a namespace without a claim mints it from the promote, so a +// chart in its own namespace restores from a capture made elsewhere. +func TestPerCapturePVCBackend_MountMintsClaimInRestoreNamespace(t *testing.T) { + kc := promotedSharedFixture() + b := &PerCapturePVCBackend{KubeClient: kc, Namespace: "nvsnap-system", StorageClass: "nvmesh-sc", Promoter: sharedPromoter(kc), Log: logrus.New()} + pm, err := b.Mount(context.Background(), xnsHash, VolumeMeta{Name: "nvsnap-cachedir", MountPath: "/opt/nvsnap", Namespace: "fn-ns"}) + if err != nil { + t.Fatal(err) + } + if pm.Volume.PersistentVolumeClaim == nil || pm.Volume.PersistentVolumeClaim.ClaimName != sharedROXName(xnsHash) { + t.Errorf("mount must name the namespace-local rox claim, got %+v", pm.Volume) + } + if _, err := kc.CoreV1().PersistentVolumeClaims("fn-ns").Get(context.Background(), sharedROXName(xnsHash), metav1.GetOptions{}); err != nil { + t.Errorf("claim must exist in fn-ns after Mount: %v", err) + } +} + +// snapshot-clone: a namespace-local clone needs a snapshot in that +// namespace, pre-provisioned from the promote's content handle. +func TestSnapshotClone_EnsureClaimPreProvisionsSnapshot(t *testing.T) { + scheme := runtime.NewScheme() + dyn := dynamicfake.NewSimpleDynamicClientWithCustomListKinds(scheme, map[schema.GroupVersionResource]string{ + volumeSnapshotGVR: "VolumeSnapshotList", volumeSnapshotContentGVR: "VolumeSnapshotContentList", + }) + ctx := context.Background() + snapName := "snap-" + ShortHash(xnsHash) + promoted := &unstructured.Unstructured{Object: map[string]any{ + "apiVersion": "snapshot.storage.k8s.io/v1", "kind": "VolumeSnapshot", + "metadata": map[string]any{"name": snapName, "namespace": "capture-ns", "labels": map[string]any{"nvsnap.io/per-capture": "true"}}, + "status": map[string]any{"readyToUse": true, "boundVolumeSnapshotContentName": "snapcontent-1", "restoreSize": "20Gi"}, + }} + content := &unstructured.Unstructured{Object: map[string]any{ + "apiVersion": "snapshot.storage.k8s.io/v1", "kind": "VolumeSnapshotContent", + "metadata": map[string]any{"name": "snapcontent-1"}, + "spec": map[string]any{"driver": "pd.csi.storage.gke.io"}, + "status": map[string]any{"snapshotHandle": "projects/p/global/snapshots/s1"}, + }} + // The namespaced copy is pre-seeded ready, standing in for the + // external snapshotter that would bind it in a real cluster. + nsSnap := &unstructured.Unstructured{Object: map[string]any{ + "apiVersion": "snapshot.storage.k8s.io/v1", "kind": "VolumeSnapshot", + "metadata": map[string]any{"name": snapName, "namespace": "fn-ns", "labels": map[string]any{"nvsnap.io/per-capture": "true", labelHashShort: ShortHash(xnsHash), labelNamespace: "fn-ns"}}, + "status": map[string]any{"readyToUse": true}, + }} + for _, o := range []*unstructured.Unstructured{promoted, content, nsSnap} { + gvr := volumeSnapshotGVR + if o.GetKind() == "VolumeSnapshotContent" { + gvr = volumeSnapshotContentGVR + } + if _, err := dyn.Resource(gvr).Namespace(o.GetNamespace()).Create(ctx, o, metav1.CreateOptions{}); err != nil { + t.Fatal(err) + } + } + kc := fake.NewSimpleClientset() + p := &SnapshotClonePromoter{KubeClient: kc, DynClient: dyn, StorageClass: "hyperdisk-ml", SnapshotClass: "hdml", ReadOnlyMany: true, SnapshotTimeout: 5 * time.Second, Log: logrus.New()} + if err := p.EnsureClaim(ctx, xnsHash, "fn-ns"); err != nil { + t.Fatal(err) + } + pre, err := dyn.Resource(volumeSnapshotContentGVR).Get(ctx, "snapcontent-1-"+nsSuffix("fn-ns"), metav1.GetOptions{}) + if err != nil { + t.Fatalf("pre-provisioned content missing: %v", err) + } + handle, _, _ := unstructured.NestedString(pre.Object, "spec", "source", "snapshotHandle") + refNS, _, _ := unstructured.NestedString(pre.Object, "spec", "volumeSnapshotRef", "namespace") + policy, _, _ := unstructured.NestedString(pre.Object, "spec", "deletionPolicy") + if handle != "projects/p/global/snapshots/s1" || refNS != "fn-ns" || policy != "Retain" { + t.Errorf("content must point at the promote's handle, bind into fn-ns and retain: handle=%q ref=%q policy=%q", handle, refNS, policy) + } + pvc, err := kc.CoreV1().PersistentVolumeClaims("fn-ns").Get(ctx, "rox-"+ShortHash(xnsHash), metav1.GetOptions{}) + if err != nil { + t.Fatalf("rox clone missing in fn-ns: %v", err) + } + if pvc.Spec.DataSource == nil || pvc.Spec.DataSource.Name != snapName || pvc.Spec.Resources.Requests[corev1.ResourceStorage] != resource.MustParse("20Gi") { + t.Errorf("clone must come from the namespaced snapshot at the promote's restore size: %+v", pvc.Spec) + } + // Per-pod-clone storage has no shared claim to reproduce. + perPod := &SnapshotClonePromoter{KubeClient: kc, DynClient: dyn, StorageClass: "gp3", ReadOnlyMany: false} + if err := perPod.EnsureClaim(ctx, xnsHash, "fn-ns"); !errors.Is(err, ErrUnsupported) { + t.Errorf("per-pod clone: want ErrUnsupported, got %v", err) + } +} diff --git a/src/compute-plane-services/nvsnap/internal/checkpointstore/promoter_shared.go b/src/compute-plane-services/nvsnap/internal/checkpointstore/promoter_shared.go index a0ac50db83..519941a4e3 100644 --- a/src/compute-plane-services/nvsnap/internal/checkpointstore/promoter_shared.go +++ b/src/compute-plane-services/nvsnap/internal/checkpointstore/promoter_shared.go @@ -43,6 +43,8 @@ package checkpointstore import ( "context" + "crypto/sha256" + "encoding/hex" "fmt" "strings" @@ -190,7 +192,7 @@ func (p *SharedVolumePromoter) Promote(ctx context.Context, in PromoteInput) (Pr return PromoteResult{SharedClaimName: roxName, ReusedWriterVolume: true}, nil } -func (p *SharedVolumePromoter) ensureSecondaryPV(ctx context.Context, primary *corev1.PersistentVolume, secName, roxName, ns string) error { +func (p *SharedVolumePromoter) ensureSecondaryPV(ctx context.Context, primary *corev1.PersistentVolume, secName, roxName, ns string, extraLabels ...map[string]string) error { if existing, err := p.KubeClient.CoreV1().PersistentVolumes().Get(ctx, secName, metav1.GetOptions{}); err == nil { return p.releaseStaleBinding(ctx, existing, roxName, ns) } else if !apierrors.IsNotFound(err) { @@ -209,6 +211,11 @@ func (p *SharedVolumePromoter) ensureSecondaryPV(ctx context.Context, primary *c "nvsnap.io/role": "reader-shared", }, } + for _, extra := range extraLabels { + for k, v := range extra { + sec.Labels[k] = v + } + } sec.Spec.AccessModes = []corev1.PersistentVolumeAccessMode{corev1.ReadOnlyMany} sec.Spec.PersistentVolumeReclaimPolicy = corev1.PersistentVolumeReclaimRetain sec.Spec.CSI = primary.Spec.CSI.DeepCopy() @@ -383,6 +390,7 @@ func (p *SharedVolumePromoter) Delete(ctx context.Context, hash, ns string) erro errs = append(errs, fmt.Sprintf("primary PV %s: %v", primaryPVName, err)) } } + errs = append(errs, p.deleteNamespacedClaims(ctx, hash)...) if len(errs) > 0 { return fmt.Errorf("shared-volume delete: %s", strings.Join(errs, "; ")) } @@ -433,3 +441,97 @@ func (p *SharedVolumePromoter) deletePV(ctx context.Context, name string) error } return nil } + +// Namespace-local claims. The capture namespace gets nvsnap-ro-pv- +// bound to rox- at promote. Any other namespace gets its own +// secondary PV, nvsnap-ro-pv--, with the handle rewritten for +// that namespace and a rox- claim bound to it: the same static-PV +// pattern NVCA uses for one model volume across tenant namespaces. Both +// carry nvsnap.io/hash-short and nvsnap.io/namespace so Delete can find +// every one of them. + +const ( + labelHashShort = "nvsnap.io/hash-short" + labelNamespace = "nvsnap.io/namespace" +) + +// nsSuffix is a short stable token for a namespace name, safe in an +// object name regardless of the namespace's length. +func nsSuffix(ns string) string { + sum := sha256.Sum256([]byte(ns)) + return hex.EncodeToString(sum[:4]) +} + +func namespacedSecondaryPVName(hash, ns string) string { + return secondaryPVName(hash) + "-" + nsSuffix(ns) +} + +// EnsureClaim mints rox- in ns from the promoted secondary PV. +func (p *SharedVolumePromoter) EnsureClaim(ctx context.Context, hash, ns string) error { + p.applyDefaults() + roxName := sharedROXName(hash) + if _, err := p.KubeClient.CoreV1().PersistentVolumeClaims(ns).Get(ctx, roxName, metav1.GetOptions{}); err == nil { + return nil + } else if !apierrors.IsNotFound(err) { + return fmt.Errorf("get rox PVC %s/%s: %w", ns, roxName, err) + } + // The promote's own secondary PV is the proof the artifact exists and + // the source of the CSI handle to rewrite. + promoted, err := p.KubeClient.CoreV1().PersistentVolumes().Get(ctx, secondaryPVName(hash), metav1.GetOptions{}) + if err != nil { + if apierrors.IsNotFound(err) { + return ErrNotFound + } + return fmt.Errorf("get promoted PV %s: %w", secondaryPVName(hash), err) + } + if promoted.Spec.CSI == nil || promoted.Spec.CSI.VolumeHandle == "" { + return fmt.Errorf("promoted PV %s has no CSI volumeHandle", promoted.Name) + } + primaryName := p.primaryPVForHandle(ctx, promoted.Spec.CSI.VolumeHandle) + if primaryName == "" { + return fmt.Errorf("no primary PV found for promoted handle %q", promoted.Spec.CSI.VolumeHandle) + } + primary, err := p.KubeClient.CoreV1().PersistentVolumes().Get(ctx, primaryName, metav1.GetOptions{}) + if err != nil { + return fmt.Errorf("get primary PV %s: %w", primaryName, err) + } + // Labels go on at create: N pods of one chart are admitted in the same + // second and each runs this, so a create-then-update would race on the + // PV's resourceVersion (seen on dev1: "the object has been modified"). + // Create is idempotent through AlreadyExists; nothing here updates. + secName := namespacedSecondaryPVName(hash, ns) + if err := p.ensureSecondaryPV(ctx, primary, secName, roxName, ns, map[string]string{labelHashShort: ShortHash(hash), labelNamespace: ns}); err != nil { + return err + } + return p.ensureSharedROXPVC(ctx, ns, hash, roxName, secName, primary) +} + +// deleteNamespacedClaims removes every rox claim and per-namespace +// secondary PV minted by EnsureClaim for hash. +func (p *SharedVolumePromoter) deleteNamespacedClaims(ctx context.Context, hash string) []string { + var errs []string + sel := labelHashShort + "=" + ShortHash(hash) + if pvcs, err := p.KubeClient.CoreV1().PersistentVolumeClaims("").List(ctx, metav1.ListOptions{LabelSelector: sel}); err == nil { + for i := range pvcs.Items { + c := &pvcs.Items[i] + if c.Name != sharedROXName(hash) { + continue + } + if derr := p.deletePVC(ctx, c.Namespace, c.Name); derr != nil { + errs = append(errs, fmt.Sprintf("rox %s/%s: %v", c.Namespace, c.Name, derr)) + } + } + } else { + errs = append(errs, fmt.Sprintf("list rox claims: %v", err)) + } + if pvs, err := p.KubeClient.CoreV1().PersistentVolumes().List(ctx, metav1.ListOptions{LabelSelector: sel}); err == nil { + for i := range pvs.Items { + if derr := p.deletePV(ctx, pvs.Items[i].Name); derr != nil { + errs = append(errs, fmt.Sprintf("secondary PV %s: %v", pvs.Items[i].Name, derr)) + } + } + } else { + errs = append(errs, fmt.Sprintf("list secondary PVs: %v", err)) + } + return errs +} diff --git a/src/compute-plane-services/nvsnap/internal/checkpointstore/promoter_snapshot.go b/src/compute-plane-services/nvsnap/internal/checkpointstore/promoter_snapshot.go index 8198c9cf22..9158272ce7 100644 --- a/src/compute-plane-services/nvsnap/internal/checkpointstore/promoter_snapshot.go +++ b/src/compute-plane-services/nvsnap/internal/checkpointstore/promoter_snapshot.go @@ -56,6 +56,15 @@ var volumeSnapshotGVR = schema.GroupVersionResource{ Resource: "volumesnapshots", } +// volumeSnapshotContentGVR is the cluster-scoped half of a snapshot. A +// namespace-local snapshot is pre-provisioned from the promote's content +// handle so a clone can be made in a namespace other than the capture's. +var volumeSnapshotContentGVR = schema.GroupVersionResource{ + Group: "snapshot.storage.k8s.io", + Version: "v1", + Resource: "volumesnapshotcontents", +} + // SnapshotClonePromoter implements Promoter via CSI VolumeSnapshot+clone. type SnapshotClonePromoter struct { KubeClient kubernetes.Interface @@ -375,6 +384,7 @@ func (p *SnapshotClonePromoter) Delete(ctx context.Context, hash, ns string) err if err := p.deleteSnapshot(ctx, ns, snapName); err != nil { errs = append(errs, fmt.Sprintf("snap %s: %v", snapName, err)) } + errs = append(errs, p.deleteNamespacedClaims(ctx, hash)...) if len(errs) > 0 { return fmt.Errorf("snapshot-clone delete: %s", strings.Join(errs, "; ")) } @@ -396,3 +406,193 @@ func (p *SnapshotClonePromoter) deleteSnapshot(ctx context.Context, ns, name str } return nil } + +// EnsureClaim mints rox- in ns as a clone of the promoted snapshot. +// VolumeSnapshots are namespaced and a clone must name a snapshot in its +// own namespace, so the promote's snapshot handle is re-exposed there +// through a pre-provisioned VolumeSnapshotContent + VolumeSnapshot pair +// (both Retain: the CSI snapshot stays owned by the capture namespace). +// Only the ReadOnlyMany strategy has a shared claim to reproduce. +func (p *SnapshotClonePromoter) EnsureClaim(ctx context.Context, hash, ns string) error { + p.applyDefaults() + if !p.ReadOnlyMany { + return ErrUnsupported + } + roxName := "rox-" + ShortHash(hash) + if _, err := p.KubeClient.CoreV1().PersistentVolumeClaims(ns).Get(ctx, roxName, metav1.GetOptions{}); err == nil { + return nil + } else if !apierrors.IsNotFound(err) { + return fmt.Errorf("get rox PVC %s/%s: %w", ns, roxName, err) + } + snapName := "snap-" + ShortHash(hash) + src, err := p.findPromotedSnapshot(ctx, snapName) + if err != nil { + return err + } + srcNS := src.GetNamespace() + if srcNS == ns { + return ErrNotFound // promote not finished in the capture namespace + } + contentName, _, _ := unstructured.NestedString(src.Object, "status", "boundVolumeSnapshotContentName") + if contentName == "" { + return ErrNotFound + } + content, err := p.DynClient.Resource(volumeSnapshotContentGVR).Get(ctx, contentName, metav1.GetOptions{}) + if err != nil { + return fmt.Errorf("get VolumeSnapshotContent %s: %w", contentName, err) + } + handle, _, _ := unstructured.NestedString(content.Object, "status", "snapshotHandle") + driver, _, _ := unstructured.NestedString(content.Object, "spec", "driver") + if handle == "" || driver == "" { + return fmt.Errorf("VolumeSnapshotContent %s has no snapshotHandle/driver yet", contentName) + } + restoreSize, _, _ := unstructured.NestedString(src.Object, "status", "restoreSize") + + nsContent := contentName + "-" + nsSuffix(ns) + pre := &unstructured.Unstructured{Object: map[string]any{ + "apiVersion": "snapshot.storage.k8s.io/v1", + "kind": "VolumeSnapshotContent", + "metadata": map[string]any{ + "name": nsContent, + "labels": map[string]any{ + "app.kubernetes.io/managed-by": "nvsnap", + "nvsnap.io/per-capture": "true", + labelHashShort: ShortHash(hash), + labelNamespace: ns, + }, + }, + "spec": map[string]any{ + "deletionPolicy": "Retain", + "driver": driver, + "volumeSnapshotClassName": p.SnapshotClass, + "source": map[string]any{"snapshotHandle": handle}, + "volumeSnapshotRef": map[string]any{"name": snapName, "namespace": ns}, + }, + }} + if _, err := p.DynClient.Resource(volumeSnapshotContentGVR).Create(ctx, pre, metav1.CreateOptions{}); err != nil && !apierrors.IsAlreadyExists(err) { + return fmt.Errorf("create pre-provisioned VolumeSnapshotContent %s: %w", nsContent, err) + } + snap := &unstructured.Unstructured{Object: map[string]any{ + "apiVersion": "snapshot.storage.k8s.io/v1", + "kind": "VolumeSnapshot", + "metadata": map[string]any{ + "name": snapName, + "namespace": ns, + "labels": map[string]any{ + "app.kubernetes.io/managed-by": "nvsnap", + "nvsnap.io/per-capture": "true", + labelHashShort: ShortHash(hash), + labelNamespace: ns, + }, + }, + "spec": map[string]any{ + "source": map[string]any{"volumeSnapshotContentName": nsContent}, + }, + }} + if _, err := p.DynClient.Resource(volumeSnapshotGVR).Namespace(ns).Create(ctx, snap, metav1.CreateOptions{}); err != nil && !apierrors.IsAlreadyExists(err) { + return fmt.Errorf("create namespaced VolumeSnapshot %s/%s: %w", ns, snapName, err) + } + if err := wait.PollUntilContextTimeout(ctx, 2*time.Second, p.SnapshotTimeout, true, func(ctx context.Context) (bool, error) { + got, err := p.DynClient.Resource(volumeSnapshotGVR).Namespace(ns).Get(ctx, snapName, metav1.GetOptions{}) + if err != nil { + return false, err + } + ready, _, _ := unstructured.NestedBool(got.Object, "status", "readyToUse") + return ready, nil + }); err != nil { + return fmt.Errorf("namespaced snapshot %s/%s not ready: %w", ns, snapName, err) + } + sizeQty := resource.MustParse("1Gi") + if restoreSize != "" { + if q, perr := resource.ParseQuantity(restoreSize); perr == nil { + sizeQty = q + } + } + return p.createROXFromSnapshot(ctx, ns, hash, roxName, snapName, sizeQty) +} + +// findPromotedSnapshot locates snap- in whichever namespace the +// capture ran in. +func (p *SnapshotClonePromoter) findPromotedSnapshot(ctx context.Context, snapName string) (*unstructured.Unstructured, error) { + list, err := p.DynClient.Resource(volumeSnapshotGVR).Namespace("").List(ctx, metav1.ListOptions{LabelSelector: "nvsnap.io/per-capture=true"}) + if err != nil { + return nil, fmt.Errorf("list VolumeSnapshots: %w", err) + } + for i := range list.Items { + it := &list.Items[i] + if it.GetName() != snapName { + continue + } + if _, isNS := it.GetLabels()[labelNamespace]; isNS { + continue // a namespaced copy, not the promote's own + } + return it, nil + } + return nil, ErrNotFound +} + +// createROXFromSnapshot is cloneROX with an explicit size, for namespaces +// that have no writer PVC to size from. +func (p *SnapshotClonePromoter) createROXFromSnapshot(ctx context.Context, ns, hash, roxName, snapName string, sizeQty resource.Quantity) error { + apiGroup := "snapshot.storage.k8s.io" + roxPVC := &corev1.PersistentVolumeClaim{ + ObjectMeta: metav1.ObjectMeta{ + Name: roxName, + Namespace: ns, + Labels: map[string]string{ + "app.kubernetes.io/managed-by": "nvsnap", + "nvsnap.io/per-capture": "true", + "nvsnap.io/role": "reader", + labelHashShort: ShortHash(hash), + labelNamespace: ns, + }, + }, + Spec: corev1.PersistentVolumeClaimSpec{ + AccessModes: []corev1.PersistentVolumeAccessMode{corev1.ReadOnlyMany}, + StorageClassName: &p.StorageClass, + Resources: corev1.VolumeResourceRequirements{Requests: corev1.ResourceList{corev1.ResourceStorage: sizeQty}}, + DataSource: &corev1.TypedLocalObjectReference{APIGroup: &apiGroup, Kind: "VolumeSnapshot", Name: snapName}, + }, + } + if _, err := p.KubeClient.CoreV1().PersistentVolumeClaims(ns).Create(ctx, roxPVC, metav1.CreateOptions{}); err != nil && !apierrors.IsAlreadyExists(err) { + return fmt.Errorf("create rox PVC %s/%s: %w", ns, roxName, err) + } + return nil +} + +// deleteNamespacedClaims removes the claims, snapshots and contents that +// EnsureClaim minted for hash in other namespaces. +func (p *SnapshotClonePromoter) deleteNamespacedClaims(ctx context.Context, hash string) []string { + var errs []string + sel := labelHashShort + "=" + ShortHash(hash) + "," + labelNamespace + if pvcs, err := p.KubeClient.CoreV1().PersistentVolumeClaims("").List(ctx, metav1.ListOptions{LabelSelector: sel}); err == nil { + for i := range pvcs.Items { + c := &pvcs.Items[i] + if derr := p.KubeClient.CoreV1().PersistentVolumeClaims(c.Namespace).Delete(ctx, c.Name, metav1.DeleteOptions{}); derr != nil && !apierrors.IsNotFound(derr) { + errs = append(errs, fmt.Sprintf("rox %s/%s: %v", c.Namespace, c.Name, derr)) + } + } + } else { + errs = append(errs, fmt.Sprintf("list rox claims: %v", err)) + } + if snaps, err := p.DynClient.Resource(volumeSnapshotGVR).Namespace("").List(ctx, metav1.ListOptions{LabelSelector: sel}); err == nil { + for i := range snaps.Items { + it := &snaps.Items[i] + if derr := p.DynClient.Resource(volumeSnapshotGVR).Namespace(it.GetNamespace()).Delete(ctx, it.GetName(), metav1.DeleteOptions{}); derr != nil && !apierrors.IsNotFound(derr) { + errs = append(errs, fmt.Sprintf("snapshot %s/%s: %v", it.GetNamespace(), it.GetName(), derr)) + } + } + } else { + errs = append(errs, fmt.Sprintf("list namespaced snapshots: %v", err)) + } + if contents, err := p.DynClient.Resource(volumeSnapshotContentGVR).List(ctx, metav1.ListOptions{LabelSelector: sel}); err == nil { + for i := range contents.Items { + if derr := p.DynClient.Resource(volumeSnapshotContentGVR).Delete(ctx, contents.Items[i].GetName(), metav1.DeleteOptions{}); derr != nil && !apierrors.IsNotFound(derr) { + errs = append(errs, fmt.Sprintf("content %s: %v", contents.Items[i].GetName(), derr)) + } + } + } else { + errs = append(errs, fmt.Sprintf("list namespaced contents: %v", err)) + } + return errs +} diff --git a/src/compute-plane-services/nvsnap/internal/election/BUILD.bazel b/src/compute-plane-services/nvsnap/internal/election/BUILD.bazel new file mode 100644 index 0000000000..66602f9272 --- /dev/null +++ b/src/compute-plane-services/nvsnap/internal/election/BUILD.bazel @@ -0,0 +1,31 @@ +load("@rules_go//go:def.bzl", "go_library", "go_test") + +go_library( + name = "election", + srcs = ["election.go"], + importpath = "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/election", + visibility = ["//src/compute-plane-services/nvsnap:__subpackages__"], + deps = [ + "//src/compute-plane-services/nvsnap/internal/checkpointstore", + "@io_k8s_api//coordination/v1:coordination", + "@io_k8s_api//core/v1:core", + "@io_k8s_apimachinery//pkg/api/errors", + "@io_k8s_apimachinery//pkg/apis/meta/v1:meta", + "@io_k8s_client_go//kubernetes", + ], +) + +go_test( + name = "election_test", + srcs = ["election_test.go"], + embed = [":election"], + deps = [ + "@io_k8s_api//coordination/v1:coordination", + "@io_k8s_api//core/v1:core", + "@io_k8s_apimachinery//pkg/apis/meta/v1:meta", + "@io_k8s_apimachinery//pkg/runtime", + "@io_k8s_apimachinery//pkg/types", + "@io_k8s_client_go//kubernetes/fake", + "@io_k8s_client_go//testing", + ], +) diff --git a/src/compute-plane-services/nvsnap/internal/election/election.go b/src/compute-plane-services/nvsnap/internal/election/election.go new file mode 100644 index 0000000000..8af07bc3b6 --- /dev/null +++ b/src/compute-plane-services/nvsnap/internal/election/election.go @@ -0,0 +1,204 @@ +/* +SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +// Package election picks one downloader per cache hash when a chart +// creates N model workers at once. The webhook calls Elect for every +// model workload it admits; exactly one admission per hash wins, because +// the election is a Kubernetes Lease create, which is atomic. The winner +// is decorated as the capture source, the losers as gated restore pods +// that nvsnap-server releases once the capture is promoted. +// +// Design: docs/proposals/helm-chart-cache-election.md. +package election + +import ( + "context" + "fmt" + "time" + + coordinationv1 "k8s.io/api/coordination/v1" + corev1 "k8s.io/api/core/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/util/uuid" + "k8s.io/client-go/kubernetes" + + "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/checkpointstore" +) + +// Pod metadata the election stamps and later selects on. +const ( + // HashAnnotation carries the full cache hash the webhook composed at + // admission. The capture watcher reads it instead of recomposing from + // the live pod, so both sides agree by construction. + HashAnnotation = "nvsnap.io/hash" + // HashLabel carries the short hash, so nvsnap-server can list every + // pod of a hash with a label selector. + HashLabel = "nvsnap.io/hash" + // RoleAnnotation records the admission decision: leader, follower or + // restore. + RoleAnnotation = "nvsnap.io/role" + // GatedLabel is "true" while a follower holds the scheduling gate. + GatedLabel = "nvsnap.io/gated" + // GateName is the schedulingGates entry a follower carries until the + // capture for its hash is promoted. + GateName = "nvsnap.io/wait-for-cache" + // ColdStartAnnotation is stamped on a follower whose leader failed; + // the pod is recreated by its controller and re-admitted. + ColdStartAnnotation = "nvsnap.io/cold-start" + + // LeaseKindLabel and LeaseKindValue mark election Leases so the + // server reconciler can list them. + LeaseKindLabel = "nvsnap.io/kind" + LeaseKindValue = "capture-election" + // ElectionIDAnnotation on the leader pod is the id the webhook minted + // at admission and used as the Lease holder. A pod has no UID or name + // yet when a mutating webhook sees its CREATE, so the webhook's own id + // is the only handle that exists on both the Lease and the pod. + ElectionIDAnnotation = "nvsnap.io/election-id" + // LeaderNamespaceAnnotation and LeaderIDAnnotation on the Lease name + // the leader pod (namespace + election id) so its liveness can be + // checked. + LeaderNamespaceAnnotation = "nvsnap.io/leader-namespace" + LeaderIDAnnotation = "nvsnap.io/leader-id" + // DeadlineAnnotation is the RFC3339 time after which the reconciler + // treats the leader as failed even if its pod is still around. + DeadlineAnnotation = "nvsnap.io/deadline" + + // DefaultDeadline bounds a cold start plus a capture. + DefaultDeadline = 60 * time.Minute +) + +// Role is the admission decision for a model workload. +type Role string + +// Roles: the leader captures, followers wait gated for its promote, and +// restore means a promoted capture already existed at admission. +const ( + RoleLeader Role = "leader" + RoleFollower Role = "follower" + RoleRestore Role = "restore" +) + +// Elector decides who downloads. Implementations must be safe for +// concurrent admissions of the same hash. +type Elector interface { + // Elect returns RoleLeader for exactly one live election per hash and + // RoleFollower for every other caller while that election stands. The + // returned id is the Lease holder; the webhook stamps it on the leader + // as ElectionIDAnnotation so the server can find the leader later. + Elect(ctx context.Context, hash string, pod *corev1.Pod) (Role, string, error) +} + +// LeaseName is the election Lease for a hash. +func LeaseName(hash string) string { return "nvsnap-capture-" + checkpointstore.ShortHash(hash) } + +// PodIdentity names a pod at admission. Deployment and DynamoGraph pods +// have neither name nor UID yet at CREATE, only generateName. +func PodIdentity(pod *corev1.Pod) string { + if pod == nil { + return "" + } + name := pod.Name + if name == "" { + name = pod.GenerateName + "*" + } + if pod.UID != "" { + return fmt.Sprintf("%s/%s(%s)", pod.Namespace, name, pod.UID) + } + return pod.Namespace + "/" + name +} + +// LeaseElector elects through a Lease create in Namespace. +type LeaseElector struct { + KubeClient kubernetes.Interface + // Namespace holds the election Leases (nvsnap-system). + Namespace string + // Deadline bounds the leader's cold start plus capture. Zero means + // DefaultDeadline. + Deadline time.Duration + // Now is a clock seam for tests. + Now func() time.Time + // NewID mints the election id; a seam for tests. nil uses a UUID. + NewID func() string +} + +func (e *LeaseElector) newID() string { + if e.NewID != nil { + return e.NewID() + } + return string(uuid.NewUUID()) +} + +func (e *LeaseElector) now() time.Time { + if e.Now != nil { + return e.Now() + } + return time.Now() +} + +func (e *LeaseElector) deadline() time.Duration { + if e.Deadline <= 0 { + return DefaultDeadline + } + return e.Deadline +} + +// Elect creates the Lease for hash with a freshly minted id as holder. +// Created means leader; AlreadyExists means follower; anything else is an +// error the caller fails open on. +func (e *LeaseElector) Elect(ctx context.Context, hash string, pod *corev1.Pod) (Role, string, error) { + if e.KubeClient == nil { + return "", "", fmt.Errorf("election: no kube client") + } + if pod == nil { + return "", "", fmt.Errorf("election: nil pod") + } + now := e.now() + secs := int32(e.deadline().Seconds()) + holder := e.newID() + lease := &coordinationv1.Lease{ + ObjectMeta: metav1.ObjectMeta{ + Name: LeaseName(hash), + Namespace: e.Namespace, + Labels: map[string]string{ + LeaseKindLabel: LeaseKindValue, + HashLabel: checkpointstore.ShortHash(hash), + }, + Annotations: map[string]string{ + HashAnnotation: hash, + LeaderNamespaceAnnotation: pod.Namespace, + LeaderIDAnnotation: holder, + DeadlineAnnotation: now.Add(e.deadline()).UTC().Format(time.RFC3339), + }, + }, + Spec: coordinationv1.LeaseSpec{ + HolderIdentity: &holder, + LeaseDurationSeconds: &secs, + AcquireTime: &metav1.MicroTime{Time: now}, + }, + } + _, err := e.KubeClient.CoordinationV1().Leases(e.Namespace).Create(ctx, lease, metav1.CreateOptions{}) + switch { + case err == nil: + return RoleLeader, holder, nil + case apierrors.IsAlreadyExists(err): + return RoleFollower, "", nil + default: + return "", "", fmt.Errorf("election: create lease %s: %w", lease.Name, err) + } +} diff --git a/src/compute-plane-services/nvsnap/internal/election/election_test.go b/src/compute-plane-services/nvsnap/internal/election/election_test.go new file mode 100644 index 0000000000..1fe43dff25 --- /dev/null +++ b/src/compute-plane-services/nvsnap/internal/election/election_test.go @@ -0,0 +1,112 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package election + +import ( + "context" + "errors" + "testing" + "time" + + coordinationv1 "k8s.io/api/coordination/v1" + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/types" + "k8s.io/client-go/kubernetes/fake" + k8stesting "k8s.io/client-go/testing" +) + +func pod(uid string) *corev1.Pod { + return &corev1.Pod{ObjectMeta: metav1.ObjectMeta{Namespace: "fn", GenerateName: "w-", UID: types.UID(uid)}} +} + +const hash = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef" + +// Exactly one of N admissions for the same hash is the leader, and it is +// the one whose UID the Lease records. +func TestLeaseElector_OneLeaderPerHash(t *testing.T) { + kc := fake.NewSimpleClientset() + now := time.Date(2026, 9, 25, 12, 0, 0, 0, time.UTC) + ids := []string{"a", "b", "c", "d"} + e := &LeaseElector{KubeClient: kc, Namespace: "nvsnap-system", Deadline: 30 * time.Minute, + Now: func() time.Time { return now }, NewID: func() string { id := ids[0]; ids = ids[1:]; return id }} + + roles := map[Role]int{} + var leaderID string + for range 4 { + r, id, err := e.Elect(context.Background(), hash, pod("")) + if err != nil { + t.Fatal(err) + } + roles[r]++ + if r == RoleLeader { + leaderID = id + } else if id != "" { + t.Errorf("follower must not receive an id, got %q", id) + } + } + if roles[RoleLeader] != 1 || roles[RoleFollower] != 3 { + t.Fatalf("roles = %v, want 1 leader 3 followers", roles) + } + lease, err := kc.CoordinationV1().Leases("nvsnap-system").Get(context.Background(), LeaseName(hash), metav1.GetOptions{}) + if err != nil { + t.Fatal(err) + } + if leaderID != "a" || *lease.Spec.HolderIdentity != "a" || lease.Annotations[LeaderIDAnnotation] != "a" || lease.Annotations[LeaderNamespaceAnnotation] != "fn" { + t.Errorf("lease must record the first admission's id as leader (got %q): %+v", leaderID, lease.ObjectMeta) + } + if lease.Annotations[DeadlineAnnotation] != now.Add(30*time.Minute).Format(time.RFC3339) { + t.Errorf("deadline = %q, want admission + 30m", lease.Annotations[DeadlineAnnotation]) + } + if lease.Labels[LeaseKindLabel] != LeaseKindValue || lease.Labels[HashLabel] == "" || lease.Annotations[HashAnnotation] != hash { + t.Errorf("lease must be selectable by kind and hash: %+v", lease.ObjectMeta) + } + // A different hash is a separate election. + other := "ffff" + hash[4:] + ids = []string{"z"} + if r, _, _ := e.Elect(context.Background(), other, pod("")); r != RoleLeader { + t.Errorf("first admission of another hash must lead, got %s", r) + } +} + +func TestLeaseElector_Errors(t *testing.T) { + e := &LeaseElector{KubeClient: fake.NewSimpleClientset(), Namespace: "nvsnap-system"} + if _, _, err := e.Elect(context.Background(), hash, nil); err == nil { + t.Error("nil pod; want error") + } + // Pods have no UID at CREATE admission; the elector must not need one. + if r, id, err := e.Elect(context.Background(), hash, pod("")); err != nil || r != RoleLeader || id == "" { + t.Errorf("uid-less pod must be electable with a minted id, got role=%s id=%q err=%v", r, id, err) + } + kc := fake.NewSimpleClientset() + kc.PrependReactor("create", "leases", func(k8stesting.Action) (bool, runtime.Object, error) { + return true, nil, errors.New("apiserver unavailable") + }) + e = &LeaseElector{KubeClient: kc, Namespace: "nvsnap-system"} + if _, _, err := e.Elect(context.Background(), hash, pod("a")); err == nil { + t.Error("a non-AlreadyExists create error must surface so the webhook fails open") + } + if (&LeaseElector{Namespace: "x"}).deadline() != DefaultDeadline { + t.Error("zero Deadline must default") + } + var _ *coordinationv1.Lease +} + +func TestPodIdentity(t *testing.T) { + if got := PodIdentity(pod("u")); got != "fn/w-*(u)" { + t.Errorf("generateName pod identity = %q", got) + } + if got := PodIdentity(pod("")); got != "fn/w-*" { + t.Errorf("admission-time pod (no uid) identity = %q", got) + } + named := pod("u") + named.Name = "w-abc" + if got := PodIdentity(named); got != "fn/w-abc(u)" { + t.Errorf("named pod identity = %q", got) + } + if PodIdentity(nil) != "" { + t.Error("nil pod must yield empty identity") + } +} diff --git a/src/compute-plane-services/nvsnap/internal/rootfsonly/composer.go b/src/compute-plane-services/nvsnap/internal/rootfsonly/composer.go index 2d8c1f7140..0f5d4d1cae 100644 --- a/src/compute-plane-services/nvsnap/internal/rootfsonly/composer.go +++ b/src/compute-plane-services/nvsnap/internal/rootfsonly/composer.go @@ -98,7 +98,7 @@ func composeEngineCompatFlags(c corev1.Container) []string { for i, cmd := range c.Command { flags = append(flags, "cmd["+itoa(i)+"]:"+cmd) } - for i, a := range c.Args { + for i, a := range stripRoleFlags(c.Args) { flags = append(flags, "arg["+itoa(i)+"]:"+a) } for _, e := range cacheRelevantEnv(c.Env) { @@ -156,7 +156,13 @@ func cacheRelevantEnv(envs []corev1.EnvVar) []corev1.EnvVar { "TORCH_DISTRIBUTED_DEBUG": {}, "VLLM_LOGGING_LEVEL": {}, } - skipPrefix := []string{"NVSNAP_"} + // Disaggregated-serving wiring (Dynamo runtime env, etcd and NATS + // endpoints) says where a worker plugs in, not what it downloads or + // compiles; see stripRoleFlags. + for _, name := range roleEnvExact { + skipExact[name] = struct{}{} + } + skipPrefix := append([]string{"NVSNAP_"}, roleEnvPrefixes...) out := make([]corev1.EnvVar, 0, len(envs)) for _, e := range envs { if _, ok := skipExact[e.Name]; ok { @@ -190,23 +196,129 @@ func cacheRelevantEnv(envs []corev1.EnvVar) []corev1.EnvVar { // engines that wrap the launch command in /bin/bash -c). var modelFlagPattern = regexp.MustCompile(`--model(?:-path)?(?:=|\s+)([^\s\\]+)`) -// inferModelID is a best-effort extractor for the human-readable model -// identifier. Empty result is fine — ModelID is for display only and -// doesn't affect hash discrimination (the args themselves are in -// EngineCompatFlags). Conventions: +// modelEnvNames are the env vars engines and NIMs read the model identity +// from when it is not on the command line, in precedence order. +var modelEnvNames = []string{"HF_MODEL_ID", "MODEL_ID", "NIM_MODEL_NAME"} + +// InferModelID is a best-effort extractor for the model identifier the +// container will download. It looks, in order, at: +// +// - command+args as a token list: "--model X", "--model-path X", +// "--model=X". Dynamo workers (python3 -m dynamo.vllm --model X) and +// any chart that lists args one per item land here. +// - each arg as free text, for the bash -lc "vllm serve --model X ..." +// wrapper convention. +// - the env vars in modelEnvNames. +// - a NIM image, whose identity is the image itself. // -// - vLLM / SGLang / TRT-LLM: --model or --model-path flag inside an -// Args[0] shell script (most nvsnap workloads use this pattern). -// - NIM: the model id is encoded in the image name -// (nvcr.io/nim//:). -func inferModelID(c corev1.Container) string { +// Empty means no model identity: the pod is not a downloader. +func InferModelID(c corev1.Container) string { + argv := append(append([]string{}, c.Command...), c.Args...) + for i, tok := range argv { + switch { + case tok == "--model" || tok == "--model-path": + if i+1 < len(argv) && !strings.HasPrefix(argv[i+1], "-") { + return argv[i+1] + } + case strings.HasPrefix(tok, "--model=") || strings.HasPrefix(tok, "--model-path="): + if v := tok[strings.IndexByte(tok, '=')+1:]; v != "" { + return v + } + } + } for _, a := range c.Args { if m := modelFlagPattern.FindStringSubmatch(a); m != nil { return m[1] } } + for _, name := range modelEnvNames { + for _, e := range c.Env { + if e.Name == name && e.Value != "" { + return e.Value + } + } + } if IsNIMImage(c.Image) { return c.Image } return "" } + +// inferModelID keeps the historical unexported name for the composer. +func inferModelID(c corev1.Container) string { return InferModelID(c) } + +// IsModelWorkload reports whether the pod is a model downloader: it asks +// for a GPU and its main container names a model. Supporting pods in a +// chart (frontend, router, planner, etcd, nats) fail one or both tests +// and are left alone by the election. +func IsModelWorkload(pod *corev1.Pod, mainContainer int) (modelID string, ok bool) { + if pod == nil || mainContainer < 0 || mainContainer >= len(pod.Spec.Containers) { + return "", false + } + if podGPURequest(pod) == 0 { + return "", false + } + modelID = InferModelID(pod.Spec.Containers[mainContainer]) + return modelID, modelID != "" +} + +// Role flags. In a disaggregated deployment (Dynamo over vLLM, SGLang or +// TRT-LLM) the prefill and decode workers run the same image, download +// the same model and differ only in a flag that names their role and in +// the KV-transfer plumbing between them. Hashing those flags would give +// each role its own cache and its own download, which is the duplication +// the election exists to remove. The model tree is identical across roles +// and mounts read-only; compile caches live in a per-pod writable shadow, +// so a role that needs different kernels recompiles into it and nothing +// is corrupted. Flags that change what is downloaded (--model, --revision, +// --tokenizer, quantization) stay in the hash. +var ( + roleFlagsNoValue = map[string]bool{ + "--is-prefill-worker": true, // dynamo.vllm + "--is-decode-worker": true, + } + roleFlagsWithValue = map[string]bool{ + "--disaggregation-mode": true, // dynamo.sglang, dynamo.trtllm + "--disaggregation-strategy": true, + "--disaggregation-bootstrap-port": true, + "--disaggregation-transfer-backend": true, + "--kv-transfer-config": true, // vLLM connector JSON, carries kv_role + "--kv-events-config": true, // KV event publishing, prefill side only + } + // GROVE_*: Grove gang-scheduling env the Dynamo operator injects, with + // the component name and the pod index in it. + roleEnvPrefixes = []string{"DYN_", "DYNAMO_", "GROVE_"} + roleEnvExact = []string{"ETCD_ENDPOINTS", "NATS_SERVER", "NATS_URL"} + + // roleFlagInString removes the same flags from a shell-script arg (the + // bash -lc "vllm serve ..." convention), value quoted or bare. + roleFlagInString = regexp.MustCompile( + `\s--(?:is-prefill-worker|is-decode-worker)\b` + + `|\s--(?:disaggregation-mode|disaggregation-strategy|disaggregation-bootstrap-port|disaggregation-transfer-backend|kv-transfer-config|kv-events-config)(?:=|\s+)(?:'[^']*'|"[^"]*"|\S+)`) +) + +// stripRoleFlags returns args without the role flags, in both the +// one-token-per-item form and the single shell-string form. +func stripRoleFlags(args []string) []string { + out := make([]string, 0, len(args)) + for i := 0; i < len(args); i++ { + a := args[i] + if roleFlagsNoValue[a] { + continue + } + if roleFlagsWithValue[a] { + if i+1 < len(args) && !strings.HasPrefix(args[i+1], "-") { + i++ + } + continue + } + if eq := strings.IndexByte(a, '='); eq > 0 && roleFlagsWithValue[a[:eq]] { + continue + } + if strings.ContainsAny(a, " \n\t") { + a = roleFlagInString.ReplaceAllString(a, "") + } + out = append(out, a) + } + return out +} diff --git a/src/compute-plane-services/nvsnap/internal/rootfsonly/composer_test.go b/src/compute-plane-services/nvsnap/internal/rootfsonly/composer_test.go index cdbc6245dd..c05415c8bd 100644 --- a/src/compute-plane-services/nvsnap/internal/rootfsonly/composer_test.go +++ b/src/compute-plane-services/nvsnap/internal/rootfsonly/composer_test.go @@ -22,6 +22,7 @@ import ( "testing" corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/api/resource" "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/checkpointstore" ) @@ -247,3 +248,155 @@ func TestCompose_CommandIncluded(t *testing.T) { t.Fatalf("Command not included: %v", in.EngineCompatFlags) } } + +// The model identity is what makes a pod a downloader, and charts spell it +// several ways: Dynamo lists "--model" and the value as separate items, +// stock charts wrap "vllm serve --model X" in one bash string, NIMs use +// env or the image. Frontend pods have none of them. +func TestInferModelID_Forms(t *testing.T) { + cases := map[string]struct { + c corev1.Container + want string + }{ + "dynamo list form": {corev1.Container{ + Command: []string{"python3", "-m", "dynamo.vllm"}, + Args: []string{"--model", "Qwen/Qwen3-0.6B", "--is-decode-worker"}, + }, "Qwen/Qwen3-0.6B"}, + "sglang list form": {corev1.Container{ + Command: []string{"python3", "-m", "dynamo.sglang"}, + Args: []string{"--model-path", "google/gemma-4-31B-it", "--tp", "2"}, + }, "google/gemma-4-31B-it"}, + "equals form": {corev1.Container{Args: []string{"--model=meta-llama/Llama-3.1-8B-Instruct"}}, "meta-llama/Llama-3.1-8B-Instruct"}, + "bash wrapper": {corev1.Container{ + Command: []string{"/bin/bash", "-lc"}, + Args: []string{"set -e\nnohup setsid vllm serve --model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --tensor-parallel-size 2 &\nwhile true; do sleep 30; done"}, + }, "TinyLlama/TinyLlama-1.1B-Chat-v1.0"}, + "env HF_MODEL_ID": {corev1.Container{Env: []corev1.EnvVar{{Name: "HF_MODEL_ID", Value: "openai/whisper-large-v3"}}}, "openai/whisper-large-v3"}, + "nim image": {corev1.Container{Image: "nvcr.io/nim/meta/llama-3.1-8b-instruct:1.8.3"}, "nvcr.io/nim/meta/llama-3.1-8b-instruct:1.8.3"}, + "dangling flag": {corev1.Container{Args: []string{"--model"}}, ""}, + "flag then flag": {corev1.Container{Args: []string{"--model", "--port"}}, ""}, + "dynamo frontend": {corev1.Container{Command: []string{"python3", "-m", "dynamo.frontend"}, Args: []string{"--router-mode", "kv"}}, ""}, + } + for name, tc := range cases { + if got := InferModelID(tc.c); got != tc.want { + t.Errorf("%s: InferModelID = %q, want %q", name, got, tc.want) + } + } +} + +func TestIsModelWorkload(t *testing.T) { + gpu := corev1.ResourceRequirements{Limits: corev1.ResourceList{"nvidia.com/gpu": resource.MustParse("2")}} + worker := &corev1.Pod{Spec: corev1.PodSpec{Containers: []corev1.Container{{ + Command: []string{"python3", "-m", "dynamo.vllm"}, Args: []string{"--model", "Qwen/Qwen3-0.6B", "--is-prefill-worker"}, Resources: gpu, + }}}} + frontend := &corev1.Pod{Spec: corev1.PodSpec{Containers: []corev1.Container{{ + Command: []string{"python3", "-m", "dynamo.frontend"}, + }}}} + gpuNoModel := &corev1.Pod{Spec: corev1.PodSpec{Containers: []corev1.Container{{ + Command: []string{"python3", "train.py"}, Resources: gpu, + }}}} + modelNoGPU := &corev1.Pod{Spec: corev1.PodSpec{Containers: []corev1.Container{{ + Args: []string{"--model", "Qwen/Qwen3-0.6B"}, + }}}} + if id, ok := IsModelWorkload(worker, 0); !ok || id != "Qwen/Qwen3-0.6B" { + t.Errorf("dynamo worker: (%q,%v), want downloader", id, ok) + } + for name, p := range map[string]*corev1.Pod{"frontend": frontend, "gpu without model": gpuNoModel, "model without gpu": modelNoGPU} { + if _, ok := IsModelWorkload(p, 0); ok { + t.Errorf("%s must not be classified as a downloader", name) + } + } + if _, ok := IsModelWorkload(worker, 3); ok { + t.Error("out-of-range main container must not classify") + } +} + +// Prefill and decode workers of one disaggregated deployment must share a +// hash: same image, same model, same download. Only the role flags and the +// Dynamo/etcd/NATS wiring differ, and none of that changes the cache tree. +func TestCompose_RoleNeutralAcrossPrefillAndDecode(t *testing.T) { + c := &HashInputComposer{CUDADriverMajor: 580} + hashOf := func(cmd, args []string, env ...corev1.EnvVar) string { + p := &corev1.Pod{Spec: corev1.PodSpec{Containers: []corev1.Container{{ + Name: "main", Image: "vllm/vllm-openai:v0.20.0", Command: cmd, Args: args, Env: env, + }}}} + return checkpointstore.ComputeHash(c.Compose(p, 0)) + } + dyn := []string{"python3", "-m", "dynamo.vllm"} + prefill := hashOf(dyn, []string{"--model", "Qwen/Qwen3-0.6B", "--is-prefill-worker"}, + corev1.EnvVar{Name: "DYN_NAMESPACE", Value: "dgd-a"}, corev1.EnvVar{Name: "ETCD_ENDPOINTS", Value: "etcd-a:2379"}) + decode := hashOf(dyn, []string{"--model", "Qwen/Qwen3-0.6B", "--is-decode-worker"}, + corev1.EnvVar{Name: "DYN_NAMESPACE", Value: "dgd-b"}, corev1.EnvVar{Name: "NATS_SERVER", Value: "nats://x:4222"}) + if prefill != decode { + t.Error("dynamo.vllm prefill and decode workers must hash the same") + } + sgl := []string{"python3", "-m", "dynamo.sglang"} + sp := hashOf(sgl, []string{"--model-path", "google/gemma-4-31B-it", "--disaggregation-mode", "prefill", "--disaggregation-bootstrap-port", "8998"}) + sd := hashOf(sgl, []string{"--model-path", "google/gemma-4-31B-it", "--disaggregation-mode=decode", "--disaggregation-transfer-backend", "nixl"}) + if sp != sd { + t.Error("dynamo.sglang prefill and decode workers must hash the same") + } + bash := []string{"/bin/bash", "-lc"} + bp := hashOf(bash, []string{`vllm serve --model Qwen/Qwen3-0.6B --is-prefill-worker --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_producer"}' > /out 2>&1`}) + bd := hashOf(bash, []string{`vllm serve --model Qwen/Qwen3-0.6B --is-decode-worker --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_consumer"}' > /out 2>&1`}) + if bp != bd { + t.Error("shell-string prefill and decode workers must hash the same") + } + // What changes the download or the engine still separates hashes. + if hashOf(dyn, []string{"--model", "Qwen/Qwen3-0.6B", "--is-prefill-worker"}) == hashOf(dyn, []string{"--model", "Qwen/Qwen3-1.7B", "--is-prefill-worker"}) { + t.Error("different models must hash differently") + } + if hashOf(dyn, []string{"--model", "Qwen/Qwen3-0.6B", "--revision", "abc"}) == hashOf(dyn, []string{"--model", "Qwen/Qwen3-0.6B", "--revision", "def"}) { + t.Error("different revisions download different files and must hash differently") + } + if hashOf(dyn, []string{"--model", "Qwen/Qwen3-0.6B"}, corev1.EnvVar{Name: "HF_HUB_OFFLINE", Value: "1"}) == hashOf(dyn, []string{"--model", "Qwen/Qwen3-0.6B"}) { + t.Error("cache-relevant env must still participate in the hash") + } +} + +func TestStripRoleFlags(t *testing.T) { + got := stripRoleFlags([]string{"--model", "m", "--is-decode-worker", "--disaggregation-mode", "decode", "--disaggregation-strategy=prefill_first", "--tp", "2", "--kv-transfer-config"}) + want := []string{"--model", "m", "--tp", "2"} + if strings.Join(got, " ") != strings.Join(want, " ") { + t.Errorf("stripRoleFlags = %v, want %v", got, want) + } + // A dangling value-flag at the end must not eat a following flag. + got = stripRoleFlags([]string{"--disaggregation-mode", "--port", "8000"}) + if strings.Join(got, " ") != "--port 8000" { + t.Errorf("value flag followed by a flag: %v", got) + } +} + +// Taken from the pods the Dynamo operator (vllm-runtime 1.1.1) created on +// dev1 for the NVCF disaggregated sample: prefill also carries +// --kv-events-config, and the operator injects Grove scheduling env whose +// values differ per component and per replica. +func TestCompose_RoleNeutral_LiveDynamoOperatorPods(t *testing.T) { + c := &HashInputComposer{CUDADriverMajor: 580} + worker := func(args []string, pclq string) *corev1.Pod { + env := []corev1.EnvVar{ + {Name: "DYN_COMPONENT", Value: pclq}, {Name: "DYN_NAMESPACE", Value: "myllm"}, {Name: "DYN_SYSTEM_PORT", Value: "9090"}, + {Name: "GROVE_PCLQ_NAME", Value: pclq}, {Name: "GROVE_PCLQ_POD_INDEX", Value: "0"}, {Name: "GROVE_PCS_NAME", Value: "myllm"}, + {Name: "NATS_SERVER", Value: "nats://dynamo-operator-nats:4222"}, + {Name: "POD_NAME", ValueFrom: &corev1.EnvVarSource{FieldRef: &corev1.ObjectFieldSelector{FieldPath: "metadata.name"}}}, + {Name: "NIXL_TELEMETRY_ENABLE", Value: "true"}, + } + return &corev1.Pod{Spec: corev1.PodSpec{Containers: []corev1.Container{{ + Name: "main", Image: "nvcr.io/nvidia/ai-dynamo/vllm-runtime:1.1.1", + Command: []string{"python3", "-m", "dynamo.vllm"}, Args: args, Env: env, + }}}} + } + prefill := worker([]string{"--model", "Qwen/Qwen3-0.6B", "--disaggregation-mode", "prefill", "--tensor-parallel-size", "1", + "--kv-transfer-config", `{"kv_connector":"NixlConnector","kv_role":"kv_both"}`, + "--kv-events-config", `{"publisher":"zmq","topic":"kv-events","endpoint":"tcp://*:20080","enable_kv_cache_events":true}`}, "myllm-0-vllmprefillworker") + decode := worker([]string{"--model", "Qwen/Qwen3-0.6B", "--disaggregation-mode", "decode", "--tensor-parallel-size", "1"}, "myllm-0-vllmdecodeworker") + replica := worker([]string{"--model", "Qwen/Qwen3-0.6B", "--disaggregation-mode", "decode", "--tensor-parallel-size", "1"}, "myllm-0-vllmdecodeworker") + replica.Spec.Containers[0].Env[4].Value = "1" // GROVE_PCLQ_POD_INDEX + hp, hd, hr := checkpointstore.ComputeHash(c.Compose(prefill, 0)), checkpointstore.ComputeHash(c.Compose(decode, 0)), checkpointstore.ComputeHash(c.Compose(replica, 0)) + if hp != hd { + t.Errorf("live prefill and decode workers must hash the same: %s vs %s", hp[:8], hd[:8]) + } + if hd != hr { + t.Errorf("two replicas of one component must hash the same: %s vs %s", hd[:8], hr[:8]) + } +} diff --git a/src/compute-plane-services/nvsnap/internal/rootfsonly/orchestrator.go b/src/compute-plane-services/nvsnap/internal/rootfsonly/orchestrator.go index 64592caea6..b4cd3cee0c 100644 --- a/src/compute-plane-services/nvsnap/internal/rootfsonly/orchestrator.go +++ b/src/compute-plane-services/nvsnap/internal/rootfsonly/orchestrator.go @@ -39,6 +39,11 @@ import ( // CaptureRequest describes a pod the agent should capture. type CaptureRequest struct { + // Hash, when set, is the content hash to store the capture under + // (stamped on the pod at admission by the election webhook). Empty + // means compute it from HashInput. + Hash string + // PodUID is metadata.uid — used by the PIDResolver and to compute // the kubelet host paths for emptyDir volumes. PodUID string @@ -207,7 +212,13 @@ func (c *Capturer) Capture(ctx context.Context, req CaptureRequest) (checkpoints return checkpointstore.Manifest{}, errors.New("rootfsonly: CaptureRequest.Spec is nil") } - hash := checkpointstore.ComputeHash(req.HashInput) + // The admission webhook stamps the hash it composed onto elected pods + // (nvsnap.io/hash) and the watcher passes it through; honour it so the + // capture lands under the hash the followers were decorated against. + hash := req.Hash + if hash == "" { + hash = checkpointstore.ComputeHash(req.HashInput) + } log := c.logger().WithFields(logrus.Fields{ "hash": checkpointstore.ShortHash(hash), "pod": req.Namespace + "/" + req.Name, diff --git a/src/compute-plane-services/nvsnap/internal/rootfsonly/orchestrator_test.go b/src/compute-plane-services/nvsnap/internal/rootfsonly/orchestrator_test.go index eb826d86a6..e4fef9b306 100644 --- a/src/compute-plane-services/nvsnap/internal/rootfsonly/orchestrator_test.go +++ b/src/compute-plane-services/nvsnap/internal/rootfsonly/orchestrator_test.go @@ -440,3 +440,39 @@ func TestCollectCacheEnv(t *testing.T) { t.Error("expected nil for empty cacheDir / nil spec / bad index") } } + +// The election webhook stamps the hash it composed at admission onto the +// pod; the capture must land under that hash, not under a recomposition +// from the live pod, or the gated followers mount a claim that never binds. +func TestCapture_StampedHashOverridesComposition(t *testing.T) { + env := newOrchTestEnv(t) + env.addProc(t, env.upperdirMountinfo()) + env.addUpperdirContent(t) + + stamped := "feedfacefeedfacefeedfacefeedfacefeedfacefeedfacefeedfacefeedface" + req := CaptureRequest{ + Hash: stamped, + PodUID: env.podUID, + Namespace: "ns", + Name: "p", + Spec: pod("vllm/x", nil, nil), + MainContainer: 0, + HashInput: checkpointstore.HashInput{ + ImageDigest: "x", ModelID: "y", CUDADriverMajor: 1, + CaptureFormatVersion: checkpointstore.CaptureFormatVersion, + }, + } + m, err := env.capturer().Capture(context.Background(), req) + if err != nil { + t.Fatal(err) + } + if m.Hash != stamped { + t.Fatalf("manifest hash = %s, want the stamped %s", m.Hash, stamped) + } + if _, err := env.backend.Stat(context.Background(), stamped); err != nil { + t.Errorf("capture not stored under the stamped hash: %v", err) + } + if _, err := env.backend.Stat(context.Background(), checkpointstore.ComputeHash(req.HashInput)); err == nil { + t.Error("capture must not also land under the recomposed hash") + } +} diff --git a/src/compute-plane-services/nvsnap/internal/rootfsonly/watcher.go b/src/compute-plane-services/nvsnap/internal/rootfsonly/watcher.go index 5eb20f5d06..d019a3c8b2 100644 --- a/src/compute-plane-services/nvsnap/internal/rootfsonly/watcher.go +++ b/src/compute-plane-services/nvsnap/internal/rootfsonly/watcher.go @@ -40,6 +40,13 @@ import ( // want nvsnap to fan out across the cluster. const DefaultCaptureLabel = "nvsnap.io/capture" +// stampedHashAnnotation is the hash the election webhook composed at +// admission (election.HashAnnotation; duplicated here so rootfsonly does +// not import election). When present it overrides the watcher's own +// composition, which would otherwise differ from admission's by the +// resolved image digest and the webhook-injected env. +const stampedHashAnnotation = "nvsnap.io/hash" + // Watcher subscribes to Pod events on the local node, detects warm pods // (label + Ready), and triggers Capturer.Capture once per pod UID via the // HashInputComposer. Idempotent at every layer: even if multiple events @@ -279,6 +286,7 @@ func (w *Watcher) runCapture(ctx context.Context, pod *corev1.Pod) { hashInput := w.Composer.Compose(pod, 0) req := CaptureRequest{ + Hash: pod.Annotations[stampedHashAnnotation], PodUID: string(pod.UID), Namespace: pod.Namespace, Name: pod.Name, diff --git a/src/compute-plane-services/nvsnap/internal/rootfsonly/watcher_test.go b/src/compute-plane-services/nvsnap/internal/rootfsonly/watcher_test.go index 394bd088a5..7226fb2200 100644 --- a/src/compute-plane-services/nvsnap/internal/rootfsonly/watcher_test.go +++ b/src/compute-plane-services/nvsnap/internal/rootfsonly/watcher_test.go @@ -493,3 +493,33 @@ func TestWatcher_CommittedCaptureStaysDeduped(t *testing.T) { t.Fatal("a committed capture released its dedup mark; the pod would be recaptured on every resync") } } + +// A pod the election webhook stamped with nvsnap.io/hash is captured under +// that hash, so the followers decorated against rox- find it. +func TestWatcher_CapturesUnderStampedHash(t *testing.T) { + env := newWatcherEnv(t) + env.addProc(t, env.upperdirMountinfo()) + env.addUpperdirContent(t) + count := &countingBackend{Backend: env.backend} + capturer := env.capturer() + capturer.Backend = count + w := env.watcher() + w.Capturer = capturer + w.sem = make(chan struct{}, w.concurrency()) + + stamped := "feedfacefeedfacefeedfacefeedfacefeedfacefeedfacefeedfacefeedface" + pod := fakePod(types.UID(env.podUID), "vllm-8b", map[string]string{DefaultCaptureLabel: "true"}, true) + pod.Annotations = map[string]string{stampedHashAnnotation: stamped} + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + w.HandlePodEvent(ctx, pod) + if !waitFor(t, 3*time.Second, func() bool { + _, err := count.Stat(ctx, stamped) + return err == nil + }) { + t.Fatalf("capture never landed under the stamped hash") + } + if _, err := count.Stat(ctx, checkpointstore.ComputeHash(w.Composer.Compose(pod, 0))); err == nil { + t.Error("watcher must not recompose the hash when one is stamped") + } +} diff --git a/src/compute-plane-services/nvsnap/internal/server/BUILD.bazel b/src/compute-plane-services/nvsnap/internal/server/BUILD.bazel index 74d28ce7ae..35bef9dd69 100644 --- a/src/compute-plane-services/nvsnap/internal/server/BUILD.bazel +++ b/src/compute-plane-services/nvsnap/internal/server/BUILD.bazel @@ -5,6 +5,7 @@ go_library( srcs = [ "agent_auth.go", "demo.go", + "election_release.go", "lookup.go", "manifests.go", "observability_proxy.go", @@ -20,6 +21,7 @@ go_library( deps = [ "//src/compute-plane-services/nvsnap/internal/checkpointstore", "//src/compute-plane-services/nvsnap/internal/db", + "//src/compute-plane-services/nvsnap/internal/election", "//src/compute-plane-services/nvsnap/internal/metrics", "@com_github_gorilla_mux//:mux", "@com_github_gorilla_websocket//:websocket", @@ -44,6 +46,7 @@ go_test( "checkpoint_status_test.go", "delete_cascade_integration_test.go", "delete_checkpoint_test.go", + "election_release_test.go", "l2_contract_test.go", "lookup_test.go", "observability_proxy_test.go", @@ -58,6 +61,7 @@ go_test( "//src/compute-plane-services/nvsnap/internal/agent", "//src/compute-plane-services/nvsnap/internal/checkpointstore", "//src/compute-plane-services/nvsnap/internal/db", + "//src/compute-plane-services/nvsnap/internal/election", "@com_github_sirupsen_logrus//:logrus", "@io_k8s_api//coordination/v1:coordination", "@io_k8s_api//core/v1:core", @@ -67,6 +71,7 @@ go_test( "@io_k8s_apimachinery//pkg/apis/meta/v1/unstructured", "@io_k8s_apimachinery//pkg/runtime", "@io_k8s_apimachinery//pkg/runtime/schema", + "@io_k8s_apimachinery//pkg/types", "@io_k8s_client_go//dynamic/fake", "@io_k8s_client_go//kubernetes/fake", "@io_k8s_client_go//testing", diff --git a/src/compute-plane-services/nvsnap/internal/server/election_release.go b/src/compute-plane-services/nvsnap/internal/server/election_release.go new file mode 100644 index 0000000000..aec4fbc00d --- /dev/null +++ b/src/compute-plane-services/nvsnap/internal/server/election_release.go @@ -0,0 +1,186 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "encoding/json" + "fmt" + "time" + + "github.com/sirupsen/logrus" + corev1 "k8s.io/api/core/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/types" + "k8s.io/client-go/kubernetes" + + "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/checkpointstore" + "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/election" +) + +// Election release: the webhook gates follower pods of a hash until the +// leader's capture is promoted (docs/proposals/helm-chart-cache-election.md). +// nvsnap-server is the one component that learns about the promote (the +// agent posts pvc-state), so it removes the gates on ready and, when the +// leader fails or overruns its deadline, evicts the followers so their +// controller recreates them and a new election runs. + +// electionReleaser holds the cluster access the release needs; both the +// HTTP handler and the reconciler use it. +type electionReleaser struct { + kube kubernetes.Interface + leaseNS string + log logrus.FieldLogger + now func() time.Time +} + +func followerSelector(hash string) string { + return fmt.Sprintf("%s=%s,%s=true", election.HashLabel, checkpointstore.ShortHash(hash), election.GatedLabel) +} + +// listGated returns the gated followers of hash across all namespaces. +func (r *electionReleaser) listGated(ctx context.Context, hash string) ([]corev1.Pod, error) { + pods, err := r.kube.CoreV1().Pods("").List(ctx, metav1.ListOptions{LabelSelector: followerSelector(hash)}) + if err != nil { + return nil, fmt.Errorf("list gated followers: %w", err) + } + return pods.Items, nil +} + +// release drops the scheduling gate on every gated follower of hash and +// deletes the election Lease; the rox is Bound at this point so the pods +// schedule straight into a warm start. Idempotent. +func (r *electionReleaser) release(ctx context.Context, hash string) (released int, err error) { + pods, err := r.listGated(ctx, hash) + if err != nil { + return 0, err + } + patch, _ := json.Marshal(map[string]any{ + "metadata": map[string]any{"labels": map[string]string{election.GatedLabel: "false"}}, + "spec": map[string]any{"schedulingGates": []any{}}, + }) + for i := range pods { + p := &pods[i] + if _, err := r.kube.CoreV1().Pods(p.Namespace).Patch(ctx, p.Name, types.MergePatchType, patch, metav1.PatchOptions{}); err != nil { + if apierrors.IsNotFound(err) { + continue + } + return released, fmt.Errorf("ungate %s/%s: %w", p.Namespace, p.Name, err) + } + released++ + } + r.deleteLease(ctx, hash) + return released, nil +} + +// evict deletes the gated followers of hash that a controller will +// recreate, and the Lease, so the recreated pods run a fresh election. A +// follower's volumes name a claim that will now never bind and pod volumes +// are immutable, so recreation is the only way to change its fate. Pods +// without a controller owner are left in place and logged. +func (r *electionReleaser) evict(ctx context.Context, hash, reason string) (evicted int, err error) { + pods, err := r.listGated(ctx, hash) + if err != nil { + return 0, err + } + for i := range pods { + p := &pods[i] + if metav1.GetControllerOf(p) == nil { + r.log.WithFields(logrus.Fields{"pod": p.Namespace + "/" + p.Name, "hash": checkpointstore.ShortHash(hash)}). + Warn("election: gated follower has no controller; left gated (delete it by hand)") + continue + } + if err := r.kube.CoreV1().Pods(p.Namespace).Delete(ctx, p.Name, metav1.DeleteOptions{}); err != nil && !apierrors.IsNotFound(err) { + return evicted, fmt.Errorf("evict %s/%s: %w", p.Namespace, p.Name, err) + } + evicted++ + } + r.deleteLease(ctx, hash) + r.log.WithFields(logrus.Fields{"hash": checkpointstore.ShortHash(hash), "evicted": evicted, "reason": reason}). + Info("election: leader gone; followers evicted for re-election") + return evicted, nil +} + +func (r *electionReleaser) deleteLease(ctx context.Context, hash string) { + if err := r.kube.CoordinationV1().Leases(r.leaseNS).Delete(ctx, election.LeaseName(hash), metav1.DeleteOptions{}); err != nil && !apierrors.IsNotFound(err) { + r.log.WithError(err).WithField("hash", checkpointstore.ShortHash(hash)).Warn("election: lease delete failed") + } +} + +// onPromoteState reacts to the agent's promote state for hash. +func (r *electionReleaser) onPromoteState(ctx context.Context, hash, state string) { + if r == nil || r.kube == nil { + return + } + switch state { + case "ready": + n, err := r.release(ctx, hash) + if err != nil { + r.log.WithError(err).WithField("hash", checkpointstore.ShortHash(hash)).Warn("election: release failed") + return + } + r.log.WithFields(logrus.Fields{"hash": checkpointstore.ShortHash(hash), "released": n}).Info("election: capture promoted; followers released") + case "failed": + n, err := r.evict(ctx, hash, "promote failed") + if err != nil { + r.log.WithError(err).WithFields(logrus.Fields{"hash": checkpointstore.ShortHash(hash), "evicted": n}).Warn("election: evict incomplete") + } + } +} + +// reconcile checks every live election: a leader pod that is gone or has +// terminated without a promote, or a Lease past its deadline, means the +// followers will never be released; evict them so a new election runs. +func (r *electionReleaser) reconcile(ctx context.Context) { + if r == nil || r.kube == nil { + return + } + leases, err := r.kube.CoordinationV1().Leases(r.leaseNS).List(ctx, metav1.ListOptions{ + LabelSelector: election.LeaseKindLabel + "=" + election.LeaseKindValue, + }) + if err != nil { + r.log.WithError(err).Warn("election: list leases failed") + return + } + now := time.Now() + if r.now != nil { + now = r.now() + } + for i := range leases.Items { + l := &leases.Items[i] + hash := l.Annotations[election.HashAnnotation] + if hash == "" { + continue + } + if dl, err := time.Parse(time.RFC3339, l.Annotations[election.DeadlineAnnotation]); err == nil && now.After(dl) { + _, _ = r.evict(ctx, hash, "deadline passed") + continue + } + if !r.leaderAlive(ctx, l.Annotations[election.LeaderNamespaceAnnotation], l.Annotations[election.LeaderIDAnnotation], hash) { + _, _ = r.evict(ctx, hash, "leader pod gone or terminated") + } + } +} + +// leaderAlive finds the leader by its election id among the pods stamped +// with the hash in its namespace and reports whether it can still capture. +func (r *electionReleaser) leaderAlive(ctx context.Context, ns, id, hash string) bool { + pods, err := r.kube.CoreV1().Pods(ns).List(ctx, metav1.ListOptions{ + LabelSelector: election.HashLabel + "=" + checkpointstore.ShortHash(hash), + }) + if err != nil { + // Cannot tell; do not evict on a transient API error. + r.log.WithError(err).Warn("election: list leader candidates failed") + return true + } + for i := range pods.Items { + p := &pods.Items[i] + if p.Annotations[election.ElectionIDAnnotation] != id { + continue + } + return p.Status.Phase != corev1.PodFailed && p.Status.Phase != corev1.PodSucceeded && p.DeletionTimestamp == nil + } + return false +} diff --git a/src/compute-plane-services/nvsnap/internal/server/election_release_test.go b/src/compute-plane-services/nvsnap/internal/server/election_release_test.go new file mode 100644 index 0000000000..a0e9b3619d --- /dev/null +++ b/src/compute-plane-services/nvsnap/internal/server/election_release_test.go @@ -0,0 +1,135 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "testing" + "time" + + "github.com/sirupsen/logrus" + coordinationv1 "k8s.io/api/coordination/v1" + corev1 "k8s.io/api/core/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/types" + "k8s.io/client-go/kubernetes/fake" + + "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/checkpointstore" + "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/election" +) + +const electTestHash = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef" + +func quietLog() logrus.FieldLogger { + l := logrus.New() + l.SetLevel(logrus.PanicLevel) + return l +} + +func gatedFollower(name string, owned bool) *corev1.Pod { + p := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Name: name, Namespace: "fn", UID: types.UID("uid-" + name), + Labels: map[string]string{election.HashLabel: checkpointstore.ShortHash(electTestHash), election.GatedLabel: "true"}}, + Spec: corev1.PodSpec{SchedulingGates: []corev1.PodSchedulingGate{{Name: election.GateName}}}, + } + if owned { + ctrl := true + p.OwnerReferences = []metav1.OwnerReference{{APIVersion: "apps/v1", Kind: "ReplicaSet", Name: "rs", UID: "rs-uid", Controller: &ctrl}} + } + return p +} + +func leaderPod(phase corev1.PodPhase) *corev1.Pod { + const id = "L" + return &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Name: "leader-" + id, Namespace: "fn", UID: types.UID("uid-" + id), + Labels: map[string]string{election.HashLabel: checkpointstore.ShortHash(electTestHash)}, + Annotations: map[string]string{election.ElectionIDAnnotation: id}}, + Status: corev1.PodStatus{Phase: phase}, + } +} + +func electionLease(leaderUID string, deadline time.Time) *coordinationv1.Lease { + return &coordinationv1.Lease{ObjectMeta: metav1.ObjectMeta{ + Name: election.LeaseName(electTestHash), Namespace: "nvsnap-system", + Labels: map[string]string{election.LeaseKindLabel: election.LeaseKindValue}, + Annotations: map[string]string{ + election.HashAnnotation: electTestHash, + election.LeaderNamespaceAnnotation: "fn", + election.LeaderIDAnnotation: leaderUID, + election.DeadlineAnnotation: deadline.UTC().Format(time.RFC3339), + }, + }} +} + +func TestElectionRelease_ReadyUngatesFollowersAndDropsLease(t *testing.T) { + kc := fake.NewSimpleClientset(gatedFollower("f1", true), gatedFollower("f2", false), electionLease("L", time.Now().Add(time.Hour))) + r := &electionReleaser{kube: kc, leaseNS: "nvsnap-system", log: quietLog()} + r.onPromoteState(context.Background(), electTestHash, "ready") + for _, name := range []string{"f1", "f2"} { + p, err := kc.CoreV1().Pods("fn").Get(context.Background(), name, metav1.GetOptions{}) + if err != nil { + t.Fatal(err) + } + if len(p.Spec.SchedulingGates) != 0 || p.Labels[election.GatedLabel] != "false" { + t.Errorf("%s: gates=%v gated=%q, want released", name, p.Spec.SchedulingGates, p.Labels[election.GatedLabel]) + } + } + if _, err := kc.CoordinationV1().Leases("nvsnap-system").Get(context.Background(), election.LeaseName(electTestHash), metav1.GetOptions{}); !apierrors.IsNotFound(err) { + t.Errorf("lease must be deleted on ready, got err=%v", err) + } +} + +func TestElectionRelease_FailedEvictsOnlyControllerOwnedFollowers(t *testing.T) { + kc := fake.NewSimpleClientset(gatedFollower("owned", true), gatedFollower("bare", false), electionLease("L", time.Now().Add(time.Hour))) + r := &electionReleaser{kube: kc, leaseNS: "nvsnap-system", log: quietLog()} + r.onPromoteState(context.Background(), electTestHash, "failed") + if _, err := kc.CoreV1().Pods("fn").Get(context.Background(), "owned", metav1.GetOptions{}); !apierrors.IsNotFound(err) { + t.Errorf("controller-owned follower must be deleted for re-election, err=%v", err) + } + if _, err := kc.CoreV1().Pods("fn").Get(context.Background(), "bare", metav1.GetOptions{}); err != nil { + t.Errorf("bare follower must be left in place, err=%v", err) + } + if _, err := kc.CoordinationV1().Leases("nvsnap-system").Get(context.Background(), election.LeaseName(electTestHash), metav1.GetOptions{}); !apierrors.IsNotFound(err) { + t.Errorf("lease must be deleted on failed, got err=%v", err) + } +} + +func TestElectionRelease_ReconcileLeaderLiveness(t *testing.T) { + now := time.Date(2026, 9, 25, 12, 0, 0, 0, time.UTC) + cases := map[string]struct { + objs []*corev1.Pod + deadline time.Time + wantEvict bool + }{ + "live leader keeps followers": {[]*corev1.Pod{leaderPod(corev1.PodRunning)}, now.Add(time.Hour), false}, + "pending leader keeps followers": {[]*corev1.Pod{leaderPod(corev1.PodPending)}, now.Add(time.Hour), false}, + "leader gone evicts": {nil, now.Add(time.Hour), true}, + "failed leader evicts": {[]*corev1.Pod{leaderPod(corev1.PodFailed)}, now.Add(time.Hour), true}, + "deadline passed evicts even live": {[]*corev1.Pod{leaderPod(corev1.PodRunning)}, now.Add(-time.Minute), true}, + } + for name, tc := range cases { + kc := fake.NewSimpleClientset(gatedFollower("f", true), electionLease("L", tc.deadline)) + for _, p := range tc.objs { + if _, err := kc.CoreV1().Pods("fn").Create(context.Background(), p, metav1.CreateOptions{}); err != nil { + t.Fatal(err) + } + } + r := &electionReleaser{kube: kc, leaseNS: "nvsnap-system", log: quietLog(), now: func() time.Time { return now }} + r.reconcile(context.Background()) + _, err := kc.CoreV1().Pods("fn").Get(context.Background(), "f", metav1.GetOptions{}) + evicted := apierrors.IsNotFound(err) + if evicted != tc.wantEvict { + t.Errorf("%s: follower evicted=%v, want %v", name, evicted, tc.wantEvict) + } + } +} + +func TestElectionRelease_NilClientIsNoop(t *testing.T) { + var r *electionReleaser + r.onPromoteState(context.Background(), electTestHash, "ready") + r.reconcile(context.Background()) + (&electionReleaser{log: quietLog()}).onPromoteState(context.Background(), electTestHash, "failed") +} diff --git a/src/compute-plane-services/nvsnap/internal/server/reconciler.go b/src/compute-plane-services/nvsnap/internal/server/reconciler.go index 52c779b436..c2c3504121 100644 --- a/src/compute-plane-services/nvsnap/internal/server/reconciler.go +++ b/src/compute-plane-services/nvsnap/internal/server/reconciler.go @@ -48,6 +48,9 @@ type Reconciler struct { Namespace string // nvsnap-system; where capture CMs live Interval time.Duration // default 30s Log logrus.FieldLogger + // Elections, when set, is checked every tick for leaders that died or + // overran their deadline (election_release.go). + Elections *electionReleaser } // Run blocks reconciling on every Interval until ctx is cancelled. Errors @@ -81,6 +84,7 @@ func (r *Reconciler) Run(ctx context.Context) { // Idempotent: existing rows by ID are skipped (CreateCheckpoint returns // an error for duplicates; we treat it as no-op). func (r *Reconciler) reconcileOnce(ctx context.Context) { + r.Elections.reconcile(ctx) selector := fmt.Sprintf("%s=%s", checkpointstore.CMLabelKind, checkpointstore.CMLabelKindCapture) cms, err := r.KubeClient.CoreV1().ConfigMaps(r.Namespace).List(ctx, metav1.ListOptions{LabelSelector: selector}) if err != nil { diff --git a/src/compute-plane-services/nvsnap/internal/server/server.go b/src/compute-plane-services/nvsnap/internal/server/server.go index f9a827bdae..2377d313f0 100644 --- a/src/compute-plane-services/nvsnap/internal/server/server.go +++ b/src/compute-plane-services/nvsnap/internal/server/server.go @@ -352,6 +352,7 @@ func (s *Server) Run(ctx context.Context) error { Namespace: captureManifestNamespace, Interval: 30 * time.Second, Log: s.log.WithField("subsys", "reconciler"), + Elections: s.electionReleaser(), } go rec.Run(ctx) @@ -2496,3 +2497,14 @@ func (s *Server) corsMiddleware(next http.Handler) http.Handler { next.ServeHTTP(w, r) }) } + +// electionReleaser is the follower release/evict helper for the +// one-downloader election. nil kube client (tests, offline) yields a +// releaser whose methods are no-ops. +func (s *Server) electionReleaser() *electionReleaser { + return &electionReleaser{ + kube: s.kubeClient, + leaseNS: captureManifestNamespace, + log: s.log.WithField("subsys", "election"), + } +} diff --git a/src/compute-plane-services/nvsnap/internal/server/sources.go b/src/compute-plane-services/nvsnap/internal/server/sources.go index 32ade970ac..842aa123e4 100644 --- a/src/compute-plane-services/nvsnap/internal/server/sources.go +++ b/src/compute-plane-services/nvsnap/internal/server/sources.go @@ -423,6 +423,9 @@ func (s *Server) updatePVCPromoteStateByHash(w http.ResponseWriter, r *http.Requ http.Error(w, "no checkpoint found for hash", http.StatusNotFound) return } + // One-downloader election: ready releases the gated followers of this + // hash, failed evicts them for re-election (election_release.go). + s.electionReleaser().onPromoteState(r.Context(), hash, req.State) _ = s.catalog.LogAudit(&db.AuditEntry{ Action: "checkpoint.pvc_promote_state", Resource: "checkpoint_hash", ResourceID: hash, Actor: "agent", diff --git a/src/compute-plane-services/nvsnap/internal/webhook/BUILD.bazel b/src/compute-plane-services/nvsnap/internal/webhook/BUILD.bazel index b7d190ea77..a90f668440 100644 --- a/src/compute-plane-services/nvsnap/internal/webhook/BUILD.bazel +++ b/src/compute-plane-services/nvsnap/internal/webhook/BUILD.bazel @@ -7,6 +7,7 @@ go_library( "auto_inject.go", "cachedir.go", "cert.go", + "election.go", "extract_coalesce.go", "l2_mount.go", "mount_prep_init.go", @@ -20,6 +21,7 @@ go_library( visibility = ["//src/compute-plane-services/nvsnap:__subpackages__"], deps = [ "//src/compute-plane-services/nvsnap/internal/checkpointstore", + "//src/compute-plane-services/nvsnap/internal/election", "//src/compute-plane-services/nvsnap/internal/rootfsonly", "//src/compute-plane-services/nvsnap/internal/tracing", "@com_github_sirupsen_logrus//:logrus", @@ -37,6 +39,7 @@ go_test( "admission_test.go", "cachedir_noshim_test.go", "cachedir_test.go", + "election_test.go", "extract_coalesce_test.go", "l2_mount_test.go", "mergeplan_test.go", @@ -49,9 +52,11 @@ go_test( embed = [":webhook"], deps = [ "//src/compute-plane-services/nvsnap/internal/checkpointstore", + "//src/compute-plane-services/nvsnap/internal/election", "//src/compute-plane-services/nvsnap/internal/rootfsonly", "@io_k8s_api//admission/v1:admission", "@io_k8s_api//core/v1:core", + "@io_k8s_apimachinery//pkg/api/resource", "@io_k8s_apimachinery//pkg/apis/meta/v1:meta", "@io_k8s_apimachinery//pkg/runtime", "@io_k8s_apimachinery//pkg/types", diff --git a/src/compute-plane-services/nvsnap/internal/webhook/cachedir.go b/src/compute-plane-services/nvsnap/internal/webhook/cachedir.go index 10f67df33c..ca822b6b00 100644 --- a/src/compute-plane-services/nvsnap/internal/webhook/cachedir.go +++ b/src/compute-plane-services/nvsnap/internal/webhook/cachedir.go @@ -202,6 +202,13 @@ func sortedEnvVars(env map[string]string) []corev1.EnvVar { // whole dir as the rox PVC root. No-op when cachedir mode is off or the // main-container index is out of range. func (m *Mutator) cacheDirCapturePatches(pod *corev1.Pod) []PatchOp { + return m.cacheDirCapturePatchesFor(pod, false) +} + +// cacheDirCapturePatchesFor is cacheDirCapturePatches with the opt-in +// label check skipped for an elected leader, which is chosen by the +// election rather than by the chart author. +func (m *Mutator) cacheDirCapturePatchesFor(pod *corev1.Pod, elected bool) []PatchOp { if m.CacheDir == "" { return nil } @@ -209,7 +216,7 @@ func (m *Mutator) cacheDirCapturePatches(pod *corev1.Pod) []PatchOp { // labeled for capture (nvsnap.io/capture: "true"), matching the rootfs // capture watcher. Un-labeled pods (system/infra, helm-chart miniservice, // anything not meant for capture) are left untouched. See CaptureLabel. - if pod.Labels[CaptureLabel] != "true" { + if !elected && pod.Labels[CaptureLabel] != "true" { return nil } if m.MainContainer < 0 || m.MainContainer >= len(pod.Spec.Containers) { @@ -287,16 +294,46 @@ func (m *Mutator) tryL2CacheDir(ctx context.Context, pod *corev1.Pod, hash strin // capture predates EntryArgv, fall through to L1. // Resolve the rox PVC (ErrNotFound = not Bound → caller falls to L1). - pm, err := m.L2Backend.Mount(ctx, hash, checkpointstore.VolumeMeta{ - Name: cacheDirVolumeName, - MountPath: m.CacheDir, - Type: "cachedir", - Namespace: pod.Namespace, - }) + pm, err := m.L2Backend.Mount(ctx, hash, cacheDirVolumeMeta(m.CacheDir, pod.Namespace)) if err != nil { return nil, err } - roxVol := pm.Volume + // Cache/model env: REPLAYED from the manifest (the per-checkpoint + // single source of truth), verbatim, so the paths match exactly what + // the capture pod ran with regardless of any later ConfigMap edit. + // Fall back to recomputing from CacheDir for pre-v0.1.0 cachedir + // captures that predate the stamped CacheEnv. NOTE: never read the + // live ConfigMap here; that would reintroduce the path-drift the + // stamp exists to prevent. + var envs []corev1.EnvVar + if len(manifest.CacheEnv) > 0 { + envs = sortedEnvVars(manifest.CacheEnv) + } else { + envs = cacheDirEnvVars(m.CacheDir) + } + return m.cacheDirRestorePatches(pod, pm.Volume, envs) +} + +// cacheDirVolumeMeta is the L2 mount request for the cachedir rox. +func cacheDirVolumeMeta(cacheDir, namespace string) checkpointstore.VolumeMeta { + return checkpointstore.VolumeMeta{ + Name: cacheDirVolumeName, + MountPath: cacheDir, + Type: "cachedir", + Namespace: namespace, + } +} + +// cacheDirRestorePatches builds the cachedir restore decoration around a +// rox volume: the rox mounted read-only at m.CacheDir, a writable emptyDir +// shadowing the cache subtree seeded from the rox, the page-cache prewarm, +// and the cache env. The volume may name a claim that does not exist yet +// (an election follower); nothing here checks the cluster. +func (m *Mutator) cacheDirRestorePatches(pod *corev1.Pod, roxVol corev1.Volume, envs []corev1.EnvVar) ([]PatchOp, error) { + if m.MainContainer < 0 || m.MainContainer >= len(pod.Spec.Containers) { + return nil, fmt.Errorf("MainContainer index %d out of range (have %d containers)", + m.MainContainer, len(pod.Spec.Containers)) + } roxVol.Name = cacheDirVolumeName if roxVol.PersistentVolumeClaim != nil { roxVol.PersistentVolumeClaim.ReadOnly = true @@ -309,7 +346,7 @@ func (m *Mutator) tryL2CacheDir(ctx context.Context, pod *corev1.Pod, hash strin } } - patches := make([]PatchOp, 0, 11+len(manifest.CacheEnv)) + patches := make([]PatchOp, 0, 11+len(envs)) if pod.Spec.Volumes == nil { patches = append(patches, PatchOp{Op: "add", Path: "/spec/volumes", Value: []any{}}) } @@ -418,19 +455,6 @@ func (m *Mutator) tryL2CacheDir(ctx context.Context, pod *corev1.Pod, hash strin patches = append(patches, PatchOp{Op: "add", Path: "/spec/initContainers/-", Value: prewarmInit}) } - // Cache/model env — REPLAYED from the manifest (the per-checkpoint - // single source of truth), verbatim, so the paths match exactly what - // the capture pod ran with regardless of any later ConfigMap edit. - // Fall back to recomputing from CacheDir for pre-v0.1.0 cachedir - // captures that predate the stamped CacheEnv. NOTE: never read the - // live ConfigMap here — that would reintroduce the path-drift the - // stamp exists to prevent. - var envs []corev1.EnvVar - if len(manifest.CacheEnv) > 0 { - envs = sortedEnvVars(manifest.CacheEnv) - } else { - envs = cacheDirEnvVars(m.CacheDir) - } // NOTE: do NOT set HF_HUB_OFFLINE here. It only suppresses benign HF // negative-cache (.no_exist) warnings, but vLLM's arg_utils keys off // HF_HUB_OFFLINE to rewrite --model from the repo-id to the resolved diff --git a/src/compute-plane-services/nvsnap/internal/webhook/election.go b/src/compute-plane-services/nvsnap/internal/webhook/election.go new file mode 100644 index 0000000000..d71dcf08e9 --- /dev/null +++ b/src/compute-plane-services/nvsnap/internal/webhook/election.go @@ -0,0 +1,167 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package webhook + +import ( + "context" + "errors" + "fmt" + "strings" + + "github.com/sirupsen/logrus" + corev1 "k8s.io/api/core/v1" + + "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/checkpointstore" + "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/election" + "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/rootfsonly" +) + +// electionPatches is the admission decision for a model workload that +// carries no explicit restore-from (docs/proposals/helm-chart-cache-election.md): +// +// - a promoted cachedir capture exists for the composed hash: restore. +// - otherwise elect. The leader gets the capture decoration and the +// capture label so the watcher picks it up; followers get the restore +// decoration against the claim the promote will bind, plus a +// scheduling gate nvsnap-server removes on ready. +// +// Returns (nil, nil) when the pod is not a model workload, the election is +// off, or the pod is better left to the existing paths. Every returned +// patch set starts with the hash stamp, so the watcher captures under the +// same hash the followers were decorated against. +func (m *Mutator) electionPatches(ctx context.Context, pod *corev1.Pod) ([]PatchOp, error) { + if m.Elector == nil || m.Composer == nil || m.CacheDir == "" || m.L2Backend == nil { + return nil, nil + } + modelID, ok := rootfsonly.IsModelWorkload(pod, m.MainContainer) + if !ok { + return nil, nil + } + hash := checkpointstore.ComputeHash(m.Composer.Compose(pod, m.MainContainer)) + log := m.logger().WithFields(logrus.Fields{ + "pod": election.PodIdentity(pod), + "hash": checkpointstore.ShortHash(hash), + "model": modelID, + }) + + mp := newMetaPatcher(pod) + manifest, statErr := m.Backend.Stat(ctx, hash) + switch { + case statErr == nil: + if manifest.CaptureMethod != "cachedir" { + log.WithField("method", manifest.CaptureMethod).Info("election: existing capture is not cachedir; leaving pod to the explicit paths") + return nil, nil + } + patches, err := m.tryL2CacheDir(ctx, pod, hash, manifest) + if err == nil && patches != nil { + log.Info("election: promoted capture exists; restoring") + return append(mp.stamp(hash, election.RoleRestore), patches...), nil + } + if err != nil && !errors.Is(err, checkpointstore.ErrNotFound) { + return nil, fmt.Errorf("restore from existing capture: %w", err) + } + // A manifest without a bound claim: promote pending or failed. A + // leader elected now would skip the capture (hash exists) and never + // promote, leaving followers gated until the deadline. Stay out. + log.Info("election: capture exists but its volume is not bound; admitting unchanged") + return nil, nil + case !errors.Is(statErr, checkpointstore.ErrNotFound): + return nil, fmt.Errorf("backend stat: %w", statErr) + } + + role, id, err := m.Elector.Elect(ctx, hash, pod) + if err != nil { + return nil, err + } + switch role { + case election.RoleLeader: + patches := mp.stamp(hash, election.RoleLeader) + patches = append(patches, mp.annotation(election.ElectionIDAnnotation, id)...) + patches = append(patches, mp.label(CaptureLabel, "true")...) + patches = append(patches, m.cacheDirCapturePatchesFor(pod, true)...) + log.WithField("election_id", id).Info("election: leader; capture decoration applied") + return patches, nil + case election.RoleFollower: + pending, ok := m.L2Backend.(checkpointstore.PendingMounter) + if !ok { + log.Info("election: follower but the L2 backend cannot name the claim ahead of promote; cold start") + return nil, nil + } + pm, ok := pending.PendingMountSpec(hash, cacheDirVolumeMeta(m.CacheDir, pod.Namespace)) + if !ok { + log.Info("election: follower but storage is per-pod clone; cold start") + return nil, nil + } + patches := mp.stamp(hash, election.RoleFollower) + patches = append(patches, mp.label(election.GatedLabel, "true")...) + restore, err := m.cacheDirRestorePatches(pod, pm.Volume, m.cacheEnvVars(m.CacheDir)) + if err != nil { + return nil, err + } + if restore == nil { + log.Info("election: follower already carries the cache mount; admitting unchanged") + return nil, nil + } + patches = append(patches, restore...) + patches = append(patches, gatePatch(pod)) + if m.L2WaitImage != "" { + // The restore decoration created /spec/initContainers if it was + // missing, so an insert at index 0 is valid here. + waitC := buildL2WaitContainer(m.L2WaitImage, m.NvSnapServerURL, hash, m.L2WaitTimeout) + patches = append(patches, PatchOp{Op: "add", Path: "/spec/initContainers/0", Value: waitC}) + } + log.WithField("claim", pm.Volume.PersistentVolumeClaim.ClaimName).Info("election: follower; gated until the capture is promoted") + return patches, nil + default: + return nil, fmt.Errorf("election returned unknown role %q", role) + } +} + +// metaPatcher sets labels and annotations on the admitted pod, creating +// each map at most once. A second "add" of the map path would replace +// the map and drop the keys set before it. +type metaPatcher struct { + pod *corev1.Pod + bootstrapped map[string]bool +} + +func newMetaPatcher(pod *corev1.Pod) *metaPatcher { + return &metaPatcher{pod: pod, bootstrapped: map[string]bool{}} +} + +func (mp *metaPatcher) set(mapPath string, existing map[string]string, key, value string) []PatchOp { + var patches []PatchOp + if existing == nil && !mp.bootstrapped[mapPath] { + mp.bootstrapped[mapPath] = true + patches = append(patches, PatchOp{Op: "add", Path: mapPath, Value: map[string]string{}}) + } + return append(patches, PatchOp{Op: "add", Path: mapPath + "/" + strings.ReplaceAll(key, "/", "~1"), Value: value}) +} + +func (mp *metaPatcher) label(key, value string) []PatchOp { + return mp.set("/metadata/labels", mp.pod.Labels, key, value) +} + +func (mp *metaPatcher) annotation(key, value string) []PatchOp { + return mp.set("/metadata/annotations", mp.pod.Annotations, key, value) +} + +// stamp records the composed hash (annotation full, label short) and the +// role on the pod. +func (mp *metaPatcher) stamp(hash string, role election.Role) []PatchOp { + patches := mp.annotation(election.HashAnnotation, hash) + patches = append(patches, mp.annotation(election.RoleAnnotation, string(role))...) + patches = append(patches, mp.label(election.HashLabel, checkpointstore.ShortHash(hash))...) + return patches +} + +// gatePatch adds the wait-for-cache scheduling gate. Gates may only be +// removed after creation, so this is the one moment to add it. +func gatePatch(pod *corev1.Pod) PatchOp { + gate := corev1.PodSchedulingGate{Name: election.GateName} + if pod.Spec.SchedulingGates == nil { + return PatchOp{Op: "add", Path: "/spec/schedulingGates", Value: []corev1.PodSchedulingGate{gate}} + } + return PatchOp{Op: "add", Path: "/spec/schedulingGates/-", Value: gate} +} diff --git a/src/compute-plane-services/nvsnap/internal/webhook/election_test.go b/src/compute-plane-services/nvsnap/internal/webhook/election_test.go new file mode 100644 index 0000000000..e258bdae76 --- /dev/null +++ b/src/compute-plane-services/nvsnap/internal/webhook/election_test.go @@ -0,0 +1,315 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package webhook + +import ( + "context" + "errors" + "strings" + "testing" + + corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/api/resource" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + + "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/checkpointstore" + "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/election" + "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/rootfsonly" +) + +// fakeElector returns a fixed role and records whether it was asked. +type fakeElector struct { + role election.Role + err error + called int + hash string +} + +func (f *fakeElector) Elect(_ context.Context, hash string, _ *corev1.Pod) (election.Role, string, error) { + f.called++ + f.hash = hash + if f.role == election.RoleLeader { + return f.role, "id-1", f.err + } + return f.role, "", f.err +} + +// pendingStub is an L2 backend that can (or cannot) name the claim ahead +// of the promote, on top of the shared stub's Mount behaviour. +type pendingStub struct { + stubL2Backend + pendingOK bool +} + +func (p *pendingStub) PendingMountSpec(hash string, vol checkpointstore.VolumeMeta) (checkpointstore.PodMount, bool) { + if !p.pendingOK { + return checkpointstore.PodMount{}, false + } + name := vol.Name + return checkpointstore.PodMount{ + Volume: corev1.Volume{Name: name, VolumeSource: corev1.VolumeSource{ + PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{ClaimName: "rox-" + checkpointstore.ShortHash(hash), ReadOnly: true}}}, + VolumeMount: corev1.VolumeMount{Name: name, MountPath: vol.MountPath, ReadOnly: true}, + }, true +} + +// dynamoWorker is a chart-shaped model worker: Deployment pod with no +// name yet, GPU request, Dynamo list-form args, no nvsnap annotations. +func dynamoWorker() *corev1.Pod { + return &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Namespace: "fn-ns", GenerateName: "dgd-worker-"}, + Spec: corev1.PodSpec{Containers: []corev1.Container{{ + Name: "main", Image: "vllm/vllm-openai:v0.20.0", + Command: []string{"python3", "-m", "dynamo.vllm"}, + Args: []string{"--model", "Qwen/Qwen3-0.6B", "--is-decode-worker"}, + Resources: corev1.ResourceRequirements{Limits: corev1.ResourceList{"nvidia.com/gpu": resource.MustParse("2")}}, + }}}, + } +} + +func electionMutator(t *testing.T, el election.Elector, l2 checkpointstore.Backend) *Mutator { + t.Helper() + return &Mutator{ + Backend: newBackend(t), + L2Backend: l2, + CacheDir: "/opt/nvsnap", + Composer: &rootfsonly.HashInputComposer{CUDADriverMajor: 580}, + Elector: el, + L2WaitImage: "nvsnap-l2-wait:test", + NvSnapServerURL: "http://nvsnap-server:8080", + } +} + +type patchView struct { + labels, annotations map[string]string + inits []string + claims []string + gates int + commandTouched bool + cacheEmptyDir bool + envNames map[string]bool +} + +func viewPatches(patches []PatchOp) patchView { + v := patchView{labels: map[string]string{}, annotations: map[string]string{}, envNames: map[string]bool{}} + for i := range patches { + p := &patches[i] + switch { + case strings.HasPrefix(p.Path, "/metadata/labels/"): + v.labels[strings.ReplaceAll(strings.TrimPrefix(p.Path, "/metadata/labels/"), "~1", "/")] = p.Value.(string) + case strings.HasPrefix(p.Path, "/metadata/annotations/"): + v.annotations[strings.ReplaceAll(strings.TrimPrefix(p.Path, "/metadata/annotations/"), "~1", "/")] = p.Value.(string) + case strings.HasSuffix(p.Path, "/command") || strings.HasSuffix(p.Path, "/args"): + v.commandTouched = true + case strings.HasPrefix(p.Path, "/spec/schedulingGates"): + v.gates++ + } + switch val := p.Value.(type) { + case corev1.Container: + if strings.HasPrefix(p.Path, "/spec/initContainers") { + if strings.HasSuffix(p.Path, "/0") { + v.inits = append([]string{val.Name}, v.inits...) + } else { + v.inits = append(v.inits, val.Name) + } + } + case []corev1.Container: + for i := range val { + v.inits = append(v.inits, val[i].Name) + } + case corev1.Volume: + if val.PersistentVolumeClaim != nil { + v.claims = append(v.claims, val.PersistentVolumeClaim.ClaimName) + } + if val.Name == cacheDirVolumeName && val.EmptyDir != nil { + v.cacheEmptyDir = true + } + case corev1.EnvVar: + v.envNames[val.Name] = true + } + } + return v +} + +func TestElection_NotAModelWorkloadIsIgnored(t *testing.T) { + el := &fakeElector{role: election.RoleLeader} + m := electionMutator(t, el, &pendingStub{pendingOK: true}) + frontend := dynamoWorker() + frontend.Spec.Containers[0].Command = []string{"python3", "-m", "dynamo.frontend"} + frontend.Spec.Containers[0].Args = []string{"--router-mode", "kv"} + frontend.Spec.Containers[0].Resources = corev1.ResourceRequirements{} + patches, err := m.Mutate(context.Background(), frontend) + if err != nil || len(patches) != 0 { + t.Fatalf("frontend must be admitted unchanged, got %d patches err=%v", len(patches), err) + } + if el.called != 0 { + t.Error("election must not run for a pod without GPU and model") + } +} + +func TestElection_LeaderGetsCaptureDecoration(t *testing.T) { + el := &fakeElector{role: election.RoleLeader} + m := electionMutator(t, el, &pendingStub{pendingOK: true}) + pod := dynamoWorker() + patches, err := m.Mutate(context.Background(), pod) + if err != nil { + t.Fatal(err) + } + v := viewPatches(patches) + if el.called != 1 { + t.Fatalf("elector called %d times, want 1", el.called) + } + if v.annotations[election.RoleAnnotation] != "leader" { + t.Errorf("role annotation = %q, want leader", v.annotations[election.RoleAnnotation]) + } + if v.annotations[election.HashAnnotation] != el.hash || len(el.hash) != 64 { + t.Errorf("hash annotation %q must be the full hash the elector saw %q", v.annotations[election.HashAnnotation], el.hash) + } + if v.labels[election.HashLabel] != checkpointstore.ShortHash(el.hash) { + t.Errorf("hash label = %q, want short hash", v.labels[election.HashLabel]) + } + if v.labels[CaptureLabel] != "true" { + t.Error("leader must carry the capture label so the watcher captures it") + } + if v.annotations[election.ElectionIDAnnotation] != "id-1" { + t.Errorf("leader must carry the election id the Lease holds, got %q", v.annotations[election.ElectionIDAnnotation]) + } + if !v.cacheEmptyDir || !v.envNames["HF_HOME"] { + t.Errorf("leader must get the capture decoration (cache emptyDir + cache env), got emptyDir=%v env=%v", v.cacheEmptyDir, v.envNames) + } + if v.gates != 0 || len(v.claims) != 0 || v.labels[election.GatedLabel] != "" { + t.Errorf("leader must not be gated or mounted on a claim: gates=%d claims=%v", v.gates, v.claims) + } + if v.commandTouched { + t.Error("leader command must be left as authored") + } +} + +func TestElection_FollowerIsGatedOnThePendingClaim(t *testing.T) { + el := &fakeElector{role: election.RoleFollower} + m := electionMutator(t, el, &pendingStub{pendingOK: true}) + pod := dynamoWorker() + patches, err := m.Mutate(context.Background(), pod) + if err != nil { + t.Fatal(err) + } + v := viewPatches(patches) + if v.annotations[election.RoleAnnotation] != "follower" || v.labels[election.GatedLabel] != "true" { + t.Errorf("follower stamp wrong: role=%q gated=%q", v.annotations[election.RoleAnnotation], v.labels[election.GatedLabel]) + } + if v.gates != 1 { + t.Errorf("follower must carry exactly one scheduling gate, got %d", v.gates) + } + want := "rox-" + checkpointstore.ShortHash(el.hash) + if len(v.claims) != 1 || v.claims[0] != want { + t.Errorf("follower must mount the pending claim %s, got %v", want, v.claims) + } + if len(v.inits) < 3 || v.inits[0] != "nvsnap-l2-wait" || v.inits[1] != "nvsnap-seed-cache" || v.inits[2] != "nvsnap-prewarm" { + t.Errorf("follower inits = %v, want [nvsnap-l2-wait nvsnap-seed-cache nvsnap-prewarm]", v.inits) + } + if v.labels[CaptureLabel] != "" { + t.Error("follower must not be a capture source") + } + if !v.envNames["HF_HOME"] { + t.Error("follower must get the cache env so the engine reads the mounted cache") + } + if v.commandTouched { + t.Error("follower command must be left as authored") + } +} + +func TestElection_FollowerOnPerPodCloneStorageStartsCold(t *testing.T) { + el := &fakeElector{role: election.RoleFollower} + m := electionMutator(t, el, &pendingStub{pendingOK: false}) + patches, err := m.Mutate(context.Background(), dynamoWorker()) + if err != nil || len(patches) != 0 { + t.Fatalf("follower without a nameable claim must be admitted unchanged, got %d patches err=%v", len(patches), err) + } +} + +func TestElection_ErrorAdmitsUnchanged(t *testing.T) { + el := &fakeElector{err: errors.New("apiserver down")} + m := electionMutator(t, el, &pendingStub{pendingOK: true}) + patches, err := m.Mutate(context.Background(), dynamoWorker()) + if err != nil || len(patches) != 0 { + t.Fatalf("election error must fail open, got %d patches err=%v", len(patches), err) + } +} + +func TestElection_PromotedCaptureRestoresWithoutElecting(t *testing.T) { + el := &fakeElector{role: election.RoleLeader} + pod := dynamoWorker() + composer := &rootfsonly.HashInputComposer{CUDADriverMajor: 580} + hash := checkpointstore.ComputeHash(composer.Compose(pod, 0)) + l2 := &pendingStub{pendingOK: true} + l2.mountResult = checkpointstore.PodMount{ + Volume: corev1.Volume{Name: "x", VolumeSource: corev1.VolumeSource{ + PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{ClaimName: "rox-" + checkpointstore.ShortHash(hash)}}}, + VolumeMount: corev1.VolumeMount{Name: "x", MountPath: "/opt/nvsnap"}, + } + m := electionMutator(t, el, l2) + manifest := checkpointstore.Manifest{CaptureMethod: "cachedir", Volumes: []checkpointstore.VolumeMeta{{Name: "v", MountPath: "/cache", Type: "emptyDir", FileCount: 1, SizeBytes: 1}}} + if _, err := m.Backend.Put(context.Background(), hash, []checkpointstore.CaptureSource{{SrcPath: srcForManifest(t, manifest)}}, manifest); err != nil { + t.Fatal(err) + } + patches, err := m.Mutate(context.Background(), pod) + if err != nil { + t.Fatal(err) + } + v := viewPatches(patches) + if el.called != 0 { + t.Error("a promoted capture must be restored without running an election") + } + if v.annotations[election.RoleAnnotation] != "restore" || v.gates != 0 || len(v.claims) != 1 { + t.Errorf("restore decoration wrong: role=%q gates=%d claims=%v", v.annotations[election.RoleAnnotation], v.gates, v.claims) + } +} + +func TestElection_CaptureWithoutBoundClaimIsLeftAlone(t *testing.T) { + el := &fakeElector{role: election.RoleLeader} + pod := dynamoWorker() + composer := &rootfsonly.HashInputComposer{CUDADriverMajor: 580} + hash := checkpointstore.ComputeHash(composer.Compose(pod, 0)) + l2 := &pendingStub{pendingOK: true} + l2.mountErr = checkpointstore.ErrNotFound + m := electionMutator(t, el, l2) + manifest := checkpointstore.Manifest{CaptureMethod: "cachedir", Volumes: []checkpointstore.VolumeMeta{{Name: "v", MountPath: "/cache", Type: "emptyDir", FileCount: 1, SizeBytes: 1}}} + if _, err := m.Backend.Put(context.Background(), hash, []checkpointstore.CaptureSource{{SrcPath: srcForManifest(t, manifest)}}, manifest); err != nil { + t.Fatal(err) + } + patches, err := m.Mutate(context.Background(), pod) + if err != nil || len(patches) != 0 || el.called != 0 { + t.Fatalf("manifest present but rox unbound: want unchanged and no election, got %d patches err=%v elected=%d", len(patches), err, el.called) + } +} + +func TestElection_ExplicitRestoreFromBypassesElection(t *testing.T) { + el := &fakeElector{role: election.RoleLeader} + m := electionMutator(t, el, &pendingStub{pendingOK: true}) + pod := dynamoWorker() + pod.Annotations = map[string]string{RestoreFromAnnotation: "deadbeef"} + _, _ = m.Mutate(context.Background(), pod) + if el.called != 0 { + t.Error("an explicit restore-from must take the old path, not the election") + } +} + +// The pod arrives with no labels and no annotations; the maps must be +// created exactly once each, or a second create would wipe earlier keys. +func TestElection_MetadataMapsBootstrappedOnce(t *testing.T) { + m := electionMutator(t, &fakeElector{role: election.RoleFollower}, &pendingStub{pendingOK: true}) + patches, err := m.Mutate(context.Background(), dynamoWorker()) + if err != nil { + t.Fatal(err) + } + creates := map[string]int{} + for _, p := range patches { + if p.Path == "/metadata/labels" || p.Path == "/metadata/annotations" { + creates[p.Path]++ + } + } + if creates["/metadata/labels"] != 1 || creates["/metadata/annotations"] != 1 { + t.Errorf("map bootstraps = %v, want exactly one each", creates) + } +} diff --git a/src/compute-plane-services/nvsnap/internal/webhook/mutate.go b/src/compute-plane-services/nvsnap/internal/webhook/mutate.go index 11ac018cd9..265ba08bc9 100644 --- a/src/compute-plane-services/nvsnap/internal/webhook/mutate.go +++ b/src/compute-plane-services/nvsnap/internal/webhook/mutate.go @@ -38,6 +38,7 @@ import ( corev1 "k8s.io/api/core/v1" "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/checkpointstore" + "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/election" "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/rootfsonly" "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/tracing" ) @@ -246,6 +247,14 @@ type Mutator struct { // own NVSNAP_PREWARM=0/1 overrides the profile either way. StorageProfile *checkpointstore.StorageProfile + // Elector, when set, turns on the one-downloader-per-hash election + // for model workloads that carry no nvsnap.io/restore-from: the + // webhook composes the hash, restores when a promoted capture exists, + // otherwise elects a leader to capture and gates the followers' + // scheduling until the promote binds. nil keeps the label-driven + // capture and explicit-hash restore paths only. See election.go. + Elector election.Elector + // L2WaitImage is the nvsnap-l2-wait init-container image ref // (nvsnap#147). When non-empty, tryL2Mount prepends a // nvsnap-l2-wait init container that polls nvsnap-server's @@ -427,6 +436,17 @@ func (m *Mutator) Mutate(ctx context.Context, pod *corev1.Pod) ([]PatchOp, error raw, ok := pod.Annotations[RestoreFromAnnotation] if !ok || raw == "" { + // One downloader per hash for chart-shaped workloads. Handles the + // pod fully (restore, leader or follower) or declines with nil so + // the label-driven capture inject below keeps its behaviour. An + // error is logged and admits the pod unchanged: the election is + // an optimisation, never a gate. + if ep, err := m.electionPatches(ctx, pod); err != nil { + m.logger().WithError(err).WithField("pod", election.PodIdentity(pod)). + Warn("election failed; admitting pod unchanged") + } else if ep != nil { + return mergePatchPlan(append(injectPatches, ep...)), nil + } // No restore-from — this is a CAPTURE/cold pod. In cachedir mode, // inject the cache/model env vars + the /opt/nvsnap emptyDir so the // engine funnels its caches there and the agent captures that dir. diff --git a/src/compute-plane-services/nvsnap/scripts/test-election-e2e.sh b/src/compute-plane-services/nvsnap/scripts/test-election-e2e.sh new file mode 100755 index 0000000000..7d243d061a --- /dev/null +++ b/src/compute-plane-services/nvsnap/scripts/test-election-e2e.sh @@ -0,0 +1,60 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# One-downloader election e2e against a stock vLLM chart +# (deploy/k8s/charts/vllm-workers): install N replicas, expect one leader +# and gated followers, the followers released after the promote and Ready +# with no downloads, then uninstall + reinstall with every pod restoring. +# Prints the cold and warm engine phases from the vLLM logs. +# +# KUBECONFIG=... ./scripts/test-election-e2e.sh [helm --set overrides...] +# NS=nvsnap-system REL=qwen MODEL=Qwen/Qwen2.5-32B-Instruct ./scripts/test-election-e2e.sh --set tensorParallel=4 +# +# Design and recorded results: docs/proposals/helm-chart-cache-election.md +set -u +HERE=$(cd "$(dirname "$0")" && pwd) +CHART=${CHART:-$HERE/../deploy/k8s/charts/vllm-workers} +NS=${NS:-nvsnap-system}; AGENT_NS=${AGENT_NS:-nvsnap-system}; REL=${REL:-vllm-workers}; MODEL=${MODEL:-Qwen/Qwen2.5-32B-Instruct}; REPLICAS=${REPLICAS:-2} +HELM_ARGS=("$@") +k() { timeout 60 kubectl "$@"; } +roles() { k get pods -n $NS -l app=$REL -o json | python3 -c " +import json,sys +for p in json.load(sys.stdin)['items']: + a=p['metadata'].get('annotations',{}); l=p['metadata'].get('labels',{}); st=p['status'] + ready=any(c['type']=='Ready' and c['status']=='True' for c in st.get('conditions',[])) + print(' %-22s role=%-8s gated=%-5s gates=%d phase=%-16s ready=%s node=%s hash=%s' % (p['metadata']['name'], a.get('nvsnap.io/role','-'), l.get('nvsnap.io/gated','-'), len(p['spec'].get('schedulingGates',[])), st.get('phase'), ready, p['spec'].get('nodeName','-'), (a.get('nvsnap.io/hash') or '-')[:8]))"; } +agentlog() { for a in $(k get pods -n $AGENT_NS -o name | grep nvsnap-agent); do k logs -n $AGENT_NS $a -c agent --since=${1:-30m} 2>/dev/null | grep -E "$2" | cut -c1-240; done; } +waitfor() { local desc=$1 secs=$2; shift 2; for i in $(seq 1 $((secs/5))); do if eval "$@" >/dev/null 2>&1; then echo " [$desc] after $((i*5))s"; return 0; fi; sleep 5; done; echo " [$desc] TIMEOUT ${secs}s"; return 1; } + +echo "=== 0. preconditions"; k get ds nvsnap-agent -n $AGENT_NS -o jsonpath='{.spec.template.spec.containers[0].image}{"\n"}'; k get deploy nvsnap-server -n $AGENT_NS -o jsonpath='{.spec.template.spec.containers[0].image}{"\n"}' +helm uninstall $REL -n $NS >/dev/null 2>&1; sleep 5 +echo "=== 1. install replicas=2 at $(date -u +%H:%M:%SZ)"; T0=$(date +%s) +helm install $REL $CHART -n $NS --set replicas=$REPLICAS --set model="$MODEL" "${HELM_ARGS[@]}" >/dev/null || exit 1 +sleep 12; roles +echo "--- webhook decisions"; agentlog 5m "election:" | tail -6 +echo "--- lease"; k get lease -n $AGENT_NS -l nvsnap.io/kind=capture-election -o custom-columns=NAME:.metadata.name,HOLDER:.spec.holderIdentity,DEADLINE:.metadata.annotations.nvsnap\\.io/deadline --no-headers +LEADER=$(k get pods -n $NS -l app=$REL -o json | python3 -c "import json,sys; print(next((p['metadata']['name'] for p in json.load(sys.stdin)['items'] if p['metadata'].get('annotations',{}).get('nvsnap.io/role')=='leader'),''))") +FOLLOWER=$(k get pods -n $NS -l app=$REL -o json | python3 -c "import json,sys; print(next((p['metadata']['name'] for p in json.load(sys.stdin)['items'] if p['metadata'].get('annotations',{}).get('nvsnap.io/role')=='follower'),''))") +H=$(k get pod $LEADER -n $NS -o jsonpath='{.metadata.annotations.nvsnap\.io/hash}' | cut -c1-8); echo "leader=$LEADER follower=$FOLLOWER hash=$H"; [ -n "$LEADER" ] && [ -n "$FOLLOWER" ] || { echo "ELECTION DID NOT HAPPEN"; exit 1; } +echo "=== 2. leader cold start"; waitfor "leader Ready" 1800 "[ \"\$(k get pod $LEADER -n $NS -o jsonpath='{.status.containerStatuses[0].ready}')\" = true ]"; echo " leader Ready at +$(( $(date +%s)-T0 ))s" +echo "=== 3. capture + promote"; waitfor "capture committed" 600 "agentlog 20m 'capture committed' | grep -q '$LEADER'"; agentlog 20m "capture committed|L2 promote complete|pvc-state" | grep "$LEADER|$H" | tail -3 +echo "=== 4. follower release"; waitfor "follower ungated" 1500 "[ \"\$(k get pod $FOLLOWER -n $NS -o jsonpath='{.spec.schedulingGates}')\" = '' ]"; echo " follower ungated at +$(( $(date +%s)-T0 ))s"; T1=$(date +%s) +k logs -n $AGENT_NS deploy/nvsnap-server --since=20m 2>/dev/null | grep -E "election" | tail -3 | cut -c1-200 +waitfor "follower Ready" 1200 "[ \"\$(k get pod $FOLLOWER -n $NS -o jsonpath='{.status.containerStatuses[0].ready}')\" = true ]"; echo " follower Ready $(( $(date +%s)-T1 ))s after ungate (leader cold: $(( T1-T0 ))s incl capture)" +roles +echo "--- follower downloaded? (0 expected)"; k logs -n $NS $FOLLOWER 2>/dev/null | grep -ciE "downloading|Fetching .* files" +echo "--- follower inits"; k get pod $FOLLOWER -n $NS -o jsonpath='{range .spec.initContainers[*]}{.name} {end}{"\n"}' +echo "--- follower serves"; k exec -n $NS $FOLLOWER -c vllm -- curl -s -m 60 http://127.0.0.1:8000/v1/completions -H 'Content-Type: application/json' -d '{"model":"'"$MODEL"'","prompt":"The capital of France is","max_tokens":4,"temperature":0}' 2>/dev/null | python3 -c "import json,sys; print(' ', repr(json.load(sys.stdin)['choices'][0]['text']))" +phases() { k logs -n $NS $1 2>/dev/null | grep -E "Starting to load model|Loading weights took|Model loading took|torch.compile takes|compiled graph|Graph capturing finished|init engine|Application startup complete" | sed -E 's/^\(([A-Za-z_0-9]+) pid=[0-9]+\) //' | cut -c1-150 | sed 's/^/ /'; } +echo "--- leader (cold) phases"; phases $LEADER; echo "--- follower (warm) phases"; phases $FOLLOWER +echo "--- capture size"; k get cm -n $NS -l nvsnap.io/kind -o json 2>/dev/null | python3 -c " +import json,sys +for cm in json.load(sys.stdin)['items']: + for v in cm['data'].values(): + try: d=json.loads(v) + except Exception: continue + if d.get('capture_method')=='cachedir' and d.get('hash','').startswith('$H'): print(' hash', d['hash'][:8], 'bytes', d.get('total_size_bytes'), 'files', d.get('file_count'))" 2>/dev/null | head -2 +echo "=== 6. reinstall (both restore)"; helm uninstall $REL -n $NS >/dev/null; sleep 20; T3=$(date +%s); helm install $REL $CHART -n $NS --set replicas=$REPLICAS --set model="$MODEL" "${HELM_ARGS[@]}" >/dev/null; sleep 12; roles +waitfor "both Ready" 1200 "[ \"\$(k get pods -n $NS -l app=$REL -o jsonpath='{.items[*].status.containerStatuses[0].ready}')\" = 'true true' ]"; echo " both Ready $(( $(date +%s)-T3 ))s after reinstall" +echo "=== done $(date -u +%H:%M:%SZ)"