feat(cuda): add a transfer-neutral CustomStorage helper core

agent/Dockerfile

@@ -15,6 +15,7 @@
 # Optional targets for CI:
 #   docker build --target linter --build-context api=../api .   # Run linting
 #   docker build --target tester --build-context api=../api .   # Run tests
+#   docker build --target cuda-helper-core-builder .             # Prove the CUDA helper builds without NIXL
 
 # =============================================================================
 # Build Arguments
@@ -33,6 +34,71 @@ ARG AGENT_BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:25.11-cuda13.0-devel-ubuntu24.0
 # but placeholder builds MUST override it with --build-arg BASE_IMAGE=<image>
 ARG BASE_IMAGE=placeholder-requires-base-image-arg
 
+# =============================================================================
+# Stage: CUDA checkpoint helper core builder (no transfer adapter)
+# =============================================================================
+FROM ${AGENT_BASE_IMAGE} AS cuda-helper-core-builder
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    build-essential \
+    && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /workspace
+
+COPY cmd/cuda-checkpoint-helper/main.cpp ./cmd/cuda-checkpoint-helper/main.cpp
+COPY cmd/cuda-checkpoint-helper/cuda_checkpoint_compat.h ./cmd/cuda-checkpoint-helper/cuda_checkpoint_compat.h
+COPY cmd/cuda-checkpoint-helper/daemon_protocol.cpp ./cmd/cuda-checkpoint-helper/daemon_protocol.cpp
+COPY cmd/cuda-checkpoint-helper/daemon_protocol.h ./cmd/cuda-checkpoint-helper/daemon_protocol.h
+COPY cmd/cuda-checkpoint-helper/daemon_protocol_test.cpp ./cmd/cuda-checkpoint-helper/daemon_protocol_test.cpp
+COPY cmd/cuda-checkpoint-helper/testdata/daemon_request_v6.hex ./cmd/cuda-checkpoint-helper/testdata/daemon_request_v6.hex
+COPY cmd/cuda-checkpoint-helper/storage_manifest.cpp ./cmd/cuda-checkpoint-helper/storage_manifest.cpp
+COPY cmd/cuda-checkpoint-helper/storage_manifest.h ./cmd/cuda-checkpoint-helper/storage_manifest.h
+COPY cmd/cuda-checkpoint-helper/storage_manifest_test.cpp ./cmd/cuda-checkpoint-helper/storage_manifest_test.cpp
+COPY cmd/cuda-checkpoint-helper/transfer_config.cpp ./cmd/cuda-checkpoint-helper/transfer_config.cpp
+COPY cmd/cuda-checkpoint-helper/transfer_config.h ./cmd/cuda-checkpoint-helper/transfer_config.h
+COPY cmd/cuda-checkpoint-helper/transfer_config_test.cpp ./cmd/cuda-checkpoint-helper/transfer_config_test.cpp
+COPY cmd/cuda-checkpoint-helper/transfer_cancellation.h ./cmd/cuda-checkpoint-helper/transfer_cancellation.h
+COPY cmd/cuda-checkpoint-helper/transfer_engine.h ./cmd/cuda-checkpoint-helper/transfer_engine.h
+COPY cmd/cuda-checkpoint-helper/transfer_engine_test.cpp ./cmd/cuda-checkpoint-helper/transfer_engine_test.cpp
+COPY cmd/cuda-checkpoint-helper/transfer_backend_unavailable.cpp ./cmd/cuda-checkpoint-helper/transfer_backend_unavailable.cpp
+
+RUN g++ -std=c++20 -O2 -Wall -Wextra -Werror -pthread \
+      -o /cuda-checkpoint-helper-no-transfer-adapter \
+      ./cmd/cuda-checkpoint-helper/main.cpp \
+      ./cmd/cuda-checkpoint-helper/daemon_protocol.cpp \
+      ./cmd/cuda-checkpoint-helper/storage_manifest.cpp \
+      ./cmd/cuda-checkpoint-helper/transfer_config.cpp \
+      ./cmd/cuda-checkpoint-helper/transfer_backend_unavailable.cpp \
+      -I/usr/local/cuda/include \
+      -L/usr/local/cuda/lib64/stubs \
+      -lcuda \
+    && ldd /cuda-checkpoint-helper-no-transfer-adapter > /tmp/helper-ldd.txt \
+    && ! grep -q nixl /tmp/helper-ldd.txt
+
+RUN g++ -std=c++20 -O2 -Wall -Wextra -Werror -pthread \
+      -o /cuda-checkpoint-helper-daemon-protocol-test \
+      ./cmd/cuda-checkpoint-helper/daemon_protocol.cpp \
+      ./cmd/cuda-checkpoint-helper/daemon_protocol_test.cpp \
+    && /cuda-checkpoint-helper-daemon-protocol-test
+
+RUN g++ -std=c++20 -O2 -Wall -Wextra -Werror \
+      -o /cuda-checkpoint-helper-storage-test \
+      ./cmd/cuda-checkpoint-helper/storage_manifest.cpp \
+      ./cmd/cuda-checkpoint-helper/storage_manifest_test.cpp \
+    && /cuda-checkpoint-helper-storage-test
+
+RUN g++ -std=c++20 -O2 -Wall -Wextra -Werror \
+      -o /cuda-checkpoint-helper-transfer-config-test \
+      ./cmd/cuda-checkpoint-helper/transfer_config.cpp \
+      ./cmd/cuda-checkpoint-helper/transfer_config_test.cpp \
+    && /cuda-checkpoint-helper-transfer-config-test
+
+RUN g++ -std=c++20 -O2 -Wall -Wextra -Werror \
+      -I/usr/local/cuda/include \
+      -o /cuda-checkpoint-helper-transfer-cancellation-test \
+      ./cmd/cuda-checkpoint-helper/transfer_engine_test.cpp \
+    && /cuda-checkpoint-helper-transfer-cancellation-test
+
 # =============================================================================
 # Stage: Go base - Common setup for Go builds
 # =============================================================================

agent/Makefile

@@ -3,15 +3,36 @@ include $(REPO_HACK_DIR)/tools.mk
 
 .DEFAULT_GOAL := build
 
-.PHONY: build test tidy lint fmt
+.PHONY: build test test-cuda-helper tidy lint fmt
 
 build:
 	GOOS=linux GOARCH=amd64 go build -o bin/agent ./cmd/agent
 	GOOS=linux GOARCH=amd64 go build -o bin/nsrestore ./cmd/nsrestore
 
-test:
+test: test-cuda-helper
 	go test ./...
 
+test-cuda-helper:
+	@mkdir -p bin
+	g++ -std=c++20 -O2 -Wall -Wextra -Werror -pthread \
+		-o bin/cuda-checkpoint-helper-daemon-protocol-test \
+		cmd/cuda-checkpoint-helper/daemon_protocol.cpp \
+		cmd/cuda-checkpoint-helper/daemon_protocol_test.cpp
+	bin/cuda-checkpoint-helper-daemon-protocol-test
+	g++ -std=c++20 -O2 -Wall -Wextra -Werror \
+		-o bin/cuda-checkpoint-helper-storage-test \
+		cmd/cuda-checkpoint-helper/storage_manifest.cpp \
+		cmd/cuda-checkpoint-helper/storage_manifest_test.cpp
+	bin/cuda-checkpoint-helper-storage-test
+	g++ -std=c++20 -O2 -Wall -Wextra -Werror \
+		-o bin/cuda-checkpoint-helper-transfer-config-test \
+		cmd/cuda-checkpoint-helper/transfer_config.cpp \
+		cmd/cuda-checkpoint-helper/transfer_config_test.cpp
+	bin/cuda-checkpoint-helper-transfer-config-test
+	g++ -std=c++20 -O2 -Wall -Wextra -Werror \
+		-o bin/cuda-checkpoint-helper-transfer-cancellation-test \
+		cmd/cuda-checkpoint-helper/transfer_engine_test.cpp
+	bin/cuda-checkpoint-helper-transfer-cancellation-test
 fmt:
 	gofmt -w .
 

agent/cmd/cuda-checkpoint-helper/README.md

@@ -0,0 +1,120 @@
+# CUDA checkpoint helper
+
+This directory contains Snapshot's node-local CUDA checkpoint helper. The
+helper isolates CUDA driver calls and CustomStorage callbacks from the Go
+agent. It is an implementation detail of the Snapshot agent, not a
+workload-facing lifecycle API and not the PageBroker transaction protocol.
+
+## Communication boundary
+
+```mermaid
+flowchart LR
+    C["Snapshot controller"] -->|"checkpoint or restore work"| A["Snapshot node agent"]
+    A -->|"resolve and validate target identity"| R["container runtime and /proc"]
+    A <-->|"versioned, bounded Unix SOCK_SEQPACKET RPC"| H["privileged CUDA helper"]
+    H <-->|"CUDA driver checkpoint/restore"| P["target process"]
+    H <-->|"transfer adapter"| B["artifact data plane"]
+```
+
+The controller decides which workload operation is running. The node agent
+owns target discovery, ordering with CRIU, durable manifest construction, and
+the final checkpoint or restore result. The helper owns CUDA driver calls,
+CustomStorage callback lifetime, per-operation transfer cancellation, and
+reporting the observed CUDA target state.
+
+The helper RPC is local to one Snapshot node agent. It does not define:
+
+- the Kubernetes Snapshot API;
+- workload quiesce or resume semantics;
+- durable `checkpoint_id` or manifest publication;
+- PageBroker transaction, commit, or abort semantics; or
+- a network API that PageBroker must implement.
+
+A PageBroker GPU engine may reuse the CUDA operation and transfer behavior
+without adopting this socket protocol. Conversely, Snapshot's local path may
+provide a transfer adapter without changing the workload lifecycle.
+
+## Running
+
+The Snapshot integration configures one privileged helper beside each node
+agent. The helper needs host PID visibility and CUDA driver access so it can
+validate and checkpoint CUDA-owning processes on that node. The operation
+socket is private to the pod and is shared with the agent through an `emptyDir`
+volume.
+
+```text
+cuda-checkpoint-helper --daemon \
+  --socket /run/cuda-checkpoint-helper/helper.sock \
+  --max-operation-seconds 3600
+```
+
+The chart checks readiness through the separate health socket derived from the
+same path:
+
+```text
+cuda-checkpoint-helper --health \
+  --socket /run/cuda-checkpoint-helper/helper.sock
+```
+
+Health succeeds only after the daemon has bound both sockets and advertised
+the deferred-CUDA capability. CustomStorage availability is reported as a
+separate capability so callers can fail before state-changing work when the
+driver or transfer adapter is unavailable.
+
+## Request and response envelope
+
+Each request contains the protocol version, action, validated node-local PID,
+PID identity, storage mode, device mapping, selected GPU UUIDs, and
+operation-specific paths. The bounded request is carried in one
+`SOCK_SEQPACKET` message so the daemon never accepts a partial request as a
+complete operation.
+
+Each response contains the protocol version, operation result, capability and
+fatal-state flags, and a bounded diagnostic payload. Health responses advertise
+capabilities before the agent starts state-changing CUDA work.
+
+The agent must revalidate PID identity before the helper signals or mutates a
+target. Raw host PIDs are node-local execution details and are never durable
+checkpoint identity.
+
+## Operation lifecycle
+
+For checkpoint, the helper locks the target, starts the CUDA checkpoint
+operation, transfers every CustomStorage extent through the selected adapter,
+and completes the CUDA operation handle before returning success. For restore,
+it restores the target from the recorded extent manifest and completes the
+handle. After all restore targets succeed, the agent sends a separate unlock
+request for each target.
+
+One helper request operates on one CUDA-owning PID. The Snapshot agent may
+issue several requests for one workload, but it retains ordering and an
+individual result for every target.
+
+The daemon retains primary contexts only for the request's selected GPU set.
+After a successful operation, it associates those references with the exact
+target PID, process start time, and cgroup. It releases them only after `/proc`
+confirms that target exited or its PID was reused, or during daemon shutdown.
+An inconclusive identity read retains the contexts and defers new work rather
+than risking release underneath a live restored target. A release failure is
+fatal because continuing would make GPU-resource ownership ambiguous.
+
+## Failure rules
+
+- Failure of any extent cancels sibling transfers for that operation.
+- The helper applies one configured cooperative watchdog, capped at one hour,
+  to extent transfers and reports an unhealthy in-flight operation after that
+  threshold. CUDA driver calls are not forcibly interruptible. The client waits
+  up to five minutes longer; an absent response is an unknown outcome and is
+  not replayed.
+- Once a state-changing request may have reached the helper, an unknown result
+  is not replayed automatically.
+- A CUDA operation handle must be completed or resolved before the helper
+  reports a reusable target. An unresolved handle is fatal to that helper
+  process.
+- Storage cleanup, including a future PageBroker abort, does not prove that the
+  CUDA target or workload is safe to resume.
+
+The no-backend build used by the first stack slice validates compilation,
+linkage, and the standalone protocol, manifest, transfer-configuration, and
+cancellation contracts without choosing a production transfer implementation.
+The Snapshot-local NIXL/POSIX adapter and its rollout are added separately.

agent/cmd/cuda-checkpoint-helper/cuda_checkpoint_compat.h

@@ -0,0 +1,112 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved. SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <cuda.h>
+
+#include <cstddef>
+#include <type_traits>
+
+namespace cuda_checkpoint_compat {
+
+#if defined(CUDA_VERSION) && CUDA_VERSION >= 13040
+
+using OperationHandle = CUcheckpointOperationHandle;
+using PerDeviceData = CUcheckpointCustomStoragePerDeviceData;
+using StorageInfo = CUcheckpointCustomStorageInfo;
+using CheckpointArgs = CUcheckpointCheckpointArgs;
+using RestoreArgs = CUcheckpointRestoreArgs;
+using OperationCompleteFn = decltype(&cuCheckpointOperationComplete);
+
+#else
+
+// Public CUDA 13.4 (13040) custom-storage ABI used while the image builds
+// against CUDA 13.0 headers. Keep these declarations local to this helper.
+struct Operation;
+using OperationHandle = Operation *;
+
+struct PerDeviceData {
+  CUdeviceptr devPtr;
+  size_t size;
+  CUstream stream;
+};
+
+struct StorageInfo {
+  OperationHandle handle;
+  PerDeviceData *perDeviceData;
+  unsigned int deviceCount;
+};
+
+struct CheckpointArgs {
+  StorageInfo **customStorageInfo_out;
+  char reserved[64 - sizeof(StorageInfo **)];
+};
+
+struct RestoreArgs {
+  CUcheckpointGpuPair *gpuPairs;
+  unsigned int gpuPairsCount;
+  unsigned int padding0;
+  StorageInfo **customStorageInfo_out;
+  char reserved[64 - sizeof(CUcheckpointGpuPair *) - 2 * sizeof(unsigned int) -
+                sizeof(StorageInfo **)];
+};
+
+using OperationCompleteFn = CUresult(CUDAAPI *)(OperationHandle);
+
+#endif
+
+inline OperationCompleteFn ResolveOperationComplete(bool *available) {
+  void *symbol = nullptr;
+  CUdriverProcAddressQueryResult query_status =
+      CU_GET_PROC_ADDRESS_SYMBOL_NOT_FOUND;
+  const CUresult status =
+      cuGetProcAddress("cuCheckpointOperationComplete", &symbol, 13040,
+                       CU_GET_PROC_ADDRESS_DEFAULT, &query_status);
+  *available = status == CUDA_SUCCESS && symbol != nullptr &&
+               query_status == CU_GET_PROC_ADDRESS_SUCCESS;
+  return *available ? reinterpret_cast<OperationCompleteFn>(symbol) : nullptr;
+}
+
+inline CUcheckpointCheckpointArgs *NativeArgs(CheckpointArgs *args) {
+  return reinterpret_cast<CUcheckpointCheckpointArgs *>(args);
+}
+
+inline CUcheckpointRestoreArgs *NativeArgs(RestoreArgs *args) {
+  return reinterpret_cast<CUcheckpointRestoreArgs *>(args);
+}
+
+static_assert(sizeof(void *) == 8,
+              "CUDA checkpoint custom storage requires a 64-bit ABI");
+static_assert(std::is_standard_layout_v<PerDeviceData>);
+static_assert(sizeof(PerDeviceData) == 24);
+static_assert(alignof(PerDeviceData) == 8);
+static_assert(offsetof(PerDeviceData, devPtr) == 0);
+static_assert(offsetof(PerDeviceData, size) == 8);
+static_assert(offsetof(PerDeviceData, stream) == 16);
+
+static_assert(std::is_standard_layout_v<StorageInfo>);
+static_assert(sizeof(StorageInfo) == 24);
+static_assert(alignof(StorageInfo) == 8);
+static_assert(offsetof(StorageInfo, handle) == 0);
+static_assert(offsetof(StorageInfo, perDeviceData) == 8);
+static_assert(offsetof(StorageInfo, deviceCount) == 16);
+
+static_assert(sizeof(CUcheckpointCheckpointArgs) == 64);
+static_assert(std::is_standard_layout_v<CheckpointArgs>);
+static_assert(sizeof(CheckpointArgs) == sizeof(CUcheckpointCheckpointArgs));
+static_assert(alignof(CheckpointArgs) == alignof(CUcheckpointCheckpointArgs));
+static_assert(offsetof(CheckpointArgs, customStorageInfo_out) == 0);
+
+static_assert(sizeof(CUcheckpointRestoreArgs) == 64);
+static_assert(std::is_standard_layout_v<RestoreArgs>);
+static_assert(sizeof(RestoreArgs) == sizeof(CUcheckpointRestoreArgs));
+static_assert(alignof(RestoreArgs) == alignof(CUcheckpointRestoreArgs));
+static_assert(offsetof(RestoreArgs, gpuPairs) == 0);
+static_assert(offsetof(RestoreArgs, gpuPairsCount) == 8);
+static_assert(offsetof(RestoreArgs, padding0) == 12);
+static_assert(offsetof(RestoreArgs, customStorageInfo_out) == 16);
+
+} // namespace cuda_checkpoint_compat