feat: round-trip test for the CUDA 13.4 CustomStorage checkpoint API

[hhzhang16]

Overview:

Adds a small standalone round-trip test for the CUDA 13.4 CustomStorage checkpoint API.

Details:

What the round-trip script does

1. Creates a clean artifact directory for the checkpoint files.

2. Creates two pipes between the parent and child processes. These pipes are for small communication + coordination messages, not for checkpoint bytes.

3. Forks before either process initializes CUDA -- this avoids inheriting partially initialized or inconsistent CUDA runtime state.

4. The child initializes CUDA, allocates an application buffer, and fills it with a deterministic pattern. Because the pattern is reproducible, the child can reconstruct the expected contents after restore without keeping a second copy.

5. The child tells the parent its application pointer and size. The default application buffer is 64 MiB.

6. The child waits while the parent performs checkpoint and restore.

7. The parent initializes CUDA and looks up cuCheckpointOperationComplete, the CustomStorage entry point.

8. The parent confirms that the child starts in RUNNING, then locks it with cuCheckpointProcessLock(child_pid).

9. The parent calls CUDA checkpoint with customStorageInfo_out enabled. CUDA returns one storage extent containing:

   • the device
   • device UUID
   • checkpoint storage pointer
   • number of bytes
   • CUDA stream

   The CustomStorage pointer is where CUDA exposes its serialized checkpoint state for our code to copy. It is not necessarily the application buffer’s pointer.

10. The parent copies the checkpoint extent from GPU memory into a file through a bounded, pinned CPU buffer. Copies are chunked so the script does not need to allocate enough pinned memory for the complete checkpoint.

11. The parent writes a temporary manifest describing the extent, calls cuCheckpointOperationComplete(handle), confirms the child is now CHECKPOINTED, and only then publishes the manifest. A failed CUDA completion cannot leave an apparently valid checkpoint manifest.

12. Before restore, the script validates the manifest, file size, device identity, and expected extent count.

13. The parent calls CUDA restore, reads the stored extent through the pinned CPU buffer, and copies each chunk back into CUDA-provided storage.

14. After restore completes, the script verifies the remaining state transitions:

    CHECKPOINTED → LOCKED → RUNNING

15. The parent unlocks the child. The child compares every restored application byte with the expected deterministic pattern and performs another CUDA operation to prove the restored process is still usable.

A 120-second watchdog prevents a driver or process hang from occupying a shared GPU indefinitely.

Tests

• Positive case: successful round-trip with validation
• Negative case: truncated artifact (delete a byte)
• Negative case: same-size corruption (overwrite bytes with incorrect data)

Results on Nscale:

• GPU: NVIDIA B200
• Driver: 595.58.03
• Build headers: CUDA 13.0
• Runtime: forward-compatible libcuda.so
• cuCheckpointOperationComplete: present

Normal round trip

• Result: passed
• Application bytes: 67,108,864
• CUDA checkpoint extent: 708,837,376 bytes
• State sequence: RUNNING → LOCKED → CHECKPOINTED → LOCKED → RUNNING
• Checkpoint time: 1.331 seconds
• Restore time: 0.404 seconds
• Post-restore CUDA operation: passed

Truncated artifact

• Artifact rejected before restore
• Result: passed

Same-size corruption

• CUDA restore completed
• Child detected restored byte mismatch at offset 0
• Result: negative test passed

Where should the reviewer start?

Related Issues

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

🔗 This PR is linked to an issue:

• Closes #XXXX

🚫 This PR is NOT linked to an issue:

• Confirmed — no related issue

[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.