diff --git a/agent/Dockerfile b/agent/Dockerfile index d6fdba85..249cfd18 100644 --- a/agent/Dockerfile +++ b/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= 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 # ============================================================================= diff --git a/agent/Makefile b/agent/Makefile index da32a629..992274a3 100644 --- a/agent/Makefile +++ b/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 . diff --git a/agent/cmd/cuda-checkpoint-helper/README.md b/agent/cmd/cuda-checkpoint-helper/README.md new file mode 100644 index 00000000..ebc08601 --- /dev/null +++ b/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. diff --git a/agent/cmd/cuda-checkpoint-helper/cuda_checkpoint_compat.h b/agent/cmd/cuda-checkpoint-helper/cuda_checkpoint_compat.h new file mode 100644 index 00000000..750718f2 --- /dev/null +++ b/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 + +#include +#include + +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(symbol) : nullptr; +} + +inline CUcheckpointCheckpointArgs *NativeArgs(CheckpointArgs *args) { + return reinterpret_cast(args); +} + +inline CUcheckpointRestoreArgs *NativeArgs(RestoreArgs *args) { + return reinterpret_cast(args); +} + +static_assert(sizeof(void *) == 8, + "CUDA checkpoint custom storage requires a 64-bit ABI"); +static_assert(std::is_standard_layout_v); +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); +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); +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); +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 diff --git a/agent/cmd/cuda-checkpoint-helper/daemon_protocol.cpp b/agent/cmd/cuda-checkpoint-helper/daemon_protocol.cpp new file mode 100644 index 00000000..ef850754 --- /dev/null +++ b/agent/cmd/cuda-checkpoint-helper/daemon_protocol.cpp @@ -0,0 +1,981 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include "daemon_protocol.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cuda_checkpoint_daemon { +namespace { + +uint16_t ReadU16(const unsigned char *data) { + return static_cast(data[0]) | static_cast(data[1]) << 8; +} + +uint32_t ReadU32(const unsigned char *data) { + return static_cast(data[0]) | static_cast(data[1]) << 8 | + static_cast(data[2]) << 16 | + static_cast(data[3]) << 24; +} + +uint64_t ReadU64(const unsigned char *data) { + return static_cast(ReadU32(data)) | + static_cast(ReadU32(data + 4)) << 32; +} + +void WriteU16(std::vector *data, size_t offset, uint16_t value) { + (*data)[offset] = value & 0xff; + (*data)[offset + 1] = value >> 8; +} + +void WriteU32(std::vector *data, size_t offset, uint32_t value) { + for (size_t index = 0; index < 4; ++index) { + (*data)[offset + index] = value >> (index * 8); + } +} + +void WriteU64(std::vector *data, size_t offset, uint64_t value) { + WriteU32(data, offset, value & 0xffffffffU); + WriteU32(data, offset + 4, value >> 32); +} + +bool ContainsNul(const std::string &value) { + return value.find('\0') != std::string::npos; +} + +bool IsAllowedStorageDirectory(const std::string &value) { + const std::filesystem::path path(value); + const std::filesystem::path root("/checkpoints"); + if (!path.is_absolute() || path.lexically_normal() != path || path == root) { + return false; + } + const auto mismatch = std::mismatch(root.begin(), root.end(), path.begin(), + path.end()); + return mismatch.first == root.end(); +} + +bool MakeSocketAddress(const std::string &path, sockaddr_un *address) { + if (path.empty() || path.front() != '/' || + path.size() >= sizeof(address->sun_path)) { + return false; + } + *address = {}; + address->sun_family = AF_UNIX; + std::memcpy(address->sun_path, path.c_str(), path.size() + 1); + return true; +} + +bool SocketConfirmedStale(const sockaddr_un &address) { + const int probe_fd = socket(AF_UNIX, SOCK_SEQPACKET | SOCK_CLOEXEC, 0); + if (probe_fd < 0) { + return false; + } + const int result = connect( + probe_fd, reinterpret_cast(&address), sizeof(address)); + const int connect_error = errno; + close(probe_fd); + return result != 0 && connect_error == ECONNREFUSED; +} + +uint64_t DurationSeconds(std::chrono::steady_clock::time_point start, + std::chrono::steady_clock::time_point end) { + return std::chrono::duration_cast(end - start).count(); +} + +int WriteWake(int fd) { + const unsigned char wake = 1; + ssize_t written; + do { + written = write(fd, &wake, sizeof(wake)); + } while (written < 0 && errno == EINTR); + if (written == static_cast(sizeof(wake))) { + return 0; + } + return written < 0 ? errno : EIO; +} + +void SetError(std::string *error, const char *operation, int error_code) { + *error = std::string(operation) + ": " + std::strerror(error_code); +} + +} // namespace + +bool ParseRequest(const unsigned char *data, size_t size, Request *request, + std::string *error) { + if (data == nullptr || request == nullptr || size < kRequestHeaderSize || + size > kMaxRequestSize) { + *error = "invalid request size"; + return false; + } + if (ReadU32(data) != kMagic || ReadU16(data + 4) != kVersion || + ReadU16(data + 6) != kRequestHeaderSize) { + *error = "invalid request protocol header"; + return false; + } + const auto action = static_cast(ReadU16(data + 8)); + const auto backend = static_cast(ReadU16(data + 10)); + if (action != Action::kHealth && action != Action::kCheckpoint && + action != Action::kRestore && action != Action::kLock && + action != Action::kUnlock) { + *error = "invalid request action"; + return false; + } + const uint32_t device_map_size = ReadU32(data + 28); + const uint32_t storage_dir_size = ReadU32(data + 32); + const uint32_t cgroup_size = ReadU32(data + 36); + const uint32_t job_file_size = ReadU32(data + 48); + const uint32_t selected_devices_size = ReadU32(data + 52); + if (device_map_size > kMaxRequestSize || storage_dir_size > kMaxRequestSize || + cgroup_size > kMaxCgroupSize || job_file_size > kMaxJobFileSize || + selected_devices_size > kMaxRequestSize || + static_cast(device_map_size) + storage_dir_size + cgroup_size + + job_file_size + selected_devices_size != + size - kRequestHeaderSize) { + *error = "invalid request payload lengths"; + return false; + } + Request parsed; + parsed.action = action; + parsed.backend = backend; + parsed.pid = ReadU32(data + 12); + parsed.transfer_buffer_count = ReadU32(data + 16); + parsed.transfer_chunk_bytes = ReadU64(data + 20); + parsed.expected_start_time_ticks = ReadU64(data + 40); + const char *payload = + reinterpret_cast(data + kRequestHeaderSize); + parsed.device_map.assign(payload, device_map_size); + parsed.storage_dir.assign(payload + device_map_size, storage_dir_size); + parsed.expected_cgroup.assign(payload + device_map_size + storage_dir_size, + cgroup_size); + parsed.job_file.assign(payload + device_map_size + storage_dir_size + + cgroup_size, + job_file_size); + parsed.selected_devices.assign( + payload + device_map_size + storage_dir_size + cgroup_size + + job_file_size, + selected_devices_size); + if (ContainsNul(parsed.device_map) || ContainsNul(parsed.storage_dir) || + ContainsNul(parsed.expected_cgroup) || ContainsNul(parsed.job_file) || + ContainsNul(parsed.selected_devices)) { + *error = "request strings contain NUL"; + return false; + } + if (action == Action::kHealth) { + if (parsed.backend != Backend::kUnspecified || parsed.pid != 0 || + parsed.transfer_buffer_count != 0 || parsed.transfer_chunk_bytes != 0 || + parsed.expected_start_time_ticks != 0 || !parsed.device_map.empty() || + !parsed.storage_dir.empty() || !parsed.expected_cgroup.empty() || + !parsed.job_file.empty() || !parsed.selected_devices.empty()) { + *error = "health request has arguments"; + return false; + } + } else if (parsed.pid == 0 || parsed.pid > INT_MAX || + parsed.expected_start_time_ticks == 0 || + parsed.expected_cgroup.empty()) { + *error = "invalid operation arguments"; + return false; + } else if (action == Action::kLock || action == Action::kUnlock) { + if ((parsed.backend != Backend::kRegular && + parsed.backend != Backend::kPosix) || + parsed.transfer_buffer_count != 0 || parsed.transfer_chunk_bytes != 0 || + !parsed.device_map.empty() || !parsed.storage_dir.empty() || + !parsed.selected_devices.empty()) { + *error = "lock/unlock request has transfer arguments"; + return false; + } + } else { + if (backend != Backend::kRegular && backend != Backend::kPosix) { + *error = "checkpoint/restore request has invalid backend"; + return false; + } + if (action == Action::kCheckpoint && !parsed.device_map.empty()) { + *error = "checkpoint request has a device map"; + return false; + } + if (backend == Backend::kRegular) { + if (parsed.transfer_buffer_count != 0 || + parsed.transfer_chunk_bytes != 0 || !parsed.storage_dir.empty() || + !parsed.selected_devices.empty()) { + *error = "regular backend request has custom-storage arguments"; + return false; + } + } else if (parsed.transfer_buffer_count == 0 || + parsed.transfer_chunk_bytes == 0 || parsed.storage_dir.empty() || + !IsAllowedStorageDirectory(parsed.storage_dir) || + parsed.selected_devices.empty()) { + *error = "POSIX backend request has invalid transfer arguments"; + return false; + } + } + if (!parsed.job_file.empty() && parsed.job_file.front() != '/') { + *error = "job file must be absolute"; + return false; + } + *request = std::move(parsed); + return true; +} + +bool EncodeRequest(const Request &request, std::vector *data, + std::string *error) { + if (request.device_map.size() > std::numeric_limits::max() || + request.storage_dir.size() > std::numeric_limits::max() || + request.expected_cgroup.size() > kMaxCgroupSize || + request.job_file.size() > kMaxJobFileSize || + request.selected_devices.size() > + std::numeric_limits::max() || + kRequestHeaderSize + request.device_map.size() + + request.storage_dir.size() + request.expected_cgroup.size() + + request.job_file.size() + request.selected_devices.size() > + kMaxRequestSize) { + *error = "request is too large"; + return false; + } + data->assign(kRequestHeaderSize + request.device_map.size() + + request.storage_dir.size() + request.expected_cgroup.size() + + request.job_file.size() + request.selected_devices.size(), + 0); + WriteU32(data, 0, kMagic); + WriteU16(data, 4, kVersion); + WriteU16(data, 6, kRequestHeaderSize); + WriteU16(data, 8, static_cast(request.action)); + WriteU16(data, 10, static_cast(request.backend)); + WriteU32(data, 12, request.pid); + WriteU32(data, 16, request.transfer_buffer_count); + WriteU64(data, 20, request.transfer_chunk_bytes); + WriteU32(data, 28, request.device_map.size()); + WriteU32(data, 32, request.storage_dir.size()); + WriteU32(data, 36, request.expected_cgroup.size()); + WriteU64(data, 40, request.expected_start_time_ticks); + WriteU32(data, 48, request.job_file.size()); + WriteU32(data, 52, request.selected_devices.size()); + std::memcpy(data->data() + kRequestHeaderSize, request.device_map.data(), + request.device_map.size()); + std::memcpy(data->data() + kRequestHeaderSize + request.device_map.size(), + request.storage_dir.data(), request.storage_dir.size()); + std::memcpy(data->data() + kRequestHeaderSize + request.device_map.size() + + request.storage_dir.size(), + request.expected_cgroup.data(), request.expected_cgroup.size()); + std::memcpy(data->data() + kRequestHeaderSize + request.device_map.size() + + request.storage_dir.size() + request.expected_cgroup.size(), + request.job_file.data(), request.job_file.size()); + std::memcpy(data->data() + kRequestHeaderSize + request.device_map.size() + + request.storage_dir.size() + request.expected_cgroup.size() + + request.job_file.size(), + request.selected_devices.data(), request.selected_devices.size()); + return true; +} + +bool ParseResponse(const unsigned char *data, size_t size, Response *response, + std::string *error) { + if (data == nullptr || response == nullptr || size < kResponseHeaderSize || + size > kMaxResponseSize) { + *error = "invalid response size"; + return false; + } + if (ReadU32(data) != kMagic || ReadU16(data + 4) != kVersion || + ReadU16(data + 6) != kResponseHeaderSize) { + *error = "invalid response protocol header"; + return false; + } + const uint32_t flags = ReadU32(data + 12); + if ((flags & ~(kResponseFatal | kResponseCapabilityDeferredCUDA | + kResponseCapabilityCustomStorage | + kResponseLockNotAcquired)) != 0) { + *error = "invalid response flags"; + return false; + } + const uint32_t output_size = ReadU32(data + 16); + const uint32_t error_size = ReadU32(data + 20); + if (static_cast(output_size) + error_size != + size - kResponseHeaderSize) { + *error = "invalid response payload lengths"; + return false; + } + response->cuda_status = static_cast(ReadU32(data + 8)); + response->flags = flags; + const char *payload = + reinterpret_cast(data + kResponseHeaderSize); + response->output.assign(payload, output_size); + response->error.assign(payload + output_size, error_size); + return true; +} + +bool EncodeResponse(const Response &response, std::vector *data, + std::string *error) { + if (response.output.size() > std::numeric_limits::max() || + response.error.size() > std::numeric_limits::max() || + kResponseHeaderSize + response.output.size() + response.error.size() > + kMaxResponseSize) { + *error = "response is too large"; + return false; + } + data->assign( + kResponseHeaderSize + response.output.size() + response.error.size(), 0); + WriteU32(data, 0, kMagic); + WriteU16(data, 4, kVersion); + WriteU16(data, 6, kResponseHeaderSize); + WriteU32(data, 8, static_cast(response.cuda_status)); + WriteU32(data, 12, response.flags); + WriteU32(data, 16, response.output.size()); + WriteU32(data, 20, response.error.size()); + std::memcpy(data->data() + kResponseHeaderSize, response.output.data(), + response.output.size()); + std::memcpy(data->data() + kResponseHeaderSize + response.output.size(), + response.error.data(), response.error.size()); + return true; +} + +ProcessExistenceState InspectProcessExistence(uint32_t pid, + const std::string &proc_root, + std::string *error) { + const std::filesystem::path process_dir = + std::filesystem::path(proc_root) / std::to_string(pid); + std::error_code exists_error; + const bool process_exists = std::filesystem::exists(process_dir, exists_error); + if (exists_error) { + *error = "failed to inspect current process directory: " + + exists_error.message(); + return ProcessExistenceState::kIndeterminate; + } + if (!process_exists) { + *error = "current process no longer exists"; + return ProcessExistenceState::kMissing; + } + return ProcessExistenceState::kExists; +} + +ProcessIdentityState InspectProcessIdentity(const Request &request, + const std::string &proc_root, + std::string *error) { + const ProcessExistenceState existence = + InspectProcessExistence(request.pid, proc_root, error); + if (existence == ProcessExistenceState::kMissing) { + return ProcessIdentityState::kExitedOrReused; + } + if (existence == ProcessExistenceState::kIndeterminate) { + return ProcessIdentityState::kIndeterminate; + } + const std::filesystem::path process_dir = + std::filesystem::path(proc_root) / std::to_string(request.pid); + std::ifstream stat(process_dir / "stat"); + std::string stat_line; + if (!std::getline(stat, stat_line)) { + *error = "failed to read current process stat"; + return ProcessIdentityState::kIndeterminate; + } + const size_t closing_paren = stat_line.rfind(')'); + if (closing_paren == std::string::npos || + closing_paren + 2 >= stat_line.size()) { + *error = "malformed current process stat"; + return ProcessIdentityState::kIndeterminate; + } + std::istringstream fields(stat_line.substr(closing_paren + 2)); + std::string field; + uint64_t start_time_ticks = 0; + for (size_t index = 0; index < 20; ++index) { + if (!(fields >> field)) { + *error = "malformed current process stat fields"; + return ProcessIdentityState::kIndeterminate; + } + if (index == 19) { + char *end = nullptr; + errno = 0; + const unsigned long long parsed = std::strtoull(field.c_str(), &end, 10); + if (errno != 0 || end == nullptr || *end != '\0') { + *error = "invalid current process start time"; + return ProcessIdentityState::kIndeterminate; + } + start_time_ticks = parsed; + } + } + std::ifstream cgroup_file(process_dir / "cgroup"); + std::ostringstream cgroup_contents; + cgroup_contents << cgroup_file.rdbuf(); + const std::string cgroup = cgroup_contents.str(); + if ((!cgroup_file.good() && !cgroup_file.eof()) || cgroup.empty() || + cgroup.size() > kMaxCgroupSize) { + *error = "failed to read valid current process cgroup"; + return ProcessIdentityState::kIndeterminate; + } + if (start_time_ticks != request.expected_start_time_ticks) { + *error = "start time mismatch"; + return ProcessIdentityState::kExitedOrReused; + } + if (cgroup != request.expected_cgroup) { + *error = "cgroup changed while PID and start time still match"; + return ProcessIdentityState::kIndeterminate; + } + return ProcessIdentityState::kMatches; +} + +bool ValidateProcessIdentity(const Request &request, + const std::string &proc_root, std::string *error) { + return InspectProcessIdentity(request, proc_root, error) == + ProcessIdentityState::kMatches; +} + +bool ExecuteValidated(const Request &request, const std::string &proc_root, + const OperationExecutor &executor, Response *response) { + if (request.action == Action::kHealth) { + response->cuda_status = 1; + response->error = "health requests must use the health socket"; + return true; + } + std::string identity_error; + if (!ValidateProcessIdentity(request, proc_root, &identity_error)) { + response->cuda_status = 1; + if (request.action == Action::kLock) { + response->flags |= kResponseLockNotAcquired; + } + response->error = + "process identity changed before CUDA operation: " + identity_error; + return true; + } + *response = executor(request); + return ResponseAllowsServerContinue(*response); +} + +bool ResponseAllowsServerContinue(const Response &response) { + return (response.flags & kResponseFatal) == 0; +} + +int32_t FinishHandledOperation(bool post_handle_succeeded, + int32_t failure_status, + const CompletionExecutor &completion, + OperationState *state) { + state->handle_returned = true; + if (!post_handle_succeeded) { + return failure_status; + } + const int32_t status = completion(); + state->completion_succeeded = status == 0; + return status; +} + +OperationHealth::OperationHealth(std::chrono::seconds max_operation_duration) + : max_operation_duration_(max_operation_duration) {} + +void OperationHealth::MarkReady(bool custom_storage_available) { + std::lock_guard lock(mutex_); + ready_ = true; + custom_storage_available_ = custom_storage_available; +} + +void OperationHealth::Begin(Action action, uint32_t pid) { + std::lock_guard lock(mutex_); + const auto now = Clock::now(); + busy_ = true; + action_ = action; + pid_ = pid; + started_ = now; +} + +void OperationHealth::End() { + std::lock_guard lock(mutex_); + busy_ = false; + action_ = Action::kHealth; + pid_ = 0; +} + +HealthSnapshot OperationHealth::Snapshot() const { + std::lock_guard lock(mutex_); + const auto now = Clock::now(); + HealthSnapshot snapshot{ + .ready = ready_, + .busy = busy_, + .healthy = ready_, + .action = action_, + .pid = pid_, + .deadline_seconds = + static_cast(max_operation_duration_.count()), + .custom_storage_available = custom_storage_available_, + }; + if (busy_) { + snapshot.elapsed_seconds = DurationSeconds(started_, now); + snapshot.healthy = ready_ && now - started_ <= max_operation_duration_; + } + return snapshot; +} + +Response HealthResponse(const OperationHealth &health) { + const HealthSnapshot snapshot = health.Snapshot(); + Response response; + response.cuda_status = snapshot.healthy ? 0 : 1; + if (snapshot.ready) { + response.flags = kResponseCapabilityDeferredCUDA; + if (snapshot.custom_storage_available) { + response.flags |= kResponseCapabilityCustomStorage; + } + } + std::ostringstream output; + output << "{\"ready\":" << (snapshot.ready ? "true" : "false") + << ",\"busy\":" << (snapshot.busy ? "true" : "false") + << ",\"healthy\":" << (snapshot.healthy ? "true" : "false") + << ",\"action\":\"" << ActionName(snapshot.action) << "\"" + << ",\"pid\":" << snapshot.pid + << ",\"elapsed_seconds\":" << snapshot.elapsed_seconds + << ",\"deadline_seconds\":" << snapshot.deadline_seconds + << ",\"custom_storage_available\":" + << (snapshot.custom_storage_available ? "true" : "false") << "}\n"; + response.output = output.str(); + if (!snapshot.healthy) { + response.error = snapshot.ready ? "operation exceeded watchdog deadline" + : "daemon is not ready"; + } + return response; +} + +Response HealthResponseAfterReap(const OperationHealth &health, + int32_t release_status, + const std::string &reap_warning) { + if (release_status != 0) { + return Response{ + .cuda_status = release_status, + .flags = kResponseFatal, + .output = "", + .error = "failed to release CUDA primary contexts for an exited " + "restore target", + }; + } + Response response = HealthResponse(health); + if (!reap_warning.empty()) { + response.error = reap_warning + "; retained contexts; reaping deferred"; + } + return response; +} + +bool OperationTimeoutMilliseconds(std::chrono::seconds duration, + unsigned int *timeout_ms, + std::string *error) { + if (timeout_ms == nullptr || duration.count() <= 0 || + static_cast(duration.count()) > + std::numeric_limits::max() / 1000ULL) { + if (error != nullptr) { + *error = "operation duration cannot be represented as a CUDA lock timeout"; + } + return false; + } + *timeout_ms = static_cast(duration.count() * 1000ULL); + return true; +} + +OwnedUnixSocket::~OwnedUnixSocket() { Close(); } + +BoundedOutputCapture::~BoundedOutputCapture() noexcept { + if (write_fd_ >= 0) { + close(write_fd_); + write_fd_ = -1; + } + if (reader_.joinable()) { + reader_.join(); + } + if (read_fd_ >= 0) { + close(read_fd_); + read_fd_ = -1; + } +} + +bool BoundedOutputCapture::Start(std::string *error) { + int pipe_fds[2]{-1, -1}; + if (pipe2(pipe_fds, O_CLOEXEC) != 0) { + if (error != nullptr) { + *error = std::string("create bounded output pipe: ") + + std::strerror(errno); + } + return false; + } + read_fd_ = pipe_fds[0]; + write_fd_ = pipe_fds[1]; + try { + reader_ = std::thread(&BoundedOutputCapture::Drain, this); + } catch (const std::system_error &e) { + close(read_fd_); + close(write_fd_); + read_fd_ = -1; + write_fd_ = -1; + if (error != nullptr) { + *error = std::string("start bounded output reader: ") + e.what(); + } + return false; + } + return true; +} + +void BoundedOutputCapture::Drain() noexcept { + std::array buffer{}; + for (;;) { + const ssize_t count = read(read_fd_, buffer.data(), buffer.size()); + if (count > 0) { + const size_t available = limit_ > output_.size() ? limit_ - output_.size() : 0; + const size_t retained = std::min(available, static_cast(count)); + output_.append(buffer.data(), retained); + truncated_ = truncated_ || retained < static_cast(count); + continue; + } + if (count == 0) { + return; + } + if (errno == EINTR) { + continue; + } + read_error_ = std::strerror(errno); + return; + } +} + +bool BoundedOutputCapture::Finish(std::string *output, bool *truncated, + std::string *error) { + if (write_fd_ >= 0) { + close(write_fd_); + write_fd_ = -1; + } + if (reader_.joinable()) { + reader_.join(); + } + if (read_fd_ >= 0) { + close(read_fd_); + read_fd_ = -1; + } + *output = std::move(output_); + *truncated = truncated_; + if (!read_error_.empty()) { + if (error != nullptr) { + *error = "read bounded output: " + read_error_; + } + return false; + } + return true; +} + +bool OwnedUnixSocket::Bind(const std::string &path, int backlog, + std::string *error) { + sockaddr_un address{}; + if (fd_ >= 0 || !MakeSocketAddress(path, &address)) { + *error = "invalid socket path"; + return false; + } + const std::string lock_path = path + ".lock"; + lock_fd_ = open(lock_path.c_str(), O_CREAT | O_CLOEXEC | O_RDWR, 0600); + if (lock_fd_ < 0 || flock(lock_fd_, LOCK_EX | LOCK_NB) != 0) { + *error = lock_fd_ < 0 ? std::strerror(errno) + : "socket path is owned by another server"; + Close(); + return false; + } + fd_ = socket(AF_UNIX, SOCK_SEQPACKET | SOCK_CLOEXEC, 0); + if (fd_ < 0) { + *error = std::strerror(errno); + Close(); + return false; + } + struct stat existing{}; + if (lstat(path.c_str(), &existing) == 0) { + if (!S_ISSOCK(existing.st_mode) || !SocketConfirmedStale(address)) { + *error = "refusing to replace non-socket or live socket path"; + Close(); + return false; + } + if (unlink(path.c_str()) != 0) { + *error = std::strerror(errno); + Close(); + return false; + } + } else if (errno != ENOENT) { + *error = std::strerror(errno); + Close(); + return false; + } + if (bind(fd_, reinterpret_cast(&address), sizeof(address)) != 0) { + *error = std::strerror(errno); + Close(); + return false; + } + path_ = path; + bound_ = true; + struct stat bound_stat{}; + if (lstat(path.c_str(), &bound_stat) != 0 || !S_ISSOCK(bound_stat.st_mode)) { + *error = "failed to inspect bound socket"; + Close(); + return false; + } + device_ = bound_stat.st_dev; + inode_ = bound_stat.st_ino; + if (backlog < 0) { + *error = "invalid backlog"; + Close(); + return false; + } + if (chmod(path.c_str(), 0600) != 0 || listen(fd_, backlog) != 0) { + *error = std::strerror(errno); + Close(); + return false; + } + return true; +} + +void OwnedUnixSocket::UnlinkIfOwned() { + if (!bound_) { + return; + } + struct stat current{}; + if (lstat(path_.c_str(), ¤t) == 0 && + static_cast(current.st_dev) == device_ && + static_cast(current.st_ino) == inode_) { + (void)unlink(path_.c_str()); + } + bound_ = false; +} + +void OwnedUnixSocket::Close() { + if (fd_ >= 0) { + close(fd_); + fd_ = -1; + } + UnlinkIfOwned(); + if (lock_fd_ >= 0) { + close(lock_fd_); + lock_fd_ = -1; + } +} + +ShutdownSignalOwner::~ShutdownSignalOwner() noexcept { + (void)StopAndJoinNoThrow(); + Close(); +} + +bool ShutdownSignalOwner::Start(std::string *error) { + if (thread_started_ || signal_fd_ >= 0) { + *error = "shutdown signal owner already started"; + return false; + } + + if (sigemptyset(&signals_) != 0 || sigaddset(&signals_, SIGTERM) != 0 || + sigaddset(&signals_, SIGINT) != 0) { + SetError(error, "build shutdown signal set", errno); + return false; + } + const int mask_error = pthread_sigmask(SIG_BLOCK, &signals_, nullptr); + if (mask_error != 0) { + SetError(error, "block shutdown signals", mask_error); + return false; + } + // Keep the mask blocked until process exit. Restoring it after the owner is + // joined would reopen a window for the default signal action during teardown. + + if (pipe2(operation_stop_pipe_, O_CLOEXEC | O_NONBLOCK) != 0) { + SetError(error, "create operation shutdown pipe", errno); + Close(); + return false; + } + if (pipe2(health_stop_pipe_, O_CLOEXEC | O_NONBLOCK) != 0) { + SetError(error, "create health shutdown pipe", errno); + Close(); + return false; + } + if (pipe2(control_pipe_, O_CLOEXEC | O_NONBLOCK) != 0) { + SetError(error, "create signal owner control pipe", errno); + Close(); + return false; + } + signal_fd_ = signalfd(-1, &signals_, SFD_CLOEXEC | SFD_NONBLOCK); + if (signal_fd_ < 0) { + SetError(error, "create shutdown signal fd", errno); + Close(); + return false; + } + + const int thread_error = pthread_create( + &thread_, nullptr, &ShutdownSignalOwner::ThreadEntry, this); + if (thread_error != 0) { + SetError(error, "create shutdown signal thread", thread_error); + Close(); + return false; + } + thread_started_ = true; + return true; +} + +bool ShutdownSignalOwner::StopAndJoin(std::string *error) { + const ShutdownResult result = StopAndJoinNoThrow(); + if (!result.ok()) { + SetError(error, result.operation, result.error_code); + return false; + } + return true; +} + +ShutdownSignalOwner::ShutdownResult +ShutdownSignalOwner::RequestShutdownNoThrow() noexcept { + if (!thread_started_) { + return {}; + } + + const int wake_error = WriteWake(control_pipe_[1]); + if (wake_error != 0) { + return { + .operation = "wake shutdown signal thread", + .error_code = wake_error, + }; + } + return {}; +} + +ShutdownSignalOwner::ShutdownResult +ShutdownSignalOwner::StopAndJoinNoThrow() noexcept { + if (!thread_started_) { + return {}; + } + + const ShutdownResult wake_result = RequestShutdownNoThrow(); + const int join_error = pthread_join(thread_, nullptr); + if (join_error != 0) { + return { + .operation = "join shutdown signal thread", + .error_code = join_error, + }; + } + thread_started_ = false; + if (!wake_result.ok()) { + return wake_result; + } + if (thread_error_ != 0) { + return { + .operation = "shutdown signal thread", + .error_code = thread_error_, + }; + } + return {}; +} + +void ShutdownSignalOwner::Close() noexcept { + if (thread_started_) { + return; + } + for (int *pipe : {operation_stop_pipe_, health_stop_pipe_, control_pipe_}) { + for (size_t index = 0; index < 2; ++index) { + if (pipe[index] >= 0) { + close(pipe[index]); + pipe[index] = -1; + } + } + } + if (signal_fd_ >= 0) { + close(signal_fd_); + signal_fd_ = -1; + } +} + +void *ShutdownSignalOwner::ThreadEntry(void *owner) noexcept { + static_cast(owner)->Run(); + return nullptr; +} + +void ShutdownSignalOwner::Run() noexcept { + pollfd descriptors[2]{ + {.fd = signal_fd_, .events = POLLIN, .revents = 0}, + {.fd = control_pipe_[0], .events = POLLIN, .revents = 0}, + }; + int poll_result; + do { + poll_result = poll(descriptors, 2, -1); + } while (poll_result < 0 && errno == EINTR); + if (poll_result < 0) { + thread_error_ = errno; + } else if ((descriptors[0].revents & + (POLLIN | POLLERR | POLLHUP | POLLNVAL)) != 0) { + signalfd_siginfo signal_info{}; + const ssize_t bytes = read(signal_fd_, &signal_info, sizeof(signal_info)); + if (bytes != static_cast(sizeof(signal_info))) { + thread_error_ = bytes < 0 ? errno : EIO; + } else if (signal_info.ssi_signo != SIGTERM && + signal_info.ssi_signo != SIGINT) { + thread_error_ = EINVAL; + } + } else if ((descriptors[1].revents & + (POLLIN | POLLERR | POLLHUP | POLLNVAL)) != 0) { + unsigned char control = 0; + const ssize_t bytes = read(control_pipe_[0], &control, sizeof(control)); + if (bytes != static_cast(sizeof(control))) { + thread_error_ = bytes < 0 ? errno : EIO; + } + } else { + thread_error_ = EIO; + } + + shutdown_requested_.store(true, std::memory_order_release); + const int operation_error = WriteWake(operation_stop_pipe_[1]); + const int health_error = WriteWake(health_stop_pipe_[1]); + if (thread_error_ == 0) { + thread_error_ = operation_error != 0 ? operation_error : health_error; + } +} + +int PollForInputOrStop(int input_fd, int stop_fd, + const std::function &before_poll, + int timeout_milliseconds) { + pollfd descriptors[2]{ + {.fd = input_fd, .events = POLLIN, .revents = 0}, + {.fd = stop_fd, .events = POLLIN, .revents = 0}, + }; + if (before_poll) { + before_poll(); + } + for (;;) { + const int result = poll(descriptors, 2, timeout_milliseconds); + if (result < 0 && errno == EINTR) { + continue; + } + if (result == 0) { + return -1; + } + if (result < 0) { + return -2; + } + if ((descriptors[1].revents & (POLLIN | POLLERR | POLLHUP | POLLNVAL)) != + 0) { + return 0; + } + if ((descriptors[0].revents & (POLLIN | POLLERR | POLLHUP | POLLNVAL)) != + 0) { + return 1; + } + } +} + +const char *ActionName(Action action) { + switch (action) { + case Action::kHealth: + return "none"; + case Action::kCheckpoint: + return "checkpoint"; + case Action::kRestore: + return "restore"; + case Action::kLock: + return "lock"; + case Action::kUnlock: + return "unlock"; + } + return "unknown"; +} + +} // namespace cuda_checkpoint_daemon diff --git a/agent/cmd/cuda-checkpoint-helper/daemon_protocol.h b/agent/cmd/cuda-checkpoint-helper/daemon_protocol.h new file mode 100644 index 00000000..323cc39f --- /dev/null +++ b/agent/cmd/cuda-checkpoint-helper/daemon_protocol.h @@ -0,0 +1,254 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cuda_checkpoint_daemon { + +constexpr uint32_t kMagic = 0x50484344; // "DCHP" in little-endian. +constexpr uint16_t kVersion = 6; +constexpr size_t kRequestHeaderSize = 56; +constexpr size_t kResponseHeaderSize = 24; +constexpr size_t kMaxRequestSize = 64 * 1024; +constexpr size_t kMaxResponseSize = 128 * 1024; +constexpr size_t kMaxCgroupSize = 4096; +constexpr size_t kMaxJobFileSize = 4096; + +constexpr uint32_t kResponseFatal = 1U << 0; +constexpr uint32_t kResponseCapabilityDeferredCUDA = 1U << 1; +constexpr uint32_t kResponseCapabilityCustomStorage = 1U << 2; +constexpr uint32_t kResponseLockNotAcquired = 1U << 3; + +enum class Action : uint16_t { + kHealth = 0, + kCheckpoint = 1, + kRestore = 2, + kLock = 3, + kUnlock = 4, +}; + +enum class Backend : uint16_t { + kUnspecified = 0, + kRegular = 1, + kPosix = 2, +}; + +// Cleanup must distinguish a confirmed exit/PID reuse from a transient +// observation failure. Releasing retained CUDA state on an inconclusive +// /proc read can break a live restored process. +enum class ProcessIdentityState : uint8_t { + kMatches, + kExitedOrReused, + kIndeterminate, +}; + +enum class ProcessExistenceState : uint8_t { + kExists, + kMissing, + kIndeterminate, +}; + +struct Request { + Action action = Action::kHealth; + Backend backend = Backend::kUnspecified; + uint32_t pid = 0; + uint32_t transfer_buffer_count = 0; + uint64_t transfer_chunk_bytes = 0; + uint64_t expected_start_time_ticks = 0; + std::string device_map; + std::string storage_dir; + std::string expected_cgroup; + std::string job_file; + std::string selected_devices; +}; + +struct Response { + int32_t cuda_status = 0; + uint32_t flags = 0; + std::string output; + std::string error; +}; + +struct OperationState { + bool handle_returned = false; + bool completion_succeeded = false; + + bool fatal() const { return handle_returned && !completion_succeeded; } +}; + +struct HealthSnapshot { + bool ready = false; + bool busy = false; + bool healthy = false; + Action action = Action::kHealth; + uint32_t pid = 0; + uint64_t elapsed_seconds = 0; + uint64_t deadline_seconds = 0; + bool custom_storage_available = false; +}; + +class OperationHealth { +public: + explicit OperationHealth(std::chrono::seconds max_operation_duration); + + void MarkReady(bool custom_storage_available); + void Begin(Action action, uint32_t pid); + void End(); + HealthSnapshot Snapshot() const; + +private: + using Clock = std::chrono::steady_clock; + + const std::chrono::seconds max_operation_duration_; + mutable std::mutex mutex_; + bool ready_ = false; + bool custom_storage_available_ = false; + bool busy_ = false; + Action action_ = Action::kHealth; + uint32_t pid_ = 0; + Clock::time_point started_{}; +}; + +class OwnedUnixSocket { +public: + OwnedUnixSocket() = default; + OwnedUnixSocket(const OwnedUnixSocket &) = delete; + OwnedUnixSocket &operator=(const OwnedUnixSocket &) = delete; + ~OwnedUnixSocket(); + + bool Bind(const std::string &path, int backlog, std::string *error); + void Close(); + int fd() const { return fd_; } + +private: + void UnlinkIfOwned(); + + int fd_ = -1; + int lock_fd_ = -1; + std::string path_; + uint64_t device_ = 0; + uint64_t inode_ = 0; + bool bound_ = false; +}; + +class ShutdownSignalOwner { +public: + struct ShutdownResult { + const char *operation = nullptr; + int error_code = 0; + + bool ok() const { return error_code == 0; } + }; + + ShutdownSignalOwner() = default; + ShutdownSignalOwner(const ShutdownSignalOwner &) = delete; + ShutdownSignalOwner &operator=(const ShutdownSignalOwner &) = delete; + ~ShutdownSignalOwner() noexcept; + + bool Start(std::string *error); + ShutdownResult RequestShutdownNoThrow() noexcept; + ShutdownResult StopAndJoinNoThrow() noexcept; + bool StopAndJoin(std::string *error); + void Close() noexcept; + bool ShutdownRequested() const { + return shutdown_requested_.load(std::memory_order_acquire); + } + int operation_stop_fd() const { return operation_stop_pipe_[0]; } + int health_stop_fd() const { return health_stop_pipe_[0]; } + +private: + static void *ThreadEntry(void *owner) noexcept; + void Run() noexcept; + + std::atomic shutdown_requested_{false}; + sigset_t signals_{}; + pthread_t thread_{}; + std::atomic thread_started_{false}; + int signal_fd_ = -1; + int operation_stop_pipe_[2]{-1, -1}; + int health_stop_pipe_[2]{-1, -1}; + int control_pipe_[2]{-1, -1}; + int thread_error_ = 0; +}; + +// BoundedOutputCapture drains a pipe continuously while retaining at most +// limit bytes. This prevents verbose CUDA diagnostics from filling ephemeral +// storage or blocking the daemon on a full pipe. +class BoundedOutputCapture { +public: + explicit BoundedOutputCapture(size_t limit) : limit_(limit) {} + BoundedOutputCapture(const BoundedOutputCapture &) = delete; + BoundedOutputCapture &operator=(const BoundedOutputCapture &) = delete; + ~BoundedOutputCapture() noexcept; + + bool Start(std::string *error); + int write_fd() const { return write_fd_; } + bool Finish(std::string *output, bool *truncated, std::string *error); + +private: + void Drain() noexcept; + + size_t limit_; + int read_fd_ = -1; + int write_fd_ = -1; + std::thread reader_; + std::string output_; + std::string read_error_; + bool truncated_ = false; +}; + +using OperationExecutor = std::function; +using CompletionExecutor = std::function; + +bool ParseRequest(const unsigned char *data, size_t size, Request *request, + std::string *error); +bool EncodeRequest(const Request &request, std::vector *data, + std::string *error); +bool ParseResponse(const unsigned char *data, size_t size, Response *response, + std::string *error); +bool EncodeResponse(const Response &response, std::vector *data, + std::string *error); +bool ValidateProcessIdentity(const Request &request, + const std::string &proc_root, std::string *error); +ProcessExistenceState InspectProcessExistence(uint32_t pid, + const std::string &proc_root, + std::string *error); +ProcessIdentityState InspectProcessIdentity(const Request &request, + const std::string &proc_root, + std::string *error); +bool ExecuteValidated(const Request &request, const std::string &proc_root, + const OperationExecutor &executor, Response *response); +bool ResponseAllowsServerContinue(const Response &response); +int32_t FinishHandledOperation(bool post_handle_succeeded, + int32_t failure_status, + const CompletionExecutor &completion, + OperationState *state); +// Returns 1 for input, 0 for stop, -1 for timeout, and -2 for poll failure +// with errno preserved for the caller. +int PollForInputOrStop(int input_fd, int stop_fd, + const std::function &before_poll = {}, + int timeout_milliseconds = -1); +const char *ActionName(Action action); +Response HealthResponse(const OperationHealth &health); +Response HealthResponseAfterReap(const OperationHealth &health, + int32_t release_status, + const std::string &reap_warning); +bool OperationTimeoutMilliseconds(std::chrono::seconds duration, + unsigned int *timeout_ms, + std::string *error); + +} // namespace cuda_checkpoint_daemon diff --git a/agent/cmd/cuda-checkpoint-helper/daemon_protocol_test.cpp b/agent/cmd/cuda-checkpoint-helper/daemon_protocol_test.cpp new file mode 100644 index 00000000..bc1c3f14 --- /dev/null +++ b/agent/cmd/cuda-checkpoint-helper/daemon_protocol_test.cpp @@ -0,0 +1,561 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include "daemon_protocol.h" + +#include +#ifdef NDEBUG +#undef NDEBUG +#endif +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +using namespace cuda_checkpoint_daemon; + +Request TestRequest(Action action) { + return Request{ + .action = action, + .backend = + action == Action::kHealth ? Backend::kUnspecified : Backend::kPosix, + .pid = 123, + .transfer_buffer_count = + action == Action::kCheckpoint || action == Action::kRestore ? 2U : 0U, + .transfer_chunk_bytes = + action == Action::kCheckpoint || action == Action::kRestore ? 4096U + : 0U, + .expected_start_time_ticks = 987654, + .device_map = action == Action::kRestore ? "GPU-a=GPU-b" : "", + .storage_dir = action == Action::kCheckpoint || action == Action::kRestore + ? "/checkpoints/cuda" + : "", + .expected_cgroup = "0::/kubepods/test\n", + .job_file = + action == Action::kHealth ? "" : "/host/proc/123/root/tmp/cuda-job", + .selected_devices = + action == Action::kCheckpoint || action == Action::kRestore + ? "GPU-12345678-1234-1234-1234-123456789abc" + : "", + }; +} + +std::vector ReadGoldenRequest() { + std::ifstream fixture( + "cmd/cuda-checkpoint-helper/testdata/daemon_request_v6.hex"); + assert(fixture.good()); + std::string encoded; + fixture >> encoded; + assert(!encoded.empty() && encoded.size() % 2 == 0); + std::vector bytes; + bytes.reserve(encoded.size() / 2); + for (size_t index = 0; index < encoded.size(); index += 2) { + const std::string byte = encoded.substr(index, 2); + char *end = nullptr; + const unsigned long value = std::strtoul(byte.c_str(), &end, 16); + assert(end == byte.c_str() + 2 && value <= 0xffUL); + bytes.push_back(static_cast(value)); + } + return bytes; +} + +void TestGoldenRequestFixture() { + const std::vector encoded = ReadGoldenRequest(); + Request parsed; + std::string error; + assert(ParseRequest(encoded.data(), encoded.size(), &parsed, &error)); + assert(parsed.action == Action::kRestore); + assert(parsed.backend == Backend::kPosix); + assert(parsed.pid == 42U); + assert(parsed.transfer_buffer_count == 2U); + assert(parsed.transfer_chunk_bytes == 8U * 1024U * 1024U); + assert(parsed.expected_start_time_ticks == 12345U); + assert(parsed.device_map == + "GPU-aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee=" + "GPU-11111111-2222-3333-4444-555555555555"); + assert(parsed.storage_dir == "/checkpoints/process-nspid-42"); + assert(parsed.expected_cgroup == "0::/kubepods/test\n"); + assert(parsed.job_file == "/host/proc/42/root/tmp/cuda-job"); + assert(parsed.selected_devices == + "GPU-12345678-1234-1234-1234-123456789abc"); + std::vector reencoded; + assert(EncodeRequest(parsed, &reencoded, &error)); + assert(reencoded == encoded); +} + +std::string CreateProcRoot(const Request &request) { + char proc_template[] = "/tmp/cuda-daemon-proc-test-XXXXXX"; + const char *proc_root = mkdtemp(proc_template); + assert(proc_root != nullptr); + const std::filesystem::path process_dir = + std::filesystem::path(proc_root) / std::to_string(request.pid); + std::filesystem::create_directories(process_dir); + { + std::ofstream stat(process_dir / "stat"); + stat << "123 (worker with spaces) S 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 " + "17 18 987654 20\n"; + std::ofstream cgroup(process_dir / "cgroup"); + cgroup << request.expected_cgroup; + } + return proc_root; +} + +void TestProtocol() { + Request request = TestRequest(Action::kRestore); + std::vector encoded; + std::string error; + assert(EncodeRequest(request, &encoded, &error)); + Request parsed; + assert(ParseRequest(encoded.data(), encoded.size(), &parsed, &error)); + assert( + parsed.pid == request.pid && parsed.storage_dir == request.storage_dir && + parsed.backend == Backend::kPosix && + parsed.expected_start_time_ticks == request.expected_start_time_ticks && + parsed.expected_cgroup == request.expected_cgroup && + parsed.job_file == request.job_file && + parsed.selected_devices == request.selected_devices); + + encoded[4] = 99; + assert(!ParseRequest(encoded.data(), encoded.size(), &parsed, &error)); + encoded.resize(kMaxRequestSize + 1); + assert(!ParseRequest(encoded.data(), encoded.size(), &parsed, &error)); + + Response response{.cuda_status = 17, + .flags = kResponseFatal | kResponseLockNotAcquired, + .output = "stdout", + .error = "stderr"}; + assert(EncodeResponse(response, &encoded, &error)); + Response parsed_response; + assert( + ParseResponse(encoded.data(), encoded.size(), &parsed_response, &error)); + assert(parsed_response.cuda_status == 17 && + parsed_response.flags == + (kResponseFatal | kResponseLockNotAcquired) && + parsed_response.output == "stdout" && + parsed_response.error == "stderr"); + encoded[16] = 200; + assert( + !ParseResponse(encoded.data(), encoded.size(), &parsed_response, &error)); + + for (const Action action : {Action::kLock, Action::kCheckpoint, + Action::kRestore, Action::kUnlock}) { + request = TestRequest(action); + assert(EncodeRequest(request, &encoded, &error)); + assert(ParseRequest(encoded.data(), encoded.size(), &parsed, &error)); + assert(parsed.action == action); + } + + request = TestRequest(Action::kCheckpoint); + request.backend = Backend::kRegular; + request.transfer_buffer_count = 0; + request.transfer_chunk_bytes = 0; + request.storage_dir.clear(); + request.selected_devices.clear(); + assert(EncodeRequest(request, &encoded, &error)); + assert(ParseRequest(encoded.data(), encoded.size(), &parsed, &error)); + assert(parsed.backend == Backend::kRegular && parsed.storage_dir.empty()); + + request.job_file = "tmp/cuda-job"; + assert(EncodeRequest(request, &encoded, &error)); + assert(!ParseRequest(encoded.data(), encoded.size(), &parsed, &error)); + + request = TestRequest(Action::kCheckpoint); + request.backend = Backend::kRegular; + request.transfer_buffer_count = 0; + request.transfer_chunk_bytes = 0; + request.storage_dir.clear(); + assert(EncodeRequest(request, &encoded, &error)); + assert(!ParseRequest(encoded.data(), encoded.size(), &parsed, &error)); + + request = TestRequest(Action::kRestore); + request.selected_devices.clear(); + assert(EncodeRequest(request, &encoded, &error)); + assert(!ParseRequest(encoded.data(), encoded.size(), &parsed, &error)); + + for (const std::string &invalid_storage : { + std::string("/checkpoints"), std::string("/checkpoints-other/cuda"), + std::string("/checkpoints/../etc"), std::string("/tmp/cuda")}) { + request = TestRequest(Action::kCheckpoint); + request.storage_dir = invalid_storage; + assert(EncodeRequest(request, &encoded, &error)); + assert(!ParseRequest(encoded.data(), encoded.size(), &parsed, &error)); + } +} + +void TestExecutionIdentityAndFatalControlFlow() { + for (const Action action : {Action::kLock, Action::kCheckpoint, + Action::kRestore, Action::kUnlock}) { + Request request = TestRequest(action); + const std::string proc_root = CreateProcRoot(request); + int executions = 0; + Response response; + assert(ExecuteValidated( + request, proc_root, + [&executions](const Request &) { + ++executions; + return Response{}; + }, + &response)); + assert(executions == 1); + request.expected_start_time_ticks++; + assert(ExecuteValidated( + request, proc_root, + [&executions](const Request &) { + ++executions; + return Response{}; + }, + &response)); + assert(executions == 1); + if (action == Action::kLock) { + assert((response.flags & kResponseLockNotAcquired) != 0); + } else { + assert((response.flags & kResponseLockNotAcquired) == 0); + } + std::filesystem::remove_all(proc_root); + } + + OperationState operation; + int completions = 0; + assert(FinishHandledOperation( + false, 17, + [&completions] { + ++completions; + return 0; + }, + &operation) == 17); + assert(completions == 0 && operation.fatal()); + operation = {}; + assert(FinishHandledOperation( + true, 17, + [&completions] { + ++completions; + return 0; + }, + &operation) == 0); + assert(completions == 1 && !operation.fatal()); + Response fatal{ + .cuda_status = 17, + .flags = kResponseFatal, + .output = "", + .error = "", + }; + assert(!ResponseAllowsServerContinue(fatal)); +} + +void TestProcessIdentityStates() { + Request request = TestRequest(Action::kRestore); + const std::string proc_root = CreateProcRoot(request); + std::string error; + assert(InspectProcessExistence(request.pid, proc_root, &error) == + ProcessExistenceState::kExists); + assert(InspectProcessIdentity(request, proc_root, &error) == + ProcessIdentityState::kMatches); + + request.expected_start_time_ticks++; + assert(InspectProcessIdentity(request, proc_root, &error) == + ProcessIdentityState::kExitedOrReused); + request.expected_start_time_ticks--; + + request.expected_cgroup = "0::/kubepods/moved\n"; + assert(InspectProcessIdentity(request, proc_root, &error) == + ProcessIdentityState::kIndeterminate); + request.expected_cgroup = "0::/kubepods/test\n"; + + const std::filesystem::path process_dir = + std::filesystem::path(proc_root) / std::to_string(request.pid); + assert(std::filesystem::remove(process_dir / "stat")); + assert(InspectProcessIdentity(request, proc_root, &error) == + ProcessIdentityState::kIndeterminate); + + std::filesystem::remove_all(process_dir); + assert(InspectProcessExistence(request.pid, proc_root, &error) == + ProcessExistenceState::kMissing); + assert(InspectProcessIdentity(request, proc_root, &error) == + ProcessIdentityState::kExitedOrReused); + + const std::filesystem::path invalid_proc_root = + std::filesystem::path(proc_root) / "symlink-loop"; + std::filesystem::create_symlink("symlink-loop", invalid_proc_root); + assert(InspectProcessExistence(request.pid, invalid_proc_root, &error) == + ProcessExistenceState::kIndeterminate); + std::filesystem::remove_all(proc_root); +} + +void TestHealthStates() { + OperationHealth regular_operation(std::chrono::seconds(1)); + regular_operation.MarkReady(false); + HealthSnapshot regular_snapshot = regular_operation.Snapshot(); + assert(regular_snapshot.ready && regular_snapshot.healthy && + !regular_snapshot.custom_storage_available); + Response regular_response = HealthResponse(regular_operation); + assert(regular_response.cuda_status == 0 && + (regular_response.flags & kResponseCapabilityDeferredCUDA) != 0 && + (regular_response.flags & kResponseCapabilityCustomStorage) == 0); + + OperationHealth no_operation(std::chrono::seconds(1)); + HealthSnapshot snapshot = no_operation.Snapshot(); + assert(!snapshot.ready && !snapshot.busy && !snapshot.healthy); + + no_operation.MarkReady(true); + snapshot = no_operation.Snapshot(); + assert(snapshot.ready && !snapshot.busy && snapshot.healthy && + snapshot.custom_storage_available); + Response response = HealthResponse(no_operation); + assert((response.flags & kResponseCapabilityDeferredCUDA) != 0 && + (response.flags & kResponseCapabilityCustomStorage) != 0); + + Response deferred_reap = + HealthResponseAfterReap(no_operation, 0, "temporary proc read failure"); + assert(deferred_reap.cuda_status == 0 && + (deferred_reap.flags & kResponseCapabilityDeferredCUDA) != 0 && + deferred_reap.error.find("retained contexts; reaping deferred") != + std::string::npos); + Response failed_reap = HealthResponseAfterReap(no_operation, 17, ""); + assert(failed_reap.cuda_status == 17 && + (failed_reap.flags & kResponseFatal) != 0); + + no_operation.Begin(Action::kCheckpoint, 123); + snapshot = no_operation.Snapshot(); + assert(snapshot.ready && snapshot.busy && snapshot.healthy); + assert(snapshot.action == Action::kCheckpoint && snapshot.pid == 123); + std::this_thread::sleep_for(std::chrono::milliseconds(1100)); + snapshot = no_operation.Snapshot(); + assert(snapshot.busy && !snapshot.healthy); + no_operation.End(); + snapshot = no_operation.Snapshot(); + assert(!snapshot.busy && snapshot.healthy); +} + +void TestSocketLifecycle() { + char socket_template[] = "/tmp/cuda-daemon-socket-test-XXXXXX"; + const char *directory = mkdtemp(socket_template); + assert(directory != nullptr); + const std::string path = std::string(directory) + "/helper.sock"; + std::string error; + + OwnedUnixSocket owner; + assert(owner.Bind(path, 1, &error)); + struct stat owned{}; + assert(lstat(path.c_str(), &owned) == 0); + OwnedUnixSocket contender; + assert(!contender.Bind(path, 1, &error)); + struct stat after_contender{}; + assert(lstat(path.c_str(), &after_contender) == 0); + assert(after_contender.st_ino == owned.st_ino && + after_contender.st_dev == owned.st_dev); + + const std::string replacement = std::string(directory) + "/replacement"; + assert(unlink(path.c_str()) == 0); + { + std::ofstream file(replacement); + file << "replacement"; + } + assert(rename(replacement.c_str(), path.c_str()) == 0); + owner.Close(); + struct stat replacement_stat{}; + assert(lstat(path.c_str(), &replacement_stat) == 0 && + S_ISREG(replacement_stat.st_mode)); + assert(unlink(path.c_str()) == 0); + + const int stale_fd = socket(AF_UNIX, SOCK_SEQPACKET | SOCK_CLOEXEC, 0); + assert(stale_fd >= 0); + sockaddr_un address{}; + address.sun_family = AF_UNIX; + std::memcpy(address.sun_path, path.c_str(), path.size() + 1); + assert(bind(stale_fd, reinterpret_cast(&address), + sizeof(address)) == 0); + close(stale_fd); + OwnedUnixSocket stale_replacer; + assert(stale_replacer.Bind(path, 1, &error)); + stale_replacer.Close(); + assert(lstat(path.c_str(), &replacement_stat) != 0 && errno == ENOENT); + + OwnedUnixSocket setup_failure; + assert(!setup_failure.Bind(path, -1, &error)); + assert(lstat(path.c_str(), &replacement_stat) != 0 && errno == ENOENT); + std::filesystem::remove_all(directory); +} + +void RunBounded(const std::function &test) { + const pid_t child = fork(); + assert(child >= 0); + if (child == 0) { + test(); + _exit(0); + } + + int status = 0; + pid_t wait_result = 0; + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(5); + while ((wait_result = waitpid(child, &status, WNOHANG)) == 0 && + std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + if (wait_result == 0) { + assert(kill(child, SIGKILL) == 0); + assert(waitpid(child, &status, 0) == child); + assert(false && "bounded shutdown test timed out"); + } + assert(wait_result == child); + assert(WIFEXITED(status) && WEXITSTATUS(status) == 0); +} + +void RunSignalOwnerScenario(bool external_signal) { + ShutdownSignalOwner signal_owner; + std::string error; + assert(signal_owner.Start(&error)); + + int silent_operation_client[2]; + int silent_health_client[2]; + assert(socketpair(AF_UNIX, SOCK_SEQPACKET | SOCK_CLOEXEC, 0, + silent_operation_client) == 0); + assert(socketpair(AF_UNIX, SOCK_SEQPACKET | SOCK_CLOEXEC, 0, + silent_health_client) == 0); + std::barrier waiters_started(3); + const auto rendezvous = [&] { waiters_started.arrive_and_wait(); }; + const auto wait_for_stop = [&](int client_fd, int stop_fd) { + assert(PollForInputOrStop(client_fd, stop_fd, rendezvous) == 0); + unsigned char wake = 0; + assert(read(stop_fd, &wake, sizeof(wake)) == sizeof(wake)); + }; + std::thread operation_waiter([&] { + wait_for_stop(silent_operation_client[0], signal_owner.operation_stop_fd()); + }); + std::thread health_waiter([&] { + wait_for_stop(silent_health_client[0], signal_owner.health_stop_fd()); + }); + waiters_started.arrive_and_wait(); + + if (external_signal) { + assert(kill(getpid(), SIGTERM) == 0); + operation_waiter.join(); + health_waiter.join(); + assert(signal_owner.ShutdownRequested()); + assert(signal_owner.StopAndJoin(&error)); + } else { + assert(signal_owner.StopAndJoin(&error)); + operation_waiter.join(); + health_waiter.join(); + assert(signal_owner.ShutdownRequested()); + } + signal_owner.Close(); + close(silent_operation_client[0]); + close(silent_operation_client[1]); + close(silent_health_client[0]); + close(silent_health_client[1]); +} + +void TestExternalSignalWakesIndependentWaiters() { + RunBounded([] { RunSignalOwnerScenario(true); }); +} + +void TestInternalStopJoinsSignalOwnerAndWakesIndependentWaiters() { + RunBounded([] { RunSignalOwnerScenario(false); }); +} + +void TestPollDoesNotConsumeStopWake() { + int stop_pipe[2]; + int silent_client[2]; + assert(pipe2(stop_pipe, O_CLOEXEC | O_NONBLOCK) == 0); + assert(socketpair(AF_UNIX, SOCK_SEQPACKET | SOCK_CLOEXEC, 0, silent_client) == + 0); + const unsigned char wake = 1; + assert(write(stop_pipe[1], &wake, sizeof(wake)) == sizeof(wake)); + assert(PollForInputOrStop(silent_client[0], stop_pipe[0]) == 0); + assert(PollForInputOrStop(silent_client[0], stop_pipe[0]) == 0); + unsigned char remaining = 0; + assert(read(stop_pipe[0], &remaining, sizeof(remaining)) == + sizeof(remaining)); + close(silent_client[0]); + close(silent_client[1]); + close(stop_pipe[0]); + close(stop_pipe[1]); +} + +void TestPollTimesOutWithoutInput() { + int stop_pipe[2]; + int silent_client[2]; + assert(pipe2(stop_pipe, O_CLOEXEC | O_NONBLOCK) == 0); + assert(socketpair(AF_UNIX, SOCK_SEQPACKET | SOCK_CLOEXEC, 0, silent_client) == + 0); + assert(PollForInputOrStop(silent_client[0], stop_pipe[0], {}, 1) == -1); + close(silent_client[0]); + close(silent_client[1]); + close(stop_pipe[0]); + close(stop_pipe[1]); +} + +void TestOperationTimeoutMilliseconds() { + unsigned int timeout_ms = 0; + std::string error; + assert(OperationTimeoutMilliseconds(std::chrono::seconds(3600), &timeout_ms, + &error)); + assert(timeout_ms == 3600000U); + assert(!OperationTimeoutMilliseconds(std::chrono::seconds(0), &timeout_ms, + &error)); + assert(!OperationTimeoutMilliseconds( + std::chrono::seconds( + static_cast(std::numeric_limits::max()) / + 1000ULL + + 1ULL), + &timeout_ms, &error)); +} + +void BoundedOutputCaptureScenario() { + BoundedOutputCapture capture(8); + std::string error; + assert(capture.Start(&error)); + const std::string payload(1024ULL * 1024ULL, 'x'); + size_t offset = 0; + while (offset < payload.size()) { + const ssize_t written = write(capture.write_fd(), payload.data() + offset, + payload.size() - offset); + assert(written > 0); + offset += static_cast(written); + } + std::string output; + bool truncated = false; + assert(capture.Finish(&output, &truncated, &error)); + assert(output == "xxxxxxxx"); + assert(truncated); +} + +void TestBoundedOutputCaptureDrainsAndTruncates() { + RunBounded([] { BoundedOutputCaptureScenario(); }); +} + +} // namespace + +int main() { + TestGoldenRequestFixture(); + TestProtocol(); + TestExecutionIdentityAndFatalControlFlow(); + TestProcessIdentityStates(); + TestHealthStates(); + TestSocketLifecycle(); + TestExternalSignalWakesIndependentWaiters(); + TestInternalStopJoinsSignalOwnerAndWakesIndependentWaiters(); + TestPollDoesNotConsumeStopWake(); + TestPollTimesOutWithoutInput(); + TestOperationTimeoutMilliseconds(); + TestBoundedOutputCaptureDrainsAndTruncates(); + return 0; +} diff --git a/agent/cmd/cuda-checkpoint-helper/main.cpp b/agent/cmd/cuda-checkpoint-helper/main.cpp new file mode 100644 index 00000000..2357e0b9 --- /dev/null +++ b/agent/cmd/cuda-checkpoint-helper/main.cpp @@ -0,0 +1,1641 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "cuda_checkpoint_compat.h" +#include "daemon_protocol.h" +#include "storage_manifest.h" +#include "transfer_config.h" +#include "transfer_engine.h" + +namespace { + +namespace storage = cuda_checkpoint_storage; +namespace transfer = cuda_checkpoint_transfer; +using Clock = std::chrono::steady_clock; +namespace daemon_protocol = cuda_checkpoint_daemon; +constexpr int kClientReceiveTimeoutMilliseconds = 5000; + +class ScopedFd { +public: + explicit ScopedFd(int fd) : fd_(fd) {} + ScopedFd(const ScopedFd &) = delete; + ScopedFd &operator=(const ScopedFd &) = delete; + ~ScopedFd() noexcept { + if (fd_ >= 0) { + close(fd_); + } + } + + int get() const { return fd_; } + +private: + int fd_; +}; + +class DaemonThreadShutdown { +public: + DaemonThreadShutdown( + daemon_protocol::ShutdownSignalOwner *signal_owner, + daemon_protocol::ShutdownSignalOwner::ShutdownResult *result) + : signal_owner_(signal_owner), result_(result) {} + DaemonThreadShutdown(const DaemonThreadShutdown &) = delete; + DaemonThreadShutdown &operator=(const DaemonThreadShutdown &) = delete; + + ~DaemonThreadShutdown() noexcept { + *result_ = signal_owner_->StopAndJoinNoThrow(); + } + +private: + daemon_protocol::ShutdownSignalOwner *signal_owner_; + daemon_protocol::ShutdownSignalOwner::ShutdownResult *result_; +}; + +constexpr uint64_t kMaxOperationSeconds = 60 * 60; + +bool ParsePositiveSeconds(const char *value, uint64_t *seconds_out) { + char *end = nullptr; + errno = 0; + const unsigned long long seconds = std::strtoull(value, &end, 10); + if (value[0] == '\0' || end == nullptr || *end != '\0' || errno != 0 || + seconds == 0 || + seconds > kMaxOperationSeconds) { + return false; + } + *seconds_out = seconds; + return true; +} + +double SecondsSince(Clock::time_point start) { + return std::chrono::duration(Clock::now() - start).count(); +} + +double SecondsBetween(Clock::time_point start, Clock::time_point end) { + return std::chrono::duration(end - start).count(); +} + +CUresult DeviceUUID(CUdevice device, std::string *uuid_out); + +class OperationContexts { +public: + OperationContexts() = default; + OperationContexts(const OperationContexts &) = delete; + OperationContexts &operator=(const OperationContexts &) = delete; + + CUresult RetainAll(int *device_count, double *enumeration_seconds, + double *retain_seconds) { + const auto enumeration_start = Clock::now(); + int count = 0; + CUresult status = cuDeviceGetCount(&count); + *enumeration_seconds = SecondsSince(enumeration_start); + *device_count = count; + if (status != CUDA_SUCCESS) { + return status; + } + + contexts_.reserve(count); + for (int ordinal = 0; ordinal < count; ++ordinal) { + const auto retain_start = Clock::now(); + CUdevice device = 0; + status = cuDeviceGet(&device, ordinal); + if (status != CUDA_SUCCESS) { + *retain_seconds += SecondsSince(retain_start); + return status; + } + CUcontext context = nullptr; + status = cuDevicePrimaryCtxRetain(&context, device); + *retain_seconds += SecondsSince(retain_start); + if (status != CUDA_SUCCESS) { + return status; + } + contexts_.push_back({device, context}); + } + return CUDA_SUCCESS; + } + + CUresult RetainSelected(const std::vector &selected_devices, + int *device_count, double *enumeration_seconds, + double *retain_seconds) { + const auto enumeration_start = Clock::now(); + int count = 0; + CUresult status = cuDeviceGetCount(&count); + *enumeration_seconds = SecondsSince(enumeration_start); + *device_count = count; + if (status != CUDA_SUCCESS) { + return status; + } + + const std::unordered_set selected(selected_devices.begin(), + selected_devices.end()); + contexts_.reserve(selected.size()); + for (int ordinal = 0; ordinal < count; ++ordinal) { + CUdevice device = 0; + status = cuDeviceGet(&device, ordinal); + if (status != CUDA_SUCCESS) { + return status; + } + std::string uuid; + status = DeviceUUID(device, &uuid); + if (status != CUDA_SUCCESS) { + return status; + } + if (!selected.contains(uuid)) { + continue; + } + const auto retain_start = Clock::now(); + CUcontext context = nullptr; + status = cuDevicePrimaryCtxRetain(&context, device); + *retain_seconds += SecondsSince(retain_start); + if (status != CUDA_SUCCESS) { + return status; + } + contexts_.push_back({device, context}); + } + if (contexts_.size() != selected.size()) { + return CUDA_ERROR_INVALID_DEVICE; + } + return CUDA_SUCCESS; + } + + std::vector DetachDevices() { + std::vector devices; + devices.reserve(contexts_.size()); + for (const auto &entry : contexts_) { + devices.push_back(entry.device); + } + contexts_.clear(); + return devices; + } + + CUresult ReleaseAll() { + CUresult first_error = CUDA_SUCCESS; + while (!contexts_.empty()) { + const CUresult status = + cuDevicePrimaryCtxRelease(contexts_.back().device); + if (first_error == CUDA_SUCCESS && status != CUDA_SUCCESS) { + first_error = status; + } + contexts_.pop_back(); + } + return first_error; + } + + CUresult ContextAndDeviceForStream(CUstream stream, CUcontext *context_out, + CUdevice *device_out) const { + CUcontext stream_context = nullptr; + CUresult status = cuStreamGetCtx(stream, &stream_context); + if (status != CUDA_SUCCESS) { + return status; + } + for (const auto &retained : contexts_) { + if (retained.context == stream_context) { + *context_out = retained.context; + *device_out = retained.device; + return CUDA_SUCCESS; + } + } + return CUDA_ERROR_INVALID_CONTEXT; + } + + ~OperationContexts() { (void)ReleaseAll(); } + + size_t size() const { return contexts_.size(); } + +private: + struct Entry { + CUdevice device; + CUcontext context; + }; + std::vector contexts_; +}; + +// CUDA 13.4 CustomStorage restore qualification found that releasing the +// helper's retained primary-context reference while the target remained alive +// could later fault that target. Keep one operation reference with the exact +// PID/start-time/cgroup identity and release it only after confirmed exit or +// PID reuse. An inconclusive /proc read must retain the reference and block new +// work rather than guessing that the target exited. +class PersistentTargetContexts { +public: + PersistentTargetContexts() = default; + PersistentTargetContexts(const PersistentTargetContexts &) = delete; + PersistentTargetContexts & + operator=(const PersistentTargetContexts &) = delete; + + CUresult Adopt(std::vector devices, + const daemon_protocol::Request &request) { + std::lock_guard lock(mutex_); + for (const auto &target : targets_) { + if (SameIdentity(target.request, request)) { + return ReleaseDevices(devices); + } + } + targets_.push_back({request, std::move(devices)}); + return CUDA_SUCCESS; + } + + CUresult ReapExited(const std::string &proc_root, + std::string *identity_error) { + std::lock_guard lock(mutex_); + CUresult first_error = CUDA_SUCCESS; + auto target = targets_.begin(); + while (target != targets_.end()) { + std::string target_error; + const auto identity_state = daemon_protocol::InspectProcessIdentity( + target->request, proc_root, &target_error); + if (identity_state == daemon_protocol::ProcessIdentityState::kMatches) { + ++target; + continue; + } + if (identity_state == + daemon_protocol::ProcessIdentityState::kIndeterminate) { + if (identity_error != nullptr && identity_error->empty()) { + *identity_error = "cannot safely determine target " + + std::to_string(target->request.pid) + + " identity: " + target_error; + } + ++target; + continue; + } + const CUresult status = ReleaseDevices(target->devices); + if (first_error == CUDA_SUCCESS && status != CUDA_SUCCESS) { + first_error = status; + } + target = targets_.erase(target); + } + return first_error; + } + + CUresult ReleaseAll() { + std::lock_guard lock(mutex_); + CUresult first_error = CUDA_SUCCESS; + for (const auto &target : targets_) { + const CUresult status = ReleaseDevices(target.devices); + if (first_error == CUDA_SUCCESS && status != CUDA_SUCCESS) { + first_error = status; + } + } + targets_.clear(); + return first_error; + } + + ~PersistentTargetContexts() { (void)ReleaseAll(); } + +private: + struct TargetContexts { + daemon_protocol::Request request; + std::vector devices; + }; + + static bool SameIdentity(const daemon_protocol::Request &left, + const daemon_protocol::Request &right) { + return left.pid == right.pid && + left.expected_start_time_ticks == right.expected_start_time_ticks && + left.expected_cgroup == right.expected_cgroup; + } + + static CUresult ReleaseDevices(const std::vector &devices) { + CUresult first_error = CUDA_SUCCESS; + for (const CUdevice device : devices) { + const CUresult status = cuDevicePrimaryCtxRelease(device); + if (first_error == CUDA_SUCCESS && status != CUDA_SUCCESS) { + first_error = status; + } + } + return first_error; + } + + std::mutex mutex_; + std::vector targets_; +}; + +int PrintUsage(FILE *stream) { + return std::fprintf(stream, + "Usage:\n" + " cuda-checkpoint-helper --get-restore-tid --pid \n" + " cuda-checkpoint-helper --daemon --socket " + " " + "[--max-operation-seconds ]\n" + " cuda-checkpoint-helper --health --socket " + "\n") < 0 + ? 1 + : 0; +} + +int PrintUsageError() { + (void)PrintUsage(stderr); + return 1; +} + +void PrintCudaError(CUresult status) { + const char *name = nullptr; + const char *message = nullptr; + (void)cuGetErrorName(status, &name); + (void)cuGetErrorString(status, &message); + std::fprintf(stderr, "%s: %s\n", + name == nullptr ? "CUDA_ERROR_UNKNOWN" : name, + message == nullptr ? "unknown CUDA error" : message); +} + +bool ParsePID(const char *value, int *pid_out) { + char *end = nullptr; + long pid = std::strtol(value, &end, 10); + if (value[0] == '\0' || end == nullptr || *end != '\0' || pid <= 0 || + pid > INT_MAX) { + return false; + } + *pid_out = static_cast(pid); + return true; +} + +bool ParseUUID(const char *value, CUuuid *uuid_out) { + if (value == nullptr || uuid_out == nullptr) { + return false; + } + std::array bytes{}; + if (!storage::ParseGPUUUID(value, &bytes)) { + return false; + } + static_assert(sizeof(uuid_out->bytes) == bytes.size()); + std::memcpy(uuid_out->bytes, bytes.data(), bytes.size()); + return true; +} + +bool ParseDeviceMap(const std::string &device_map, + std::vector *pairs, + std::vector *storage_pairs = nullptr) { + if (device_map.empty()) { + return true; + } + std::unordered_set source_uuids; + std::unordered_set destination_uuids; + std::istringstream input(device_map); + std::string pair; + while (std::getline(input, pair, ',')) { + size_t separator = pair.find('='); + if (separator == std::string::npos || + pair.find('=', separator + 1) != std::string::npos) { + return false; + } + CUcheckpointGpuPair parsed{}; + const std::string source_input = pair.substr(0, separator); + const std::string destination_input = pair.substr(separator + 1); + std::string source; + std::string destination; + if (!ParseUUID(source_input.c_str(), &parsed.oldUuid) || + !ParseUUID(destination_input.c_str(), &parsed.newUuid) || + !storage::CanonicalizeGPUUUID(source_input, &source) || + !storage::CanonicalizeGPUUUID(destination_input, &destination) || + !source_uuids.insert(source).second || + !destination_uuids.insert(destination).second) { + return false; + } + pairs->push_back(parsed); + if (storage_pairs != nullptr) { + storage_pairs->push_back({std::move(source), std::move(destination)}); + } + } + return !pairs->empty(); +} + +bool ParseDeviceSelection(const std::string &selected_devices, + std::vector *devices) { + if (selected_devices.empty()) { + return false; + } + std::unordered_set seen; + std::istringstream input(selected_devices); + std::string value; + while (std::getline(input, value, ',')) { + std::string canonical; + if (!storage::CanonicalizeGPUUUID(value, &canonical) || + !seen.insert(canonical).second) { + return false; + } + devices->push_back(std::move(canonical)); + } + return !devices->empty(); +} + +CUresult DeviceUUID(CUdevice device, std::string *uuid_out) { + CUuuid uuid{}; + CUresult status = cuDeviceGetUuid(&uuid, device); + if (status != CUDA_SUCCESS) { + return status; + } + std::array bytes{}; + static_assert(sizeof(uuid.bytes) == bytes.size()); + std::memcpy(bytes.data(), uuid.bytes, bytes.size()); + *uuid_out = storage::FormatGPUUUID(bytes); + return CUDA_SUCCESS; +} + +struct CustomStorageResult { + CUresult status = CUDA_SUCCESS; + daemon_protocol::OperationState operation; + bool fatal = false; +}; + +CustomStorageResult +DoCustomStorage(int pid, bool checkpoint, const std::string &device_map, + const std::filesystem::path &storage_dir, + const transfer::TransferOptions &transfer_options, + Clock::time_point operation_deadline, + Clock::time_point helper_main_start, + cuda_checkpoint_compat::OperationCompleteFn operation_complete, + const daemon_protocol::Request *daemon_request, + PersistentTargetContexts *persistent_contexts) { + const auto custom_storage_start = Clock::now(); + if (operation_complete == nullptr) { + std::fprintf(stderr, "CUDA custom storage unavailable\n"); + return {CUDA_ERROR_NOT_SUPPORTED, {}}; + } + const auto storage_directory_start = Clock::now(); + if (!storage_dir.is_absolute()) { + std::fprintf(stderr, "custom storage directory must be absolute\n"); + return {CUDA_ERROR_INVALID_VALUE, {}}; + } + if (checkpoint) { + std::error_code filesystem_error; + std::filesystem::create_directories(storage_dir, filesystem_error); + struct stat directory_stat{}; + if (filesystem_error || lstat(storage_dir.c_str(), &directory_stat) != 0 || + !S_ISDIR(directory_stat.st_mode) || + chmod(storage_dir.c_str(), 0700) != 0) { + std::fprintf(stderr, "failed to create custom storage directory\n"); + return {CUDA_ERROR_OPERATING_SYSTEM, {}}; + } + std::string remove_error; + if (!storage::RemoveManifest(storage_dir, &remove_error)) { + std::fprintf(stderr, + "failed to clear stale custom storage manifest: %s\n", + remove_error.c_str()); + return {CUDA_ERROR_OPERATING_SYSTEM, {}}; + } + } else { + struct stat directory_stat{}; + if (lstat(storage_dir.c_str(), &directory_stat) != 0 || + !S_ISDIR(directory_stat.st_mode) || + (directory_stat.st_mode & 0022) != 0) { + std::fprintf(stderr, "custom storage directory is missing or invalid\n"); + return {CUDA_ERROR_INVALID_VALUE, {}}; + } + } + const double storage_directory_validation_seconds = + SecondsSince(storage_directory_start); + + int visible_cuda_device_count = 0; + double device_enumeration_seconds = 0.0; + double target_context_discovery_seconds = 0.0; + double primary_context_retain_seconds = 0.0; + double primary_context_release_seconds = 0.0; + OperationContexts operation_contexts; + std::vector selected_devices; + if (daemon_request != nullptr && + !ParseDeviceSelection(daemon_request->selected_devices, + &selected_devices)) { + std::fprintf(stderr, "invalid selected CUDA devices\n"); + return {CUDA_ERROR_INVALID_VALUE, {}}; + } + // cuCheckpointProcessCheckpoint/Restore requires the helper to retain the + // primary contexts before it can return the CustomStorage streams. Retain a + // fresh operation reference for only the target's selected devices. A + // successful daemon operation transfers that reference to the target- + // identity cache; direct CLI use releases it at invocation end. + CUresult status = CUDA_SUCCESS; + if (selected_devices.empty()) { + status = operation_contexts.RetainAll( + &visible_cuda_device_count, &device_enumeration_seconds, + &primary_context_retain_seconds); + } else { + status = operation_contexts.RetainSelected( + selected_devices, &visible_cuda_device_count, + &device_enumeration_seconds, &primary_context_retain_seconds); + } + if (status != CUDA_SUCCESS) { + return {status, {}}; + } + + const auto manifest_validation_start = Clock::now(); + std::vector manifest; + std::string manifest_error; + if (!checkpoint && + (!storage::ReadManifest(storage_dir, &manifest, &manifest_error) || + !storage::ValidateExtentFiles(storage_dir, manifest, &manifest_error))) { + std::fprintf(stderr, "custom storage manifest validation failed: %s\n", + manifest_error.c_str()); + return {CUDA_ERROR_INVALID_VALUE, {}}; + } + const double manifest_validation_seconds = + SecondsSince(manifest_validation_start); + + cuda_checkpoint_compat::StorageInfo *info = nullptr; + std::vector gpu_pairs; + std::vector storage_pairs; + const auto device_map_preparation_start = Clock::now(); + if (!checkpoint && !ParseDeviceMap(device_map, &gpu_pairs, &storage_pairs)) { + return {CUDA_ERROR_INVALID_VALUE, {}}; + } + const double device_map_preparation_seconds = + SecondsSince(device_map_preparation_start); + if (daemon_request != nullptr) { + std::string identity_error; + if (!daemon_protocol::ValidateProcessIdentity(*daemon_request, "/host/proc", + &identity_error)) { + std::fprintf(stderr, + "process identity changed before CUDA operation: %s\n", + identity_error.c_str()); + return {CUDA_ERROR_INVALID_VALUE, {}}; + } + } + const auto cuda_process_api_start = Clock::now(); + if (checkpoint) { + cuda_checkpoint_compat::CheckpointArgs args{}; + args.customStorageInfo_out = &info; + status = cuCheckpointProcessCheckpoint( + pid, cuda_checkpoint_compat::NativeArgs(&args)); + } else { + cuda_checkpoint_compat::RestoreArgs args{}; + args.gpuPairs = gpu_pairs.empty() ? nullptr : gpu_pairs.data(); + args.gpuPairsCount = gpu_pairs.size(); + args.customStorageInfo_out = &info; + status = cuCheckpointProcessRestore( + pid, cuda_checkpoint_compat::NativeArgs(&args)); + } + const double cuda_process_api_seconds = SecondsSince(cuda_process_api_start); + if (status != CUDA_SUCCESS) { + return {status, {}}; + } + daemon_protocol::OperationState operation{.handle_returned = true}; + const auto post_handle_failure = [&operation, &operation_contexts]( + CUresult failure) { + const CUresult status = + static_cast(daemon_protocol::FinishHandledOperation( + false, failure, + [] { return static_cast(CUDA_SUCCESS); }, &operation)); + const CUresult release_status = operation_contexts.ReleaseAll(); + if (release_status != CUDA_SUCCESS) { + std::fprintf( + stderr, + "failed to release operation CUDA contexts with status %d while " + "handling status %d\n", + static_cast(release_status), static_cast(status)); + } + return CustomStorageResult{.status = status, + .operation = operation, + .fatal = operation.fatal() || + release_status != CUDA_SUCCESS}; + }; + const auto metadata_job_construction_start = Clock::now(); + if (info == nullptr || info->handle == nullptr || + info->deviceCount > + static_cast(visible_cuda_device_count) || + (info->deviceCount > 0 && info->perDeviceData == nullptr)) { + std::fprintf(stderr, "CUDA returned invalid custom storage information\n"); + return post_handle_failure(CUDA_ERROR_INVALID_VALUE); + } + + size_t pinned_bytes = 0; + std::string transfer_config_error; + if (!transfer::CalculatePinnedBytes(info->deviceCount, transfer_options, + &pinned_bytes, &transfer_config_error)) { + std::fprintf(stderr, "custom storage transfer configuration invalid: %s\n", + transfer_config_error.c_str()); + return post_handle_failure(CUDA_ERROR_INVALID_VALUE); + } + + std::vector contexts(info->deviceCount); + std::vector devices(info->deviceCount); + std::vector device_extents; + device_extents.reserve(info->deviceCount); + for (unsigned int index = 0; index < info->deviceCount; ++index) { + const auto target_context_discovery_start = Clock::now(); + status = operation_contexts.ContextAndDeviceForStream( + info->perDeviceData[index].stream, &contexts[index], &devices[index]); + target_context_discovery_seconds += + SecondsSince(target_context_discovery_start); + if (status != CUDA_SUCCESS) { + return post_handle_failure(status); + } + std::string uuid; + status = DeviceUUID(devices[index], &uuid); + if (status != CUDA_SUCCESS) { + return post_handle_failure(status); + } + device_extents.push_back( + {std::move(uuid), info->perDeviceData[index].size}); + } + + if (checkpoint && !storage::BuildCheckpointManifest(device_extents, &manifest, + &manifest_error)) { + std::fprintf(stderr, "invalid checkpoint custom storage mapping: %s\n", + manifest_error.c_str()); + return post_handle_failure(CUDA_ERROR_INVALID_VALUE); + } + std::vector transfer_jobs; + if (!storage::BuildTransferJobs( + manifest, device_extents, + checkpoint ? std::vector{} : storage_pairs, + &transfer_jobs, &manifest_error)) { + std::fprintf(stderr, "invalid restore custom storage mapping: %s\n", + manifest_error.c_str()); + return post_handle_failure(CUDA_ERROR_INVALID_VALUE); + } + + size_t total_bytes = 0; + for (const auto &extent : manifest) { + if (extent.size > std::numeric_limits::max() - total_bytes) { + std::fprintf(stderr, "custom storage byte count overflow\n"); + return post_handle_failure(CUDA_ERROR_INVALID_VALUE); + } + total_bytes += extent.size; + } + const double metadata_job_construction_seconds = + SecondsSince(metadata_job_construction_start); + + const auto start = Clock::now(); + const auto worker_orchestration_start = Clock::now(); + std::vector workers; + std::vector worker_success(transfer_jobs.size(), 0); + std::vector worker_errors(transfer_jobs.size()); + std::vector worker_metrics(transfer_jobs.size()); + transfer::TransferCancellation cancellation(operation_deadline); + std::string worker_start_error; + try { + workers.reserve(transfer_jobs.size()); + for (size_t job_index = 0; job_index < transfer_jobs.size(); ++job_index) { + workers.emplace_back([&, job_index] { + const auto &job = transfer_jobs[job_index]; + try { + const auto &device_data = info->perDeviceData[job.device_index]; + const transfer::StorageLayout layout{ + {{storage_dir / manifest[job.extent_index].filename, + device_data.size}}, + {{0, device_data.size, 0, 0}}, + }; + const bool transferred = transfer::TransferExtent( + device_data.devPtr, device_data.size, device_data.stream, + contexts[job.device_index], layout, + checkpoint ? transfer::TransferOperation::kCheckpoint + : transfer::TransferOperation::kRestore, + transfer_options, &cancellation, &worker_metrics[job_index], + &worker_errors[job_index]); + worker_success[job_index] = transferred; + if (!transferred) { + cancellation.Cancel(); + } + } catch (const std::exception &exception) { + cancellation.Cancel(); + worker_errors[job_index] = exception.what(); + } catch (...) { + cancellation.Cancel(); + worker_errors[job_index] = "unknown worker exception"; + } + }); + } + } catch (const std::exception &exception) { + cancellation.Cancel(); + worker_start_error = exception.what(); + } catch (...) { + cancellation.Cancel(); + worker_start_error = "unknown thread creation exception"; + } + for (auto &worker : workers) { + worker.join(); + } + const double worker_orchestration_seconds = + SecondsSince(worker_orchestration_start); + if (!worker_start_error.empty()) { + std::fprintf(stderr, "failed to start custom storage worker: %s\n", + worker_start_error.c_str()); + return post_handle_failure(CUDA_ERROR_OPERATING_SYSTEM); + } + for (size_t job_index = 0; job_index < transfer_jobs.size(); ++job_index) { + if (!worker_success[job_index]) { + std::fprintf(stderr, + "custom storage transfer failed for device index %zu: %s\n", + transfer_jobs[job_index].device_index, + worker_errors[job_index].c_str()); + return post_handle_failure(CUDA_ERROR_OPERATING_SYSTEM); + } + } + + size_t transferred_bytes = 0; + double setup_service_seconds = 0.0; + double pipeline_service_seconds = 0.0; + double storage_service_seconds = 0.0; + double cuda_wait_service_seconds = 0.0; + double fsync_service_seconds = 0.0; + double cleanup_service_seconds = 0.0; + for (const auto &metrics : worker_metrics) { + if (metrics.bytes > + std::numeric_limits::max() - transferred_bytes) { + std::fprintf(stderr, "custom storage transferred byte count overflow\n"); + return post_handle_failure(CUDA_ERROR_OPERATING_SYSTEM); + } + transferred_bytes += metrics.bytes; + setup_service_seconds += metrics.setup_seconds; + pipeline_service_seconds += metrics.pipeline_seconds; + storage_service_seconds += metrics.storage_seconds; + cuda_wait_service_seconds += metrics.cuda_wait_seconds; + fsync_service_seconds += metrics.fsync_seconds; + cleanup_service_seconds += metrics.cleanup_seconds; + } + if (transferred_bytes != total_bytes) { + std::fprintf(stderr, + "custom storage transfer coverage mismatch: transferred=%zu " + "expected=%zu\n", + transferred_bytes, total_bytes); + return post_handle_failure(CUDA_ERROR_OPERATING_SYSTEM); + } + + const auto post_transfer_validation_start = Clock::now(); + if (checkpoint) { + if (!storage::ValidateExtentFiles(storage_dir, manifest, &manifest_error)) { + std::fprintf(stderr, "custom storage extent validation failed: %s\n", + manifest_error.c_str()); + return post_handle_failure(CUDA_ERROR_OPERATING_SYSTEM); + } + if (!storage::WriteManifest(storage_dir, manifest, &manifest_error)) { + std::fprintf(stderr, "custom storage manifest write failed: %s\n", + manifest_error.c_str()); + return post_handle_failure(CUDA_ERROR_OPERATING_SYSTEM); + } + } + const double post_transfer_validation_seconds = + SecondsSince(post_transfer_validation_start); + + // This is the sole acknowledgment point; CUDA exposes no public abort for + // failures above. + const auto operation_complete_start = Clock::now(); + status = static_cast(daemon_protocol::FinishHandledOperation( + true, CUDA_SUCCESS, + [operation_complete, info] { + return static_cast(operation_complete(info->handle)); + }, + &operation)); + const double cuda_operation_complete_seconds = + SecondsSince(operation_complete_start); + if (status != CUDA_SUCCESS) { + if (checkpoint && !storage::RemoveManifest(storage_dir, &manifest_error)) { + std::fprintf(stderr, + "failed to remove custom storage manifest after CUDA " + "completion failure: %s\n", + manifest_error.c_str()); + } + return post_handle_failure(status); + } + + // Preserve the original transfer interval: worker setup through CUDA + // acknowledgment. + const double seconds = SecondsSince(start); + const double gib_per_second = seconds == 0.0 + ? 0.0 + : static_cast(total_bytes) / + (1024.0 * 1024.0 * 1024.0) / seconds; + const size_t retained_context_count = operation_contexts.size(); + const auto primary_context_release_start = Clock::now(); + const bool persist_for_target = + daemon_request != nullptr && persistent_contexts != nullptr; + const CUresult primary_context_release_status = + persist_for_target + ? persistent_contexts->Adopt(operation_contexts.DetachDevices(), + *daemon_request) + : operation_contexts.ReleaseAll(); + primary_context_release_seconds += + SecondsSince(primary_context_release_start); + const char *primary_context_release_state = + persist_for_target ? "deferred_until_target_exit" : "completed"; + const char *context_lifecycle = + persist_for_target ? "target_identity" : "invocation"; + const auto telemetry_end = Clock::now(); + const double custom_storage_total_seconds = + SecondsBetween(custom_storage_start, telemetry_end); + const double helper_main_to_telemetry_seconds = + SecondsBetween(helper_main_start, telemetry_end); + std::fprintf( + stdout, + "{\"event\":\"cuda_custom_storage_transfer\",\"schema_version\":1," + "\"operation\":\"%s\",\"devices\":%zu,\"bytes\":%zu," + "\"duration_seconds\":%.6f,\"effective_gib_per_second\":%.6f," + "\"transfer_buffer_count\":%zu,\"transfer_chunk_bytes\":%zu," + "\"pinned_bytes\":%zu,\"setup_service_seconds\":%.6f," + "\"pipeline_service_seconds\":%.6f,\"storage_service_seconds\":%.6f," + "\"cuda_wait_service_seconds\":%.6f,\"fsync_service_seconds\":%.6f," + "\"cleanup_service_seconds\":%.6f," + "\"timing_scope\":\"monotonic_wall;totals_contain_subphases;" + "service_seconds_are_cross_worker_sums_and_may_overlap\"," + "\"helper_main_to_telemetry_seconds\":%.6f," + "\"custom_storage_total_seconds\":%.6f," + "\"storage_directory_validation_seconds\":%.6f," + "\"cuda_device_count\":%d," + "\"retained_context_count\":%zu," + "\"device_enumeration_seconds\":%.6f," + "\"target_context_discovery_seconds\":%.6f," + "\"primary_context_retain_seconds\":%.6f," + "\"manifest_validation_seconds\":%.6f," + "\"device_map_preparation_seconds\":%.6f," + "\"cuda_process_api_seconds\":%.6f," + "\"metadata_job_construction_seconds\":%.6f," + "\"worker_orchestration_seconds\":%.6f," + "\"post_transfer_validation_seconds\":%.6f," + "\"cuda_operation_complete_seconds\":%.6f," + "\"primary_context_release_seconds\":%.6f," + "\"primary_context_release_state\":\"%s\"," + "\"primary_context_release_success\":%s," + "\"primary_context_release_status\":%d," + "\"context_lifecycle\":\"%s\"}\n", + checkpoint ? "checkpoint" : "restore", manifest.size(), total_bytes, + seconds, gib_per_second, transfer_options.buffer_count, + transfer_options.chunk_bytes, pinned_bytes, setup_service_seconds, + pipeline_service_seconds, storage_service_seconds, + cuda_wait_service_seconds, fsync_service_seconds, cleanup_service_seconds, + helper_main_to_telemetry_seconds, custom_storage_total_seconds, + storage_directory_validation_seconds, visible_cuda_device_count, + retained_context_count, + device_enumeration_seconds, + target_context_discovery_seconds, primary_context_retain_seconds, + manifest_validation_seconds, + device_map_preparation_seconds, cuda_process_api_seconds, + metadata_job_construction_seconds, worker_orchestration_seconds, + post_transfer_validation_seconds, cuda_operation_complete_seconds, + primary_context_release_seconds, primary_context_release_state, + primary_context_release_status == CUDA_SUCCESS ? "true" : "false", + static_cast(primary_context_release_status), context_lifecycle); + if (primary_context_release_status != CUDA_SUCCESS) { + std::fprintf( + stderr, + "warning: retained CUDA primary context release failed with status %d " + "after operation acknowledgment\n", + static_cast(primary_context_release_status)); + } + if (primary_context_release_status != CUDA_SUCCESS) { + return {primary_context_release_status, operation}; + } + return {CUDA_SUCCESS, operation, false}; +} + +CUresult DoRegularCheckpoint(int pid) { + CUcheckpointCheckpointArgs args{}; + return cuCheckpointProcessCheckpoint(pid, &args); +} + +CUresult DoLegacyRestore(int pid, const std::string &device_map) { + std::vector pairs; + if (!ParseDeviceMap(device_map, &pairs)) { + return CUDA_ERROR_INVALID_VALUE; + } + CUcheckpointRestoreArgs args{}; + args.gpuPairs = pairs.empty() ? nullptr : pairs.data(); + args.gpuPairsCount = pairs.size(); + return cuCheckpointProcessRestore(pid, &args); +} + +daemon_protocol::Response RunDaemonOperation( + const daemon_protocol::Request &request, + cuda_checkpoint_compat::OperationCompleteFn operation_complete, + std::chrono::seconds max_operation_duration, + PersistentTargetContexts *persistent_contexts) { + daemon_protocol::Response response; + // Until the driver lock call begins, every failure leaves the source process + // running. Clear this immediately before the call so callers only treat + // failures that are known to precede CUDA mutation as safe to leave alive. + if (request.action == daemon_protocol::Action::kLock) { + response.flags |= daemon_protocol::kResponseLockNotAcquired; + } + std::string reap_error; + const CUresult release_status = + persistent_contexts->ReapExited("/host/proc", &reap_error); + if (release_status != CUDA_SUCCESS) { + response.cuda_status = release_status; + response.flags |= daemon_protocol::kResponseFatal; + response.error = + "failed to release CUDA primary contexts for an exited target"; + return response; + } + if (!reap_error.empty()) { + response.cuda_status = CUDA_ERROR_OPERATING_SYSTEM; + response.error = reap_error + "; retained contexts and deferred operation"; + return response; + } + constexpr size_t kPerStreamCaptureLimit = + daemon_protocol::kMaxResponseSize / 2 - 256; + daemon_protocol::BoundedOutputCapture output_capture( + kPerStreamCaptureLimit); + daemon_protocol::BoundedOutputCapture error_capture(kPerStreamCaptureLimit); + std::string capture_setup_error; + if (!output_capture.Start(&capture_setup_error) || + !error_capture.Start(&capture_setup_error)) { + response.cuda_status = CUDA_ERROR_OPERATING_SYSTEM; + response.flags |= daemon_protocol::kResponseFatal; + response.error = capture_setup_error; + return response; + } + (void)std::fflush(stdout); + (void)std::fflush(stderr); + const int saved_stdout = dup(STDOUT_FILENO); + const int saved_stderr = dup(STDERR_FILENO); + if (saved_stdout < 0 || saved_stderr < 0 || + dup2(output_capture.write_fd(), STDOUT_FILENO) < 0 || + dup2(error_capture.write_fd(), STDERR_FILENO) < 0) { + response.cuda_status = CUDA_ERROR_OPERATING_SYSTEM; + response.flags |= daemon_protocol::kResponseFatal; + response.error = "failed to redirect daemon operation output"; + } else { + if ((!request.job_file.empty() && + setenv("CUDA_CHECKPOINT_JOB_FILE", request.job_file.c_str(), 1) != + 0) || + (request.job_file.empty() && + unsetenv("CUDA_CHECKPOINT_JOB_FILE") != 0)) { + response.cuda_status = CUDA_ERROR_OPERATING_SYSTEM; + std::perror("configure CUDA_CHECKPOINT_JOB_FILE"); + } else if (request.backend == daemon_protocol::Backend::kPosix && + operation_complete == nullptr) { + response.cuda_status = CUDA_ERROR_NOT_SUPPORTED; + std::fprintf( + stderr, + "CUDA POSIX CustomStorage backend requested but the CUDA 13.4 " + "driver API or transfer adapter is unavailable\n"); + } else if (request.action == daemon_protocol::Action::kLock || + request.action == daemon_protocol::Action::kUnlock) { + std::string identity_error; + if (!daemon_protocol::ValidateProcessIdentity(request, "/host/proc", + &identity_error)) { + response.cuda_status = CUDA_ERROR_INVALID_VALUE; + std::fprintf( + stderr, "process identity changed immediately before CUDA %s: %s\n", + daemon_protocol::ActionName(request.action), + identity_error.c_str()); + } else if (request.action == daemon_protocol::Action::kLock) { + CUcheckpointLockArgs lock_args{}; + std::string timeout_error; + if (!daemon_protocol::OperationTimeoutMilliseconds( + max_operation_duration, &lock_args.timeoutMs, + &timeout_error)) { + response.cuda_status = CUDA_ERROR_INVALID_VALUE; + response.flags |= daemon_protocol::kResponseFatal; + std::fprintf(stderr, "%s\n", timeout_error.c_str()); + } else { + response.flags &= ~daemon_protocol::kResponseLockNotAcquired; + response.cuda_status = + cuCheckpointProcessLock(request.pid, &lock_args); + if (response.cuda_status == CUDA_ERROR_NOT_READY) { + // CUDA guarantees a timed-out lock leaves the process RUNNING. + response.flags |= daemon_protocol::kResponseLockNotAcquired; + } + } + } else { + CUcheckpointUnlockArgs args{}; + response.cuda_status = cuCheckpointProcessUnlock(request.pid, &args); + } + } else if (request.backend == daemon_protocol::Backend::kRegular) { + std::string identity_error; + if (!daemon_protocol::ValidateProcessIdentity(request, "/host/proc", + &identity_error)) { + response.cuda_status = CUDA_ERROR_INVALID_VALUE; + std::fprintf(stderr, + "process identity changed immediately before regular " + "CUDA %s: %s\n", + daemon_protocol::ActionName(request.action), + identity_error.c_str()); + } else if (request.action == daemon_protocol::Action::kCheckpoint) { + response.cuda_status = DoRegularCheckpoint(request.pid); + } else { + response.cuda_status = DoLegacyRestore(request.pid, request.device_map); + } + } else { + transfer::TransferOptions options{ + .buffer_count = request.transfer_buffer_count, + .chunk_bytes = static_cast(request.transfer_chunk_bytes), + }; + std::string validation_error; + if (!transfer::ValidateTransferOptions(options, &validation_error)) { + response.cuda_status = CUDA_ERROR_INVALID_VALUE; + std::fprintf(stderr, "invalid transfer configuration: %s\n", + validation_error.c_str()); + } else { + const auto operation_start = Clock::now(); + if (max_operation_duration > + Clock::time_point::max() - operation_start) { + response.cuda_status = CUDA_ERROR_INVALID_VALUE; + std::fprintf(stderr, + "configured operation duration exceeds the steady " + "clock range\n"); + } else { + const CustomStorageResult result = DoCustomStorage( + request.pid, + request.action == daemon_protocol::Action::kCheckpoint, + request.device_map, request.storage_dir, options, + operation_start + max_operation_duration, operation_start, + operation_complete, &request, persistent_contexts); + response.cuda_status = result.status; + if (result.operation.fatal() || result.fatal) { + response.flags |= daemon_protocol::kResponseFatal; + } + } + } + } + if (response.cuda_status != CUDA_SUCCESS) { + PrintCudaError(static_cast(response.cuda_status)); + } + } + (void)std::fflush(stdout); + (void)std::fflush(stderr); + bool output_restore_failed = false; + if (saved_stdout >= 0) { + output_restore_failed = dup2(saved_stdout, STDOUT_FILENO) < 0; + close(saved_stdout); + } + if (saved_stderr >= 0) { + output_restore_failed = + dup2(saved_stderr, STDERR_FILENO) < 0 || output_restore_failed; + close(saved_stderr); + } + if (output_restore_failed) { + response.cuda_status = CUDA_ERROR_OPERATING_SYSTEM; + response.flags |= daemon_protocol::kResponseFatal; + } + bool output_truncated = false; + bool error_truncated = false; + std::string captured_output; + std::string captured_error; + std::string output_capture_error; + std::string error_capture_error; + const bool output_finished = output_capture.Finish( + &captured_output, &output_truncated, &output_capture_error); + const bool error_finished = error_capture.Finish( + &captured_error, &error_truncated, &error_capture_error); + if (!output_finished || !error_finished) { + response.cuda_status = CUDA_ERROR_OPERATING_SYSTEM; + response.flags |= daemon_protocol::kResponseFatal; + for (const std::string *capture_error : {&output_capture_error, + &error_capture_error}) { + if (capture_error->empty()) { + continue; + } + if (!response.error.empty()) { + response.error += '\n'; + } + response.error += *capture_error; + } + } + if (!captured_output.empty()) { + if (!response.output.empty()) { + response.output += '\n'; + } + response.output += captured_output; + } + if (!captured_error.empty()) { + if (!response.error.empty()) { + response.error += '\n'; + } + response.error += captured_error; + } + if (output_truncated) { + response.output += "\n[stdout truncated at daemon response limit]\n"; + } + if (error_truncated) { + response.error += "\n[stderr truncated at daemon response limit]\n"; + } + if (output_restore_failed) { + if (!response.error.empty()) { + response.error += '\n'; + } + response.error += "failed to restore daemon output descriptors"; + } + return response; +} + +bool ValidSocketPath(const std::string &path) { + const std::filesystem::path socket_path(path); + const std::string filename = socket_path.filename(); + const bool clean_filename = + !filename.empty() && + std::all_of(filename.begin(), filename.end(), [](unsigned char c) { + return std::isalnum(c) || c == '.' || c == '_' || c == '-'; + }); + return !path.empty() && path.front() == '/' && + path.size() + sizeof(".health") <= sizeof(sockaddr_un::sun_path) && + socket_path.lexically_normal() == socket_path && + socket_path.parent_path() == + std::filesystem::path("/run/cuda-checkpoint-helper") && + clean_filename; +} + +int RunHealthClient(const std::string &socket_path) { + sockaddr_un address{}; + const std::string health_socket_path = socket_path + ".health"; + if (!ValidSocketPath(socket_path) || + health_socket_path.size() >= sizeof(address.sun_path)) { + std::fprintf(stderr, "invalid daemon socket path\n"); + return 1; + } + address.sun_family = AF_UNIX; + std::memcpy(address.sun_path, health_socket_path.c_str(), + health_socket_path.size() + 1); + const int socket_fd = socket(AF_UNIX, SOCK_SEQPACKET | SOCK_CLOEXEC, 0); + if (socket_fd < 0 || + connect(socket_fd, reinterpret_cast(&address), + sizeof(address)) != 0) { + if (socket_fd >= 0) { + close(socket_fd); + } + return 1; + } + timeval timeout{.tv_sec = kClientReceiveTimeoutMilliseconds / 1000, + .tv_usec = 0}; + if (setsockopt(socket_fd, SOL_SOCKET, SO_SNDTIMEO, &timeout, + sizeof(timeout)) != 0 || + setsockopt(socket_fd, SOL_SOCKET, SO_RCVTIMEO, &timeout, + sizeof(timeout)) != 0) { + close(socket_fd); + return 1; + } + daemon_protocol::Request request; + std::vector packet; + std::string error; + if (!daemon_protocol::EncodeRequest(request, &packet, &error) || + send(socket_fd, packet.data(), packet.size(), MSG_NOSIGNAL) != + static_cast(packet.size())) { + close(socket_fd); + return 1; + } + packet.resize(daemon_protocol::kMaxResponseSize + 1); + const ssize_t received = + recv(socket_fd, packet.data(), packet.size(), MSG_TRUNC); + close(socket_fd); + daemon_protocol::Response response; + if (received <= 0 || + static_cast(received) > daemon_protocol::kMaxResponseSize || + !daemon_protocol::ParseResponse(packet.data(), received, &response, + &error) || + response.cuda_status != CUDA_SUCCESS || + (response.flags & daemon_protocol::kResponseCapabilityDeferredCUDA) == + 0) { + return 1; + } + return 0; +} + +bool RunHealthServer(daemon_protocol::OwnedUnixSocket *socket, int shutdown_fd, + int log_fd, + const daemon_protocol::OperationHealth *health, + PersistentTargetContexts *persistent_contexts) { + std::vector packet(daemon_protocol::kMaxRequestSize + 1); + for (;;) { + const int server_poll = + daemon_protocol::PollForInputOrStop(socket->fd(), shutdown_fd); + if (server_poll == -2) { + dprintf(log_fd, "health socket poll failed: %s\n", std::strerror(errno)); + return false; + } + if (server_poll == 0) { + return true; + } + const int accepted_fd = + accept4(socket->fd(), nullptr, nullptr, SOCK_CLOEXEC); + if (accepted_fd < 0) { + if (errno == EINTR || errno == EAGAIN || errno == ECONNABORTED) { + continue; + } + if (errno == EMFILE || errno == ENFILE || errno == ENOBUFS || + errno == ENOMEM) { + dprintf(log_fd, "health socket accept temporarily failed: %s\n", + std::strerror(errno)); + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + continue; + } + dprintf(log_fd, "health socket accept failed: %s\n", + std::strerror(errno)); + return false; + } + ScopedFd client_fd(accepted_fd); + const int client_poll = daemon_protocol::PollForInputOrStop( + client_fd.get(), shutdown_fd, {}, kClientReceiveTimeoutMilliseconds); + if (client_poll == -2) { + dprintf(log_fd, "health client poll failed: %s\n", + std::strerror(errno)); + return false; + } + if (client_poll == 0) { + return true; + } + if (client_poll < 0) { + continue; + } + const ssize_t received = + recv(client_fd.get(), packet.data(), packet.size(), MSG_TRUNC); + daemon_protocol::Request request; + daemon_protocol::Response response; + std::string error; + if (received <= 0 || + static_cast(received) > daemon_protocol::kMaxRequestSize || + !daemon_protocol::ParseRequest(packet.data(), received, &request, + &error)) { + response.cuda_status = CUDA_ERROR_INVALID_VALUE; + response.error = + received <= 0 ? "failed to receive health request" : error; + } else if (request.action != daemon_protocol::Action::kHealth) { + response.cuda_status = CUDA_ERROR_INVALID_VALUE; + response.error = "health socket accepts only health requests"; + } else { + std::string reap_error; + const CUresult release_status = + persistent_contexts->ReapExited("/host/proc", &reap_error); + response = daemon_protocol::HealthResponseAfterReap( + *health, static_cast(release_status), reap_error); + if (!reap_error.empty()) { + // An unreadable /proc identity is not proof that the target exited. + // Keep liveness successful so kubelet does not restart the helper and + // release retained contexts during shutdown. Operation requests remain + // fail-closed until identity can be established again. + dprintf(log_fd, "target-context reaping deferred: %s\n", + reap_error.c_str()); + } + } + std::vector encoded; + if (daemon_protocol::EncodeResponse(response, &encoded, &error)) { + (void)send(client_fd.get(), encoded.data(), encoded.size(), MSG_NOSIGNAL); + } + } + return true; +} + +int RunDaemon(const std::string &socket_path, uint64_t max_operation_seconds) { + daemon_protocol::ShutdownSignalOwner signal_owner; + std::string setup_error; + if (!signal_owner.Start(&setup_error)) { + std::fprintf(stderr, "daemon shutdown setup failed: %s\n", + setup_error.c_str()); + return 1; + } + + const std::filesystem::path path(socket_path); + if (!ValidSocketPath(socket_path)) { + std::fprintf(stderr, "invalid daemon socket path\n"); + return 1; + } + std::error_code filesystem_error; + std::filesystem::create_directories(path.parent_path(), filesystem_error); + if (filesystem_error || chmod(path.parent_path().c_str(), 0700) != 0) { + std::fprintf(stderr, "failed to create private daemon socket directory\n"); + return 1; + } + + const auto init_start = Clock::now(); + CUresult status = cuInit(0); + const double init_seconds = SecondsSince(init_start); + if (status != CUDA_SUCCESS) { + PrintCudaError(status); + return 1; + } + int device_count = 0; + double enumeration_seconds = 0.0; + const auto enumeration_start = Clock::now(); + status = cuDeviceGetCount(&device_count); + enumeration_seconds = SecondsSince(enumeration_start); + if (status != CUDA_SUCCESS) { + PrintCudaError(status); + return 1; + } + int driver_version = 0; + (void)cuDriverGetVersion(&driver_version); + bool custom_storage_driver_api_available = false; + const cuda_checkpoint_compat::OperationCompleteFn operation_complete = + cuda_checkpoint_compat::ResolveOperationComplete( + &custom_storage_driver_api_available); + const bool custom_storage_transfer_backend_available = + transfer::TransferBackendAvailable(); + const bool custom_storage_available = + custom_storage_driver_api_available && + custom_storage_transfer_backend_available; + const cuda_checkpoint_compat::OperationCompleteFn + custom_storage_operation_complete = + custom_storage_available ? operation_complete : nullptr; + PersistentTargetContexts persistent_contexts; + + daemon_protocol::OwnedUnixSocket operation_socket; + daemon_protocol::OwnedUnixSocket health_socket; + std::string socket_error; + if (!operation_socket.Bind(socket_path, 16, &socket_error)) { + std::fprintf(stderr, "daemon operation socket setup failed: %s\n", + socket_error.c_str()); + return 1; + } + if (!health_socket.Bind(socket_path + ".health", 4, &socket_error)) { + std::fprintf(stderr, "daemon health socket setup failed: %s\n", + socket_error.c_str()); + return 1; + } + daemon_protocol::OperationHealth operation_health{ + std::chrono::seconds(max_operation_seconds)}; + operation_health.MarkReady(custom_storage_available); + // Operation capture redirects process-wide stderr. Keep the health thread on + // the original container-log descriptor so its diagnostics cannot leak into + // an unrelated operation response. + ScopedFd health_log_fd(dup(STDERR_FILENO)); + if (health_log_fd.get() < 0) { + std::fprintf(stderr, "duplicate daemon health log descriptor failed: %s\n", + std::strerror(errno)); + return 1; + } + daemon_protocol::ShutdownSignalOwner::ShutdownResult shutdown_result; + daemon_protocol::ShutdownSignalOwner::ShutdownResult health_shutdown_result; + std::atomic health_thread_failed{false}; + bool daemon_fatal = false; + { + // The guard is destroyed before the jthread: it stops and joins the + // signal owner, which wakes the health server, and then jthread joins the + // server. + std::jthread health_thread; + try { + health_thread = std::jthread([&]() noexcept { + try { + if (!RunHealthServer(&health_socket, signal_owner.health_stop_fd(), + health_log_fd.get(), &operation_health, + &persistent_contexts)) { + health_thread_failed.store(true, std::memory_order_release); + health_shutdown_result = signal_owner.RequestShutdownNoThrow(); + } + } catch (...) { + health_shutdown_result = signal_owner.RequestShutdownNoThrow(); + health_thread_failed.store(true, std::memory_order_release); + } + }); + } catch (const std::system_error &exception) { + std::fprintf(stderr, "create daemon health thread failed: %s\n", + exception.what()); + return 1; + } + DaemonThreadShutdown shutdown_threads(&signal_owner, &shutdown_result); + try { + std::fprintf( + stdout, + "{\"event\":\"cuda_checkpoint_daemon_ready\",\"schema_version\":1," + "\"cuda_init_seconds\":%.6f," + "\"cuda_device_count\":%d,\"device_enumeration_seconds\":%.6f," + "\"primary_context_retain_seconds\":%.6f,\"cuda_driver_version\":%" + "d," + "\"custom_storage_driver_api_available\":%s," + "\"custom_storage_transfer_backend_available\":%s," + "\"custom_storage_available\":%s," + "\"context_lifecycle\":\"target_identity\"}\n", + init_seconds, device_count, enumeration_seconds, 0.0, + driver_version, + custom_storage_driver_api_available ? "true" : "false", + custom_storage_transfer_backend_available ? "true" : "false", + custom_storage_available ? "true" : "false"); + std::fflush(stdout); + + std::vector packet(daemon_protocol::kMaxRequestSize + 1); + while (!signal_owner.ShutdownRequested() && !daemon_fatal) { + const int server_poll = daemon_protocol::PollForInputOrStop( + operation_socket.fd(), signal_owner.operation_stop_fd()); + if (server_poll == -2) { + std::fprintf(stderr, "operation socket poll failed: %s\n", + std::strerror(errno)); + daemon_fatal = true; + break; + } + if (server_poll == 0) { + break; + } + const int accepted_fd = + accept4(operation_socket.fd(), nullptr, nullptr, SOCK_CLOEXEC); + if (accepted_fd < 0) { + if (errno == EINTR || errno == EAGAIN || errno == ECONNABORTED) { + continue; + } + if (errno == EMFILE || errno == ENFILE || errno == ENOBUFS || + errno == ENOMEM) { + std::fprintf(stderr, + "operation socket accept temporarily failed: %s\n", + std::strerror(errno)); + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + continue; + } + if (signal_owner.ShutdownRequested()) { + break; + } + std::perror("operation socket accept"); + daemon_fatal = true; + break; + } + ScopedFd client_fd(accepted_fd); + const int client_poll = daemon_protocol::PollForInputOrStop( + client_fd.get(), signal_owner.operation_stop_fd(), {}, + kClientReceiveTimeoutMilliseconds); + if (client_poll == -2) { + std::fprintf(stderr, "operation client poll failed: %s\n", + std::strerror(errno)); + daemon_fatal = true; + break; + } + if (client_poll == 0) { + break; + } + if (client_poll < 0) { + continue; + } + const ssize_t received = + recv(client_fd.get(), packet.data(), packet.size(), MSG_TRUNC); + daemon_protocol::Response response; + daemon_protocol::Request request; + std::string protocol_error; + if (received <= 0 || + static_cast(received) > daemon_protocol::kMaxRequestSize || + !daemon_protocol::ParseRequest(packet.data(), received, &request, + &protocol_error)) { + response.cuda_status = CUDA_ERROR_INVALID_VALUE; + response.error = + received <= 0 ? "failed to receive request" : protocol_error; + } else if (request.action == daemon_protocol::Action::kHealth) { + response.cuda_status = CUDA_ERROR_INVALID_VALUE; + response.error = "health requests must use the health socket"; + } else { + const auto rpc_start = Clock::now(); + operation_health.Begin(request.action, request.pid); + daemon_fatal = !daemon_protocol::ExecuteValidated( + request, "/host/proc", + [custom_storage_operation_complete, max_operation_seconds, + &persistent_contexts]( + const daemon_protocol::Request &validated) { + return RunDaemonOperation( + validated, custom_storage_operation_complete, + std::chrono::seconds(max_operation_seconds), + &persistent_contexts); + }, + &response); + operation_health.End(); + std::fprintf(stdout, + "{\"event\":\"cuda_checkpoint_daemon_operation\"," + "\"schema_version\":1,\"action\":\"%s\"," + "\"pid\":%u,\"cuda_status\":%d,\"fatal\":%s,\"rpc_" + "service_seconds\":%.6f}\n", + daemon_protocol::ActionName(request.action), request.pid, + response.cuda_status, + (response.flags & daemon_protocol::kResponseFatal) != 0 + ? "true" + : "false", + SecondsSince(rpc_start)); + std::fflush(stdout); + } + std::vector encoded; + if (!daemon_protocol::EncodeResponse(response, &encoded, + &protocol_error)) { + daemon_protocol::Response bounded{ + .cuda_status = CUDA_ERROR_OPERATING_SYSTEM, + .flags = response.flags & daemon_protocol::kResponseFatal, + .output = "", + .error = "daemon response exceeded protocol limit", + }; + (void)daemon_protocol::EncodeResponse(bounded, &encoded, + &protocol_error); + } + (void)send(client_fd.get(), encoded.data(), encoded.size(), + MSG_NOSIGNAL); + } + } catch (const std::exception &exception) { + std::fprintf(stderr, "daemon processing failed: %s\n", exception.what()); + daemon_fatal = true; + } catch (...) { + std::fprintf(stderr, "daemon processing failed: unknown exception\n"); + daemon_fatal = true; + } + } + if (health_thread_failed.load(std::memory_order_acquire)) { + std::fprintf(stderr, "daemon health thread failed\n"); + daemon_fatal = true; + if (!health_shutdown_result.ok()) { + shutdown_result = health_shutdown_result; + } + } + if (!shutdown_result.ok()) { + std::fprintf(stderr, "daemon shutdown failed: %s: %s\n", + shutdown_result.operation, + std::strerror(shutdown_result.error_code)); + daemon_fatal = true; + } + signal_owner.Close(); + health_socket.Close(); + operation_socket.Close(); + const auto release_start = Clock::now(); + status = persistent_contexts.ReleaseAll(); + std::fprintf( + stdout, + "{\"event\":\"cuda_checkpoint_daemon_stopped\",\"schema_version\":1," + "\"primary_context_release_seconds\":%.6f,\"primary_context_release_" + "status\":%d,\"context_lifecycle\":\"target_identity\"}\n", + SecondsSince(release_start), static_cast(status)); + return status == CUDA_SUCCESS && !daemon_fatal ? 0 : 1; +} + +} // namespace + +int main(int argc, char **argv) { + int pid = 0; + bool have_pid = false; + bool get_restore_tid = false; + bool daemon = false; + bool health = false; + std::string socket_path; + uint64_t max_operation_seconds = kMaxOperationSeconds; + + if (argc == 1) { + return PrintUsageError(); + } + for (int index = 1; index < argc; ++index) { + std::string argument = argv[index]; + if (argument == "--daemon") { + daemon = true; + } else if (argument == "--health") { + health = true; + } else if (argument == "--socket" && ++index < argc) { + socket_path = argv[index]; + } else if (argument == "--max-operation-seconds" && ++index < argc && + ParsePositiveSeconds(argv[index], &max_operation_seconds)) { + } else if (argument == "--get-restore-tid") { + get_restore_tid = true; + } else if ((argument == "--pid" || argument == "-p") && ++index < argc && + ParsePID(argv[index], &pid)) { + have_pid = true; + } else if (argument == "--help" || argument == "-h") { + return PrintUsage(stdout); + } else { + return PrintUsageError(); + } + } + + if (daemon || health) { + if (static_cast(daemon) + static_cast(health) != 1 || + socket_path.empty() || have_pid || get_restore_tid || + (health && max_operation_seconds != kMaxOperationSeconds)) { + return PrintUsageError(); + } + return daemon ? RunDaemon(socket_path, max_operation_seconds) + : RunHealthClient(socket_path); + } + + if (!get_restore_tid || !have_pid) { + return PrintUsageError(); + } + + CUresult status = cuInit(0); + if (status != CUDA_SUCCESS) { + PrintCudaError(status); + return 1; + } + int tid = 0; + status = cuCheckpointProcessGetRestoreThreadId(pid, &tid); + if (status == CUDA_ERROR_INVALID_VALUE) { + std::string process_error; + const daemon_protocol::ProcessExistenceState existence = + daemon_protocol::InspectProcessExistence(pid, "/proc", &process_error); + if (existence == daemon_protocol::ProcessExistenceState::kExists) { + // The output pointer is valid and the candidate PID still exists. The + // driver uses INVALID_VALUE for a live process without CUDA checkpoint + // state. Keep that negative result distinct from helper, driver, and + // raced-PID failures so the agent can fail closed on the latter. + return std::fprintf(stdout, "none\n") < 0 ? 1 : 0; + } + std::fprintf(stderr, "CUDA restore-tid candidate validation failed: %s\n", + process_error.c_str()); + return 1; + } + if (status != CUDA_SUCCESS) { + PrintCudaError(status); + return 1; + } + return std::fprintf(stdout, "%d\n", tid) < 0 ? 1 : 0; +} diff --git a/agent/cmd/cuda-checkpoint-helper/storage_manifest.cpp b/agent/cmd/cuda-checkpoint-helper/storage_manifest.cpp new file mode 100644 index 00000000..8c6ec2dd --- /dev/null +++ b/agent/cmd/cuda-checkpoint-helper/storage_manifest.cpp @@ -0,0 +1,554 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include "storage_manifest.h" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cuda_checkpoint_storage { +namespace { + +constexpr size_t kMaximumDeviceCount = 1024; +std::atomic kTemporaryManifestSequence{0}; + +class FileDescriptor { +public: + explicit FileDescriptor(int fd = -1) : fd_(fd) {} + FileDescriptor(const FileDescriptor &) = delete; + FileDescriptor &operator=(const FileDescriptor &) = delete; + + ~FileDescriptor() { + if (fd_ >= 0) { + (void)close(fd_); + } + } + + int get() const { return fd_; } + + bool Close() { + if (fd_ < 0) { + return true; + } + const int fd = fd_; + fd_ = -1; + return close(fd) == 0; + } + +private: + int fd_; +}; + +int HexValue(char value) { + if (value >= '0' && value <= '9') { + return value - '0'; + } + if (value >= 'a' && value <= 'f') { + return value - 'a' + 10; + } + if (value >= 'A' && value <= 'F') { + return value - 'A' + 10; + } + return -1; +} + +bool NormalizeExtent(const ManifestExtent &extent, size_t index, + ManifestExtent *normalized, std::string *error) { + if (!CanonicalizeGPUUUID(extent.source_uuid, &normalized->source_uuid)) { + *error = "invalid source GPU UUID in helper manifest"; + return false; + } + if (extent.size == 0) { + *error = "zero-sized extent in helper manifest"; + return false; + } + normalized->size = extent.size; + normalized->filename = extent.filename; + if (normalized->filename != DeviceFilename(index)) { + *error = "invalid deterministic extent filename in helper manifest"; + return false; + } + return true; +} + +bool NormalizeManifest(const std::vector &extents, + std::vector *normalized, + std::string *error) { + if (extents.size() > kMaximumDeviceCount) { + *error = "helper manifest has too many device extents"; + return false; + } + normalized->clear(); + normalized->reserve(extents.size()); + std::unordered_set source_uuids; + for (size_t index = 0; index < extents.size(); ++index) { + ManifestExtent extent; + if (!NormalizeExtent(extents[index], index, &extent, error)) { + return false; + } + if (!source_uuids.insert(extent.source_uuid).second) { + *error = "duplicate source GPU UUID in helper manifest"; + return false; + } + normalized->push_back(std::move(extent)); + } + return true; +} + +void RemoveTemporaryManifest(const std::filesystem::path &temporary) { + std::error_code ignored; + std::filesystem::remove(temporary, ignored); +} + +bool RemoveTemporaryManifests(const std::filesystem::path &directory, + bool *removed, std::string *error) { + std::error_code iterator_error; + std::filesystem::directory_iterator iterator(directory, iterator_error); + if (iterator_error) { + *error = "scan helper directory for temporary manifests: " + + iterator_error.message(); + return false; + } + const std::filesystem::directory_iterator end; + while (iterator != end) { + const auto entry = *iterator; + const std::string name = entry.path().filename().string(); + if (name != kLegacyTemporaryManifestName && + name.rfind(kTemporaryManifestPrefix, 0) != 0) { + iterator.increment(iterator_error); + if (iterator_error) { + *error = "scan helper directory for temporary manifests: " + + iterator_error.message(); + return false; + } + continue; + } + if (unlink(entry.path().c_str()) != 0 && errno != ENOENT) { + *error = "remove temporary helper manifest: " + + std::string(std::strerror(errno)); + return false; + } + *removed = true; + iterator.increment(iterator_error); + if (iterator_error) { + *error = "scan helper directory for temporary manifests: " + + iterator_error.message(); + return false; + } + } + return true; +} + +bool WriteAll(int fd, const std::string &contents) { + size_t offset = 0; + while (offset < contents.size()) { + const ssize_t written = + write(fd, contents.data() + offset, contents.size() - offset); + if (written < 0 && errno == EINTR) { + continue; + } + if (written <= 0) { + return false; + } + offset += static_cast(written); + } + return true; +} + +} // namespace + +bool ParseGPUUUID(std::string_view value, + std::array *bytes_out) { + if (bytes_out == nullptr) { + return false; + } + if (value.size() == 40) { + if (value.substr(0, 4) != "GPU-") { + return false; + } + value.remove_prefix(4); + } + if (value.size() != 36 || value[8] != '-' || value[13] != '-' || + value[18] != '-' || value[23] != '-') { + return false; + } + + size_t input_index = 0; + for (size_t byte_index = 0; byte_index < bytes_out->size(); ++byte_index) { + if (input_index == 8 || input_index == 13 || input_index == 18 || + input_index == 23) { + ++input_index; + } + const int high = HexValue(value[input_index]); + const int low = HexValue(value[input_index + 1]); + if (high < 0 || low < 0) { + return false; + } + (*bytes_out)[byte_index] = static_cast((high << 4) | low); + input_index += 2; + } + return input_index == value.size(); +} + +std::string FormatGPUUUID(const std::array &bytes) { + static constexpr char kHex[] = "0123456789abcdef"; + std::string result = "GPU-"; + result.reserve(40); + for (size_t index = 0; index < bytes.size(); ++index) { + if (index == 4 || index == 6 || index == 8 || index == 10) { + result.push_back('-'); + } + result.push_back(kHex[bytes[index] >> 4]); + result.push_back(kHex[bytes[index] & 0x0f]); + } + return result; +} + +bool CanonicalizeGPUUUID(std::string_view value, std::string *canonical_out) { + if (canonical_out == nullptr) { + return false; + } + std::array bytes{}; + if (!ParseGPUUUID(value, &bytes)) { + return false; + } + *canonical_out = FormatGPUUUID(bytes); + return true; +} + +std::string DeviceFilename(size_t index) { + char filename[32]; + std::snprintf(filename, sizeof(filename), "device-%04zu.bin", index); + return filename; +} + +bool BuildCheckpointManifest(const std::vector &devices, + std::vector *extents, + std::string *error) { + if (extents == nullptr || error == nullptr) { + return false; + } + if (devices.size() > kMaximumDeviceCount) { + *error = "CUDA returned too many custom storage devices"; + return false; + } + + extents->clear(); + extents->reserve(devices.size()); + std::unordered_set source_uuids; + for (size_t index = 0; index < devices.size(); ++index) { + std::string source_uuid; + if (!CanonicalizeGPUUUID(devices[index].uuid, &source_uuid)) { + *error = "CUDA returned an invalid source GPU UUID"; + return false; + } + if (!source_uuids.insert(source_uuid).second) { + *error = "CUDA returned duplicate source GPU UUIDs"; + return false; + } + if (devices[index].size == 0) { + *error = "CUDA returned a zero-sized custom storage extent"; + return false; + } + extents->push_back( + {std::move(source_uuid), devices[index].size, DeviceFilename(index)}); + } + return true; +} + +bool BuildTransferJobs(const std::vector &extents, + const std::vector &devices, + const std::vector &device_pairs, + std::vector *jobs, std::string *error) { + if (jobs == nullptr || error == nullptr) { + return false; + } + + std::vector normalized_extents; + if (!NormalizeManifest(extents, &normalized_extents, error)) { + return false; + } + + std::unordered_map extent_by_source; + for (size_t index = 0; index < normalized_extents.size(); ++index) { + extent_by_source.emplace(normalized_extents[index].source_uuid, index); + } + + std::unordered_map source_to_destination; + std::unordered_map destination_to_source; + for (const auto &pair : device_pairs) { + std::string source; + std::string destination; + if (!CanonicalizeGPUUUID(pair.source_uuid, &source) || + !CanonicalizeGPUUUID(pair.destination_uuid, &destination)) { + *error = "invalid GPU UUID in CUDA device map"; + return false; + } + if (!source_to_destination.emplace(source, destination).second) { + *error = "duplicate source GPU UUID in CUDA device map"; + return false; + } + if (!destination_to_source.emplace(destination, source).second) { + *error = "ambiguous destination GPU UUID in CUDA device map"; + return false; + } + } + + jobs->clear(); + jobs->reserve(devices.size()); + std::unordered_set destination_uuids; + std::unordered_set consumed_extents; + for (size_t device_index = 0; device_index < devices.size(); ++device_index) { + std::string destination; + if (!CanonicalizeGPUUUID(devices[device_index].uuid, &destination)) { + *error = "CUDA returned an invalid destination GPU UUID"; + return false; + } + if (!destination_uuids.insert(destination).second) { + *error = "CUDA returned duplicate destination GPU UUIDs"; + return false; + } + if (devices[device_index].size == 0) { + *error = "CUDA returned a zero-sized custom storage extent"; + return false; + } + + std::string source = destination; + if (!device_pairs.empty()) { + const auto source_it = destination_to_source.find(destination); + if (source_it == destination_to_source.end()) { + *error = "destination GPU UUID is missing from the CUDA device map"; + return false; + } + source = source_it->second; + } + + const auto extent_it = extent_by_source.find(source); + if (extent_it == extent_by_source.end()) { + *error = "no saved extent matches the source GPU UUID"; + return false; + } + const size_t extent_index = extent_it->second; + if (!consumed_extents.insert(extent_index).second) { + *error = "saved source GPU extent was matched more than once"; + return false; + } + if (normalized_extents[extent_index].size != devices[device_index].size) { + *error = "saved extent size does not match the source GPU UUID"; + return false; + } + jobs->push_back({device_index, extent_index}); + } + + if (consumed_extents.size() != normalized_extents.size()) { + *error = "one or more saved GPU extents were not consumed"; + return false; + } + return true; +} + +bool WriteManifest(const std::filesystem::path &directory, + const std::vector &extents, + std::string *error) { + if (error == nullptr) { + return false; + } + std::vector normalized; + if (!NormalizeManifest(extents, &normalized, error)) { + return false; + } + + bool removed_temporary = false; + if (!RemoveTemporaryManifests(directory, &removed_temporary, error)) { + return false; + } + const std::string temporary_name = + std::string(kTemporaryManifestPrefix) + std::to_string(getpid()) + "." + + std::to_string(kTemporaryManifestSequence.fetch_add(1)); + const auto temporary = directory / temporary_name; + const auto manifest_path = directory / kManifestName; + std::ostringstream serialized; + serialized << "version 2\n"; + serialized << "device_count " << normalized.size() << "\n"; + for (size_t index = 0; index < normalized.size(); ++index) { + serialized << "device " << index << " " << normalized[index].source_uuid + << " " << normalized[index].size << " " + << normalized[index].filename << "\n"; + } + + FileDescriptor fd(open(temporary.c_str(), + O_WRONLY | O_CREAT | O_EXCL | O_CLOEXEC, 0600)); + if (fd.get() < 0) { + *error = "open temporary helper manifest: " + + std::string(std::strerror(errno)); + return false; + } + if (!WriteAll(fd.get(), serialized.str()) || fsync(fd.get()) != 0 || + !fd.Close()) { + const int write_error = errno == 0 ? EIO : errno; + RemoveTemporaryManifest(temporary); + *error = "write temporary helper manifest: " + + std::string(std::strerror(write_error)); + return false; + } + if (rename(temporary.c_str(), manifest_path.c_str()) != 0) { + const int rename_error = errno; + RemoveTemporaryManifest(temporary); + *error = "commit helper manifest: " + + std::string(std::strerror(rename_error)); + return false; + } + + FileDescriptor directory_fd( + open(directory.c_str(), O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + if (directory_fd.get() < 0) { + std::error_code ignored; + std::filesystem::remove(manifest_path, ignored); + *error = "open helper directory for fsync"; + return false; + } + if (fsync(directory_fd.get()) != 0 || !directory_fd.Close()) { + std::error_code ignored; + std::filesystem::remove(manifest_path, ignored); + *error = "fsync helper directory"; + return false; + } + return true; +} + +bool ReadManifest(const std::filesystem::path &directory, + std::vector *extents, std::string *error) { + if (extents == nullptr || error == nullptr) { + return false; + } + const auto manifest_path = directory / kManifestName; + struct stat manifest_stat{}; + if (lstat(manifest_path.c_str(), &manifest_stat) != 0 || + !S_ISREG(manifest_stat.st_mode)) { + *error = "helper manifest is missing or not a regular file"; + return false; + } + + std::ifstream input(manifest_path); + std::string key; + unsigned int version = 0; + if (!(input >> key >> version) || key != "version") { + *error = "invalid helper manifest version header"; + return false; + } + if (version == 1) { + *error = "unsafe helper manifest version 1 is not supported"; + return false; + } + if (version != 2) { + *error = "unsupported helper manifest version"; + return false; + } + + size_t device_count = 0; + if (!(input >> key >> device_count) || key != "device_count" || + device_count > kMaximumDeviceCount) { + *error = "invalid helper manifest device count"; + return false; + } + + std::vector parsed; + parsed.reserve(device_count); + for (size_t expected_index = 0; expected_index < device_count; + ++expected_index) { + size_t index = 0; + ManifestExtent extent; + if (!(input >> key >> index >> extent.source_uuid >> extent.size >> + extent.filename) || + key != "device" || index != expected_index) { + *error = "invalid helper manifest device entry"; + return false; + } + std::string canonical; + if (!CanonicalizeGPUUUID(extent.source_uuid, &canonical) || + canonical != extent.source_uuid) { + *error = "helper manifest source GPU UUID is not canonical"; + return false; + } + parsed.push_back(std::move(extent)); + } + std::string trailing; + if (input >> trailing) { + *error = "unexpected helper manifest data"; + return false; + } + + std::vector normalized; + if (!NormalizeManifest(parsed, &normalized, error)) { + return false; + } + *extents = std::move(normalized); + return true; +} + +bool ValidateExtentFiles(const std::filesystem::path &directory, + const std::vector &extents, + std::string *error) { + if (error == nullptr) { + return false; + } + std::vector normalized; + if (!NormalizeManifest(extents, &normalized, error)) { + return false; + } + for (const auto &extent : normalized) { + struct stat extent_stat{}; + const auto path = directory / extent.filename; + if (lstat(path.c_str(), &extent_stat) != 0 || + !S_ISREG(extent_stat.st_mode) || extent_stat.st_size < 0 || + static_cast(extent_stat.st_size) != extent.size) { + *error = "extent file is missing, invalid, or has the wrong size"; + return false; + } + } + return true; +} + +bool RemoveManifest(const std::filesystem::path &directory, + std::string *error) { + if (error == nullptr) { + return false; + } + bool removed = false; + if (unlink((directory / kManifestName).c_str()) == 0) { + removed = true; + } else if (errno != ENOENT) { + *error = + "remove stale helper manifest: " + std::string(std::strerror(errno)); + return false; + } + if (!RemoveTemporaryManifests(directory, &removed, error)) { + return false; + } + if (removed) { + FileDescriptor directory_fd( + open(directory.c_str(), O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + if (directory_fd.get() < 0 || fsync(directory_fd.get()) != 0 || + !directory_fd.Close()) { + *error = "fsync helper directory after removing manifest"; + return false; + } + } + return true; +} + +} // namespace cuda_checkpoint_storage diff --git a/agent/cmd/cuda-checkpoint-helper/storage_manifest.h b/agent/cmd/cuda-checkpoint-helper/storage_manifest.h new file mode 100644 index 00000000..58bdcc7b --- /dev/null +++ b/agent/cmd/cuda-checkpoint-helper/storage_manifest.h @@ -0,0 +1,72 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace cuda_checkpoint_storage { + +constexpr const char *kManifestName = "manifest.txt"; +constexpr const char *kLegacyTemporaryManifestName = "manifest.txt.tmp"; +constexpr const char *kTemporaryManifestPrefix = "manifest.txt.tmp."; + +struct ManifestExtent { + std::string source_uuid; + size_t size; + std::string filename; +}; + +struct DeviceExtent { + std::string uuid; + size_t size; +}; + +struct DevicePair { + std::string source_uuid; + std::string destination_uuid; +}; + +struct TransferJob { + size_t device_index; + size_t extent_index; +}; + +bool ParseGPUUUID(std::string_view value, + std::array *bytes_out); +std::string FormatGPUUUID(const std::array &bytes); +bool CanonicalizeGPUUUID(std::string_view value, std::string *canonical_out); + +std::string DeviceFilename(size_t index); + +bool BuildCheckpointManifest(const std::vector &devices, + std::vector *extents, + std::string *error); + +// Builds jobs in current local-device order. A nonempty device map is +// interpreted as source->destination and reversed to recover each destination's +// source UUID. Extra pairs are permitted because a process may export only a +// subset of the GPUs assigned to its container. +bool BuildTransferJobs(const std::vector &extents, + const std::vector &devices, + const std::vector &device_pairs, + std::vector *jobs, std::string *error); + +bool WriteManifest(const std::filesystem::path &directory, + const std::vector &extents, + std::string *error); +bool ReadManifest(const std::filesystem::path &directory, + std::vector *extents, std::string *error); +bool ValidateExtentFiles(const std::filesystem::path &directory, + const std::vector &extents, + std::string *error); +bool RemoveManifest(const std::filesystem::path &directory, std::string *error); + +} // namespace cuda_checkpoint_storage diff --git a/agent/cmd/cuda-checkpoint-helper/storage_manifest_test.cpp b/agent/cmd/cuda-checkpoint-helper/storage_manifest_test.cpp new file mode 100644 index 00000000..e2fe3b69 --- /dev/null +++ b/agent/cmd/cuda-checkpoint-helper/storage_manifest_test.cpp @@ -0,0 +1,336 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include "storage_manifest.h" + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace storage = cuda_checkpoint_storage; + +namespace { + +constexpr const char *kSourceA = "GPU-00000000-0000-0000-0000-00000000000a"; +constexpr const char *kSourceB = "GPU-00000000-0000-0000-0000-00000000000b"; +constexpr const char *kSourceFallback = + "GPU-00000000-0000-0000-0000-00000000000c"; +constexpr const char *kDestinationA = + "GPU-10000000-0000-0000-0000-00000000000a"; +constexpr const char *kDestinationB = + "GPU-10000000-0000-0000-0000-00000000000b"; +constexpr const char *kDestinationFallback = + "GPU-10000000-0000-0000-0000-00000000000c"; + +bool Check(bool condition, const std::string &message) { + if (!condition) { + std::cerr << message << "\n"; + } + return condition; +} + +bool TestGPUUUIDParsing() { + std::array parsed{}; + std::string canonical; + return Check(storage::ParseGPUUUID(kSourceA, &parsed), + "canonical GPU UUID was rejected") && + Check(storage::FormatGPUUUID(parsed) == kSourceA, + "GPU UUID did not round-trip") && + Check(storage::CanonicalizeGPUUUID( + "00000000-0000-0000-0000-00000000000A", &canonical) && + canonical == kSourceA, + "bare uppercase GPU UUID was not canonicalized") && + Check(!storage::ParseGPUUUID( + "GPU-00000000-0000-0000-0000-00000000000g", &parsed), + "non-hex GPU UUID was accepted") && + Check(!storage::ParseGPUUUID("GPU-0000", &parsed), + "short GPU UUID was accepted"); +} + +bool TestEqualSizeNonOrderPreservingMap() { + const std::vector extents{ + {kSourceA, 4096, storage::DeviceFilename(0)}, + {kSourceB, 4096, storage::DeviceFilename(1)}, + }; + // CUDA returns destination B first. Equal sizes must not permit an + // index-based A/B swap. + const std::vector destinations{ + {kDestinationB, 4096}, + {kDestinationA, 4096}, + }; + const std::vector pairs{ + {kSourceA, kDestinationA}, + {kSourceB, kDestinationB}, + // This GPU is assigned to the container but is not exported by this + // process. Its explicit fallback pair must not invalidate the subset. + {kSourceFallback, kDestinationFallback}, + }; + + std::vector jobs; + std::string error; + return Check(storage::BuildTransferJobs(extents, destinations, pairs, &jobs, + &error), + error) && + Check(jobs.size() == 2, "expected two transfer jobs") && + Check(jobs[0].device_index == 0 && jobs[0].extent_index == 1, + "destination B was not matched to source B's deterministic " + "file") && + Check( + jobs[1].device_index == 1 && jobs[1].extent_index == 0, + "destination A was not matched to source A's deterministic file"); +} + +bool TestEmptyV2Manifest() { + char path[] = "/tmp/cuda-storage-manifest-test-XXXXXX"; + const char *directory = mkdtemp(path); + if (!Check(directory != nullptr, "mkdtemp failed")) { + return false; + } + + std::string error; + std::vector loaded; + std::vector jobs; + const bool result = + Check(storage::WriteManifest(directory, {}, &error), error) && + Check(storage::ReadManifest(directory, &loaded, &error), error) && + Check(loaded.empty(), "empty v2 manifest did not round-trip") && + Check(storage::BuildTransferJobs(loaded, {}, {}, &jobs, &error), error) && + Check(jobs.empty(), "zero-device restore produced transfer jobs"); + std::error_code ignored; + std::filesystem::remove_all(directory, ignored); + return result; +} + +bool TestNonemptyV2ManifestRoundTrip() { + char path[] = "/tmp/cuda-storage-manifest-roundtrip-test-XXXXXX"; + const char *directory = mkdtemp(path); + if (!Check(directory != nullptr, "mkdtemp failed")) { + return false; + } + const std::vector extents{ + {kSourceA, 4096, storage::DeviceFilename(0)}, + {kSourceB, 8192, storage::DeviceFilename(1)}, + }; + + std::string error; + std::vector loaded; + const auto manifest_path = + std::filesystem::path(directory) / storage::kManifestName; + const bool result = + Check(storage::WriteManifest(directory, extents, &error), error) && + Check((std::filesystem::status(manifest_path).permissions() & + std::filesystem::perms::all) == + (std::filesystem::perms::owner_read | + std::filesystem::perms::owner_write), + "committed manifest permissions are not 0600") && + Check(storage::ReadManifest(directory, &loaded, &error), error) && + Check(loaded.size() == 2 && loaded[0].source_uuid == kSourceA && + loaded[0].size == 4096 && + loaded[0].filename == "device-0000.bin" && + loaded[1].source_uuid == kSourceB && loaded[1].size == 8192 && + loaded[1].filename == "device-0001.bin", + "nonempty v2 manifest did not preserve UUID, size, and filename"); + std::error_code ignored; + std::filesystem::remove_all(directory, ignored); + return result; +} + +bool TestV1Rejected() { + char path[] = "/tmp/cuda-storage-manifest-v1-test-XXXXXX"; + const char *directory = mkdtemp(path); + if (!Check(directory != nullptr, "mkdtemp failed")) { + return false; + } + { + std::ofstream output(std::filesystem::path(directory) / + storage::kManifestName); + output << "version 1\n" + "device_count 1\n" + "device 0 4096 device-0000.bin\n"; + } + + std::vector extents; + std::string error; + const bool result = Check(!storage::ReadManifest(directory, &extents, &error), + "unsafe v1 manifest was accepted") && + Check(error.find("version 1") != std::string::npos, + "v1 rejection did not identify the unsafe version"); + std::error_code ignored; + std::filesystem::remove_all(directory, ignored); + return result; +} + +bool TestUnconsumedExtentRejected() { + const std::vector extents{ + {kSourceA, 4096, storage::DeviceFilename(0)}, + {kSourceB, 4096, storage::DeviceFilename(1)}, + }; + const std::vector destinations{{kSourceA, 4096}}; + std::vector jobs; + std::string error; + return Check( + !storage::BuildTransferJobs(extents, destinations, {}, &jobs, &error), + "restore accepted an unconsumed saved extent"); +} + +bool TestUnsafeMappingsRejected() { + const std::vector extents{ + {kSourceA, 4096, storage::DeviceFilename(0)}, + {kSourceB, 4096, storage::DeviceFilename(1)}, + }; + std::vector jobs; + std::string error; + + if (!Check(!storage::BuildTransferJobs( + extents, {{kDestinationA, 4096}, {kDestinationB, 4096}}, + {{kSourceA, kDestinationA}}, &jobs, &error), + "restore accepted a destination missing from the device map")) { + return false; + } + if (!Check(!storage::BuildTransferJobs( + extents, {{kDestinationA, 4096}, {kDestinationB, 4096}}, + {{kSourceA, kDestinationA}, {kSourceB, kDestinationA}}, &jobs, + &error), + "restore accepted an ambiguous destination UUID")) { + return false; + } + if (!Check(!storage::BuildTransferJobs( + extents, {{kDestinationA, 4096}, {kDestinationB, 8192}}, + {{kSourceA, kDestinationA}, {kSourceB, kDestinationB}}, &jobs, + &error), + "restore accepted a UUID-matched extent with the wrong size")) { + return false; + } + if (!Check(!storage::BuildTransferJobs( + {{kSourceA, 4096, storage::DeviceFilename(0)}, + {kSourceA, 4096, storage::DeviceFilename(1)}}, + {{kSourceA, 4096}, {kSourceB, 4096}}, {}, &jobs, &error), + "restore accepted duplicate saved source UUIDs")) { + return false; + } + return Check(!storage::BuildTransferJobs(extents, + {{kSourceA, 4096}, {kSourceA, 4096}}, + {}, &jobs, &error), + "restore accepted duplicate destination UUIDs"); +} + +bool TestDuplicateCheckpointUUIDRejected() { + std::vector extents; + std::string error; + return Check(!storage::BuildCheckpointManifest( + {{kSourceA, 4096}, {kSourceA, 4096}}, &extents, &error), + "checkpoint accepted duplicate source UUIDs"); +} + +bool TestWrongDeterministicFilenameRejected() { + std::vector jobs; + std::string error; + return Check(!storage::BuildTransferJobs( + {{kSourceA, 4096, "device-0001.bin"}}, + {{kSourceA, 4096}}, {}, &jobs, &error), + "manifest accepted an extent with the wrong deterministic " + "filename"); +} + +bool TestValidateExtentFiles() { + char path[] = "/tmp/cuda-storage-extent-test-XXXXXX"; + const char *directory = mkdtemp(path); + if (!Check(directory != nullptr, "mkdtemp failed")) { + return false; + } + const std::filesystem::path extent_path = + std::filesystem::path(directory) / storage::DeviceFilename(0); + { + std::ofstream extent(extent_path, std::ios::binary); + extent << "bad"; + } + const std::vector extents{ + {kSourceA, 4, storage::DeviceFilename(0)}, + }; + std::string error; + const bool rejected_wrong_size = + !storage::ValidateExtentFiles(directory, extents, &error); + const bool resized = truncate(extent_path.c_str(), 4) == 0; + const bool accepted_exact_size = + resized && storage::ValidateExtentFiles(directory, extents, &error); + std::error_code ignored; + std::filesystem::remove_all(directory, ignored); + return Check(rejected_wrong_size, + "ValidateExtentFiles accepted an incorrect extent size") && + Check(resized, "failed to resize extent fixture") && + Check(accepted_exact_size, + "ValidateExtentFiles rejected the exact extent size"); +} + +bool TestRemoveManifest() { + char path[] = "/tmp/cuda-storage-remove-test-XXXXXX"; + const char *directory = mkdtemp(path); + if (!Check(directory != nullptr, "mkdtemp failed")) { + return false; + } + const auto manifest = std::filesystem::path(directory) / "manifest.txt"; + const auto temporary = std::filesystem::path(directory) / + storage::kLegacyTemporaryManifestName; + const auto unique_temporary = std::filesystem::path(directory) / + (std::string(storage::kTemporaryManifestPrefix) + + "123.456"); + { + std::ofstream(manifest) << "manifest"; + std::ofstream(temporary) << "temporary"; + std::ofstream(unique_temporary) << "temporary"; + } + std::string error; + const bool first = storage::RemoveManifest(directory, &error); + const bool removed = !std::filesystem::exists(manifest) && + !std::filesystem::exists(temporary) && + !std::filesystem::exists(unique_temporary); + const bool second = storage::RemoveManifest(directory, &error); + std::error_code ignored; + std::filesystem::remove_all(directory, ignored); + return Check(first, error) && Check(removed, "manifest files remain") && + Check(second, "repeated RemoveManifest failed"); +} + +bool TestStaleTemporaryManifestDoesNotBlockWrite() { + char path[] = "/tmp/cuda-storage-stale-temporary-test-XXXXXX"; + const char *directory = mkdtemp(path); + if (!Check(directory != nullptr, "mkdtemp failed")) { + return false; + } + const auto stale = std::filesystem::path(directory) / + (std::string(storage::kTemporaryManifestPrefix) + + "111.222"); + std::ofstream(stale) << "stale"; + std::string error; + const bool wrote = storage::WriteManifest(directory, {}, &error); + const bool cleaned = !std::filesystem::exists(stale); + std::error_code ignored; + std::filesystem::remove_all(directory, ignored); + return Check(wrote, error) && + Check(cleaned, "stale temporary manifest was not removed"); +} + +} // namespace + +int main() { + if (!TestGPUUUIDParsing() || !TestEqualSizeNonOrderPreservingMap() || + !TestEmptyV2Manifest() || + !TestNonemptyV2ManifestRoundTrip() || !TestV1Rejected() || + !TestUnconsumedExtentRejected() || !TestUnsafeMappingsRejected() || + !TestDuplicateCheckpointUUIDRejected() || + !TestWrongDeterministicFilenameRejected() || + !TestValidateExtentFiles() || !TestRemoveManifest() || + !TestStaleTemporaryManifestDoesNotBlockWrite()) { + return 1; + } + std::cout << "cuda checkpoint storage manifest tests passed\n"; + return 0; +} diff --git a/agent/cmd/cuda-checkpoint-helper/testdata/daemon_request_v6.hex b/agent/cmd/cuda-checkpoint-helper/testdata/daemon_request_v6.hex new file mode 100644 index 00000000..ab737472 --- /dev/null +++ b/agent/cmd/cuda-checkpoint-helper/testdata/daemon_request_v6.hex @@ -0,0 +1 @@ +4443485006003800020002002a000000020000000000800000000000510000001d0000001200000039300000000000001f000000280000004750552d61616161616161612d626262622d636363632d646464642d6565656565656565656565653d4750552d31313131313131312d323232322d333333332d343434342d3535353535353535353535352f636865636b706f696e74732f70726f636573732d6e737069642d3432303a3a2f6b756265706f64732f746573740a2f686f73742f70726f632f34322f726f6f742f746d702f637564612d6a6f624750552d31323334353637382d313233342d313233342d313233342d313233343536373839616263 diff --git a/agent/cmd/cuda-checkpoint-helper/transfer_backend_unavailable.cpp b/agent/cmd/cuda-checkpoint-helper/transfer_backend_unavailable.cpp new file mode 100644 index 00000000..1f94cf2c --- /dev/null +++ b/agent/cmd/cuda-checkpoint-helper/transfer_backend_unavailable.cpp @@ -0,0 +1,28 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include "transfer_engine.h" + +namespace cuda_checkpoint_transfer { + +bool TransferBackendAvailable() { return false; } + +bool TransferExtent(CUdeviceptr, size_t, CUstream, CUcontext, + const StorageLayout &, TransferOperation, + const TransferOptions &, TransferCancellation *cancellation, + TransferMetrics *metrics, std::string *error) { + if (metrics != nullptr) { + *metrics = {}; + } + if (error != nullptr) { + *error = "no CustomStorage transfer backend is linked"; + } + if (cancellation != nullptr) { + cancellation->Cancel(); + } + return false; +} + +} // namespace cuda_checkpoint_transfer diff --git a/agent/cmd/cuda-checkpoint-helper/transfer_cancellation.h b/agent/cmd/cuda-checkpoint-helper/transfer_cancellation.h new file mode 100644 index 00000000..9de92dac --- /dev/null +++ b/agent/cmd/cuda-checkpoint-helper/transfer_cancellation.h @@ -0,0 +1,32 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include + +namespace cuda_checkpoint_transfer { + +class TransferCancellation { +public: + using Clock = std::chrono::steady_clock; + + TransferCancellation() = default; + explicit TransferCancellation(Clock::time_point deadline) + : deadline_(deadline) {} + + void Cancel() { cancelled_.store(true, std::memory_order_relaxed); } + bool DeadlineExceeded() const { return Clock::now() >= deadline_; } + bool IsCancelled() const { + return cancelled_.load(std::memory_order_relaxed) || DeadlineExceeded(); + } + +private: + std::atomic cancelled_{false}; + Clock::time_point deadline_ = Clock::time_point::max(); +}; + +} // namespace cuda_checkpoint_transfer diff --git a/agent/cmd/cuda-checkpoint-helper/transfer_config.cpp b/agent/cmd/cuda-checkpoint-helper/transfer_config.cpp new file mode 100644 index 00000000..f136daed --- /dev/null +++ b/agent/cmd/cuda-checkpoint-helper/transfer_config.cpp @@ -0,0 +1,284 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include "transfer_config.h" + +#include + +#include +#include +#include +#include +#include +#include + +namespace cuda_checkpoint_transfer { +namespace { + +bool CheckedAdd(size_t left, size_t right, size_t *result) { + if (right > std::numeric_limits::max() - left) { + return false; + } + *result = left + right; + return true; +} + +bool CheckedMultiply(size_t left, size_t right, size_t *result) { + if (left != 0 && right > std::numeric_limits::max() / left) { + return false; + } + *result = left * right; + return true; +} + +} // namespace + +bool ParseSize(std::string_view value, size_t *parsed) { + if (parsed == nullptr || value.empty()) { + return false; + } + size_t result = 0; + const auto conversion = + std::from_chars(value.data(), value.data() + value.size(), result); + if (conversion.ec != std::errc{} || + conversion.ptr != value.data() + value.size()) { + return false; + } + *parsed = result; + return true; +} + +bool ValidateTransferOptions(const TransferOptions &options, + std::string *error) { + if (error == nullptr) { + return false; + } + if (options.buffer_count == 0 || options.buffer_count > kMaximumBufferCount) { + *error = "transfer buffer count must be between 1 and " + + std::to_string(kMaximumBufferCount); + return false; + } + if (options.chunk_bytes < kMinimumChunkBytes || + options.chunk_bytes > kMaximumChunkBytes || + options.chunk_bytes % kBufferAlignment != 0) { + *error = "transfer chunk bytes must be a 4096-byte multiple between " + + std::to_string(kMinimumChunkBytes) + " and " + + std::to_string(kMaximumChunkBytes); + return false; + } + size_t pinned_bytes = 0; + if (!CheckedMultiply(options.buffer_count, options.chunk_bytes, + &pinned_bytes) || + pinned_bytes > kMaximumPinnedBytesPerDevice) { + *error = "transfer buffers exceed the 1 GiB per-device pinned-memory limit"; + return false; + } + return true; +} + +bool CalculatePinnedBytes(size_t device_count, const TransferOptions &options, + size_t *bytes, std::string *error) { + if (bytes == nullptr || error == nullptr || + !ValidateTransferOptions(options, error)) { + return false; + } + size_t per_device = 0; + if (!CheckedMultiply(options.buffer_count, options.chunk_bytes, + &per_device) || + !CheckedMultiply(device_count, per_device, bytes)) { + *error = "pinned-memory size calculation overflow"; + return false; + } + if (*bytes > kMaximumPinnedBytesPerOperation) { + *error = + "transfer buffers exceed the 2 GiB per-operation pinned-memory limit"; + return false; + } + return true; +} + +int StorageFileOpenFlags(TransferOperation operation) { + const int access = operation == TransferOperation::kCheckpoint + ? O_RDWR | O_CREAT | O_TRUNC + : O_RDONLY; + return access | O_CLOEXEC | O_NOFOLLOW; +} + +bool BuildTransferChunks(size_t extent_size, const StorageLayout &storage, + const TransferOptions &options, + std::vector *chunks, + std::string *error) { + if (chunks == nullptr || error == nullptr || + !ValidateTransferOptions(options, error)) { + return false; + } + if (extent_size == 0 || storage.files.empty() || storage.ranges.empty()) { + *error = "transfer extent and storage layout must be nonempty"; + return false; + } + + std::set paths; + for (const auto &file : storage.files) { + if (file.path.empty() || !file.path.is_absolute() || file.size == 0) { + *error = "storage files must have absolute paths and nonzero sizes"; + return false; + } + if (!paths.insert(file.path.lexically_normal()).second) { + *error = "storage file paths must be unique"; + return false; + } + } + + chunks->clear(); + size_t expected_logical_offset = 0; + size_t chunk_index = 0; + std::vector>> physical_ranges( + storage.files.size()); + for (const auto &range : storage.ranges) { + if (range.size == 0 || range.logical_offset != expected_logical_offset || + range.file_index >= storage.files.size()) { + *error = "storage ranges must be nonempty and exactly cover the logical " + "extent in order"; + return false; + } + size_t logical_end = 0; + size_t file_end = 0; + if (!CheckedAdd(range.logical_offset, range.size, &logical_end) || + !CheckedAdd(range.file_offset, range.size, &file_end)) { + *error = "storage range offset calculation overflow"; + return false; + } + if (logical_end > extent_size || + file_end > storage.files[range.file_index].size) { + *error = "storage range exceeds its logical extent or physical file"; + return false; + } + physical_ranges[range.file_index].push_back({range.file_offset, file_end}); + + const size_t range_chunk_count = + range.size / options.chunk_bytes + + static_cast(range.size % options.chunk_bytes != 0); + if (range_chunk_count > kMaximumTransferChunkCount - chunks->size()) { + *error = "transfer layout has too many chunks"; + return false; + } + for (size_t range_offset = 0; range_offset < range.size;) { + const size_t length = + std::min(options.chunk_bytes, range.size - range_offset); + chunks->push_back({range.logical_offset + range_offset, length, + range.file_index, range.file_offset + range_offset, + chunk_index % options.buffer_count}); + range_offset += length; + ++chunk_index; + } + expected_logical_offset = logical_end; + } + if (expected_logical_offset != extent_size) { + *error = "storage ranges do not cover the complete logical extent"; + return false; + } + for (auto &ranges : physical_ranges) { + std::sort(ranges.begin(), ranges.end()); + } + for (size_t file_index = 0; file_index < physical_ranges.size(); + ++file_index) { + const auto &ranges = physical_ranges[file_index]; + if (ranges.empty() || ranges.front().first != 0) { + *error = "storage ranges must exactly cover every physical file"; + return false; + } + for (size_t index = 1; index < ranges.size(); ++index) { + if (ranges[index].first != ranges[index - 1].second) { + *error = "storage ranges must not overlap or leave gaps within a " + "physical file"; + return false; + } + } + if (ranges.back().second != storage.files[file_index].size) { + *error = "storage ranges must exactly cover every physical file"; + return false; + } + } + return true; +} + +bool BuildContiguousStorageLayout(const std::filesystem::path &base_path, + size_t extent_size, size_t file_count, + StorageLayout *storage, std::string *error) { + if (storage == nullptr || error == nullptr) { + return false; + } + if (base_path.empty() || !base_path.is_absolute() || extent_size == 0 || + file_count == 0 || file_count > 64 || file_count > extent_size) { + *error = "storage base path must be absolute and file count must be " + "between 1 and 64 without exceeding the transfer byte count"; + return false; + } + + storage->files.clear(); + storage->ranges.clear(); + storage->files.reserve(file_count); + storage->ranges.reserve(file_count); + const size_t base_size = extent_size / file_count; + const size_t remainder = extent_size % file_count; + size_t logical_offset = 0; + for (size_t index = 0; index < file_count; ++index) { + const size_t size = base_size + (index < remainder ? 1 : 0); + std::filesystem::path path = base_path; + if (file_count > 1) { + char suffix[32]; + std::snprintf(suffix, sizeof(suffix), ".part-%04zu", index); + path += suffix; + } + storage->files.push_back({std::move(path), size}); + storage->ranges.push_back({logical_offset, size, index, 0}); + logical_offset += size; + } + return true; +} + +std::string JsonEscape(std::string_view value) { + static constexpr char kHex[] = "0123456789abcdef"; + std::string result; + result.reserve(value.size() + 2); + result.push_back('"'); + for (const unsigned char byte : value) { + switch (byte) { + case '"': + result += "\\\""; + break; + case '\\': + result += "\\\\"; + break; + case '\b': + result += "\\b"; + break; + case '\f': + result += "\\f"; + break; + case '\n': + result += "\\n"; + break; + case '\r': + result += "\\r"; + break; + case '\t': + result += "\\t"; + break; + default: + if (byte < 0x20 || byte >= 0x80) { + result += "\\u00"; + result.push_back(kHex[byte >> 4]); + result.push_back(kHex[byte & 0x0f]); + } else { + result.push_back(static_cast(byte)); + } + } + } + result.push_back('"'); + return result; +} + +} // namespace cuda_checkpoint_transfer diff --git a/agent/cmd/cuda-checkpoint-helper/transfer_config.h b/agent/cmd/cuda-checkpoint-helper/transfer_config.h new file mode 100644 index 00000000..04e4847b --- /dev/null +++ b/agent/cmd/cuda-checkpoint-helper/transfer_config.h @@ -0,0 +1,78 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include +#include +#include + +namespace cuda_checkpoint_transfer { + +enum class TransferOperation { + kCheckpoint, + kRestore, +}; + +constexpr size_t kDefaultBufferCount = 1; +constexpr size_t kDefaultChunkBytes = 64ULL * 1024ULL * 1024ULL; +constexpr size_t kMinimumChunkBytes = 1ULL * 1024ULL * 1024ULL; +constexpr size_t kMaximumChunkBytes = 256ULL * 1024ULL * 1024ULL; +constexpr size_t kMaximumBufferCount = 8; +constexpr size_t kMaximumPinnedBytesPerDevice = + 1ULL * 1024ULL * 1024ULL * 1024ULL; +constexpr size_t kMaximumPinnedBytesPerOperation = + 2ULL * 1024ULL * 1024ULL * 1024ULL; +constexpr size_t kMaximumTransferChunkCount = 1024ULL * 1024ULL; +constexpr size_t kBufferAlignment = 4096; + +struct TransferOptions { + size_t buffer_count = kDefaultBufferCount; + size_t chunk_bytes = kDefaultChunkBytes; +}; + +struct StorageFile { + std::filesystem::path path; + size_t size = 0; +}; + +struct StorageRange { + size_t logical_offset = 0; + size_t size = 0; + size_t file_index = 0; + size_t file_offset = 0; +}; + +struct StorageLayout { + std::vector files; + std::vector ranges; +}; + +struct TransferChunk { + size_t logical_offset; + size_t size; + size_t file_index; + size_t file_offset; + size_t slot_index; +}; + +bool ParseSize(std::string_view value, size_t *parsed); +bool ValidateTransferOptions(const TransferOptions &options, + std::string *error); +bool CalculatePinnedBytes(size_t device_count, const TransferOptions &options, + size_t *bytes, std::string *error); +int StorageFileOpenFlags(TransferOperation operation); +bool BuildTransferChunks(size_t extent_size, const StorageLayout &storage, + const TransferOptions &options, + std::vector *chunks, + std::string *error); +bool BuildContiguousStorageLayout(const std::filesystem::path &base_path, + size_t extent_size, size_t file_count, + StorageLayout *storage, std::string *error); +std::string JsonEscape(std::string_view value); + +} // namespace cuda_checkpoint_transfer diff --git a/agent/cmd/cuda-checkpoint-helper/transfer_config_test.cpp b/agent/cmd/cuda-checkpoint-helper/transfer_config_test.cpp new file mode 100644 index 00000000..331dffcc --- /dev/null +++ b/agent/cmd/cuda-checkpoint-helper/transfer_config_test.cpp @@ -0,0 +1,208 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include "transfer_config.h" + +#include + +#include +#include +#include +#include + +namespace transfer = cuda_checkpoint_transfer; + +namespace { + +bool Check(bool condition, const std::string &message) { + if (!condition) { + std::cerr << message << "\n"; + } + return condition; +} + +bool TestStorageOpenModePolicy() { + const int restore_flags = + transfer::StorageFileOpenFlags(transfer::TransferOperation::kRestore); + const int checkpoint_flags = + transfer::StorageFileOpenFlags(transfer::TransferOperation::kCheckpoint); + return Check((restore_flags & O_ACCMODE) == O_RDONLY, + "restore storage is not opened read-only") && + Check((restore_flags & (O_CREAT | O_TRUNC)) == 0, + "restore storage can be created or truncated") && + Check((restore_flags & (O_CLOEXEC | O_NOFOLLOW)) == + (O_CLOEXEC | O_NOFOLLOW), + "restore storage lacks descriptor safeguards") && + Check((checkpoint_flags & O_ACCMODE) == O_RDWR, + "checkpoint storage is not opened writable") && + Check((checkpoint_flags & (O_CREAT | O_TRUNC)) == (O_CREAT | O_TRUNC), + "checkpoint storage is not created and truncated") && + Check((checkpoint_flags & (O_CLOEXEC | O_NOFOLLOW)) == + (O_CLOEXEC | O_NOFOLLOW), + "checkpoint storage lacks descriptor safeguards"); +} + +bool TestOptionParsingAndBounds() { + size_t parsed = 0; + std::string error; + return Check(transfer::ParseSize("67108864", &parsed) && parsed == 67108864, + "valid size was not parsed") && + Check(!transfer::ParseSize("-1", &parsed), + "negative size was accepted") && + Check(!transfer::ParseSize("64MiB", &parsed), + "size with trailing data was accepted") && + Check(transfer::ValidateTransferOptions({transfer::kDefaultBufferCount, + transfer::kDefaultChunkBytes}, + &error), + error) && + Check(!transfer::ValidateTransferOptions( + {0, transfer::kDefaultChunkBytes}, &error), + "zero slots accepted") && + Check(!transfer::ValidateTransferOptions( + {1, transfer::kMinimumChunkBytes + 1}, &error), + "unaligned chunk accepted") && + Check( + !transfer::ValidateTransferOptions( + {transfer::kMaximumBufferCount, transfer::kMaximumChunkBytes}, + &error), + "more than 1 GiB per device was accepted"); +} + +bool TestPinnedMemoryCalculation() { + size_t bytes = 0; + std::string error; + return Check(transfer::CalculatePinnedBytes( + 8, {1, transfer::kDefaultChunkBytes}, &bytes, &error), + error) && + Check(bytes == 512ULL * 1024ULL * 1024ULL, + "default eight-device pinned bytes are wrong") && + Check(!transfer::CalculatePinnedBytes( + 3, + {transfer::kMaximumBufferCount, 128ULL * 1024ULL * 1024ULL}, + &bytes, &error), + "operation pinned-memory cap was not enforced") && + Check(!transfer::CalculatePinnedBytes( + std::numeric_limits::max(), + {1, transfer::kDefaultChunkBytes}, &bytes, &error), + "pinned-memory overflow was accepted"); +} + +bool TestChunkRingAndShardedLayout() { + transfer::StorageLayout storage; + std::vector chunks; + std::string error; + constexpr size_t kMiB = 1024ULL * 1024ULL; + if (!Check(transfer::BuildContiguousStorageLayout("/tmp/extent", 130 * kMiB, + 2, &storage, &error), + error) || + !Check(storage.files.size() == 2 && storage.ranges.size() == 2, + "expected two files and ranges") || + !Check(storage.files[0].size == 65 * kMiB && + storage.files[0].path == "/tmp/extent.part-0000" && + storage.files[1].size == 65 * kMiB && + storage.files[1].path == "/tmp/extent.part-0001" && + storage.ranges[0].logical_offset == 0 && + storage.ranges[0].size == 65 * kMiB && + storage.ranges[0].file_index == 0 && + storage.ranges[1].logical_offset == 65 * kMiB && + storage.ranges[1].size == 65 * kMiB && + storage.ranges[1].file_index == 1, + "contiguous storage range mapping is wrong") || + !Check(transfer::BuildTransferChunks(130 * kMiB, storage, {2, 64 * kMiB}, + &chunks, &error), + error) || + !Check(chunks.size() == 4, + "expected shard boundaries to split the four chunks")) { + return false; + } + + return Check(chunks[0].logical_offset == 0 && chunks[0].size == 64 * kMiB && + chunks[0].file_index == 0 && chunks[0].file_offset == 0 && + chunks[0].slot_index == 0, + "first chunk is wrong") && + Check(chunks[1].logical_offset == 64 * kMiB && + chunks[1].size == kMiB && chunks[1].file_index == 0 && + chunks[1].file_offset == 64 * kMiB && + chunks[1].slot_index == 1, + "first shard tail is wrong") && + Check(chunks[2].logical_offset == 65 * kMiB && + chunks[2].size == 64 * kMiB && chunks[2].file_index == 1 && + chunks[2].file_offset == 0 && chunks[2].slot_index == 0, + "second shard first chunk is wrong") && + Check(chunks[3].logical_offset == 129 * kMiB && + chunks[3].size == kMiB && chunks[3].file_index == 1 && + chunks[3].file_offset == 64 * kMiB && + chunks[3].slot_index == 1, + "second shard tail is wrong"); +} + +bool TestRelativeStorageLayoutRejected() { + transfer::StorageLayout storage; + std::string error; + return Check(!transfer::BuildContiguousStorageLayout( + "relative/extent", 4096, 1, &storage, &error), + "relative storage base path was accepted"); +} + +bool TestLayoutGapsAndOverflowRejected() { + std::vector chunks; + std::string error; + const transfer::StorageLayout gap_layout{ + {{"/tmp/a", 4096}, {"/tmp/b", 4095}}, + {{0, 4096, 0, 0}, {4097, 4095, 1, 0}}, + }; + const transfer::StorageLayout overflow_layout{ + {{"/tmp/a", std::numeric_limits::max()}}, + {{0, std::numeric_limits::max(), 0, 1}}, + }; + const size_t excessive_size = + (transfer::kMaximumTransferChunkCount + 1) * transfer::kMinimumChunkBytes; + const transfer::StorageLayout excessive_layout{ + {{"/tmp/a", excessive_size}}, + {{0, excessive_size, 0, 0}}, + }; + const transfer::StorageLayout duplicate_file_layout{ + {{"/tmp/a", 4096}, {"/tmp/../tmp/a", 4096}}, + {{0, 4096, 0, 0}, {4096, 4096, 1, 0}}, + }; + return Check(!transfer::BuildTransferChunks(8192, gap_layout, + {1, transfer::kMinimumChunkBytes}, + &chunks, &error), + "logical layout gap was accepted") && + Check(!transfer::BuildTransferChunks( + std::numeric_limits::max(), overflow_layout, + {1, transfer::kMinimumChunkBytes}, &chunks, &error), + "file offset overflow was accepted") && + Check(!transfer::BuildTransferChunks(excessive_size, excessive_layout, + {1, transfer::kMinimumChunkBytes}, + &chunks, &error), + "excessive transfer chunk count was accepted") && + Check(!transfer::BuildTransferChunks(8192, duplicate_file_layout, + {1, transfer::kMinimumChunkBytes}, + &chunks, &error), + "duplicate physical file path was accepted"); +} + +bool TestJsonEscaping() { + return Check(transfer::JsonEscape("file\\name\n\"value\"") == + "\"file\\\\name\\n\\\"value\\\"\"", + "JSON escaping is wrong") && + Check(transfer::JsonEscape(std::string("bad\xff", 4)) == + "\"bad\\u00ff\"", + "non-UTF-8 byte was emitted into JSON"); +} + +} // namespace + +int main() { + if (!TestOptionParsingAndBounds() || !TestPinnedMemoryCalculation() || + !TestStorageOpenModePolicy() || !TestChunkRingAndShardedLayout() || + !TestRelativeStorageLayoutRejected() || + !TestLayoutGapsAndOverflowRejected() || !TestJsonEscaping()) { + return 1; + } + std::cout << "CUDA transfer configuration tests passed\n"; + return 0; +} diff --git a/agent/cmd/cuda-checkpoint-helper/transfer_engine.h b/agent/cmd/cuda-checkpoint-helper/transfer_engine.h new file mode 100644 index 00000000..9a446b47 --- /dev/null +++ b/agent/cmd/cuda-checkpoint-helper/transfer_engine.h @@ -0,0 +1,47 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +#include +#include +#include + +#include "transfer_cancellation.h" +#include "transfer_config.h" + +namespace cuda_checkpoint_transfer { + +struct StorageFileMetrics { + size_t bytes = 0; + double storage_seconds = 0.0; + double fsync_seconds = 0.0; +}; + +struct TransferMetrics { + size_t bytes = 0; + double setup_seconds = 0.0; + double pipeline_seconds = 0.0; + double storage_seconds = 0.0; + double cuda_wait_seconds = 0.0; + double fsync_seconds = 0.0; + double cleanup_seconds = 0.0; + double total_seconds = 0.0; + std::vector files; +}; + +// TransferBackendAvailable reports whether this helper binary was linked with +// a transfer adapter that can service CustomStorage extents. +bool TransferBackendAvailable(); + +bool TransferExtent(CUdeviceptr device_ptr, size_t extent_size, CUstream stream, + CUcontext context, const StorageLayout &storage, + TransferOperation operation, const TransferOptions &options, + TransferCancellation *cancellation, + TransferMetrics *metrics, std::string *error); + +} // namespace cuda_checkpoint_transfer diff --git a/agent/cmd/cuda-checkpoint-helper/transfer_engine_test.cpp b/agent/cmd/cuda-checkpoint-helper/transfer_engine_test.cpp new file mode 100644 index 00000000..c1a39abd --- /dev/null +++ b/agent/cmd/cuda-checkpoint-helper/transfer_engine_test.cpp @@ -0,0 +1,30 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include "transfer_cancellation.h" + +#include +#include + +int main() { + using namespace std::chrono_literals; + using cuda_checkpoint_transfer::TransferCancellation; + + TransferCancellation unlimited; + assert(!unlimited.DeadlineExceeded()); + assert(!unlimited.IsCancelled()); + + TransferCancellation active(TransferCancellation::Clock::now() + 1h); + assert(!active.DeadlineExceeded()); + assert(!active.IsCancelled()); + + TransferCancellation expired(TransferCancellation::Clock::now() - 1ms); + assert(expired.DeadlineExceeded()); + assert(expired.IsCancelled()); + + active.Cancel(); + assert(active.IsCancelled()); + return 0; +}