From 189b87af88486ecbd953dc7d924398691a8a5966 Mon Sep 17 00:00:00 2001 From: Balaji Ganesan Date: Thu, 24 Sep 2026 15:41:49 -0700 Subject: [PATCH 1/6] fix(nvsnap): stop the capture watcher silently retiring a pod it never captured A pod meeting every documented precondition was not captured on dev1: label nvsnap.io/capture=true, PodReady True, 2 GPUs, right node, right namespace, watcher confirmed running with --rootfs-capture. No manifest appeared and no rootfsonly log line was emitted in 15 minutes, across creation with the label and a later remove and re-add. Two defects, and the second is why the first was invisible. The dedup set marks a pod UID as SCHEDULED, and handlePodEvent treats a marked UID as nothing to do. runCapture has five exits and only the capture-error path released the mark. A warmup cancelled by context, a cancelled wait for a capture slot, or a pod refresh that abandons the capture all returned with the UID still marked, so the pod was never retried and every later event returned at the alreadyScheduled check. The mark is now released by defer on every path except a committed capture, which is the only outcome that should retire a pod. The one Info line in handlePodEvent was gated on gpus < 2, so a multi-GPU pod -- the case the rootfs path exists to serve -- produced no output on any branch. A capture that silently never happened looked exactly like one never triggered. Every decision now says what it did: skipped for label, skipped for not Ready, skipped as already scheduled, or scheduling, with the gpu count and warmup. The two previously silent aborts in runCapture log as well. Three tests. A cancelled warmup and an abandoned pod replacement both release the mark so a retry can fire, and a committed capture keeps it so the watcher does not recapture the same pod on every resync, which is the case the release must not break. Mutation-checked by never releasing: four tests turn red. The first attempt reported zero failures because the edit did not match and nothing was mutated, so the run now asserts the target was found and the mutant compiles. A mutation that did not apply proves nothing, which is the second time that trap has come up. Relates to #2099 Co-Authored-By: Balaji Ganesan --- .../nvsnap/internal/rootfsonly/watcher.go | 37 +++++++-- .../internal/rootfsonly/watcher_test.go | 79 +++++++++++++++++++ 2 files changed, 108 insertions(+), 8 deletions(-) diff --git a/src/compute-plane-services/nvsnap/internal/rootfsonly/watcher.go b/src/compute-plane-services/nvsnap/internal/rootfsonly/watcher.go index 72b76e8204..5eb20f5d06 100644 --- a/src/compute-plane-services/nvsnap/internal/rootfsonly/watcher.go +++ b/src/compute-plane-services/nvsnap/internal/rootfsonly/watcher.go @@ -145,10 +145,13 @@ func (w *Watcher) handlePodEvent(ctx context.Context, obj any) { if !ok || pod == nil { return } + plog := w.logger().WithField("pod", pod.Namespace+"/"+pod.Name) if !w.isLabeledForCapture(pod) { + plog.Debug("watcher: skipping, not labeled for capture") return } if !IsPodReady(pod) { + plog.Debug("watcher: skipping, pod not Ready yet") return } // Rootfs-only capture is the multi-GPU fallback (cuda-checkpoint can't @@ -165,15 +168,18 @@ func (w *Watcher) handlePodEvent(ctx context.Context, obj any) { // signal (the watch already filters by it at line 130), so we honor it // even for single-GPU pods and let rootfs replace CRIU. gpus := podGPURequest(pod) - if gpus < 2 { - w.logger().WithFields(logrus.Fields{ - "pod": pod.Namespace + "/" + pod.Name, - "gpus": gpus, - }).Info("rootfs-only watcher capturing single-GPU pod (explicit nvsnap.io/capture=true opt-in)") - } if _, alreadyScheduled := w.captured.LoadOrStore(pod.UID, struct{}{}); alreadyScheduled { + plog.Debug("watcher: skipping, a capture is already scheduled or committed for this pod UID") return } + // Logged for every pod, not only single-GPU ones. This used to fire only + // when gpus < 2, so a multi-GPU pod -- the case this path exists to serve + // -- produced no output on any branch, and a capture that silently never + // happened was indistinguishable from one never triggered. + plog.WithFields(logrus.Fields{ + "gpus": gpus, + "warmup": w.WarmupDelay, + }).Info("watcher: scheduling rootfs capture (nvsnap.io/capture=true)") go w.runCapture(ctx, pod.DeepCopy()) } @@ -228,6 +234,19 @@ func (w *Watcher) refreshPodForCapture(ctx context.Context, pod *corev1.Pod, log // retried by re-firing on subsequent Pod Update events (we clear captured // on persistent error so the next event re-tries). func (w *Watcher) runCapture(ctx context.Context, pod *corev1.Pod) { + // captured marks a pod UID as scheduled, and handlePodEvent treats a + // marked UID as nothing-to-do. Only a committed capture may keep the + // mark: every other exit here (warmup cancelled, semaphore wait + // cancelled, pod refresh failed, capture failed) has to release it or + // the pod is never retried and every later event returns silently. + // Previously only the capture-error path released it, so an abort left + // the UID poisoned for the life of the agent. + committed := false + defer func() { + if !committed { + w.captured.Delete(pod.UID) + } + }() log := w.logger().WithFields(logrus.Fields{ "pod": pod.Namespace + "/" + pod.Name, "pod_uid": string(pod.UID), @@ -236,6 +255,7 @@ func (w *Watcher) runCapture(ctx context.Context, pod *corev1.Pod) { select { case <-time.After(w.WarmupDelay): case <-ctx.Done(): + log.Debug("watcher: warmup cancelled before capture; releasing the pod for retry") return } } @@ -243,6 +263,7 @@ func (w *Watcher) runCapture(ctx context.Context, pod *corev1.Pod) { case w.sem <- struct{}{}: defer func() { <-w.sem }() case <-ctx.Done(): + log.Debug("watcher: cancelled waiting for a capture slot; releasing the pod for retry") return } @@ -282,8 +303,7 @@ func (w *Watcher) runCapture(ctx context.Context, pod *corev1.Pod) { m, err := w.Capturer.Capture(captureCtx, req) if err != nil { log.WithError(err).Warn("capture failed; will retry on next Update event") - // Allow retry on subsequent events. - w.captured.Delete(pod.UID) + // The deferred release above re-arms the pod; no explicit Delete. return } log.WithFields(logrus.Fields{ @@ -291,6 +311,7 @@ func (w *Watcher) runCapture(ctx context.Context, pod *corev1.Pod) { "size_mib": m.TotalSizeBytes / 1024 / 1024, "files": m.FileCount, }).Info("capture committed for pod") + committed = true } func (w *Watcher) isLabeledForCapture(pod *corev1.Pod) bool { 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 f345dd3944..394bd088a5 100644 --- a/src/compute-plane-services/nvsnap/internal/rootfsonly/watcher_test.go +++ b/src/compute-plane-services/nvsnap/internal/rootfsonly/watcher_test.go @@ -414,3 +414,82 @@ func TestRefreshPodForCapture_FallsBackWhenReadFails(t *testing.T) { t.Error("failed re-read should fall back to the pre-warmup copy") } } + +// The dedup set marks a pod UID as SCHEDULED, and handlePodEvent treats a +// marked UID as nothing-to-do. Every abort therefore has to release it, or +// the pod is never retried and every later event returns silently. Only the +// capture-error path used to release, so the aborts below poisoned the UID +// for the life of the agent -- and because the multi-GPU branch logged +// nothing, the result was invisible. + +func TestWatcher_CancelledWarmupIsRetryable(t *testing.T) { + env := newWatcherEnv(t) + w := env.watcher() + w.Capturer = env.capturer() + w.sem = make(chan struct{}, w.concurrency()) + // Long enough that cancellation lands inside the warmup wait. + w.WarmupDelay = 5 * time.Second + + pod := fakePod(types.UID(env.podUID), "p", map[string]string{DefaultCaptureLabel: "true"}, true) + ctx, cancel := context.WithCancel(context.Background()) + w.HandlePodEvent(ctx, pod) + if !waitFor(t, time.Second, func() bool { _, ok := w.captured.Load(pod.UID); return ok }) { + t.Fatal("expected the UID to be marked while the capture is scheduled") + } + cancel() + + if !waitFor(t, 2*time.Second, func() bool { _, ok := w.captured.Load(pod.UID); return !ok }) { + t.Fatal("warmup was cancelled but the UID stayed marked; the pod can never be retried") + } +} + +func TestWatcher_AbandonedPodReplacementIsRetryable(t *testing.T) { + env := newWatcherEnv(t) + w := env.watcher() + w.Capturer = env.capturer() + w.sem = make(chan struct{}, w.concurrency()) + w.WarmupDelay = 0 + + // refreshPodForCapture re-reads the pod and abandons the capture when the + // UID changed, meaning ours was deleted and recreated. The client has no + // such pod at all, which drives the same abandon path. + pod := fakePod(types.UID("uid-that-is-not-in-the-fake-client"), "ghost", + map[string]string{DefaultCaptureLabel: "true"}, true) + + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + w.HandlePodEvent(ctx, pod) + + if !waitFor(t, 2*time.Second, func() bool { _, ok := w.captured.Load(pod.UID); return !ok }) { + t.Fatal("capture was abandoned but the UID stayed marked; the pod can never be retried") + } +} + +// A committed capture must KEEP the mark, or the watcher recaptures the same +// pod on every resync. This is the case the release must not break. +func TestWatcher_CommittedCaptureStaysDeduped(t *testing.T) { + env := newWatcherEnv(t) + // Same fixtures the happy-path test uses, so the capture actually + // commits; without them Capture fails and the release correctly fires, + // which would make this assert the opposite of what it means to. + env.addProc(t, env.upperdirMountinfo()) + env.addUpperdirContent(t) + w := env.watcher() + w.Capturer = env.capturer() + w.sem = make(chan struct{}, w.concurrency()) + w.WarmupDelay = 0 + + pod := fakePod(types.UID(env.podUID), "p", map[string]string{DefaultCaptureLabel: "true"}, true) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + w.HandlePodEvent(ctx, pod) + + // Once the capture commits the mark must persist. + if !waitFor(t, 3*time.Second, func() bool { _, ok := w.captured.Load(pod.UID); return ok }) { + t.Fatal("expected the UID to remain marked after a committed capture") + } + time.Sleep(200 * time.Millisecond) + if _, ok := w.captured.Load(pod.UID); !ok { + t.Fatal("a committed capture released its dedup mark; the pod would be recaptured on every resync") + } +} From 67d15d834f95fb7a84ca250dd16afe8e02a8e557 Mon Sep 17 00:00:00 2001 From: Balaji Ganesan Date: Thu, 24 Sep 2026 17:00:12 -0700 Subject: [PATCH 2/6] fix(nvsnap): make the cachedir restore shim exec the restoring pod's own command A cachedir restore seeded 2.21GB from the capture with no download, then the pod exited without serving. The prewarm shim had exec'd ["sleep","30"]. tryL2CacheDir handed the shim the CAPTURED pod's recorded entry argv. That is wrong in principle: a cachedir restore is a warm cold-start of a fresh pod that carries its own manifest command, not a resurrection of the captured process. It was also wrong in practice, because for the bash-wrapper convention nvsnap's own manifests use (nohup setsid & ... while true; do sleep 30; done) the pid resolver landed on the idle sleep and recorded that as EntryArgv. The shim now execs the restoring pod's own command and args, which the webhook is about to overwrite with the shim path anyway. The recorded EntryArgv is kept only as a fallback for pods that declare neither and so rely on the image entrypoint. The hard "no recorded EntryArgv" guard becomes "nothing to exec from either source". Four tests: the pod's command wins over the manifest, image-entrypoint pods fall back to the manifest, neither is an error rather than a silent empty exec, and a composition test through the real tryL2CacheDir with a stubbed L2 backend asserts the NVSNAP_ORIG_COMMAND patch carries the pod's argv and not the captured [sleep 30] -- the exact shape that failed on dev1. Mutation-checked by reverting to the manifest argv: the composition test turns red. The first attempt did not compile and the harness reported it compiled, because `cmd | head && echo` reports head's status; the check now tests the build's own exit code. Third time today a non-compiling mutant almost passed as evidence. Relates to #2099 Co-Authored-By: Balaji Ganesan --- .../nvsnap/internal/webhook/BUILD.bazel | 1 + .../nvsnap/internal/webhook/cachedir.go | 35 +++++++- .../webhook/cachedir_entryargv_test.go | 89 +++++++++++++++++++ 3 files changed, 121 insertions(+), 4 deletions(-) create mode 100644 src/compute-plane-services/nvsnap/internal/webhook/cachedir_entryargv_test.go diff --git a/src/compute-plane-services/nvsnap/internal/webhook/BUILD.bazel b/src/compute-plane-services/nvsnap/internal/webhook/BUILD.bazel index bd2d0ef45d..9f6fc58e7e 100644 --- a/src/compute-plane-services/nvsnap/internal/webhook/BUILD.bazel +++ b/src/compute-plane-services/nvsnap/internal/webhook/BUILD.bazel @@ -35,6 +35,7 @@ go_test( name = "webhook_test", srcs = [ "admission_test.go", + "cachedir_entryargv_test.go", "cachedir_test.go", "extract_coalesce_test.go", "l2_mount_test.go", diff --git a/src/compute-plane-services/nvsnap/internal/webhook/cachedir.go b/src/compute-plane-services/nvsnap/internal/webhook/cachedir.go index 699ad879f6..5065a308e1 100644 --- a/src/compute-plane-services/nvsnap/internal/webhook/cachedir.go +++ b/src/compute-plane-services/nvsnap/internal/webhook/cachedir.go @@ -262,6 +262,33 @@ func (m *Mutator) cacheDirCapturePatches(pod *corev1.Pod) []PatchOp { return patches } +// restoreEntryArgv decides what the prewarm shim execs once the cache is +// seeded: the restoring pod's OWN command and args, which the webhook is +// about to overwrite with the shim, falling back to the capture's recorded +// entry argv only when the pod declares neither and so relies on the image +// entrypoint. +// +// Replaying the captured argv unconditionally was wrong twice over. A fresh +// pod carries its own manifest command, and a cachedir restore is a warm +// cold-start of THAT pod, not a resurrection of the captured process. And for +// the bash-wrapper convention (nohup setsid & ... while true; do +// sleep 30; done), the pid resolver landed on the idle sleep, so EntryArgv +// was recorded as ["sleep","30"]; the restored pod prewarmed 2.2GB perfectly, +// exec'd sleep 30, exited, and never served (dev1, 2026-09-24). +func restoreEntryArgv(main corev1.Container, manifest checkpointstore.Manifest) ([]string, error) { + own := make([]string, 0, len(main.Command)+len(main.Args)) + own = append(own, main.Command...) + own = append(own, main.Args...) + if len(own) > 0 { + return own, nil + } + if len(manifest.EntryArgv) > 0 { + return manifest.EntryArgv, nil + } + return nil, fmt.Errorf("cachedir restore: pod declares no command or args and capture %s has no recorded EntryArgv (re-capture needed)", manifest.Hash) +} + + // tryL2CacheDir injects a cachedir RESTORE: the rox PVC mounted // read-only at m.CacheDir directly (no overlayfs), the cache/model env // vars set identically to capture, and nvsnap-rootfs-restore as the @@ -285,9 +312,9 @@ func (m *Mutator) tryL2CacheDir(ctx context.Context, pod *corev1.Pod, hash strin // prewarm. Same rule as the overlay path: never fall back to the // pod's command/args (ENTRYPOINT-only images carry neither). If the // capture predates EntryArgv, fall through to L1. - if len(manifest.EntryArgv) == 0 { - return nil, fmt.Errorf("cachedir restore: capture %s has no recorded EntryArgv (re-capture needed): %w", - checkpointstore.ShortHash(hash), checkpointstore.ErrNotFound) + entryArgv, err := restoreEntryArgv(pod.Spec.Containers[m.MainContainer], manifest) + if err != nil { + return nil, err } // Resolve the rox PVC (ErrNotFound = not Bound → caller falls to L1). @@ -314,7 +341,7 @@ func (m *Mutator) tryL2CacheDir(ctx context.Context, pod *corev1.Pod, hash strin } } - argvJSON, err := json.Marshal(manifest.EntryArgv) + argvJSON, err := json.Marshal(entryArgv) if err != nil { return nil, fmt.Errorf("marshal entrypoint argv: %w", err) } diff --git a/src/compute-plane-services/nvsnap/internal/webhook/cachedir_entryargv_test.go b/src/compute-plane-services/nvsnap/internal/webhook/cachedir_entryargv_test.go new file mode 100644 index 0000000000..add1eb4252 --- /dev/null +++ b/src/compute-plane-services/nvsnap/internal/webhook/cachedir_entryargv_test.go @@ -0,0 +1,89 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package webhook + +import ( + "context" + "encoding/json" + "strings" + "testing" + + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + + "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/checkpointstore" +) + +// The captured pod used the bash-wrapper convention, so the pid resolver +// recorded the wrapper's idle loop as the entry command. A restoring pod +// must run ITS OWN command, never the captured one. +var wrapperArgv = []string{"/bin/bash", "-lc", "nohup setsid vllm serve --model x & tail -F /vllm.out & while true; do sleep 30; done"} + +func TestRestoreEntryArgv_PrefersPodOwnCommand(t *testing.T) { + main := corev1.Container{Command: wrapperArgv[:2], Args: wrapperArgv[2:]} + got, err := restoreEntryArgv(main, checkpointstore.Manifest{EntryArgv: []string{"sleep", "30"}}) + if err != nil { + t.Fatal(err) + } + if strings.Join(got, "\x00") != strings.Join(wrapperArgv, "\x00") { + t.Fatalf("shim would exec %v, want the pod's own command %v", got, wrapperArgv) + } +} + +func TestRestoreEntryArgv_FallsBackToManifestForImageEntrypoint(t *testing.T) { + got, err := restoreEntryArgv(corev1.Container{}, checkpointstore.Manifest{EntryArgv: []string{"vllm", "serve"}}) + if err != nil || strings.Join(got, " ") != "vllm serve" { + t.Fatalf("pod with no command should fall back to the manifest, got %v err=%v", got, err) + } +} + +func TestRestoreEntryArgv_ErrorsWhenNothingToExec(t *testing.T) { + if _, err := restoreEntryArgv(corev1.Container{}, checkpointstore.Manifest{Hash: "abc"}); err == nil { + t.Fatal("no pod command and no recorded EntryArgv must be an error, not a silent empty exec") + } +} + +// Composition: the whole restore patch set, through the real tryL2CacheDir +// with a stubbed L2 backend, must hand the shim the pod's own argv. This is +// the exact shape that exited on dev1 with the captured ["sleep","30"]. +func TestTryL2CacheDir_ShimExecsPodOwnCommand(t *testing.T) { + m := &Mutator{ + CacheDir: "/opt/nvsnap", + MainContainer: 0, + L2Backend: &stubL2Backend{mountResult: checkpointstore.PodMount{ + Volume: corev1.Volume{Name: "x", VolumeSource: corev1.VolumeSource{ + PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{ClaimName: "rox-abc"}}}, + VolumeMount: corev1.VolumeMount{Name: "x", MountPath: "/opt/nvsnap"}, + }}, + } + pod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Name: "reuse", Namespace: "ns"}, + Spec: corev1.PodSpec{Containers: []corev1.Container{{ + Name: "vllm", Image: "vllm/vllm-openai:v0.20.0", + Command: wrapperArgv[:2], Args: wrapperArgv[2:], + }}}, + } + manifest := checkpointstore.Manifest{Hash: "abc", EntryArgv: []string{"sleep", "30"}, CaptureMethod: "cachedir"} + + patches, err := m.tryL2CacheDir(context.Background(), pod, "abc", manifest) + if err != nil { + t.Fatal(err) + } + var orig string + for _, p := range patches { + if e, ok := p.Value.(corev1.EnvVar); ok && e.Name == "NVSNAP_ORIG_COMMAND" { + orig = e.Value + } + } + if orig == "" { + t.Fatal("no NVSNAP_ORIG_COMMAND env patch emitted") + } + var argv []string + if err := json.Unmarshal([]byte(orig), &argv); err != nil { + t.Fatal(err) + } + if strings.Join(argv, "\x00") != strings.Join(wrapperArgv, "\x00") { + t.Fatalf("shim would exec %v; the captured [sleep 30] leaked through instead of the pod's own command", argv) + } +} From 56a3d7b21ed63533f818717708745acc3bca79a9 Mon Sep 17 00:00:00 2001 From: Balaji Ganesan Date: Thu, 24 Sep 2026 17:43:12 -0700 Subject: [PATCH 3/6] fix(nvsnap): cachedir restore mounts the cache and runs the pod's own command, no shim The previous commit made the prewarm shim exec the restoring pod's own argv instead of the captured one. That fixed the symptom at the wrong layer. A cachedir restore is a volume mount plus a seeded cache shadow; the pod is a fresh container from its own image and its entrypoint should run exactly as authored. Nothing in front of it is needed, and the design comment already described the restore that way before the shim was bolted on. What the shim added, and why none of it is load-bearing here: - Page-cache prewarm, 8s for 2.2GB on dev1. The kernel warms on demand. - Recreating EntryRuntimeDirs. That exists for rootfs restores, where a captured filesystem lands in a pristine container whose entrypoint will not run again. For this capture the recorded dirs were host and image state (/run/systemd/*, /run/nvidia*, /run/lock) that any fresh container already has, plus /run/vllm, which all fourteen cachedir workload manifests create themselves with mkdir -p. - chdir and exec of an argv that had to be chosen correctly. Choosing it wrong is how a restored pod seeded 2.2GB perfectly and then ran sleep 30. Removed from tryL2CacheDir: the command rewrite and args removal, the five shim env vars (NVSNAP_NO_OVERLAY, NVSNAP_PREWARM_DIR, NVSNAP_ORIG_COMMAND, NVSNAP_ORIG_CWD, NVSNAP_RUNTIME_DIRS), the nvsnap-tools hostPath volume and mount that existed only to carry the shim binary, and restoreEntryArgv, whose job no longer exists. The rox mount, the writable cache shadow, the seed init container and the replayed cache env are unchanged. Consequence worth stating: nvsnap-rootfs-restore is now reachable only from the whole-rootfs restore branch, so the scoping of that path's removal, which this commit's predecessor had contradicted, holds again. One composition test through the real tryL2CacheDir with a stubbed L2 backend: no patch touches command or args, no shim env or tools volume is emitted, and the rox mount, seed init and cache env are. Mutation-checked by re-adding a command rewrite: the mutant compiles and the test turns red. Relates to #2099 Co-Authored-By: Balaji Ganesan --- .../nvsnap/internal/webhook/BUILD.bazel | 2 +- .../nvsnap/internal/webhook/cachedir.go | 88 +++--------------- .../webhook/cachedir_entryargv_test.go | 89 ------------------ .../internal/webhook/cachedir_noshim_test.go | 91 +++++++++++++++++++ 4 files changed, 103 insertions(+), 167 deletions(-) delete mode 100644 src/compute-plane-services/nvsnap/internal/webhook/cachedir_entryargv_test.go create mode 100644 src/compute-plane-services/nvsnap/internal/webhook/cachedir_noshim_test.go diff --git a/src/compute-plane-services/nvsnap/internal/webhook/BUILD.bazel b/src/compute-plane-services/nvsnap/internal/webhook/BUILD.bazel index 9f6fc58e7e..b7d190ea77 100644 --- a/src/compute-plane-services/nvsnap/internal/webhook/BUILD.bazel +++ b/src/compute-plane-services/nvsnap/internal/webhook/BUILD.bazel @@ -35,7 +35,7 @@ go_test( name = "webhook_test", srcs = [ "admission_test.go", - "cachedir_entryargv_test.go", + "cachedir_noshim_test.go", "cachedir_test.go", "extract_coalesce_test.go", "l2_mount_test.go", diff --git a/src/compute-plane-services/nvsnap/internal/webhook/cachedir.go b/src/compute-plane-services/nvsnap/internal/webhook/cachedir.go index 5065a308e1..c2361145f3 100644 --- a/src/compute-plane-services/nvsnap/internal/webhook/cachedir.go +++ b/src/compute-plane-services/nvsnap/internal/webhook/cachedir.go @@ -32,9 +32,9 @@ limitations under the License. // - Restore pod (CaptureMethod=="cachedir"): mount the rox read-only // at m.CacheDir (model stays RO — the big part, never copied), shadow // the cache subtree /cache with a writable emptyDir seeded -// by a nvsnap-seed-cache init container (cp from the rox), then run -// nvsnap-rootfs-restore in no-overlay mode to prewarm the tree into -// page cache (same cgroup as the engine) and exec the entrypoint. +// by a nvsnap-seed-cache init container (cp from the rox), and run the +// pod's own command untouched. No shim: a cachedir restore is a warm +// cold-start of THIS pod, so its entrypoint runs exactly as authored. // // Why the writable cache shadow (ember rule #3, verified): engines write // JIT/log/lock files into the cache at startup — flashinfer opens @@ -50,7 +50,6 @@ package webhook import ( "context" - "encoding/json" "fmt" "os" "path/filepath" @@ -262,32 +261,6 @@ func (m *Mutator) cacheDirCapturePatches(pod *corev1.Pod) []PatchOp { return patches } -// restoreEntryArgv decides what the prewarm shim execs once the cache is -// seeded: the restoring pod's OWN command and args, which the webhook is -// about to overwrite with the shim, falling back to the capture's recorded -// entry argv only when the pod declares neither and so relies on the image -// entrypoint. -// -// Replaying the captured argv unconditionally was wrong twice over. A fresh -// pod carries its own manifest command, and a cachedir restore is a warm -// cold-start of THAT pod, not a resurrection of the captured process. And for -// the bash-wrapper convention (nohup setsid & ... while true; do -// sleep 30; done), the pid resolver landed on the idle sleep, so EntryArgv -// was recorded as ["sleep","30"]; the restored pod prewarmed 2.2GB perfectly, -// exec'd sleep 30, exited, and never served (dev1, 2026-09-24). -func restoreEntryArgv(main corev1.Container, manifest checkpointstore.Manifest) ([]string, error) { - own := make([]string, 0, len(main.Command)+len(main.Args)) - own = append(own, main.Command...) - own = append(own, main.Args...) - if len(own) > 0 { - return own, nil - } - if len(manifest.EntryArgv) > 0 { - return manifest.EntryArgv, nil - } - return nil, fmt.Errorf("cachedir restore: pod declares no command or args and capture %s has no recorded EntryArgv (re-capture needed)", manifest.Hash) -} - // tryL2CacheDir injects a cachedir RESTORE: the rox PVC mounted // read-only at m.CacheDir directly (no overlayfs), the cache/model env @@ -312,10 +285,6 @@ func (m *Mutator) tryL2CacheDir(ctx context.Context, pod *corev1.Pod, hash strin // prewarm. Same rule as the overlay path: never fall back to the // pod's command/args (ENTRYPOINT-only images carry neither). If the // capture predates EntryArgv, fall through to L1. - entryArgv, err := restoreEntryArgv(pod.Spec.Containers[m.MainContainer], manifest) - if err != nil { - return nil, err - } // Resolve the rox PVC (ErrNotFound = not Bound → caller falls to L1). pm, err := m.L2Backend.Mount(ctx, hash, checkpointstore.VolumeMeta{ @@ -335,22 +304,16 @@ func (m *Mutator) tryL2CacheDir(ctx context.Context, pod *corev1.Pod, hash strin main := pod.Spec.Containers[m.MainContainer] for _, vm := range main.VolumeMounts { - if vm.Name == cacheDirVolumeName || vm.MountPath == m.CacheDir || - vm.Name == nvsnapToolsVolumeName || vm.MountPath == nvsnapToolsMountPath { + if vm.Name == cacheDirVolumeName || vm.MountPath == m.CacheDir { return nil, nil // already wired } } - argvJSON, err := json.Marshal(entryArgv) - if err != nil { - return nil, fmt.Errorf("marshal entrypoint argv: %w", err) - } root := m.HostBundleRoot if root == "" { root = DefaultHostBundleRoot } - hostPathDir := corev1.HostPathDirectory patches := make([]PatchOp, 0, 11+len(manifest.CacheEnv)) if pod.Spec.Volumes == nil { @@ -373,15 +336,8 @@ func (m *Mutator) tryL2CacheDir(ctx context.Context, pod *corev1.Pod, hash strin // Volumes: // - roxVol: rox PVC (RO) — the captured cache+model tree. - // - toolsVol: hostPath bundle (the nvsnap-rootfs-restore shim binary). - // - cacheRW: writable emptyDir that shadows /cache so the - // engine's JIT/log/lock writes don't hit EROFS on the RO rox. - toolsVol := corev1.Volume{Name: nvsnapToolsVolumeName, VolumeSource: corev1.VolumeSource{HostPath: &corev1.HostPathVolumeSource{ - Path: root + "/nvsnap", - Type: &hostPathDir, - }}} cacheRWVol := corev1.Volume{Name: cacheRWVolumeName, VolumeSource: corev1.VolumeSource{EmptyDir: &corev1.EmptyDirVolumeSource{}}} - vols := []corev1.Volume{roxVol, toolsVol, cacheRWVol} + vols := []corev1.Volume{roxVol, cacheRWVol} for i := range vols { patches = append(patches, PatchOp{Op: "add", Path: "/spec/volumes/-", Value: vols[i]}) } @@ -393,7 +349,6 @@ func (m *Mutator) tryL2CacheDir(ctx context.Context, pod *corev1.Pod, hash strin for _, vm := range []corev1.VolumeMount{ {Name: cacheDirVolumeName, MountPath: m.CacheDir, ReadOnly: true}, {Name: cacheRWVolumeName, MountPath: cacheSub}, - {Name: nvsnapToolsVolumeName, MountPath: nvsnapToolsMountPath, ReadOnly: true}, } { patches = append(patches, PatchOp{ Op: "add", @@ -450,11 +405,6 @@ func (m *Mutator) tryL2CacheDir(ctx context.Context, pod *corev1.Pod, hash strin envs = cacheDirEnvVars(m.CacheDir) } envs = append(envs, - corev1.EnvVar{Name: "NVSNAP_NO_OVERLAY", Value: "1"}, - corev1.EnvVar{Name: "NVSNAP_PREWARM_DIR", Value: m.CacheDir}, - corev1.EnvVar{Name: "NVSNAP_ORIG_COMMAND", Value: string(argvJSON)}, - corev1.EnvVar{Name: envRuntimeDirs, Value: runtimeDirsJSON(manifest.EntryRuntimeDirs)}, - corev1.EnvVar{Name: "NVSNAP_ORIG_CWD", Value: manifest.EntryCwd}, // 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 @@ -470,27 +420,11 @@ func (m *Mutator) tryL2CacheDir(ctx context.Context, pod *corev1.Pod, hash strin patches = append(patches, appendEnv(m.MainContainer, e)) } - // Override command to the shim; clear args (re-applied from - // NVSNAP_ORIG_COMMAND by the shim after prewarm). - cmdOp := "replace" - if main.Command == nil { - cmdOp = "add" - } - patches = append(patches, PatchOp{ - Op: cmdOp, - Path: fmt.Sprintf("/spec/containers/%d/command", m.MainContainer), - Value: []string{nvsnapToolsMountPath + "/nvsnap-rootfs-restore"}, - }) - if main.Args != nil { - patches = append(patches, PatchOp{ - Op: "remove", - Path: fmt.Sprintf("/spec/containers/%d/args", m.MainContainer), - }) - } - - // NOTE: unlike the overlay path, NO securityContext / SYS_ADMIN / - // seccomp-unconfined changes — the no-overlay shim only reads files - // (prewarm) and exec's; it issues no mount(2)/pivot_root. Device - // isolation and the default profiles stay fully intact. + // The command is left exactly as authored. The old shim rewrite exec'd + // the CAPTURED pod's argv, which for the bash-wrapper convention was the + // idle sleep, so the restored pod seeded 2.2GB and then exited without + // serving (dev1, 2026-09-24). Nothing here needs a wrapper. + // No securityContext, SYS_ADMIN or seccomp changes: this path only adds + // volumes, mounts, env and a copying init container. return patches, nil } diff --git a/src/compute-plane-services/nvsnap/internal/webhook/cachedir_entryargv_test.go b/src/compute-plane-services/nvsnap/internal/webhook/cachedir_entryargv_test.go deleted file mode 100644 index add1eb4252..0000000000 --- a/src/compute-plane-services/nvsnap/internal/webhook/cachedir_entryargv_test.go +++ /dev/null @@ -1,89 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -package webhook - -import ( - "context" - "encoding/json" - "strings" - "testing" - - corev1 "k8s.io/api/core/v1" - metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" - - "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/checkpointstore" -) - -// The captured pod used the bash-wrapper convention, so the pid resolver -// recorded the wrapper's idle loop as the entry command. A restoring pod -// must run ITS OWN command, never the captured one. -var wrapperArgv = []string{"/bin/bash", "-lc", "nohup setsid vllm serve --model x & tail -F /vllm.out & while true; do sleep 30; done"} - -func TestRestoreEntryArgv_PrefersPodOwnCommand(t *testing.T) { - main := corev1.Container{Command: wrapperArgv[:2], Args: wrapperArgv[2:]} - got, err := restoreEntryArgv(main, checkpointstore.Manifest{EntryArgv: []string{"sleep", "30"}}) - if err != nil { - t.Fatal(err) - } - if strings.Join(got, "\x00") != strings.Join(wrapperArgv, "\x00") { - t.Fatalf("shim would exec %v, want the pod's own command %v", got, wrapperArgv) - } -} - -func TestRestoreEntryArgv_FallsBackToManifestForImageEntrypoint(t *testing.T) { - got, err := restoreEntryArgv(corev1.Container{}, checkpointstore.Manifest{EntryArgv: []string{"vllm", "serve"}}) - if err != nil || strings.Join(got, " ") != "vllm serve" { - t.Fatalf("pod with no command should fall back to the manifest, got %v err=%v", got, err) - } -} - -func TestRestoreEntryArgv_ErrorsWhenNothingToExec(t *testing.T) { - if _, err := restoreEntryArgv(corev1.Container{}, checkpointstore.Manifest{Hash: "abc"}); err == nil { - t.Fatal("no pod command and no recorded EntryArgv must be an error, not a silent empty exec") - } -} - -// Composition: the whole restore patch set, through the real tryL2CacheDir -// with a stubbed L2 backend, must hand the shim the pod's own argv. This is -// the exact shape that exited on dev1 with the captured ["sleep","30"]. -func TestTryL2CacheDir_ShimExecsPodOwnCommand(t *testing.T) { - m := &Mutator{ - CacheDir: "/opt/nvsnap", - MainContainer: 0, - L2Backend: &stubL2Backend{mountResult: checkpointstore.PodMount{ - Volume: corev1.Volume{Name: "x", VolumeSource: corev1.VolumeSource{ - PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{ClaimName: "rox-abc"}}}, - VolumeMount: corev1.VolumeMount{Name: "x", MountPath: "/opt/nvsnap"}, - }}, - } - pod := &corev1.Pod{ - ObjectMeta: metav1.ObjectMeta{Name: "reuse", Namespace: "ns"}, - Spec: corev1.PodSpec{Containers: []corev1.Container{{ - Name: "vllm", Image: "vllm/vllm-openai:v0.20.0", - Command: wrapperArgv[:2], Args: wrapperArgv[2:], - }}}, - } - manifest := checkpointstore.Manifest{Hash: "abc", EntryArgv: []string{"sleep", "30"}, CaptureMethod: "cachedir"} - - patches, err := m.tryL2CacheDir(context.Background(), pod, "abc", manifest) - if err != nil { - t.Fatal(err) - } - var orig string - for _, p := range patches { - if e, ok := p.Value.(corev1.EnvVar); ok && e.Name == "NVSNAP_ORIG_COMMAND" { - orig = e.Value - } - } - if orig == "" { - t.Fatal("no NVSNAP_ORIG_COMMAND env patch emitted") - } - var argv []string - if err := json.Unmarshal([]byte(orig), &argv); err != nil { - t.Fatal(err) - } - if strings.Join(argv, "\x00") != strings.Join(wrapperArgv, "\x00") { - t.Fatalf("shim would exec %v; the captured [sleep 30] leaked through instead of the pod's own command", argv) - } -} diff --git a/src/compute-plane-services/nvsnap/internal/webhook/cachedir_noshim_test.go b/src/compute-plane-services/nvsnap/internal/webhook/cachedir_noshim_test.go new file mode 100644 index 0000000000..18dc739166 --- /dev/null +++ b/src/compute-plane-services/nvsnap/internal/webhook/cachedir_noshim_test.go @@ -0,0 +1,91 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package webhook + +import ( + "context" + "strings" + "testing" + + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + + "github.com/NVIDIA/nvcf/src/compute-plane-services/nvsnap/internal/checkpointstore" +) + +// A cachedir restore is a volume mount plus a seeded cache shadow, and the +// pod runs its own command untouched. The shim that used to sit in front of +// the entrypoint exec'd the CAPTURED pod's argv -- for the bash-wrapper +// convention that was the idle sleep, so the restored pod seeded 2.2GB and +// exited without serving. This pins the no-shim contract end to end through +// the real patch builder. +func TestTryL2CacheDir_NoShim_PodCommandUntouched(t *testing.T) { + m := &Mutator{ + CacheDir: "/opt/nvsnap", + MainContainer: 0, + L2Backend: &stubL2Backend{mountResult: checkpointstore.PodMount{ + Volume: corev1.Volume{Name: "x", VolumeSource: corev1.VolumeSource{ + PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{ClaimName: "rox-abc"}}}, + VolumeMount: corev1.VolumeMount{Name: "x", MountPath: "/opt/nvsnap"}, + }}, + } + wrapper := []string{"/bin/bash", "-lc", "nohup setsid vllm serve & while true; do sleep 30; done"} + pod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Name: "reuse", Namespace: "ns"}, + Spec: corev1.PodSpec{Containers: []corev1.Container{{ + Name: "vllm", Image: "vllm/vllm-openai:v0.20.0", Command: wrapper[:2], Args: wrapper[2:], + }}}, + } + manifest := checkpointstore.Manifest{Hash: "abc", EntryArgv: []string{"sleep", "30"}, CaptureMethod: "cachedir"} + + patches, err := m.tryL2CacheDir(context.Background(), pod, "abc", manifest) + if err != nil { + t.Fatal(err) + } + + var sawRoxMount, sawSeedInit, sawCacheEnv bool + for _, p := range patches { + // Nothing may touch the entrypoint. + if strings.HasSuffix(p.Path, "/command") || strings.HasSuffix(p.Path, "/args") { + t.Fatalf("cachedir restore must not rewrite the pod command, got patch %s %s", p.Op, p.Path) + } + switch v := p.Value.(type) { + case corev1.EnvVar: + switch v.Name { + case "NVSNAP_ORIG_COMMAND", "NVSNAP_NO_OVERLAY", "NVSNAP_PREWARM_DIR", "NVSNAP_ORIG_CWD", envRuntimeDirs: + t.Fatalf("shim env %s leaked into a cachedir restore", v.Name) + case "HF_HOME": + sawCacheEnv = true + } + case corev1.Volume: + if v.Name == nvsnapToolsVolumeName { + t.Fatal("the nvsnap-tools shim bundle must not be mounted into a cachedir restore") + } + case corev1.VolumeMount: + if v.MountPath == "/opt/nvsnap" { + sawRoxMount = true + } + case corev1.Container: + if v.Name == "nvsnap-seed-cache" { + sawSeedInit = true + } + } + if vols, ok := p.Value.([]corev1.Volume); ok { + for _, v := range vols { + if v.Name == nvsnapToolsVolumeName { + t.Fatal("the nvsnap-tools shim bundle must not be mounted into a cachedir restore") + } + } + } + } + if !sawRoxMount { + t.Error("expected the rox cache mounted at /opt/nvsnap") + } + if !sawSeedInit { + t.Error("expected the nvsnap-seed-cache init container") + } + if !sawCacheEnv { + t.Error("expected the cache env (HF_HOME) replayed from the manifest") + } +} From 82dcc791b70ef98cf9351650d42159018aaee124 Mon Sep 17 00:00:00 2001 From: Balaji Ganesan Date: Thu, 24 Sep 2026 17:56:00 -0700 Subject: [PATCH 4/6] fix(nvsnap): bring the page-cache prewarm back as an init container Removing the entrypoint shim also removed its page-cache prewarm, on the evidence that time-to-Ready was 59s with it and 59s without. That was measured on a 2GB model and does not generalise. The prewarm exists for large models on network-attached rox storage, where the engine faults a safetensors set in by mmap as small random reads and a parallel sequential read-ahead beats that badly. It measurably helped large models on vLLM, and the seed init copies only {cache}, so {model}, the big part, was starting cold. Back as a nvsnap-prewarm init container after nvsnap-seed-cache: same node so same page cache, same pod cgroup so the same memory accounting, and the pod's own command stays untouched, which is the whole point of retiring the shim. It reads the rox tree read-only as root with six parallel workers, matching the retired Go prewarmer, and ends in || true because a read error must never fail a restore. NVSNAP_PREWARM=0 on the workload skips it, the same knob the shim honoured. Tests: the restore now emits exactly [nvsnap-seed-cache, nvsnap-prewarm] in that order, the prewarm reuses the workload image, mounts only the rox read-only, runs as root and is best-effort; NVSNAP_PREWARM=0 omits it and still leaves the command alone. Mutation-checked by dropping the prewarm append: the mutant compiles and the test turns red. Relates to #2099 Co-Authored-By: Balaji Ganesan --- .../nvsnap/internal/webhook/cachedir.go | 45 ++++++++++++++- .../internal/webhook/cachedir_noshim_test.go | 57 +++++++++++++++++++ 2 files changed, 100 insertions(+), 2 deletions(-) diff --git a/src/compute-plane-services/nvsnap/internal/webhook/cachedir.go b/src/compute-plane-services/nvsnap/internal/webhook/cachedir.go index c2361145f3..729e99c2e5 100644 --- a/src/compute-plane-services/nvsnap/internal/webhook/cachedir.go +++ b/src/compute-plane-services/nvsnap/internal/webhook/cachedir.go @@ -32,8 +32,9 @@ limitations under the License. // - Restore pod (CaptureMethod=="cachedir"): mount the rox read-only // at m.CacheDir (model stays RO — the big part, never copied), shadow // the cache subtree /cache with a writable emptyDir seeded -// by a nvsnap-seed-cache init container (cp from the rox), and run the -// pod's own command untouched. No shim: a cachedir restore is a warm +// by a nvsnap-seed-cache init container (cp from the rox), prewarm the +// rox tree into page cache with a nvsnap-prewarm init container, and run +// the pod's own command untouched. No shim: a cachedir restore is a warm // cold-start of THIS pod, so its entrypoint runs exactly as authored. // // Why the writable cache shadow (ember rule #3, verified): engines write @@ -391,6 +392,35 @@ func (m *Mutator) tryL2CacheDir(ctx context.Context, pod *corev1.Pod, hash strin } patches = append(patches, PatchOp{Op: "add", Path: "/spec/initContainers/-", Value: seedInit}) + // Page-cache prewarm of the rox tree, model included, AHEAD of the engine. + // This is the one thing the retired entrypoint shim did that earns its + // keep: on network-attached rox storage (NVMesh, EBS) a large safetensors + // set is faulted in by mmap as small random reads, and a parallel + // sequential read-ahead beats that badly for large models on vLLM. The + // seed init above deliberately does not touch {model}, so without this the + // biggest part of the tree starts cold. + // + // An init container, not a shim: same node so same page cache, same pod + // cgroup so the same memory accounting, and the pod's own command runs + // untouched. Best-effort by construction: a read error must never fail a + // restore, so the pipeline ends in || true. Six workers, matching the + // retired Go prewarmer. NVSNAP_PREWARM=0 on the workload skips it, the + // same knob the shim honoured. + if !podDisablesPrewarm(main) { + prewarmInit := corev1.Container{ + Name: "nvsnap-prewarm", + Image: main.Image, + Command: []string{"sh", "-c", fmt.Sprintf( + "find %s -type f -print0 2>/dev/null | xargs -0 -r -P 6 -n 16 cat > /dev/null 2>&1 || true", + cacheSeedSrcPath)}, + VolumeMounts: []corev1.VolumeMount{ + {Name: cacheDirVolumeName, MountPath: cacheSeedSrcPath, ReadOnly: true}, + }, + SecurityContext: &corev1.SecurityContext{RunAsUser: &seedRoot}, + } + 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. @@ -428,3 +458,14 @@ func (m *Mutator) tryL2CacheDir(ctx context.Context, pod *corev1.Pod, hash strin // volumes, mounts, env and a copying init container. return patches, nil } + +// podDisablesPrewarm honours NVSNAP_PREWARM=0 on the workload container, the +// same opt-out the retired entrypoint shim read. +func podDisablesPrewarm(main corev1.Container) bool { + for _, e := range main.Env { + if e.Name == "NVSNAP_PREWARM" && e.Value == "0" { + return true + } + } + return false +} diff --git a/src/compute-plane-services/nvsnap/internal/webhook/cachedir_noshim_test.go b/src/compute-plane-services/nvsnap/internal/webhook/cachedir_noshim_test.go index 18dc739166..c445143caf 100644 --- a/src/compute-plane-services/nvsnap/internal/webhook/cachedir_noshim_test.go +++ b/src/compute-plane-services/nvsnap/internal/webhook/cachedir_noshim_test.go @@ -45,6 +45,7 @@ func TestTryL2CacheDir_NoShim_PodCommandUntouched(t *testing.T) { } var sawRoxMount, sawSeedInit, sawCacheEnv bool + var inits []corev1.Container for _, p := range patches { // Nothing may touch the entrypoint. if strings.HasSuffix(p.Path, "/command") || strings.HasSuffix(p.Path, "/args") { @@ -67,6 +68,7 @@ func TestTryL2CacheDir_NoShim_PodCommandUntouched(t *testing.T) { sawRoxMount = true } case corev1.Container: + inits = append(inits, v) if v.Name == "nvsnap-seed-cache" { sawSeedInit = true } @@ -88,4 +90,59 @@ func TestTryL2CacheDir_NoShim_PodCommandUntouched(t *testing.T) { if !sawCacheEnv { t.Error("expected the cache env (HF_HOME) replayed from the manifest") } + // The prewarm must come back as an init container, after the seed, reading + // the rox read-only as root, and never as an entrypoint wrapper. It is the + // one piece of the retired shim that measurably helps large models. + if len(inits) != 2 || inits[0].Name != "nvsnap-seed-cache" || inits[1].Name != "nvsnap-prewarm" { + names := []string{} + for _, c := range inits { + names = append(names, c.Name) + } + t.Fatalf("want init containers [nvsnap-seed-cache nvsnap-prewarm] in that order, got %v", names) + } + pw := inits[1] + if pw.Image != pod.Spec.Containers[0].Image { + t.Errorf("prewarm should reuse the workload image (already pulled), got %q", pw.Image) + } + if len(pw.VolumeMounts) != 1 || !pw.VolumeMounts[0].ReadOnly || pw.VolumeMounts[0].Name != cacheDirVolumeName { + t.Errorf("prewarm must mount only the rox, read-only, got %+v", pw.VolumeMounts) + } + if pw.SecurityContext == nil || pw.SecurityContext.RunAsUser == nil || *pw.SecurityContext.RunAsUser != 0 { + t.Error("prewarm must run as root: the rox files are root-owned") + } + cmd := strings.Join(pw.Command, " ") + if !strings.Contains(cmd, cacheSeedSrcPath) || !strings.Contains(cmd, "|| true") { + t.Errorf("prewarm must read the rox tree and be best-effort, got %q", cmd) + } +} + +// NVSNAP_PREWARM=0 on the workload skips the prewarm init, matching the knob +// the retired shim honoured, and touches nothing else. +func TestTryL2CacheDir_PrewarmOptOut(t *testing.T) { + m := &Mutator{ + CacheDir: "/opt/nvsnap", MainContainer: 0, + L2Backend: &stubL2Backend{mountResult: checkpointstore.PodMount{ + Volume: corev1.Volume{Name: "x", VolumeSource: corev1.VolumeSource{PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{ClaimName: "rox-abc"}}}, + VolumeMount: corev1.VolumeMount{Name: "x", MountPath: "/opt/nvsnap"}, + }}, + } + pod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Name: "reuse", Namespace: "ns"}, + Spec: corev1.PodSpec{Containers: []corev1.Container{{ + Name: "vllm", Image: "img", Command: []string{"serve"}, + Env: []corev1.EnvVar{{Name: "NVSNAP_PREWARM", Value: "0"}}, + }}}, + } + patches, err := m.tryL2CacheDir(context.Background(), pod, "abc", checkpointstore.Manifest{Hash: "abc", CaptureMethod: "cachedir"}) + if err != nil { + t.Fatal(err) + } + for _, p := range patches { + if c, ok := p.Value.(corev1.Container); ok && c.Name == "nvsnap-prewarm" { + t.Fatal("NVSNAP_PREWARM=0 must skip the prewarm init container") + } + if strings.HasSuffix(p.Path, "/command") { + t.Fatal("opt-out must not rewrite the command either") + } + } } From 4f203c713538f36ca124395aff869748d48504c4 Mon Sep 17 00:00:00 2001 From: Balaji Ganesan Date: Thu, 24 Sep 2026 19:20:35 -0700 Subject: [PATCH 5/6] docs(nvsnap): record the page-cache prewarm A/B on a 70B cachedir restore Four restores of Llama-3.1-70B (vLLM TP=4) from the same 131.6 GB rox on a ~1 GB/s network block volume, page cache dropped before each: prewarm 245 s and 249 s, no prewarm 263 s and 251 s, cold 632 s. The prewarm is neutral when a single reader already saturates the volume; its value is a property of the storage class (latency-bound single streams on a high-throughput volume such as Hyperdisk ML), so the doc says so instead of attributing it to model size. Co-Authored-By: Balaji Ganesan --- .../nvsnap/docs/BENCHMARK.md | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/src/compute-plane-services/nvsnap/docs/BENCHMARK.md b/src/compute-plane-services/nvsnap/docs/BENCHMARK.md index ba5eb7c492..fb723ba625 100644 --- a/src/compute-plane-services/nvsnap/docs/BENCHMARK.md +++ b/src/compute-plane-services/nvsnap/docs/BENCHMARK.md @@ -69,6 +69,28 @@ bring-up floor (scheduling, framework import, CUDA init, tensor-parallel worker spawn) that both cold and warm starts pay; cachedir compresses the disk work, not live process initialization. +### Page-cache prewarm A/B, September 2026 + +The cachedir restore adds an `nvsnap-prewarm` init container that reads the +whole rox tree (`find | xargs -P 6 cat`) before the engine starts, so the +engine's weight load hits page cache instead of storage. Opt out per pod with +`NVSNAP_PREWARM=0`. Measured on Llama-3.1-70B-Instruct, vLLM v0.20.0 TP=4, +131.6 GB rox on a network block volume (NVMesh, ~1 GB/s per volume), same node, +page cache dropped before every run: + +| Restore | Ready | Prewarm init | +|---|---:|---:| +| prewarm | 245 s, 249 s | 138 s, 137 s | +| no prewarm | 263 s, 251 s | -- | + +Cold start on the same node: 632 s. The prewarm is neutral here: a single +reader already saturates the volume, so the prewarm only moves the same +storage-bound read ahead of the engine instead of overlapping it with process +bring-up. The prewarm pays off on volumes whose single-stream reads are +latency-bound but whose aggregate throughput is high (Hyperdisk ML), where the +engine's per-fault mmap reads leave the volume idle and the parallel sweep does +not. Treat the prewarm value as a property of the storage class, not the model. + ## Environment - **Cluster**: example-gpu-cluster (GKE) From b5129215f8ca893af1d6cad9f669a2f82f7f7c7d Mon Sep 17 00:00:00 2001 From: Balaji Ganesan Date: Thu, 24 Sep 2026 19:49:48 -0700 Subject: [PATCH 6/6] feat(nvsnap): make the cachedir prewarm a storage-profile decision The page-cache prewarm on a cachedir restore helps or not depending on the volume, not the model: on Hyperdisk ML the engine's per-fault mmap reads leave a high-throughput volume idle and the parallel sweep wins; on a volume one reader already saturates (NVMesh, 70B TP=4: 247 s vs 257 s) it is neutral. A per-pod env var was the only knob, so an operator could not set the right default for a cluster's storage. StorageProfile gains `prewarm` (default on) and `prewarmParallelism` (default 6), settable per provisioner through the existing nvsnap-storage-profiles ConfigMap. The agent resolves the profile once at startup, keeps it next to the L2 backend and hands it to the webhook, whose cachedir restore now takes the init container's presence and reader count from it. NVSNAP_PREWARM=0/1 on the pod still overrides the profile either way, so the existing opt-out keeps working. Tests cover the profile defaults and ConfigMap parsing, the webhook's on/off/parallelism/override matrix through the real patch builder, and the resolver returning the ConfigMap policy intact; the webhook tests were mutation-checked against a resolver that ignores the profile. Also drops two leftovers of the removed entrypoint shim in cachedir.go (an ineffectual hostPath root assignment and an empty append that vet rejected) and gofmt debt on the branch. Co-Authored-By: Balaji Ganesan --- .../nvsnap/docs/BENCHMARK.md | 5 +- .../design/STORAGE-AGNOSTIC-L2-PROMOTION.md | 15 ++++ .../nvsnap/internal/agent/agent.go | 5 ++ .../nvsnap/internal/agent/checkpoint_v2.go | 1 - .../nvsnap/internal/agent/l2_integration.go | 18 +++-- .../internal/agent/l2_profile_prewarm_test.go | 62 ++++++++++++++++ .../internal/agent/webhook_integration.go | 3 + .../checkpointstore/storage_profile.go | 29 ++++++++ .../checkpointstore/storage_profile_test.go | 52 ++++++++++++++ .../nvsnap/internal/webhook/cachedir.go | 71 +++++++++++-------- .../internal/webhook/cachedir_noshim_test.go | 63 ++++++++++++++++ .../nvsnap/internal/webhook/mutate.go | 7 ++ 12 files changed, 292 insertions(+), 39 deletions(-) create mode 100644 src/compute-plane-services/nvsnap/internal/agent/l2_profile_prewarm_test.go diff --git a/src/compute-plane-services/nvsnap/docs/BENCHMARK.md b/src/compute-plane-services/nvsnap/docs/BENCHMARK.md index fb723ba625..a5fc566765 100644 --- a/src/compute-plane-services/nvsnap/docs/BENCHMARK.md +++ b/src/compute-plane-services/nvsnap/docs/BENCHMARK.md @@ -89,7 +89,10 @@ storage-bound read ahead of the engine instead of overlapping it with process bring-up. The prewarm pays off on volumes whose single-stream reads are latency-bound but whose aggregate throughput is high (Hyperdisk ML), where the engine's per-fault mmap reads leave the volume idle and the parallel sweep does -not. Treat the prewarm value as a property of the storage class, not the model. +not. Treat the prewarm value as a property of the storage class, not the model: +the default and the reader count live in the storage profile (`prewarm`, +`prewarmParallelism` in the `nvsnap-storage-profiles` ConfigMap), see +`docs/design/STORAGE-AGNOSTIC-L2-PROMOTION.md`. ## Environment diff --git a/src/compute-plane-services/nvsnap/docs/design/STORAGE-AGNOSTIC-L2-PROMOTION.md b/src/compute-plane-services/nvsnap/docs/design/STORAGE-AGNOSTIC-L2-PROMOTION.md index dc6d72d405..64add87b34 100644 --- a/src/compute-plane-services/nvsnap/docs/design/STORAGE-AGNOSTIC-L2-PROMOTION.md +++ b/src/compute-plane-services/nvsnap/docs/design/STORAGE-AGNOSTIC-L2-PROMOTION.md @@ -215,6 +215,7 @@ data: strategy: shared-volume volumeHandleTransform: nvmesh mountOptions: [ro, norecovery, nouuid] # xfs RO multi-mount needs nouuid + prewarm: false # one reader saturates the volume efs.csi.aws.com: # EFS — shared RWX filesystem strategy: shared-volume volumeHandleTransform: none @@ -230,6 +231,20 @@ disables L2. The built-in table carries these same defaults so a stock install needs no ConfigMap; the ConfigMap exists for new/3rd-party backends and overrides. +Two optional keys tune the cachedir restore's page-cache prewarm per storage +class, because whether the parallel read-ahead helps depends on the volume, +not the model (see the 70B A/B in `docs/BENCHMARK.md`): + +- `prewarm` (default `true`): add the `nvsnap-prewarm` init container that + reads the whole rox tree before the engine starts. Set `false` where a + single reader already saturates the volume. +- `prewarmParallelism` (default 6): concurrent readers in the sweep. Raise it + on volumes whose throughput scales with parallel readers (Hyperdisk ML). + +A pod's own `NVSNAP_PREWARM=0` or `NVSNAP_PREWARM=1` overrides the profile. +The agent resolves the profile once at startup and hands it to the webhook, +so a ConfigMap edit takes effect on the next agent restart. + ## What stays put (no change) - Hash-keyed Lease serialization, single-copy mount-holder write phase, the diff --git a/src/compute-plane-services/nvsnap/internal/agent/agent.go b/src/compute-plane-services/nvsnap/internal/agent/agent.go index b46ca1b285..1c9f5ee23d 100644 --- a/src/compute-plane-services/nvsnap/internal/agent/agent.go +++ b/src/compute-plane-services/nvsnap/internal/agent/agent.go @@ -269,6 +269,11 @@ type Agent struct { // → CRIU dumps stay on L1 hostpath + L4 blobstore; multi-node // restore falls back to L3 peer cascade. l2Backend checkpointstore.Backend + // l2Profile is the StorageProfile resolved for the L2 StorageClass at + // startup (nil when L2 is off or nothing matched). The webhook reads + // its prewarm policy; the promoter strategy is already baked into + // l2Backend. + l2Profile *checkpointstore.StorageProfile // 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/checkpoint_v2.go b/src/compute-plane-services/nvsnap/internal/agent/checkpoint_v2.go index 51f50a690e..e06a819ec0 100644 --- a/src/compute-plane-services/nvsnap/internal/agent/checkpoint_v2.go +++ b/src/compute-plane-services/nvsnap/internal/agent/checkpoint_v2.go @@ -503,4 +503,3 @@ func tailOfFile(path string, n int) string { } return strings.Join(lines, " | ") } - 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 3a6291ef13..8f3ba19ef9 100644 --- a/src/compute-plane-services/nvsnap/internal/agent/l2_integration.go +++ b/src/compute-plane-services/nvsnap/internal/agent/l2_integration.go @@ -51,12 +51,16 @@ const storageProfilesConfigMap = "nvsnap-storage-profiles" // bad ConfigMap, construct error) — the backend then falls back to its // default snapshot-clone ROX promoter. Always logs the resolved strategy // (or the reason for fallback) so an operator can see what L2 will do. -func resolveL2Promoter(ctx context.Context, kc kubernetes.Interface, dyn dynamic.Interface, scName, namespace string, log logrus.FieldLogger) checkpointstore.Promoter { +// +// The resolved profile is returned alongside the promoter because the +// webhook consumes it too (cachedir prewarm policy); nil when nothing +// matched, so callers fall back to the profile's zero-value defaults. +func resolveL2Promoter(ctx context.Context, kc kubernetes.Interface, dyn dynamic.Interface, scName, namespace string, log logrus.FieldLogger) (checkpointstore.Promoter, *checkpointstore.StorageProfile) { sc, err := kc.StorageV1().StorageClasses().Get(ctx, scName, metav1.GetOptions{}) if err != nil { log.WithError(err).WithField("sc", scName). Warn("L2 storage profile: cannot read StorageClass; falling back to default snapshot-clone ROX promoter") - return nil + return nil, nil } provisioner := sc.Provisioner volType := sc.Parameters["type"] // disambiguates pd.csi (hyperdisk-ml vs pd-ssd) @@ -76,20 +80,21 @@ func resolveL2Promoter(ctx context.Context, kc kubernetes.Interface, dyn dynamic if !ok { log.WithFields(logrus.Fields{"sc": scName, "provisioner": provisioner, "type": volType}). Warn("L2 storage profile: no profile for provisioner[/type]; falling back to default snapshot-clone ROX promoter (add an entry to the nvsnap-storage-profiles ConfigMap to support this backend)") - return nil + return nil, nil } promoter, err := checkpointstore.NewPromoterFromProfile(profile, scName, kc, dyn, log) if err != nil { log.WithError(err).WithField("strategy", profile.Strategy). Warn("L2 storage profile: cannot construct promoter; falling back to default") - return nil + return nil, &profile } log.WithFields(logrus.Fields{ "sc": scName, "provisioner": provisioner, "type": volType, "profile_key": key, "strategy": profile.Strategy, "read_only_many": profile.ReadOnlyMany, "vh_transform": profile.VolumeHandleTransform, + "prewarm": profile.PrewarmEnabled(), "prewarm_workers": profile.PrewarmWorkers(), }).Info("L2 storage profile resolved") - return promoter + return promoter, &profile } // startL2Backend constructs the PerCapturePVCBackend if L2 is enabled @@ -169,7 +174,8 @@ func (a *Agent) startL2Backend(_ context.Context, cfg L2BackendConfig) (checkpoi // parameters.type (nvsnap#171). nil ⇒ the backend's applyDefaults // falls back to the snapshot-clone ROX promoter (Hyperdisk-ML // behavior) — back-compat for clusters with no profile match. - promoter := resolveL2Promoter(context.Background(), kc, dyn, cfg.StorageClass, cfg.Namespace, log) + promoter, profile := resolveL2Promoter(context.Background(), kc, dyn, cfg.StorageClass, cfg.Namespace, log) + a.l2Profile = profile // 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/l2_profile_prewarm_test.go b/src/compute-plane-services/nvsnap/internal/agent/l2_profile_prewarm_test.go new file mode 100644 index 0000000000..a48dd27b0f --- /dev/null +++ b/src/compute-plane-services/nvsnap/internal/agent/l2_profile_prewarm_test.go @@ -0,0 +1,62 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package agent + +import ( + "context" + "testing" + + "github.com/sirupsen/logrus" + corev1 "k8s.io/api/core/v1" + storagev1 "k8s.io/api/storage/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/client-go/kubernetes/fake" +) + +// The resolver returns the profile next to the promoter so the webhook can +// read the prewarm policy, and a nvsnap-storage-profiles ConfigMap entry +// that turns the prewarm off for a provisioner reaches it unchanged. +func TestResolveL2Promoter_ReturnsProfileWithPrewarmPolicy(t *testing.T) { + log := logrus.New() + log.SetLevel(logrus.PanicLevel) + sc := &storagev1.StorageClass{ + ObjectMeta: metav1.ObjectMeta{Name: "nvcf-sc"}, + Provisioner: "nvmesh-csi.excelero.com", + } + + kc := fake.NewSimpleClientset(sc) + promoter, profile := resolveL2Promoter(context.Background(), kc, nil, "nvcf-sc", "nvsnap-system", log) + if promoter == nil || profile == nil { + t.Fatalf("built-in NVMesh profile: promoter=%v profile=%v, want both", promoter, profile) + } + if !profile.PrewarmEnabled() || profile.PrewarmWorkers() != 6 { + t.Errorf("built-in profile ships prewarm on with 6 readers, got enabled=%v workers=%d", profile.PrewarmEnabled(), profile.PrewarmWorkers()) + } + + cm := &corev1.ConfigMap{ + ObjectMeta: metav1.ObjectMeta{Name: storageProfilesConfigMap, Namespace: "nvsnap-system"}, + Data: map[string]string{"profiles.yaml": ` +nvmesh-csi.excelero.com: + strategy: shared-volume + volumeHandleTransform: nvmesh + mountOptions: [ro, norecovery, nouuid] + prewarm: false + prewarmParallelism: 2 +`}, + } + kc = fake.NewSimpleClientset(sc, cm) + promoter, profile = resolveL2Promoter(context.Background(), kc, nil, "nvcf-sc", "nvsnap-system", log) + if promoter == nil || profile == nil { + t.Fatalf("ConfigMap overlay: promoter=%v profile=%v, want both", promoter, profile) + } + if profile.PrewarmEnabled() || profile.PrewarmWorkers() != 2 { + t.Errorf("ConfigMap prewarm policy lost on the way to the webhook: enabled=%v workers=%d", profile.PrewarmEnabled(), profile.PrewarmWorkers()) + } + + // No match: nothing to hand the webhook, it falls back to defaults. + unknown := &storagev1.StorageClass{ObjectMeta: metav1.ObjectMeta{Name: "weka"}, Provisioner: "csi.weka.io"} + if _, profile := resolveL2Promoter(context.Background(), fake.NewSimpleClientset(unknown), nil, "weka", "nvsnap-system", log); profile != nil { + t.Errorf("unmatched provisioner must yield a nil profile, got %+v", profile) + } +} 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 273ae84f73..dbc0fabd01 100644 --- a/src/compute-plane-services/nvsnap/internal/agent/webhook_integration.go +++ b/src/compute-plane-services/nvsnap/internal/agent/webhook_integration.go @@ -149,6 +149,9 @@ func (a *Agent) startWebhook(ctx context.Context, cfg WebhookConfig, backend che // restore-side RO mount agree. Empty = standard paths. CacheDir: a.config.RootfsCapture.PodCacheDir, CacheEnvFile: a.config.RootfsCapture.PodCacheEnvFile, + // Storage profile of the L2 StorageClass: the cachedir restore + // reads its prewarm policy (on/off, reader count). nil = defaults. + StorageProfile: a.l2Profile, Composer: &rootfsonly.HashInputComposer{ CUDADriverMajor: a.config.RootfsCapture.CUDADriverMajor, }, diff --git a/src/compute-plane-services/nvsnap/internal/checkpointstore/storage_profile.go b/src/compute-plane-services/nvsnap/internal/checkpointstore/storage_profile.go index 49c6173ce0..c71b491c4f 100644 --- a/src/compute-plane-services/nvsnap/internal/checkpointstore/storage_profile.go +++ b/src/compute-plane-services/nvsnap/internal/checkpointstore/storage_profile.go @@ -69,6 +69,35 @@ type StorageProfile struct { // (RO mount of a possibly-dirty log). Empty ⇒ inherit the primary PV's // options unchanged (fine for NFS-like backends: EFS, Filestore). MountOptions []string `json:"mountOptions,omitempty"` + // Prewarm controls the nvsnap-prewarm init container the webhook adds + // to a cachedir restore: a parallel sequential read of the whole rox + // tree so the engine's weight load hits page cache. nil (absent) means + // on. Whether it helps is a property of the volume, not the model: on + // Hyperdisk ML the engine's per-fault mmap reads leave a high-throughput + // volume idle and the sweep wins; on a volume a single reader already + // saturates (NVMesh, measured 70B TP=4: 247 s vs 257 s) it is neutral. + // A pod's own NVSNAP_PREWARM=0/1 overrides this either way. + Prewarm *bool `json:"prewarm,omitempty"` + // PrewarmParallelism is the number of concurrent readers in the sweep. + // 0 means DefaultPrewarmParallelism. + PrewarmParallelism int `json:"prewarmParallelism,omitempty"` +} + +// DefaultPrewarmParallelism is the reader count when a profile does not set +// one; matches the retired Go prewarmer. +const DefaultPrewarmParallelism = 6 + +// PrewarmEnabled reports the profile's prewarm default (on when unset). +func (p StorageProfile) PrewarmEnabled() bool { + return p.Prewarm == nil || *p.Prewarm +} + +// PrewarmWorkers returns the reader count, defaulted and clamped to >= 1. +func (p StorageProfile) PrewarmWorkers() int { + if p.PrewarmParallelism < 1 { + return DefaultPrewarmParallelism + } + return p.PrewarmParallelism } // builtinProfiles is the compiled-in provider table. Keyed by diff --git a/src/compute-plane-services/nvsnap/internal/checkpointstore/storage_profile_test.go b/src/compute-plane-services/nvsnap/internal/checkpointstore/storage_profile_test.go index 46d28331e4..aa1441bb8e 100644 --- a/src/compute-plane-services/nvsnap/internal/checkpointstore/storage_profile_test.go +++ b/src/compute-plane-services/nvsnap/internal/checkpointstore/storage_profile_test.go @@ -122,3 +122,55 @@ func TestNewPromoterFromProfile(t *testing.T) { t.Error("unknown strategy should error") } } + +// The prewarm policy defaults to on with six readers when a profile says +// nothing, and a ConfigMap entry can turn it off or resize it. +func TestStorageProfile_PrewarmPolicy(t *testing.T) { + var unset StorageProfile + if !unset.PrewarmEnabled() || unset.PrewarmWorkers() != DefaultPrewarmParallelism { + t.Errorf("zero profile: enabled=%v workers=%d, want on/%d", unset.PrewarmEnabled(), unset.PrewarmWorkers(), DefaultPrewarmParallelism) + } + off := false + if (StorageProfile{Prewarm: &off}).PrewarmEnabled() { + t.Error("prewarm: false must disable") + } + on := true + if !(StorageProfile{Prewarm: &on}).PrewarmEnabled() { + t.Error("prewarm: true must enable") + } + if got := (StorageProfile{PrewarmParallelism: 12}).PrewarmWorkers(); got != 12 { + t.Errorf("parallelism 12 -> %d", got) + } + if got := (StorageProfile{PrewarmParallelism: -3}).PrewarmWorkers(); got != DefaultPrewarmParallelism { + t.Errorf("negative parallelism must fall back to the default, got %d", got) + } + + m, err := ParseConfigMapProfiles(` +nvmesh-csi.excelero.com: + strategy: shared-volume + volumeHandleTransform: nvmesh + prewarm: false +pd.csi.storage.gke.io/hyperdisk-ml: + strategy: snapshot-clone + snapshotClass: hdml-images-snapshot-class + readOnlyMany: true + prewarmParallelism: 16 +`) + if err != nil { + t.Fatalf("parse: %v", err) + } + if m["nvmesh-csi.excelero.com"].PrewarmEnabled() { + t.Error("ConfigMap prewarm: false not honoured") + } + hd := m["pd.csi.storage.gke.io/hyperdisk-ml"] + if !hd.PrewarmEnabled() || hd.PrewarmWorkers() != 16 { + t.Errorf("hyperdisk entry: enabled=%v workers=%d, want on/16", hd.PrewarmEnabled(), hd.PrewarmWorkers()) + } + // Every built-in ships with the prewarm on: nothing measured so far + // justifies turning it off by default anywhere. + for k, p := range builtinProfiles { + if !p.PrewarmEnabled() { + t.Errorf("built-in %s ships with prewarm off", k) + } + } +} diff --git a/src/compute-plane-services/nvsnap/internal/webhook/cachedir.go b/src/compute-plane-services/nvsnap/internal/webhook/cachedir.go index 729e99c2e5..10f67df33c 100644 --- a/src/compute-plane-services/nvsnap/internal/webhook/cachedir.go +++ b/src/compute-plane-services/nvsnap/internal/webhook/cachedir.go @@ -262,7 +262,6 @@ func (m *Mutator) cacheDirCapturePatches(pod *corev1.Pod) []PatchOp { return patches } - // tryL2CacheDir injects a cachedir RESTORE: the rox PVC mounted // read-only at m.CacheDir directly (no overlayfs), the cache/model env // vars set identically to capture, and nvsnap-rootfs-restore as the @@ -310,12 +309,6 @@ func (m *Mutator) tryL2CacheDir(ctx context.Context, pod *corev1.Pod, hash strin } } - - root := m.HostBundleRoot - if root == "" { - root = DefaultHostBundleRoot - } - patches := make([]PatchOp, 0, 11+len(manifest.CacheEnv)) if pod.Spec.Volumes == nil { patches = append(patches, PatchOp{Op: "add", Path: "/spec/volumes", Value: []any{}}) @@ -403,16 +396,20 @@ func (m *Mutator) tryL2CacheDir(ctx context.Context, pod *corev1.Pod, hash strin // An init container, not a shim: same node so same page cache, same pod // cgroup so the same memory accounting, and the pod's own command runs // untouched. Best-effort by construction: a read error must never fail a - // restore, so the pipeline ends in || true. Six workers, matching the - // retired Go prewarmer. NVSNAP_PREWARM=0 on the workload skips it, the - // same knob the shim honoured. - if !podDisablesPrewarm(main) { + // restore, so the pipeline ends in || true. + // + // Whether the sweep pays off is a property of the volume, not the + // model (70B A/B in docs/BENCHMARK.md), so the default and the reader + // count come from the L2 StorageClass's StorageProfile, editable per + // cluster through the nvsnap-storage-profiles ConfigMap. A pod's own + // NVSNAP_PREWARM=0/1 still wins, the same knob the shim honoured. + if m.prewarmWanted(main) { prewarmInit := corev1.Container{ Name: "nvsnap-prewarm", Image: main.Image, Command: []string{"sh", "-c", fmt.Sprintf( - "find %s -type f -print0 2>/dev/null | xargs -0 -r -P 6 -n 16 cat > /dev/null 2>&1 || true", - cacheSeedSrcPath)}, + "find %s -type f -print0 2>/dev/null | xargs -0 -r -P %d -n 16 cat > /dev/null 2>&1 || true", + cacheSeedSrcPath, m.prewarmWorkers())}, VolumeMounts: []corev1.VolumeMount{ {Name: cacheDirVolumeName, MountPath: cacheSeedSrcPath, ReadOnly: true}, }, @@ -434,18 +431,16 @@ func (m *Mutator) tryL2CacheDir(ctx context.Context, pod *corev1.Pod, hash strin } else { envs = cacheDirEnvVars(m.CacheDir) } - envs = append(envs, - // 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 - // local snapshot path (engine/arg_utils.py: "when use hf offline, - // replace model ... to local model path"). Capture (cold, online) keeps - // the repo-id, so offline-at-restore changes the model string -> - // different vLLM torch.compile config_hash -> compile-cache MISS -> - // ~20s recompile every restore (gpt-oss-120b, 2026-06-19). The warnings - // are harmless; the recompile is not. Leave offline unset so capture and - // restore compute the same config_hash and the compile cache is reused. - ) + // 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 + // local snapshot path (engine/arg_utils.py: "when use hf offline, + // replace model ... to local model path"). Capture (cold, online) keeps + // the repo-id, so offline-at-restore changes the model string -> + // different vLLM torch.compile config_hash -> compile-cache MISS -> + // ~20s recompile every restore (gpt-oss-120b, 2026-06-19). The warnings + // are harmless; the recompile is not. Leave offline unset so capture and + // restore compute the same config_hash and the compile cache is reused. for _, e := range envs { patches = append(patches, appendEnv(m.MainContainer, e)) } @@ -459,13 +454,27 @@ func (m *Mutator) tryL2CacheDir(ctx context.Context, pod *corev1.Pod, hash strin return patches, nil } -// podDisablesPrewarm honours NVSNAP_PREWARM=0 on the workload container, the -// same opt-out the retired entrypoint shim read. -func podDisablesPrewarm(main corev1.Container) bool { +// prewarmWanted decides whether the cachedir restore gets the nvsnap-prewarm +// init container. An explicit NVSNAP_PREWARM on the workload container wins +// ("0" off, anything else on); otherwise the storage profile decides, and +// with no profile the answer is on. +func (m *Mutator) prewarmWanted(main corev1.Container) bool { for _, e := range main.Env { - if e.Name == "NVSNAP_PREWARM" && e.Value == "0" { - return true + if e.Name == "NVSNAP_PREWARM" { + return e.Value != "0" } } - return false + if m.StorageProfile == nil { + return true + } + return m.StorageProfile.PrewarmEnabled() +} + +// prewarmWorkers is the sweep's reader count from the storage profile, or +// the default without one. +func (m *Mutator) prewarmWorkers() int { + if m.StorageProfile == nil { + return checkpointstore.DefaultPrewarmParallelism + } + return m.StorageProfile.PrewarmWorkers() } diff --git a/src/compute-plane-services/nvsnap/internal/webhook/cachedir_noshim_test.go b/src/compute-plane-services/nvsnap/internal/webhook/cachedir_noshim_test.go index c445143caf..230de33906 100644 --- a/src/compute-plane-services/nvsnap/internal/webhook/cachedir_noshim_test.go +++ b/src/compute-plane-services/nvsnap/internal/webhook/cachedir_noshim_test.go @@ -146,3 +146,66 @@ func TestTryL2CacheDir_PrewarmOptOut(t *testing.T) { } } } + +// The storage profile of the L2 StorageClass decides the prewarm default +// and its reader count; the pod's own NVSNAP_PREWARM still overrides it. +func TestTryL2CacheDir_PrewarmFollowsStorageProfile(t *testing.T) { + newMutator := func(p *checkpointstore.StorageProfile) *Mutator { + return &Mutator{ + CacheDir: "/opt/nvsnap", MainContainer: 0, StorageProfile: p, + L2Backend: &stubL2Backend{mountResult: checkpointstore.PodMount{ + Volume: corev1.Volume{Name: "x", VolumeSource: corev1.VolumeSource{PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{ClaimName: "rox-abc"}}}, + VolumeMount: corev1.VolumeMount{Name: "x", MountPath: "/opt/nvsnap"}, + }}, + } + } + newPod := func(env ...corev1.EnvVar) *corev1.Pod { + return &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Name: "reuse", Namespace: "ns"}, + Spec: corev1.PodSpec{Containers: []corev1.Container{{ + Name: "vllm", Image: "img", Command: []string{"serve"}, Env: env, + }}}, + } + } + prewarmCmd := func(t *testing.T, m *Mutator, pod *corev1.Pod) (string, bool) { + t.Helper() + patches, err := m.tryL2CacheDir(context.Background(), pod, "abc", checkpointstore.Manifest{Hash: "abc", CaptureMethod: "cachedir"}) + if err != nil { + t.Fatal(err) + } + var sawSeed bool + for _, p := range patches { + c, ok := p.Value.(corev1.Container) + if !ok { + continue + } + switch c.Name { + case "nvsnap-seed-cache": + sawSeed = true + case "nvsnap-prewarm": + return strings.Join(c.Command, " "), true + } + } + if !sawSeed { + t.Fatal("the seed init must be present regardless of the prewarm policy") + } + return "", false + } + off, on := false, true + + if cmd, ok := prewarmCmd(t, newMutator(nil), newPod()); !ok || !strings.Contains(cmd, "-P 6 ") { + t.Errorf("no profile: want the prewarm with the default 6 readers, got ok=%v cmd=%q", ok, cmd) + } + if cmd, ok := prewarmCmd(t, newMutator(&checkpointstore.StorageProfile{PrewarmParallelism: 16}), newPod()); !ok || !strings.Contains(cmd, "-P 16 ") { + t.Errorf("profile parallelism 16 not applied, got ok=%v cmd=%q", ok, cmd) + } + if _, ok := prewarmCmd(t, newMutator(&checkpointstore.StorageProfile{Prewarm: &off}), newPod()); ok { + t.Error("profile prewarm: false must drop the prewarm init") + } + if _, ok := prewarmCmd(t, newMutator(&checkpointstore.StorageProfile{Prewarm: &off}), newPod(corev1.EnvVar{Name: "NVSNAP_PREWARM", Value: "1"})); !ok { + t.Error("NVSNAP_PREWARM=1 on the pod must override a profile that turns the prewarm off") + } + if _, ok := prewarmCmd(t, newMutator(&checkpointstore.StorageProfile{Prewarm: &on}), newPod(corev1.EnvVar{Name: "NVSNAP_PREWARM", Value: "0"})); ok { + t.Error("NVSNAP_PREWARM=0 on the pod must override a profile that turns the prewarm on") + } +} diff --git a/src/compute-plane-services/nvsnap/internal/webhook/mutate.go b/src/compute-plane-services/nvsnap/internal/webhook/mutate.go index 889254fc32..11ac018cd9 100644 --- a/src/compute-plane-services/nvsnap/internal/webhook/mutate.go +++ b/src/compute-plane-services/nvsnap/internal/webhook/mutate.go @@ -239,6 +239,13 @@ type Mutator struct { // the agent's --cachedir-env-file flag (a mounted ConfigMap). CacheEnvFile string + // StorageProfile is the profile resolved for the L2 StorageClass + // (nil when L2 is off or nothing matched). The cachedir restore takes + // its page-cache prewarm policy from here: whether to add the + // nvsnap-prewarm init container and how many readers it runs. A pod's + // own NVSNAP_PREWARM=0/1 overrides the profile either way. + StorageProfile *checkpointstore.StorageProfile + // 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