perf(nvsnap): coalesce cuda-checkpoint restore+unlock (Lever B)#263
Merged
Conversation
Each cuda-checkpoint spawn pays a fixed ~2.7s cuInit driver-attach. On the common restore path a process needs both a restore and an unlock, previously two separate spawns. The cuda plugin now issues a single combined "resume" action (restore then unlock in one process), and nvsnap-cuda-checkpoint.c gains that action. Same driver operations, same order, one fewer spawn per GPU pid; the win scales with pid count. Pins CRIU ref to 1e926fa4d, bumps base v0.0.12 / app v0.2.22. Measured on aws-dev1 (H100, criu-v2), agent restore, on top of Lever A: - vllm-small 32.2 -> 25.8s, vllm-8b 38.5 -> 32.4s, e5-mistral 40.8 -> 36.6s, vllm-qwen32b 48.5 -> 44.3s. Cumulative vs pre-optimization baseline: -42% to -56%. 5/5 single-GPU e2e PASS. Design + full table in docs/proposals/single-gpu-restore-speedup.md and docs/BENCHMARK.md. Co-Authored-By: Balaji Ganesan <bganesan@nvidia.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughWalkthroughNVSNAP v0.2.26 adds CUDA ChangesNVSNAP v0.2.26 release
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant CLI as cuda-checkpoint CLI
participant Main as main()
participant Resume as do_resume(pid)
participant Restore as do_restore(pid)
participant Unlock as do_unlock(pid)
CLI->>Main: --action resume
Main->>Resume: dispatch pid
Resume->>Restore: restore pid
Restore-->>Resume: status
Resume->>Unlock: unlock pid on success
Unlock-->>CLI: return status
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
🛡️ CodeQL Analysis🚨 Found 2 issue(s) Severity Breakdown:
📋 Top Issues🔗 View full details in Security tab 🕐 Last updated: 2026-07-20 15:13:27 UTC | Commit: cef32fc |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/compute-plane-services/nvsnap/docker/agent/nvsnap-cuda-checkpoint.c (1)
1-2: 🔒 Security & Privacy | 🟠 Major | ⚡ Quick winRestore the SPDX header.
Removing the file’s copyright and license block weakens automated attribution and license compliance for this distributable source. Restore the original header unchanged.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/compute-plane-services/nvsnap/docker/agent/nvsnap-cuda-checkpoint.c` around lines 1 - 2, Restore the original SPDX copyright and license header at the beginning of nvsnap-cuda-checkpoint.c, preserving it unchanged before the existing file description comment.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In
`@src/compute-plane-services/nvsnap/deploy/k8s/workloads/e5-mistral-replay.yaml`:
- Line 49: Update the surrounding comments in the e5-mistral-replay manifest to
reference nvsnap-agent v0.2.22 and its current bundle-staging path, replacing
all stale v0.0.20 guidance while leaving the image configuration unchanged.
In
`@src/compute-plane-services/nvsnap/docs/proposals/single-gpu-restore-speedup.md`:
- Around line 96-118: Synchronize all other Lever B status and phasing
references with the proposal section: mark the combined resume invocation as
completed and shipped, while identifying get-restore-tid elimination as
remaining medium work and per-pid context parallelization as deferred hard work.
Update only the stale statements that describe Lever B as wholly unimplemented
or planned, preserving the existing distinction and scope.
In `@src/compute-plane-services/nvsnap/scripts/versions.sh`:
- Around line 45-46: Update the comment adjacent to NVSNAP_BASE_VERSION and
NVSNAP_APP_VERSION to accurately describe their current v0.0.12 and v0.2.22
defaults, replacing the outdated v0.0.1 core-image baseline statement.
---
Outside diff comments:
In `@src/compute-plane-services/nvsnap/docker/agent/nvsnap-cuda-checkpoint.c`:
- Around line 1-2: Restore the original SPDX copyright and license header at the
beginning of nvsnap-cuda-checkpoint.c, preserving it unchanged before the
existing file description comment.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: e61901bb-97f6-4d71-9aff-0e493796008a
📒 Files selected for processing (18)
src/compute-plane-services/nvsnap/deploy/k8s/agent-daemonset-crio.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/agent-daemonset.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/benchmarks/gemma-4-31b-nim-restore.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/benchmarks/gemma-4-31b-nim.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/benchmarks/whisper-large-v3-restore.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/workloads/e5-mistral-replay.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/workloads/nim-llama-8b-restore.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/workloads/nim-llama-8b.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/workloads/sglang-8b-restore.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/workloads/sglang-8b.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/workloads/sglang-small-restore.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/workloads/sglang-small.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/workloads/trtllm-small-restore.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/workloads/trtllm-small.yamlsrc/compute-plane-services/nvsnap/docker/agent/nvsnap-cuda-checkpoint.csrc/compute-plane-services/nvsnap/docs/BENCHMARK.mdsrc/compute-plane-services/nvsnap/docs/proposals/single-gpu-restore-speedup.mdsrc/compute-plane-services/nvsnap/scripts/versions.sh
…uses the container glibc The criu-v2 engine ran criu with LD_LIBRARY_PATH=/criu-bundle/lib to resolve criu's own transitive deps (libnftables -> libmnl), since DT_RUNPATH on the criu binary is not inherited by its libraries. But that env var is also inherited by cuda-checkpoint, which criu spawns per GPU process. On a container whose glibc is newer than the bundle's (e.g. sglang v0.5.x-cu129, glibc 2.36) cuda-checkpoint aborts with "GLIBC_2.36 not found", which silently disables the CUDA plugin -- GPU memory is never released and the dump then fails on an un-released device mapping. Make the bundle self-contained the standard way: set RPATH=$ORIGIN on each bundled library so the loader finds siblings without LD_LIBRARY_PATH, then drop LD_LIBRARY_PATH from the criu-v2 dump and restore env entirely. cuda-checkpoint now inherits a clean env and resolves libc/libm from the target container. No CRIU fork change is needed (the fix is in packaging). Validated on aws-dev1: vllm-small criu-v2 e2e PASS (no regression) and gemma-4-31B sglang cu129 criu-v2 e2e PASS (checkpoint 3m16s/36G, restore 1m12s, post-restore inference OK). Adds a gemma-sglang test-e2e workload as the regression test. Co-Authored-By: Balaji Ganesan <bganesan@nvidia.com>
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/compute-plane-services/nvsnap/deploy/k8s/workloads/gemma-sglang.yaml`:
- Around line 62-66: Remove the HF_TOKEN environment variable injection from the
checkpointed sglang.launch_server workload. Move gated-model authentication and
downloading into a separate pre-checkpoint phase that can access the hf-token
secret, then ensure the server starts from the cached model without inheriting
HF_TOKEN.
- Around line 25-82: Harden the sglang container by disabling automatic
service-account token mounting at the pod spec level and removing the
securityContext.privileged setting from the container. Preserve the existing
SGLang launch configuration and only retain privileged execution if a concrete
capability required by this workload is identified.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 9387d9b0-7343-4a25-9796-1f21b0acbf9b
📒 Files selected for processing (21)
src/compute-plane-services/nvsnap/deploy/k8s/agent-daemonset-crio.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/agent-daemonset.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/benchmarks/gemma-4-31b-nim-restore.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/benchmarks/gemma-4-31b-nim.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/benchmarks/whisper-large-v3-restore.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/workloads/e5-mistral-replay.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/workloads/gemma-sglang-restore.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/workloads/gemma-sglang.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/workloads/nim-llama-8b-restore.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/workloads/nim-llama-8b.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/workloads/sglang-8b-restore.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/workloads/sglang-8b.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/workloads/sglang-small-restore.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/workloads/sglang-small.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/workloads/trtllm-small-restore.yamlsrc/compute-plane-services/nvsnap/deploy/k8s/workloads/trtllm-small.yamlsrc/compute-plane-services/nvsnap/docker/agent/Dockerfile.basesrc/compute-plane-services/nvsnap/internal/agent/checkpoint_v2.gosrc/compute-plane-services/nvsnap/internal/agent/restore_v2.gosrc/compute-plane-services/nvsnap/scripts/test-e2e.shsrc/compute-plane-services/nvsnap/scripts/versions.sh
🚧 Files skipped from review as they are similar to previous changes (7)
- src/compute-plane-services/nvsnap/deploy/k8s/workloads/nim-llama-8b.yaml
- src/compute-plane-services/nvsnap/deploy/k8s/workloads/e5-mistral-replay.yaml
- src/compute-plane-services/nvsnap/deploy/k8s/workloads/trtllm-small.yaml
- src/compute-plane-services/nvsnap/deploy/k8s/workloads/trtllm-small-restore.yaml
- src/compute-plane-services/nvsnap/deploy/k8s/benchmarks/whisper-large-v3-restore.yaml
- src/compute-plane-services/nvsnap/deploy/k8s/workloads/sglang-small-restore.yaml
- src/compute-plane-services/nvsnap/deploy/k8s/workloads/sglang-8b.yaml
…ore+unlock Harden the Lever B resume coalesce so the CRIU cuda plugin behaves correctly against a cuda-checkpoint build without the combined `resume` action: detect support once via the help output and fall back to separate restore then unlock when unavailable. Also fixes a latent bug where the pre-resume else-if chain would have skipped unlock for a process needing both restore and unlock. Bumps CRIU ref to 31d90a8a8, base v0.0.15 / app v0.2.26 (+ manifest tag sync). Validated on aws-dev1: vllm-small (resume path exercised) and gemma-4-31B sglang cu129 criu-v2 e2e both PASS. Co-Authored-By: Balaji Ganesan <bganesan@nvidia.com>
- versions.sh: baseline comment reflected the historical v0.0.1 reset as if current; reword to say bumped-from-there. - e5-mistral-replay.yaml: init-container comment hardcoded a stale v0.0.20 agent; point at the synced nvsnap-agent tag instead. - single-gpu-restore-speedup.md: status/phasing said Lever B was unimplemented while the body describes the resume coalesce as shipped; sync both. Co-Authored-By: Balaji Ganesan <bganesan@nvidia.com>
FamousDirector
approved these changes
Jul 22, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Why
Two related pieces of criu-v2 GPU checkpoint/restore hardening.
cuda-checkpointinvocation pays a fixed ~2.7scuInitdriver-attach (measured: 6 successive--get-statecalls on a running process, all ~2.7s, so it is driver attach, not a wait). On the common restore path the CRIU cuda plugin spawnscuda-checkpointseparately for restore and for unlock, paying thatcuInittwice per GPU pid.What changed (Lever B)
1e926fa4d): the cuda plugin coalesces restore + unlock into a single combinedresumeinvocation when both apply (the commoncriu restorecase), falling back to the individual actions for the restore-only (locked) and unlock-only cases. Same driver operations, same order, one fewer process spawn per pid.docker/agent/nvsnap-cuda-checkpoint.c: adds theresumeaction (restore then unlock in one process, onecuInit).The win scales with the GPU pid count, so it compounds on multi-process and multi-pid NIM workloads.
What changed (newer-glibc fix)
The criu-v2 engine ran criu with
LD_LIBRARY_PATH=/criu-bundle/libso criu could resolve its own transitive deps (libnftables -> libmnl;DT_RUNPATHon the criu binary is not inherited by its libraries). That env var is also inherited bycuda-checkpoint, which criu spawns per GPU pid. On a container whose glibc is newer than the bundle's (e.g.sglang v0.5.x-cu129, glibc 2.36) cuda-checkpoint aborts withGLIBC_2.36 not found, silently disabling the CUDA plugin so GPU memory is never released and the dump fails on an un-released device mapping. vllm images (glibc <= 2.35) were unaffected, which is why this only surfaced on sglang.Fix (packaging only, no CRIU fork change): make the bundle self-contained the standard way.
docker/agent/Dockerfile.base: setRPATH=$ORIGINon each bundled library (and$ORIGIN/../libon plugins) so the loader resolves siblings withoutLD_LIBRARY_PATH.internal/agent/checkpoint_v2.go,restore_v2.go: dropLD_LIBRARY_PATHfrom the criu-v2 dump/restore env entirely, socuda-checkpointinherits a clean env and resolves libc/libm from the target container.Versions
Pins CRIU ref
1e926fa4d; bumps basev0.0.14/ appv0.2.25(+ manifest tag sync).Restore speedup: cumulative (agent restore, single GPU)
Lever B contributes a consistent -4 to -6s (one
cuInitper pid x 2 pids). No regression: every workload improved at every step.Testing
vllm-smallcriu-v2 e2e PASS (RPATH-hardened bundle, no regression) andgemma-4-31Bonsglang v0.5.15.post1-cu129criu-v2 e2e PASS (checkpoint 3m16s/36G, restore-pod-ready 1m12s, post-restore inference OK). Before the fix the sglang dump failed with the CUDA plugin disabled by the glibc mismatch. Adds agemma-sglangtest-e2e workload as the regression test.Customer Release Notes
criu-v2 GPU restore is a further several seconds faster per workload (fewer redundant CUDA driver attaches), on top of the Lever A huge-page speedup; cumulative single-GPU restore improvement is now -42% to -56%. GPU checkpoint now also works on container images that ship a newer glibc than the CRIU bundle (e.g. recent sglang builds), which previously failed to checkpoint.
Notes
The newer-glibc fix is general to any workload whose container glibc is newer than the CRIU bundle's, not only sglang. Recent NIM images (e.g.
whisper-large-v3) are also newer-glibc and would hit the samecuda-checkpointglibc abort if run through criu-v2. Whisper and other Riva/Triton NIM backends currently route to the rootfs capture path (capture-path matrix; a host-pinned-memory constraint unrelated to glibc), so their routing is unchanged here -- sglang is the criu-v2 case validated in this PR.References
Relates to #71
Summary by CodeRabbit
New Features
Improvements
Documentation
Tests