feat(nvsnap): checkpoint and restore tensor-parallel workloads through gpushare

[balajinvda]

Why

#2300 imported gpushare (libnvsnap_gpushare.so and nvsnap-gpu-suspend), which lets processes that share GPU memory (NCCL P2P and NVLS, CUDA IPC, page-locked host memory) be checkpointed with their default settings. Nothing in nvsnap used it yet. This wires it into the agent's criu-v2 capture and restore and the webhook, so a tensor-parallel vLLM pod can be captured and restored by nvsnap itself. A manual rehearsal of exactly this path on GB300 (7B, TP=4, cross-node) passed before this change.

What changed

• Webhook: pods annotated nvsnap.io/gpushare: "true" get, on every container requesting GPUs, the node bundle at /nvsnap (read-only), the shim appended to LD_PRELOAD (keeping the pod's value; containers whose LD_PRELOAD comes from a reference are skipped with a warning), and a per-pod chunk store from the node-local checkpoint disk at /var/run/nvsnap/gpushare (subPathExpr on the pod uid).
• Capture (checkpoint_v2.go, gpushare.go): selects the dumped session's processes that load the shim and refuses a GPU process without it; records the GPU map; suspend and stop inside the pod's namespaces; dumps without the CUDA plugin and with --image-io-mode direct when the bundled CRIU supports it; resumes the source on --leave-running or after a failed dump; moves the store into <checkpoint>/gpushare; records pids, store and shim paths in the metadata. Multi-GPU CRIU is accepted when the shim is loaded. One capture per container at a time (ErrCaptureInProgress).
• Restore (restore_v2.go): checks the placeholder provides the shim and the store, restores without the CUDA plugin, then resume --gpu-map limited to the placeholder's allocated GPUs (NVIDIA_VISIBLE_DEVICES).
• Placeholders: the agent-generated one and the e2e templates mount the shim and the checkpoint's store. Also fixes the agent-generated placeholder naming the agent's in-container checkpoint path as its hostPath (that path does not exist on the node) and requests the checkpoint's GPU count.
• e2e: vllm-small-gpushare (TP=1) and vllm-tp4-gpushare (TP=4); annotated workloads stay on criu-v2 at any GPU count (harness and conformance test).
• Docs: docs/GPUSHARE.md gains "Use through nvsnap" with the flow and the integration's limits.

Customer Release Notes

Not customer visible yet (opt-in annotation, not enabled anywhere).

Plan Summary

No chart changes. New hostPath volumes on opted-in pods and placeholders only: <checkpoint root>/gpushare-pods (DirectoryOrCreate) and the node bundle.

Usage

Annotate the workload pod nvsnap.io/gpushare: "true"; capture and restore as for any criu-v2 workload. See docs/GPUSHARE.md.

Testing

• Unit tests: webhook placement through applied JSON patches (store, bundle, LD_PRELOAD append and idempotence, reference skip, bundle reuse, through Mutate); session process selection on a fake procfs (leader first, helpers ignored, GPU-without-shim refused, no-shim plain); dump and restore argv with and without gpushare and direct I/O; generated placeholder for gpushare and plain checkpoints (node paths, GPU count); GPU UUID parsing from environ.
• go test ./..., golangci-lint --new-from-rev clean, check-gazelle clean.
• Manual rehearsal of this path on GB300, 7B TP=4 cross-node: CRIU restore 24 GB in 6.4 s with direct I/O, resume 9.9 s, output matches.
• e2e runs with an image from this branch (TP=1 regression, TP=4 on GB300) to follow on this PR.

Notes

Stacked on #2304 (#2300, #2261, #2106, #2104). Known limits listed in the doc: an image-only LD_PRELOAD (Dockerfile ENV) is replaced; per-pod store directories are emptied on capture, not deleted with the pod; single-pod workloads only.

While wiring this, Bazel drift across the stack was found and fixed on each branch (#2104, #2106, #2261): BUILD files had not been regenerated, and #2261 left a pierrec/lz4/v4 repo import in MODULE.bazel after its last user was removed. Those checks only run on PRs to main.

References

Relates to #2299

Related Pull Requests

#2300, #2304, #2261

Dependencies

None.

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features
  • Added GPU-sharing support for checkpointing and restoring GPU workloads, including workloads using multiple GPUs.
  • Added sample vLLM workloads for single-GPU and four-GPU serving, with matching restore configurations.
• Documentation
  • Added guidance on GPU-sharing setup, checkpoint and restore requirements, and configuration limitations.

[balajinvda]

@coderabbitai review

[coderabbitai]

⚠️ Action not completed

Pull request base or head changed.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

🗂️ Base branches to auto review (1)

• main

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/nvcf/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: aac81634-6c86-4660-a4a7-6cf39f89f03e

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

📝 Generate docstrings

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

• Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

---

_{Comment @coderabbitai help to get the list of available commands.}