feat(snapshot): add shared CUDA CustomStorage operation layer

[hhzhang16]

Overview:

• Adds a benchmark-local CUDA 13.4 CustomStorage operation wrapper shared by target-PID checkpoint and restore.
• Adds a standalone single-GPU harness validating checkpoint, POSIX extent storage, restore, exact byte correctness, corruption handling, and restore-failure cleanup.

Details:

• Requires Linux and CUDA 13.4+ toolkit headers and driver support.

Where should the reviewer start?

• benchmarks/cuda_custom_storage/custom_storage_operation.{hpp,cpp} — CUDA operation ownership and completion
• benchmarks/cuda_custom_storage/storage.cpp — POSIX transfer and checkpoint/restore sequencing
• benchmarks/cuda_custom_storage/roundtrip.cpp and workload.cpp — controller orchestration and workload verification

Related Issues

⚠️ This section is required. Choose one path below and delete the other.

🔗 This PR is linked to an issue:

• Closes feat: round-trip test for the CUDA 13.4 CustomStorage checkpoint API #12042

---

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[devin-ai-integration]

Devin Review found 4 potential issues.

[coderabbitai]

Walkthrough

This change adds CUDA checkpoint compatibility types, custom-storage operation management, and a standalone CUDA 13.4 round-trip benchmark. The benchmark coordinates a child workload, transfers checkpoint extents through files, validates restore behavior, and tests truncated and same-size corrupted artifacts.

Changes

CUDA CustomStorage round-trip

Layer / File(s)	Summary
Checkpoint API compatibility deploy/snapshot/cmd/cuda-checkpoint-helper/cuda_checkpoint_compat.h	Defines CUDA 13.4 checkpoint types or compatible fallback structures. Resolves cuCheckpointOperationComplete dynamically and validates ABI layout compatibility.
Custom-storage operation lifecycle deploy/snapshot/cmd/cuda-checkpoint-helper/custom_storage_operation.h, deploy/snapshot/cmd/cuda-checkpoint-helper/custom_storage_operation.cpp	Adds Operation state management for checkpoint and restore calls, storage metadata validation, move construction, one-shot completion, and fatal-state tracking.
Round-trip benchmark and verification benchmarks/cuda_custom_storage/roundtrip.cpp, benchmarks/cuda_custom_storage/Makefile, benchmarks/cuda_custom_storage/README.md	Adds the executable, build targets, and documentation. The harness coordinates CUDA processes, transfers extents through checkpoint artifacts, validates state transitions, handles timeouts, and checks truncated and corrupted artifacts.

Estimated code review effort: 5 (Critical) | ~120 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	The PR implements all required features from issue #12042: parent-child process workflow, checkpoint/restore cycle, checkpoint extent handling, file persistence, state validation, and negative test cases.
Out of Scope Changes check	✅ Passed	All changes directly support the stated objective of adding a shared CUDA CustomStorage operation layer and the round-trip test harness specified in issue #12042.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Title check	✅ Passed	The title clearly summarizes the primary change: adding a shared CUDA CustomStorage operation layer.
Description check	✅ Passed	The description includes all required sections, explains the scope, identifies review entry points, and links issue #12042.

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Fix failing CI checks

---

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

[coderabbitai]

Actionable comments posted: 5

🧹 Nitpick comments (6)

benchmarks/cuda_custom_storage/Makefile (2)

12-25: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Consider adding a test target.

checkmake reports that the required test target is missing. The README already documents the exact round-trip command. A test target would encode that command once and remove the duplication between the README and manual invocation. The target needs a GPU and the CUDA 13.4 driver, so document that requirement in the recipe.

♻️ Proposed addition

-.PHONY: all clean
+.PHONY: all clean test
 
 all: $(TARGET)
+
+# Requires a GPU and a CUDA driver that exposes the CUDA 13.4 CustomStorage API.
+ARTIFACT_DIR ?= $(shell mktemp -d)/artifact
+BYTES ?= 67108864
+
+test: $(TARGET)
+	timeout 130s ./$(TARGET) --artifact-dir "$(ARTIFACT_DIR)" --bytes $(BYTES)

🤖 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 `@benchmarks/cuda_custom_storage/Makefile` around lines 12 - 25, Add a `test`
target to the Makefile following the existing pattern of the `all` and `clean`
targets. First, add `test` to the `.PHONY` declaration alongside `all` and
`clean`. Then, create a `test:` recipe that executes the round-trip command
currently documented in the README, removing the duplication between
documentation and the Makefile. Include a comment in the recipe documenting that
the target requires a GPU and CUDA 13.4 driver to run successfully.

Source: Linters/SAST tools

---

9-10: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Anchor SHARED_DIR to the Makefile location.

SHARED_DIR is relative to the working directory. The README documents make -C benchmarks/cuda_custom_storage, which works. make -f benchmarks/cuda_custom_storage/Makefile from the repository root fails, because the relative path no longer resolves. Derive the directory from MAKEFILE_LIST so both invocations work.

♻️ Proposed refactor

 TARGET := cuda-custom-storage-roundtrip
-SHARED_DIR := ../../deploy/snapshot/cmd/cuda-checkpoint-helper
-SOURCES := roundtrip.cpp $(SHARED_DIR)/custom_storage_operation.cpp
+HERE := $(dir $(lastword $(MAKEFILE_LIST)))
+SHARED_DIR := $(HERE)../../deploy/snapshot/cmd/cuda-checkpoint-helper
+SOURCES := $(HERE)roundtrip.cpp $(SHARED_DIR)/custom_storage_operation.cpp

🤖 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 `@benchmarks/cuda_custom_storage/Makefile` around lines 9 - 10, Update
SHARED_DIR in the Makefile to derive its base directory from MAKEFILE_LIST
rather than the caller’s working directory, while preserving the existing path
to cuda-checkpoint-helper and SOURCES usage so both make invocation forms
resolve correctly.

benchmarks/cuda_custom_storage/roundtrip.cpp (3)

547-548: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Replace the bool checkpoint parameter with a direction enum.

CopyExtent selects two different behaviors from one boolean. The call sites read CopyExtent(..., true) at Line 624 and CopyExtent(..., false) at Line 649, so the direction is not visible without opening the function. A small enum removes that ambiguity and matches the existing CorruptionMode style in this file.

♻️ Proposed refactor

+enum class CopyDirection {
+  kDeviceToFile,
+  kFileToDevice,
+};
+
 void
-CopyExtent(const compat::PerDeviceData& device_data, CUcontext context, const fs::path& path, bool checkpoint)
+CopyExtent(const compat::PerDeviceData& device_data, CUcontext context, const fs::path& path, CopyDirection direction)
 {
   CheckCUDA(cuCtxSetCurrent(context), "cuCtxSetCurrent");
+  const bool checkpoint = direction == CopyDirection::kDeviceToFile;

Then update the call sites:

CopyExtent(info->perDeviceData[0], context, checkpoint_file, CopyDirection::kDeviceToFile);
CopyExtent(info->perDeviceData[0], context, directory / kExtentFilename, CopyDirection::kFileToDevice);

🤖 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 `@benchmarks/cuda_custom_storage/roundtrip.cpp` around lines 547 - 548, Replace
the boolean checkpoint parameter in CopyExtent with a direction enum, such as
CopyDirection, defining explicit device-to-file and file-to-device values.
Update CopyExtent’s branching to use the enum and change both call sites to pass
the corresponding direction values instead of true or false.

---

768-769: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Rewrite the --bytes branch with a real body.

This branch has an empty body and performs the assignment plus the validation inside the condition. ParseUnsigned writes into options.bytes before options.bytes > 0 is evaluated, so an invalid value mutates the option and then falls through to the usage error. The behavior is correct, but the control flow is hard to follow and it is easy to break during a later edit. Parse into a local, then validate.

♻️ Proposed refactor

-    } else if (
-        argument == "--bytes" && ++index < argc && ParseUnsigned(argv[index], &options.bytes) && options.bytes > 0) {
+    } else if (argument == "--bytes" && index + 1 < argc) {
+      ++index;
+      size_t bytes = 0;
+      if (!ParseUnsigned(argv[index], &bytes) || bytes == 0) {
+        throw std::runtime_error("--bytes must be a positive byte count");
+      }
+      options.bytes = bytes;
     } else if (argument == "--device" && ++index < argc) {

🤖 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 `@benchmarks/cuda_custom_storage/roundtrip.cpp` around lines 768 - 769, Rewrite
the --bytes branch in the argument-parsing logic to use a real block body: parse
the argument into a local unsigned value, validate that parsing succeeds and the
value is greater than zero, then assign the validated value to options.bytes.
Preserve the existing usage-error path for missing or invalid values.

---

789-796: 🔒 Security & Privacy | 🔵 Trivial | 💤 Low value

Close the permission window on the artifact directory.

fs::create_directory applies the process umask, which is commonly 0022, so the directory is briefly group- and world-readable. The harness narrows it to owner_all only on the next statement. The README targets shared GPU hosts, so another local user can open the directory during that window and later read checkpoint.bin. Set the umask around the creation, or create the directory and immediately verify ownership before writing.

🤖 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 `@benchmarks/cuda_custom_storage/roundtrip.cpp` around lines 789 - 796, The
permission window between fs::create_directory and fs::permissions allows other
users to access the artifact directory. Close this window by setting the process
umask to a restrictive value (such as 0077) before calling fs::create_directory
on options.artifact_dir, then restore the previous umask immediately after. This
ensures the directory is created with owner-only permissions from the start,
eliminating the brief exposure period where the default umask (commonly 0022)
would make it group- and world-readable.

deploy/snapshot/cmd/cuda-checkpoint-helper/custom_storage_operation.h (1)

19-25: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win

Add a debug assertion in the destructor to enforce the documented "must exit" contract.

The comment above the class states that a caller must check fatal() and exit instead of relying on destruction to complete the operation. Nothing in the class enforces this today. Add a destructor that asserts !fatal(). This catches contract violations during testing without changing release behavior, since assert is compiled out under NDEBUG.

The move constructor already resets the moved-from object's state via std::exchange, so a moved-from Operation reports fatal() == false and will not trigger a false assertion.

🛡️ Proposed destructor addition

   Operation() = default;
   Operation(const Operation&) = delete;
   Operation& operator=(const Operation&) = delete;
   Operation(Operation&& other) noexcept;
   Operation& operator=(Operation&&) = delete;
+  ~Operation();

🤖 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 `@deploy/snapshot/cmd/cuda-checkpoint-helper/custom_storage_operation.h` around
lines 19 - 25, Update the Operation class by adding a destructor that asserts
!fatal(), enforcing the documented requirement that callers check fatal() and
exit before destruction. Keep the existing move-constructor behavior so
moved-from instances remain non-fatal, and rely on assert’s normal NDEBUG
behavior for release builds.

🤖 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 `@benchmarks/cuda_custom_storage/roundtrip.cpp`:
- Around line 1-4: Add benchmarks/cuda_custom_storage/** to
.github/codeowners/areas.yaml and assign it to the same owning team used for
deploy/snapshot/cmd/cuda-checkpoint-helper. No direct changes are needed in
benchmarks/cuda_custom_storage/roundtrip.cpp (lines 1-4),
benchmarks/cuda_custom_storage/Makefile (lines 1-2), or
benchmarks/cuda_custom_storage/README.md (lines 1-4); they are covered by the
new glob.
- Around line 623-633: Move the CopyExtent call into the existing try block in
the checkpoint flow, keeping checkpoint_file available to the catch cleanup.
Ensure failures during copying or checkpoint completion remove the partial
checkpoint_file before rethrowing, while preserving the current completion and
process-state checks.
- Around line 862-876: Update the kTruncate branch to record the checkpoint
file’s actual size after resize_file and assert it equals artifact.extent.size -
1 before calling ValidateExtentFile; remove reliance on matching “wrong size” in
the caught exception, while preserving ValidateExtentFile as the rejection check
and existing success output.
- Around line 644-648: Add validation of the restored extent's device identity
in the Run function. After the existing ValidateStandaloneStorageInfo call,
extract the GPU device from the stream field in info->perDeviceData[0] and
compare it with artifact.extent.device to ensure they match. If the devices
differ, throw a std::runtime_error similar to the existing size mismatch error.
Place this device comparison check before any CopyExtent operation to validate
device identity in addition to the existing size check.

In `@deploy/snapshot/cmd/cuda-checkpoint-helper/custom_storage_operation.cpp`:
- Around line 29-38: Persist storage-shape validity so invalid adopted
checkpoints cannot be reported as successful: in
deploy/snapshot/cmd/cuda-checkpoint-helper/custom_storage_operation.cpp lines
29-38, have Operation::Adopt store the HasValidStorageShape result in
storage_shape_valid_ and return the corresponding status; in lines 68-81,
require storage_shape_valid_ before invoking complete(handle) and transfer it in
the move constructor; in
deploy/snapshot/cmd/cuda-checkpoint-helper/custom_storage_operation.h lines
32-42, add the private bool storage_shape_valid_ initialized to false.

---

Nitpick comments:
In `@benchmarks/cuda_custom_storage/Makefile`:
- Around line 12-25: Add a `test` target to the Makefile following the existing
pattern of the `all` and `clean` targets. First, add `test` to the `.PHONY`
declaration alongside `all` and `clean`. Then, create a `test:` recipe that
executes the round-trip command currently documented in the README, removing the
duplication between documentation and the Makefile. Include a comment in the
recipe documenting that the target requires a GPU and CUDA 13.4 driver to run
successfully.
- Around line 9-10: Update SHARED_DIR in the Makefile to derive its base
directory from MAKEFILE_LIST rather than the caller’s working directory, while
preserving the existing path to cuda-checkpoint-helper and SOURCES usage so both
make invocation forms resolve correctly.

In `@benchmarks/cuda_custom_storage/roundtrip.cpp`:
- Around line 547-548: Replace the boolean checkpoint parameter in CopyExtent
with a direction enum, such as CopyDirection, defining explicit device-to-file
and file-to-device values. Update CopyExtent’s branching to use the enum and
change both call sites to pass the corresponding direction values instead of
true or false.
- Around line 768-769: Rewrite the --bytes branch in the argument-parsing logic
to use a real block body: parse the argument into a local unsigned value,
validate that parsing succeeds and the value is greater than zero, then assign
the validated value to options.bytes. Preserve the existing usage-error path for
missing or invalid values.
- Around line 789-796: The permission window between fs::create_directory and
fs::permissions allows other users to access the artifact directory. Close this
window by setting the process umask to a restrictive value (such as 0077) before
calling fs::create_directory on options.artifact_dir, then restore the previous
umask immediately after. This ensures the directory is created with owner-only
permissions from the start, eliminating the brief exposure period where the
default umask (commonly 0022) would make it group- and world-readable.

In `@deploy/snapshot/cmd/cuda-checkpoint-helper/custom_storage_operation.h`:
- Around line 19-25: Update the Operation class by adding a destructor that
asserts !fatal(), enforcing the documented requirement that callers check
fatal() and exit before destruction. Keep the existing move-constructor behavior
so moved-from instances remain non-fatal, and rely on assert’s normal NDEBUG
behavior for release builds.

🪄 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: 45f0134f-4e5f-429d-b082-b2f5c876d276

📥 Commits

Reviewing files that changed from the base of the PR and between 00ceeec and 143138e.

📒 Files selected for processing (6)

• benchmarks/cuda_custom_storage/Makefile
• benchmarks/cuda_custom_storage/README.md
• benchmarks/cuda_custom_storage/roundtrip.cpp
• deploy/snapshot/cmd/cuda-checkpoint-helper/cuda_checkpoint_compat.h
• deploy/snapshot/cmd/cuda-checkpoint-helper/custom_storage_operation.cpp
• deploy/snapshot/cmd/cuda-checkpoint-helper/custom_storage_operation.h

[hhzhang16]

/ok to test 3ac6bd3

[dfeigin-nv]

The PR Description is stale. Besides the comments - LGTM

[dmitry-tokarev-nv]

Approve from DevOps process only.

[hhzhang16]

/ok to test 614e1c1

[hhzhang16]

/ok to test c296a30

[github-actions]

🌿 Fern Docs Preview: https://nvidia-preview-019fdd3c-2865-7675-9702-f80269665ddf.docs.buildwithfern.com/dynamo/dev