diff --git a/agent/Dockerfile b/agent/Dockerfile index 249cfd18..c0a710fd 100644 --- a/agent/Dockerfile +++ b/agent/Dockerfile @@ -28,6 +28,9 @@ ARG GO_VERSION=1.26.6 # --build-arg CRIU_REF= ARG CRIU_REPO=https://github.com/checkpoint-restore/criu.git ARG CRIU_REF=criu-dev +ARG NIXL_REPO=https://github.com/ai-dynamo/nixl.git +# Immutable pin containing POSIX short-write/error recovery used by #11584. +ARG NIXL_REF=71ca3cb249a481ae892fcc042a456e1a9b87463d ARG AGENT_BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:25.11-cuda13.0-devel-ubuntu24.04@sha256:8315e2455736c4f9b597f15c5fb4d31f834e798e0c6b66bbdbdbac491ce26bd1 # For placeholder target only - this default allows agent builds to succeed, @@ -153,6 +156,64 @@ RUN CGO_ENABLED=0 GOOS=${TARGETOS} GOARCH=${TARGETARCH} go build -ldflags="-w -s # Corresponding source for the Go modules compiled into the binaries above. RUN mkdir -p /go-src && go mod vendor -o /go-src/vendor +# ============================================================================= +# Stage: NIXL SDK/runtime builder (POSIX backend only) +# ============================================================================= +FROM ${AGENT_BASE_IMAGE} AS nixl-builder + +ARG NIXL_REPO +ARG NIXL_REF + +RUN apt-get update && apt-get install -y --no-install-recommends \ + build-essential \ + ca-certificates \ + git \ + libaio-dev \ + meson \ + ninja-build \ + pkg-config \ + pybind11-dev \ + python3 \ + && rm -rf /var/lib/apt/lists/* + +RUN git init /tmp/nixl \ + && cd /tmp/nixl \ + && git remote add origin ${NIXL_REPO} \ + && git fetch --depth 1 origin ${NIXL_REF} \ + && test "$(git rev-parse FETCH_HEAD)" = "${NIXL_REF}" \ + && git checkout --detach ${NIXL_REF} \ + && CXXFLAGS="-Wno-error=maybe-uninitialized -Wno-maybe-uninitialized" meson setup build-test \ + --buildtype=debugoptimized \ + -Dbuild_tests=true \ + -Dbuild_examples=false \ + -Denable_plugins=POSIX \ + && meson compile -C build-test \ + && NIXL_PLUGIN_DIR=/tmp/nixl/build-test/src/plugins/posix \ + ./build-test/test/unit/plugins/posix/nixl_posix_test -n 32 \ + && meson setup build \ + --prefix=/usr/local \ + --libdir=lib \ + --buildtype=release \ + -Dbuild_tests=false \ + -Dbuild_examples=false \ + -Denable_plugins=POSIX \ + && meson compile -C build \ + && DESTDIR=/nixl-install meson install -C build --strip \ + && mkdir -p /nixl-runtime/lib/plugins \ + && printf '%s\n' "${NIXL_REF}" > /nixl-runtime/NIXL_COMMIT \ + && cp /tmp/nixl/subprojects/liburing-liburing-2.14/LICENSE /nixl-runtime/liburing-LICENSE \ + && cp \ + /nixl-install/usr/local/lib/libnixl.so \ + /nixl-install/usr/local/lib/libnixl_build.so \ + /nixl-install/usr/local/lib/libnixl_common.so \ + /nixl-install/usr/local/lib/libserdes.so \ + /nixl-install/usr/local/lib/libstream.so \ + /nixl-install/usr/local/lib/libfile_utils.so \ + /nixl-runtime/lib/ \ + && cp -a /nixl-install/usr/local/lib/liburing.so* /nixl-runtime/lib/ \ + && cp /nixl-install/usr/local/lib/plugins/libplugin_POSIX.so /nixl-runtime/lib/plugins/ \ + && tar -czf /nixl-src.tar.gz -C /tmp --exclude-vcs nixl + # ============================================================================= # Stage: CUDA checkpoint helper builder # ============================================================================= @@ -164,14 +225,61 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ WORKDIR /workspace -COPY cmd/cuda-checkpoint-helper/main.c ./cmd/cuda-checkpoint-helper/main.c +COPY --from=nixl-builder /nixl-install/usr/local /usr/local +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.cpp ./cmd/cuda-checkpoint-helper/transfer_engine.cpp +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/ns-bind-mount/main.c ./cmd/ns-bind-mount/main.c -RUN gcc -O2 -Wall -Wextra -o /cuda-checkpoint-helper \ - ./cmd/cuda-checkpoint-helper/main.c \ +RUN g++ -std=c++20 -O2 -Wall -Wextra -Werror -pthread -o /cuda-checkpoint-helper \ + ./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_engine.cpp \ -I/usr/local/cuda/include \ + -I/usr/local/include \ -L/usr/local/cuda/lib64/stubs \ - -lcuda + -L/usr/local/lib \ + -Wl,-rpath,/usr/local/lib \ + -lnixl -lnixl_build -lcuda + +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-test \ + ./cmd/cuda-checkpoint-helper/transfer_config.cpp \ + ./cmd/cuda-checkpoint-helper/transfer_config_test.cpp \ + && /cuda-checkpoint-helper-transfer-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 # ns-bind-mount bind-mounts the snapshot binaries into a target container's mount # namespace. It needs no CUDA headers, only mount_setattr (Linux 5.12+). @@ -250,6 +358,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ libprotobuf-c1 \ libgnutls30t64 \ libnftables1 \ + libaio1t64 \ iproute2 \ iptables \ procps \ @@ -272,6 +381,7 @@ RUN sh /tmp/collect-sources.sh /tmp/base-packages.tsv /sources/dpkg # CRIU source, at the exact ref this image was built from. COPY --from=criu-builder /criu-src.tar.gz /sources/criu/criu-src.tar.gz +COPY --from=nixl-builder /nixl-src.tar.gz /sources/nixl/nixl-src.tar.gz # Go module source for the statically linked binaries. COPY --from=builder /go-src/vendor /sources/go/vendor @@ -292,6 +402,12 @@ RUN criu --version COPY --from=criu-builder /tmp/cuda-checkpoint/bin/x86_64_Linux/cuda-checkpoint /usr/local/sbin/cuda-checkpoint COPY --from=criu-builder /tmp/cuda-checkpoint/LICENSE /legal/cuda-checkpoint/LICENSE COPY --from=cuda-helper-builder /cuda-checkpoint-helper /usr/local/bin/cuda-checkpoint-helper +COPY --from=nixl-builder /nixl-runtime/lib /usr/local/lib +COPY --from=nixl-builder /nixl-runtime/NIXL_COMMIT /usr/local/share/nixl/NIXL_COMMIT +COPY --from=nixl-builder /tmp/nixl/LICENSE /legal/NIXL/LICENSE +COPY --from=nixl-builder /nixl-runtime/liburing-LICENSE /legal/liburing/LICENSE +ENV NIXL_PLUGIN_DIR=/usr/local/lib/plugins +RUN ldconfig # nsmount resolves this helper at /usr/local/sbin/ns-bind-mount (see # agent/internal/nsmount/mount.go defaultBinaryPath). COPY --from=cuda-helper-builder /ns-bind-mount /usr/local/sbin/ns-bind-mount diff --git a/agent/cmd/cuda-checkpoint-helper/README.md b/agent/cmd/cuda-checkpoint-helper/README.md index ebc08601..bad1a7bb 100644 --- a/agent/cmd/cuda-checkpoint-helper/README.md +++ b/agent/cmd/cuda-checkpoint-helper/README.md @@ -88,7 +88,12 @@ 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. +individual result for every target. V1 holds one process-local operation slot +across the whole multi-PID sequence within one target container, so sequences +handled by one agent do not interleave. Separate agent DaemonSets on the same +node are not coordinated; +deployments must avoid that topology. Host-scoped and per-GPU scheduling are +follow-ups. The daemon retains primary contexts only for the request's selected GPU set. After a successful operation, it associates those references with the exact @@ -114,7 +119,9 @@ fatal because continuing would make GPU-resource ownership ambiguous. - 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. +The helper has two link-time transfer variants. The first stack slice links the +no-backend implementation to validate compilation, linkage, and the standalone +protocol, manifest, transfer-configuration, and cancellation contracts without +NIXL. This Snapshot integration links the NIXL-backed POSIX adapter. +`custom_storage_available` is true only when both the CUDA CustomStorage driver +API and the linked transfer adapter are available. diff --git a/agent/cmd/cuda-checkpoint-helper/main.c b/agent/cmd/cuda-checkpoint-helper/main.c deleted file mode 100644 index d0bedbcb..00000000 --- a/agent/cmd/cuda-checkpoint-helper/main.c +++ /dev/null @@ -1,405 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include -#include -#include -#include -#include -#include - -static int -print_usage(FILE* stream) -{ - return fprintf( - stream, - "Usage:\n" - " cuda-checkpoint-helper --get-state --pid [--job-file ]\n" - " cuda-checkpoint-helper --get-restore-tid --pid [--job-file ]\n" - " cuda-checkpoint-helper --action lock|checkpoint|restore|unlock --pid [--timeout ] " - "[--device-map ] [--job-file ]\n") < 0 - ? 1 - : 0; -} - -static void -print_cuda_error(CUresult status) -{ - const char* name = NULL; - const char* msg = NULL; - - (void)cuGetErrorName(status, &name); - (void)cuGetErrorString(status, &msg); - - if (name == NULL) { - name = "CUDA_ERROR_UNKNOWN"; - } - if (msg == NULL) { - msg = "unknown CUDA error"; - } - - fprintf(stderr, "%s: %s\n", name, msg); -} - -static int -parse_pid(const char* pid_str, int* pid_out) -{ - char* end = NULL; - long pid = strtol(pid_str, &end, 10); - - if (pid_str[0] == '\0' || end == NULL || *end != '\0' || pid <= 0 || pid > INT_MAX) { - return -1; - } - - *pid_out = (int)pid; - return 0; -} - -static int -parse_timeout_ms(const char* timeout_str, unsigned int* timeout_ms_out) -{ - char* end = NULL; - unsigned long timeout_ms = strtoul(timeout_str, &end, 10); - - if (timeout_str[0] == '\0' || end == NULL || *end != '\0' || timeout_ms > UINT_MAX) { - return -1; - } - - *timeout_ms_out = (unsigned int)timeout_ms; - return 0; -} - -static int -parse_hex_byte(const char* src, unsigned char* byte_out) -{ - char tmp[3]; - char* end = NULL; - long value; - - tmp[0] = src[0]; - tmp[1] = src[1]; - tmp[2] = '\0'; - - value = strtol(tmp, &end, 16); - if (end == NULL || *end != '\0' || value < 0 || value > 255) { - return -1; - } - - *byte_out = (unsigned char)value; - return 0; -} - -static int -parse_uuid(const char* uuid_str, CUuuid* uuid_out) -{ - size_t len; - int i; - - if (uuid_str == NULL || uuid_out == NULL) { - return -1; - } - - len = strlen(uuid_str); - if (len == 40) { - if (strncmp(uuid_str, "GPU-", 4) != 0) { - return -1; - } - uuid_str += 4; - len -= 4; - } - - if (len != 36) { - return -1; - } - - for (i = 0; i < 16; ++i) { - if (*uuid_str == '-') { - ++uuid_str; - } - if (!isxdigit((unsigned char)uuid_str[0]) || !isxdigit((unsigned char)uuid_str[1])) { - return -1; - } - if (parse_hex_byte(uuid_str, (unsigned char*)&uuid_out->bytes[i]) != 0) { - return -1; - } - uuid_str += 2; - } - - return *uuid_str == '\0' ? 0 : -1; -} - -static int -parse_device_map(const char* device_map, CUcheckpointGpuPair** pairs_out, unsigned int* count_out) -{ - char* copy = NULL; - char* pair = NULL; - char* pair_save = NULL; - unsigned int count = 0; - CUcheckpointGpuPair* pairs = NULL; - - *pairs_out = NULL; - *count_out = 0; - - if (device_map == NULL || device_map[0] == '\0') { - return 0; - } - - copy = strdup(device_map); - if (copy == NULL) { - return -1; - } - - for (pair = copy; *pair != '\0'; ++pair) { - if (*pair == ',') { - ++count; - } - } - ++count; - - pairs = calloc(count, sizeof(*pairs)); - if (pairs == NULL) { - free(copy); - return -1; - } - - count = 0; - pair = strtok_r(copy, ",", &pair_save); - while (pair != NULL) { - char* uuid_save = NULL; - char* old_uuid = strtok_r(pair, "=", &uuid_save); - char* new_uuid = strtok_r(NULL, "=", &uuid_save); - - if (old_uuid == NULL || new_uuid == NULL || strtok_r(NULL, "=", &uuid_save) != NULL) { - free(copy); - free(pairs); - return -1; - } - if (parse_uuid(old_uuid, &pairs[count].oldUuid) != 0 || parse_uuid(new_uuid, &pairs[count].newUuid) != 0) { - free(copy); - free(pairs); - return -1; - } - - ++count; - pair = strtok_r(NULL, ",", &pair_save); - } - - free(copy); - *pairs_out = pairs; - *count_out = count; - return 0; -} - -static const char* -process_state_string(CUprocessState state) -{ - switch (state) { - case CU_PROCESS_STATE_RUNNING: - return "running"; - case CU_PROCESS_STATE_LOCKED: - return "locked"; - case CU_PROCESS_STATE_CHECKPOINTED: - return "checkpointed"; - case CU_PROCESS_STATE_FAILED: - return "failed"; - default: - return "unknown"; - } -} - -static CUresult -do_lock(int pid, unsigned int timeout_ms) -{ - CUcheckpointLockArgs args; - - memset(&args, 0, sizeof(args)); - args.timeoutMs = timeout_ms; - return cuCheckpointProcessLock(pid, &args); -} - -static CUresult -do_checkpoint(int pid) -{ - CUcheckpointCheckpointArgs args; - - memset(&args, 0, sizeof(args)); - return cuCheckpointProcessCheckpoint(pid, &args); -} - -static CUresult -do_restore(int pid, const char* device_map) -{ - CUcheckpointRestoreArgs args; - CUcheckpointGpuPair* pairs = NULL; - unsigned int pair_count = 0; - CUresult status; - - memset(&args, 0, sizeof(args)); - if (parse_device_map(device_map, &pairs, &pair_count) != 0) { - return CUDA_ERROR_INVALID_VALUE; - } - - args.gpuPairs = pairs; - args.gpuPairsCount = pair_count; - status = cuCheckpointProcessRestore(pid, &args); - free(pairs); - return status; -} - -static CUresult -do_unlock(int pid) -{ - CUcheckpointUnlockArgs args; - - memset(&args, 0, sizeof(args)); - return cuCheckpointProcessUnlock(pid, &args); -} - -static CUresult -do_get_state(int pid, CUprocessState* state_out) -{ - return cuCheckpointProcessGetState(pid, state_out); -} - -static CUresult -do_get_restore_tid(int pid, int* tid_out) -{ - return cuCheckpointProcessGetRestoreThreadId(pid, tid_out); -} - -int -main(int argc, char** argv) -{ - const char* action = NULL; - const char* device_map = ""; - const char* job_file = ""; - int pid = 0; - int have_pid = 0; - int do_get_state_flag = 0; - int do_get_restore_tid_flag = 0; - unsigned int timeout_ms = 0; - int i; - CUresult status; - - if (argc == 1) { - return print_usage(stderr); - } - - for (i = 1; i < argc; ++i) { - if (strcmp(argv[i], "--get-state") == 0) { - do_get_state_flag = 1; - continue; - } - if (strcmp(argv[i], "--get-restore-tid") == 0) { - do_get_restore_tid_flag = 1; - continue; - } - if (strcmp(argv[i], "--action") == 0) { - if (++i >= argc) { - return print_usage(stderr); - } - action = argv[i]; - continue; - } - if (strcmp(argv[i], "--pid") == 0 || strcmp(argv[i], "-p") == 0) { - if (++i >= argc || parse_pid(argv[i], &pid) != 0) { - return print_usage(stderr); - } - have_pid = 1; - continue; - } - if (strcmp(argv[i], "--timeout") == 0 || strcmp(argv[i], "-t") == 0) { - if (++i >= argc || parse_timeout_ms(argv[i], &timeout_ms) != 0) { - return print_usage(stderr); - } - continue; - } - if (strcmp(argv[i], "--device-map") == 0 || strcmp(argv[i], "-d") == 0) { - if (++i >= argc) { - return print_usage(stderr); - } - device_map = argv[i]; - continue; - } - if (strcmp(argv[i], "--job-file") == 0) { - if (++i >= argc || argv[i][0] == '\0') { - return print_usage(stderr); - } - job_file = argv[i]; - continue; - } - if (strcmp(argv[i], "--help") == 0 || strcmp(argv[i], "-h") == 0) { - return print_usage(stdout); - } - return print_usage(stderr); - } - - if ((do_get_state_flag + do_get_restore_tid_flag + (action != NULL ? 1 : 0)) != 1) { - return print_usage(stderr); - } - if (!have_pid) { - return print_usage(stderr); - } - - if (job_file[0] != '\0' && setenv("CUDA_CHECKPOINT_JOB_FILE", job_file, 1) != 0) { - perror("setenv CUDA_CHECKPOINT_JOB_FILE"); - return 1; - } - - if (do_get_state_flag) { - CUprocessState state; - - if (timeout_ms != 0 || device_map[0] != '\0') { - return print_usage(stderr); - } - status = do_get_state(pid, &state); - if (status != CUDA_SUCCESS) { - print_cuda_error(status); - return 1; - } - return fprintf(stdout, "%s\n", process_state_string(state)) < 0 ? 1 : 0; - } - - if (do_get_restore_tid_flag) { - int tid = 0; - - if (timeout_ms != 0 || device_map[0] != '\0') { - return print_usage(stderr); - } - status = do_get_restore_tid(pid, &tid); - if (status != CUDA_SUCCESS) { - print_cuda_error(status); - return 1; - } - return fprintf(stdout, "%d\n", tid) < 0 ? 1 : 0; - } - - if (strcmp(action, "lock") == 0) { - status = do_lock(pid, timeout_ms); - } else if (strcmp(action, "checkpoint") == 0) { - if (timeout_ms != 0 || device_map[0] != '\0') { - return print_usage(stderr); - } - status = do_checkpoint(pid); - } else if (strcmp(action, "restore") == 0) { - if (timeout_ms != 0) { - return print_usage(stderr); - } - status = do_restore(pid, device_map); - } else if (strcmp(action, "unlock") == 0) { - if (timeout_ms != 0 || device_map[0] != '\0') { - return print_usage(stderr); - } - status = do_unlock(pid); - } else { - return print_usage(stderr); - } - - if (status != CUDA_SUCCESS) { - print_cuda_error(status); - return 1; - } - return 0; -} diff --git a/agent/cmd/cuda-checkpoint-helper/transfer_engine.cpp b/agent/cmd/cuda-checkpoint-helper/transfer_engine.cpp new file mode 100644 index 00000000..4e607a6e --- /dev/null +++ b/agent/cmd/cuda-checkpoint-helper/transfer_engine.cpp @@ -0,0 +1,849 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include "transfer_engine.h" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cuda_checkpoint_transfer { +namespace { + +using Clock = std::chrono::steady_clock; +constexpr auto kNixlTransferTimeout = std::chrono::minutes(30); + +double ElapsedSeconds(Clock::time_point start) { + return std::chrono::duration(Clock::now() - start).count(); +} + +void AppendError(std::string *error, const std::string &detail) { + if (!error->empty()) { + *error += "; "; + } + *error += detail; +} + +std::string CudaError(CUresult status) { + const char *name = nullptr; + const char *message = nullptr; + (void)cuGetErrorName(status, &name); + (void)cuGetErrorString(status, &message); + return std::string(name == nullptr ? "CUDA_ERROR_UNKNOWN" : name) + ": " + + (message == nullptr ? "unknown CUDA error" : message); +} + +class FileDescriptor { +public: + explicit FileDescriptor(int fd = -1) : fd_(fd) {} + FileDescriptor(const FileDescriptor &) = delete; + FileDescriptor &operator=(const FileDescriptor &) = delete; + FileDescriptor(FileDescriptor &&other) noexcept + : fd_(std::exchange(other.fd_, -1)) {} + FileDescriptor &operator=(FileDescriptor &&) = delete; + + ~FileDescriptor() { + if (fd_ >= 0) { + (void)close(fd_); + } + } + + int get() const { return fd_; } + + bool Close(std::string *error) { + if (fd_ < 0) { + return true; + } + const int fd = std::exchange(fd_, -1); + if (close(fd) == 0) { + return true; + } + AppendError(error, "close storage file failed: " + + std::string(std::strerror(errno))); + return false; + } + +private: + int fd_; +}; + +class TransferSlot { +public: + TransferSlot() = default; + TransferSlot(const TransferSlot &) = delete; + TransferSlot &operator=(const TransferSlot &) = delete; + + ~TransferSlot() { + if (cuda_may_access_) { + return; + } + if (event_ != nullptr) { + (void)cuEventDestroy(event_); + } + if (data_ != nullptr) { + if (cuMemHostUnregister(data_) == CUDA_SUCCESS) { + free(data_); + } + } + } + + CUresult Allocate(size_t size) { + if (posix_memalign(&data_, kBufferAlignment, size) != 0) { + return CUDA_ERROR_OUT_OF_MEMORY; + } + CUresult status = cuMemHostRegister(data_, size, 0); + if (status != CUDA_SUCCESS) { + free(data_); + data_ = nullptr; + return status; + } + status = cuEventCreate(&event_, CU_EVENT_DISABLE_TIMING); + if (status != CUDA_SUCCESS) { + if (cuMemHostUnregister(data_) == CUDA_SUCCESS) { + free(data_); + data_ = nullptr; + } + } + return status; + } + + void *data() const { return data_; } + CUevent event() const { return event_; } + bool pending() const { return pending_; } + void set_pending(bool pending) { pending_ = pending; } + void set_cuda_may_access(bool cuda_may_access) { + cuda_may_access_ = cuda_may_access; + } + + void MarkCUDAComplete() { + pending_ = false; + cuda_may_access_ = false; + } + + bool Close(std::string *error) { + if (cuda_may_access_) { + AppendError(error, "CUDA transfer buffer retained because stream drain " + "did not complete"); + return false; + } + bool success = true; + if (event_ != nullptr) { + const CUresult status = cuEventDestroy(event_); + event_ = nullptr; + if (status != CUDA_SUCCESS) { + success = false; + AppendError(error, + "destroy CUDA transfer event failed: " + CudaError(status)); + } + } + if (data_ != nullptr) { + const CUresult status = cuMemHostUnregister(data_); + if (status == CUDA_SUCCESS) { + free(data_); + data_ = nullptr; + } else { + success = false; + AppendError(error, "unregister CUDA transfer buffer failed: " + + CudaError(status)); + } + } + return success; + } + +private: + void *data_ = nullptr; + CUevent event_ = nullptr; + bool pending_ = false; + bool cuda_may_access_ = false; +}; + +class StreamDrainGuard { +public: + StreamDrainGuard(CUstream stream, + std::vector> *slots) + : stream_(stream), slots_(slots) {} + StreamDrainGuard(const StreamDrainGuard &) = delete; + StreamDrainGuard &operator=(const StreamDrainGuard &) = delete; + + ~StreamDrainGuard() { + if (armed_ && cuStreamSynchronize(stream_) == CUDA_SUCCESS) { + for (auto &slot : *slots_) { + slot->MarkCUDAComplete(); + } + } + } + + void Disarm() { armed_ = false; } + +private: + CUstream stream_; + std::vector> *slots_; + bool armed_ = true; +}; + +class NixlRegistrationGuard { +public: + NixlRegistrationGuard(nixlAgent *agent, nixl_reg_dlist_t *dram_registration, + nixl_reg_dlist_t *file_registration) + : agent_(agent), dram_registration_(dram_registration), + file_registration_(file_registration) {} + NixlRegistrationGuard(const NixlRegistrationGuard &) = delete; + NixlRegistrationGuard &operator=(const NixlRegistrationGuard &) = delete; + + ~NixlRegistrationGuard() { + if (armed_) { + (void)agent_->deregisterMem(*file_registration_); + (void)agent_->deregisterMem(*dram_registration_); + } + } + + void Disarm() { armed_ = false; } + +private: + nixlAgent *agent_; + nixl_reg_dlist_t *dram_registration_; + nixl_reg_dlist_t *file_registration_; + bool armed_ = true; +}; + +bool OpenStorageFiles(const StorageLayout &storage, TransferOperation operation, + std::vector *files, std::string *error) { + files->clear(); + files->reserve(storage.files.size()); + for (size_t index = 0; index < storage.files.size(); ++index) { + const auto &file = storage.files[index]; + if (file.size > static_cast(std::numeric_limits::max())) { + *error = "storage file " + std::to_string(index) + + " is too large for POSIX offsets"; + return false; + } + const std::filesystem::path normalized = file.path.lexically_normal(); + if (!normalized.is_absolute() || normalized.filename().empty()) { + *error = "storage file " + std::to_string(index) + + " does not have a valid absolute path"; + return false; + } + + // Resolve every parent component through directory descriptors. O_NOFOLLOW + // on the final open alone does not prevent a writable parent directory from + // being replaced by a symlink between validation and use. + FileDescriptor root(open("/", O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + if (root.get() < 0) { + *error = "open storage root failed: " + std::string(std::strerror(errno)); + return false; + } + std::vector parents; + int parent_fd = root.get(); + for (const auto &component : normalized.relative_path().parent_path()) { + const std::string name = component.string(); + if (name.empty() || name == "." || name == "..") { + *error = "storage file " + std::to_string(index) + + " contains an unsafe path component"; + return false; + } + const int next_fd = openat(parent_fd, name.c_str(), + O_RDONLY | O_DIRECTORY | O_NOFOLLOW | + O_CLOEXEC); + if (next_fd < 0) { + *error = "open storage parent for file " + std::to_string(index) + + " failed: " + std::strerror(errno); + return false; + } + parents.emplace_back(next_fd); + parent_fd = parents.back().get(); + } + FileDescriptor descriptor(openat(parent_fd, normalized.filename().c_str(), + StorageFileOpenFlags(operation), 0600)); + if (descriptor.get() < 0) { + *error = "open storage file " + std::to_string(index) + + " failed: " + std::strerror(errno); + return false; + } + if (operation == TransferOperation::kCheckpoint) { + if (fchmod(descriptor.get(), 0600) != 0 || + ftruncate(descriptor.get(), static_cast(file.size)) != 0) { + *error = "prepare checkpoint storage file " + std::to_string(index) + + " failed: " + std::strerror(errno); + return false; + } + } else { + struct stat file_stat{}; + if (fstat(descriptor.get(), &file_stat) != 0 || + !S_ISREG(file_stat.st_mode) || file_stat.st_size < 0 || + static_cast(file_stat.st_size) != file.size) { + *error = "restore storage file " + std::to_string(index) + + " is not regular or has the wrong size"; + return false; + } + } + files->push_back(std::move(descriptor)); + } + return true; +} + +bool NixlTransfer(nixlAgent *agent, const std::string &agent_name, + nixl_xfer_op_t operation, void *buffer, int file_fd, + size_t file_offset, size_t length, + TransferCancellation *cancellation, std::string *error) { + nixl_xfer_dlist_t dram(DRAM_SEG); + nixl_xfer_dlist_t file(FILE_SEG); + dram.addDesc(nixlBlobDesc(reinterpret_cast(buffer), length, 0)); + file.addDesc(nixlBlobDesc(file_offset, length, file_fd)); + nixlXferReqH *request = nullptr; + nixl_status_t status = + agent->createXferReq(operation, dram, file, agent_name, request); + if (status != NIXL_SUCCESS) { + *error = "NIXL createXferReq failed with status " + std::to_string(status); + return false; + } + status = agent->postXferReq(request); + const auto fallback_deadline = Clock::now() + kNixlTransferTimeout; + bool cancellation_requested = false; + std::chrono::microseconds poll_delay(50); + while (status == NIXL_IN_PROG) { + const bool configured_deadline_exceeded = + cancellation != nullptr && cancellation->DeadlineExceeded(); + const bool fallback_deadline_exceeded = + cancellation == nullptr && Clock::now() >= fallback_deadline; + if (!cancellation_requested && + ((cancellation != nullptr && cancellation->IsCancelled()) || + fallback_deadline_exceeded)) { + cancellation_requested = true; + if (configured_deadline_exceeded) { + *error = "NIXL transfer exceeded the configured operation deadline"; + } else if (fallback_deadline_exceeded) { + *error = "NIXL transfer exceeded the 30-minute fallback deadline"; + } else { + *error = "NIXL transfer canceled after another extent failed"; + } + // releaseXferReq is the NIXL cancellation API. A successful release + // transfers ownership back to NIXL and frees the handle. If cancellation + // cannot complete immediately, retain the request and keep polling until + // it reaches a terminal state; abandoning it would leave NIXL using the + // registered buffer and file descriptor after their owners are destroyed. + const nixl_status_t cancel_status = agent->releaseXferReq(request); + if (cancel_status == NIXL_SUCCESS) { + return false; + } + AppendError(error, "NIXL cancellation is pending with status " + + std::to_string(cancel_status)); + } + status = agent->getXferStatus(request); + if (status == NIXL_IN_PROG) { + std::this_thread::sleep_for(poll_delay); + poll_delay = std::min(poll_delay * 2, std::chrono::microseconds(5000)); + } + } + const nixl_status_t release_status = agent->releaseXferReq(request); + if (cancellation_requested) { + if (release_status != NIXL_SUCCESS) { + AppendError(error, "releaseXferReq after cancellation failed with status " + + std::to_string(release_status)); + } + return false; + } + if (status != NIXL_SUCCESS) { + *error = "NIXL transfer failed with status " + std::to_string(status); + if (release_status != NIXL_SUCCESS) { + AppendError(error, "releaseXferReq also failed with status " + + std::to_string(release_status)); + } + return false; + } + if (release_status != NIXL_SUCCESS) { + *error = "NIXL releaseXferReq failed with status " + + std::to_string(release_status); + return false; + } + return true; +} + +bool WaitForSlot(TransferSlot *slot, TransferMetrics *metrics, + std::string *error) { + if (!slot->pending()) { + return true; + } + const auto start = Clock::now(); + const CUresult status = cuEventSynchronize(slot->event()); + metrics->cuda_wait_seconds += ElapsedSeconds(start); + if (status != CUDA_SUCCESS) { + *error = "CUDA event synchronization failed: " + CudaError(status); + return false; + } + slot->MarkCUDAComplete(); + return true; +} + +bool EnqueueCopy(TransferOperation operation, const TransferChunk &chunk, + TransferSlot *slot, CUdeviceptr device_ptr, CUstream stream, + bool *cuda_work_posted, std::string *error) { + CUresult status = CUDA_SUCCESS; + if (operation == TransferOperation::kCheckpoint) { + status = cuMemcpyDtoHAsync(slot->data(), device_ptr + chunk.logical_offset, + chunk.size, stream); + } else { + status = cuMemcpyHtoDAsync(device_ptr + chunk.logical_offset, slot->data(), + chunk.size, stream); + } + if (status != CUDA_SUCCESS) { + *error = "CUDA asynchronous copy failed at logical offset " + + std::to_string(chunk.logical_offset) + ": " + CudaError(status); + return false; + } + *cuda_work_posted = true; + slot->set_cuda_may_access(true); + status = cuEventRecord(slot->event(), stream); + if (status != CUDA_SUCCESS) { + *error = "CUDA event record failed at logical offset " + + std::to_string(chunk.logical_offset) + ": " + CudaError(status); + return false; + } + slot->set_pending(true); + return true; +} + +bool DrainCUDA(std::vector> *slots, + CUstream stream, bool force_stream_sync, bool cuda_work_posted, + TransferMetrics *metrics, std::string *error) { + bool success = true; + for (auto &slot : *slots) { + std::string wait_error; + if (!WaitForSlot(slot.get(), metrics, &wait_error)) { + success = false; + AppendError(error, wait_error); + } + } + if (cuda_work_posted && (force_stream_sync || !success)) { + const auto start = Clock::now(); + const CUresult status = cuStreamSynchronize(stream); + metrics->cuda_wait_seconds += ElapsedSeconds(start); + if (status != CUDA_SUCCESS) { + success = false; + AppendError(error, "CUDA stream drain failed: " + CudaError(status)); + } else { + for (auto &slot : *slots) { + slot->MarkCUDAComplete(); + } + } + } + return success; +} + +bool TransferPipeline(const std::vector &chunks, + const std::vector &files, + std::vector> *slots, + nixlAgent *agent, const std::string &agent_name, + CUdeviceptr device_ptr, CUstream stream, + TransferOperation operation, TransferMetrics *metrics, + TransferCancellation *cancellation, std::string *error) { + bool success = true; + bool cuda_work_posted = false; + if (operation == TransferOperation::kRestore) { + for (const auto &chunk : chunks) { + if (cancellation != nullptr && cancellation->IsCancelled()) { + *error = "transfer canceled after another extent failed"; + success = false; + break; + } + TransferSlot *slot = (*slots)[chunk.slot_index].get(); + if (!WaitForSlot(slot, metrics, error)) { + success = false; + if (cancellation != nullptr) { + cancellation->Cancel(); + } + break; + } + if (cancellation != nullptr && cancellation->IsCancelled()) { + *error = "transfer canceled after another extent failed"; + success = false; + break; + } + const auto storage_start = Clock::now(); + std::string transfer_error; + const bool transferred = + NixlTransfer(agent, agent_name, NIXL_READ, slot->data(), + files[chunk.file_index].get(), chunk.file_offset, + chunk.size, cancellation, &transfer_error); + const double storage_seconds = ElapsedSeconds(storage_start); + metrics->storage_seconds += storage_seconds; + metrics->files[chunk.file_index].storage_seconds += storage_seconds; + if (!transferred) { + *error = "storage read failed at logical offset " + + std::to_string(chunk.logical_offset) + ": " + transfer_error; + success = false; + if (cancellation != nullptr) { + cancellation->Cancel(); + } + break; + } + metrics->files[chunk.file_index].bytes += chunk.size; + if (cancellation != nullptr && cancellation->IsCancelled()) { + *error = "transfer canceled after another extent failed"; + success = false; + break; + } + if (!EnqueueCopy(operation, chunk, slot, device_ptr, stream, + &cuda_work_posted, error)) { + success = false; + if (cancellation != nullptr) { + cancellation->Cancel(); + } + break; + } + } + } else { + size_t next_chunk = 0; + const size_t initial_count = std::min(chunks.size(), slots->size()); + for (; next_chunk < initial_count; ++next_chunk) { + if (cancellation != nullptr && cancellation->IsCancelled()) { + *error = "transfer canceled after another extent failed"; + success = false; + break; + } + const auto &chunk = chunks[next_chunk]; + if (!EnqueueCopy(operation, chunk, (*slots)[chunk.slot_index].get(), + device_ptr, stream, &cuda_work_posted, error)) { + success = false; + if (cancellation != nullptr) { + cancellation->Cancel(); + } + break; + } + } + size_t completed = 0; + for (; success && completed < next_chunk; ++completed) { + if (cancellation != nullptr && cancellation->IsCancelled()) { + *error = "transfer canceled after another extent failed"; + success = false; + break; + } + const auto &chunk = chunks[completed]; + TransferSlot *slot = (*slots)[chunk.slot_index].get(); + if (!WaitForSlot(slot, metrics, error)) { + success = false; + if (cancellation != nullptr) { + cancellation->Cancel(); + } + break; + } + if (cancellation != nullptr && cancellation->IsCancelled()) { + *error = "transfer canceled after another extent failed"; + success = false; + break; + } + const auto storage_start = Clock::now(); + std::string transfer_error; + const bool transferred = + NixlTransfer(agent, agent_name, NIXL_WRITE, slot->data(), + files[chunk.file_index].get(), chunk.file_offset, + chunk.size, cancellation, &transfer_error); + const double storage_seconds = ElapsedSeconds(storage_start); + metrics->storage_seconds += storage_seconds; + metrics->files[chunk.file_index].storage_seconds += storage_seconds; + if (!transferred) { + *error = "storage write failed at logical offset " + + std::to_string(chunk.logical_offset) + ": " + transfer_error; + success = false; + if (cancellation != nullptr) { + cancellation->Cancel(); + } + break; + } + metrics->files[chunk.file_index].bytes += chunk.size; + if (next_chunk < chunks.size()) { + const auto &next = chunks[next_chunk]; + if (next.slot_index != chunk.slot_index) { + *error = "internal transfer ring scheduling error"; + success = false; + if (cancellation != nullptr) { + cancellation->Cancel(); + } + break; + } + if (cancellation != nullptr && cancellation->IsCancelled()) { + *error = "transfer canceled after another extent failed"; + success = false; + break; + } + if (!EnqueueCopy(operation, next, slot, device_ptr, stream, + &cuda_work_posted, error)) { + success = false; + if (cancellation != nullptr) { + cancellation->Cancel(); + } + break; + } + ++next_chunk; + } + } + if (success && completed != chunks.size()) { + *error = "internal checkpoint transfer coverage error"; + success = false; + if (cancellation != nullptr) { + cancellation->Cancel(); + } + } + } + + if (!DrainCUDA(slots, stream, !success, cuda_work_posted, metrics, error)) { + success = false; + if (cancellation != nullptr) { + cancellation->Cancel(); + } + } + return success; +} + +std::string MakeAgentName() { + static std::atomic sequence{0}; + return "cuda-custom-storage-" + std::to_string(getpid()) + "-" + + std::to_string(sequence.fetch_add(1)); +} + +} // namespace + +bool TransferBackendAvailable() { return true; } + +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) { + if (metrics == nullptr || error == nullptr) { + return false; + } + *metrics = {}; + metrics->files.resize(storage.files.size()); + error->clear(); + const auto total_start = Clock::now(); + + std::vector chunks; + if (!BuildTransferChunks(extent_size, storage, options, &chunks, error)) { + if (cancellation != nullptr) { + cancellation->Cancel(); + } + metrics->total_seconds = ElapsedSeconds(total_start); + return false; + } + if (extent_size > std::numeric_limits::max() - device_ptr) { + *error = "CUDA device extent address calculation overflow"; + if (cancellation != nullptr) { + cancellation->Cancel(); + } + metrics->total_seconds = ElapsedSeconds(total_start); + return false; + } + if (device_ptr == 0 || context == nullptr) { + *error = "CUDA device pointer and context must be valid"; + if (cancellation != nullptr) { + cancellation->Cancel(); + } + metrics->total_seconds = ElapsedSeconds(total_start); + return false; + } + if (cancellation != nullptr && cancellation->IsCancelled()) { + *error = "transfer canceled before setup"; + metrics->total_seconds = ElapsedSeconds(total_start); + return false; + } + + CUresult cuda_status = cuCtxSetCurrent(context); + if (cuda_status != CUDA_SUCCESS) { + *error = "set CUDA context failed: " + CudaError(cuda_status); + if (cancellation != nullptr) { + cancellation->Cancel(); + } + metrics->total_seconds = ElapsedSeconds(total_start); + return false; + } + + const auto setup_start = Clock::now(); + std::vector> slots; + slots.reserve(options.buffer_count); + for (size_t index = 0; index < options.buffer_count; ++index) { + auto slot = std::make_unique(); + cuda_status = slot->Allocate(options.chunk_bytes); + if (cuda_status != CUDA_SUCCESS) { + *error = "allocate CUDA-registered transfer slot " + + std::to_string(index) + " failed: " + CudaError(cuda_status); + if (cancellation != nullptr) { + cancellation->Cancel(); + } + metrics->setup_seconds = ElapsedSeconds(setup_start); + metrics->total_seconds = ElapsedSeconds(total_start); + return false; + } + slots.push_back(std::move(slot)); + } + + std::vector files; + if (!OpenStorageFiles(storage, operation, &files, error)) { + if (cancellation != nullptr) { + cancellation->Cancel(); + } + metrics->setup_seconds = ElapsedSeconds(setup_start); + metrics->total_seconds = ElapsedSeconds(total_start); + return false; + } + + const std::string agent_name = MakeAgentName(); + nixlAgentConfig config; + config.useProgThread = true; + auto agent = std::make_unique(agent_name, config); + nixl_b_params_t params; + params["use_aio"] = "true"; + nixlBackendH *backend = nullptr; + nixl_status_t nixl_status = agent->createBackend("POSIX", params, backend); + if (nixl_status != NIXL_SUCCESS) { + *error = "create NIXL POSIX backend failed with status " + + std::to_string(nixl_status); + if (cancellation != nullptr) { + cancellation->Cancel(); + } + metrics->setup_seconds = ElapsedSeconds(setup_start); + metrics->total_seconds = ElapsedSeconds(total_start); + return false; + } + + nixl_reg_dlist_t dram_registration(DRAM_SEG); + nixl_reg_dlist_t file_registration(FILE_SEG); + for (const auto &slot : slots) { + dram_registration.addDesc(nixlBlobDesc( + reinterpret_cast(slot->data()), options.chunk_bytes, 0)); + } + for (size_t index = 0; index < storage.files.size(); ++index) { + file_registration.addDesc( + nixlBlobDesc(0, storage.files[index].size, files[index].get())); + } + + bool dram_registered = false; + bool file_registered = false; + nixl_status = agent->registerMem(dram_registration); + if (nixl_status == NIXL_SUCCESS) { + dram_registered = true; + nixl_status = agent->registerMem(file_registration); + file_registered = nixl_status == NIXL_SUCCESS; + } + if (!dram_registered || !file_registered) { + *error = std::string("register NIXL ") + + (dram_registered ? "file" : "DRAM") + + " memory failed with status " + std::to_string(nixl_status); + if (dram_registered) { + const nixl_status_t cleanup_status = + agent->deregisterMem(dram_registration); + if (cleanup_status != NIXL_SUCCESS) { + AppendError(error, "NIXL DRAM cleanup failed with status " + + std::to_string(cleanup_status)); + } + } + metrics->setup_seconds = ElapsedSeconds(setup_start); + metrics->total_seconds = ElapsedSeconds(total_start); + return false; + } + metrics->setup_seconds = ElapsedSeconds(setup_start); + NixlRegistrationGuard registration_guard(agent.get(), &dram_registration, + &file_registration); + + const auto pipeline_start = Clock::now(); + bool success = false; + { + StreamDrainGuard stream_drain_guard(stream, &slots); + success = TransferPipeline(chunks, files, &slots, agent.get(), agent_name, + device_ptr, stream, operation, metrics, + cancellation, error); + if (success) { + stream_drain_guard.Disarm(); + } + } + metrics->pipeline_seconds = ElapsedSeconds(pipeline_start); + + if (success && cancellation != nullptr && cancellation->IsCancelled()) { + success = false; + *error = "transfer canceled after another extent failed"; + } + if (success && operation == TransferOperation::kCheckpoint) { + for (size_t index = 0; index < files.size(); ++index) { + const auto fsync_start = Clock::now(); + const int result = fsync(files[index].get()); + const double fsync_seconds = ElapsedSeconds(fsync_start); + metrics->fsync_seconds += fsync_seconds; + metrics->files[index].fsync_seconds += fsync_seconds; + if (result != 0) { + success = false; + if (cancellation != nullptr) { + cancellation->Cancel(); + } + *error = "fsync storage file " + std::to_string(index) + + " failed: " + std::strerror(errno); + break; + } + } + } + + const auto cleanup_start = Clock::now(); + const nixl_status_t file_deregister_status = + agent->deregisterMem(file_registration); + const nixl_status_t dram_deregister_status = + agent->deregisterMem(dram_registration); + registration_guard.Disarm(); + if (file_deregister_status != NIXL_SUCCESS) { + success = false; + if (cancellation != nullptr) { + cancellation->Cancel(); + } + AppendError(error, "NIXL file deregistration failed with status " + + std::to_string(file_deregister_status)); + } + if (dram_deregister_status != NIXL_SUCCESS) { + success = false; + if (cancellation != nullptr) { + cancellation->Cancel(); + } + AppendError(error, "NIXL DRAM deregistration failed with status " + + std::to_string(dram_deregister_status)); + } + agent.reset(); + for (auto &file : files) { + if (!file.Close(error)) { + success = false; + if (cancellation != nullptr) { + cancellation->Cancel(); + } + } + } + for (auto &slot : slots) { + if (!slot->Close(error)) { + success = false; + if (cancellation != nullptr) { + cancellation->Cancel(); + } + } + } + slots.clear(); + metrics->cleanup_seconds = ElapsedSeconds(cleanup_start); + metrics->total_seconds = ElapsedSeconds(total_start); + if (success) { + metrics->bytes = extent_size; + } + return success; +} + +} // namespace cuda_checkpoint_transfer diff --git a/agent/internal/controller/controller.go b/agent/internal/controller/controller.go index 1ccbbb0c..f8e0466f 100644 --- a/agent/internal/controller/controller.go +++ b/agent/internal/controller/controller.go @@ -705,10 +705,10 @@ func (w *NodeController) runRestore(ctx context.Context, pod *corev1.Pod, artifa return err != nil, err } -// recoverCompletedRestore finalizes an interrupted status update without -// replaying CRIU. RestoreInProgress is normally retryable after an agent -// restart, but the restore-complete sentinel proves that execution already -// finished and only the snapshot/Restored=True status write remains. +// recoverCompletedRestore resolves a prior in-progress attempt without +// replaying CRIU or CUDA. The completion sentinel proves execution finished; +// without it the prior state-changing outcome is unknown, so V1 terminates the +// immutable placeholder and requires a fresh restore pod. func (op *restoreOperation) recoverCompletedRestore(ctx context.Context) (bool, error) { condition := findRestoredCondition(op.pod) if condition == nil || condition.Status != corev1.ConditionFalse || condition.Reason != restoreInProgressReason { @@ -724,7 +724,11 @@ func (op *restoreOperation) recoverCompletedRestore(ctx context.Context) (bool, return false, fmt.Errorf("check restore completion sentinel: %w", err) } if !exists { - return false, nil + interruptedErr := errors.New("restore was interrupted with an unknown CRIU/CUDA outcome; refusing to replay into the existing placeholder") + if err := op.failRestore(ctx, interruptedErr); err != nil { + return true, err + } + return true, nil } if err := op.markRestoreSucceeded(ctx); err != nil { return true, fmt.Errorf("finalize completed restore: %w", err) @@ -768,6 +772,7 @@ func (op *restoreOperation) executeRestore(ctx context.Context) (int, error) { TargetPodIP: op.pod.Status.PodIP, ContainerName: op.artifact.ContainerName, Clientset: w.clientset, + CUDATransfer: w.config.CUDACheckpoint.TransferSettings(), } return w.restoreFn(ctx, w.runtime, op.log, req, w.injector) } @@ -775,15 +780,17 @@ func (op *restoreOperation) executeRestore(ctx context.Context) (int, error) { func (op *restoreOperation) failRestore(ctx context.Context, restoreErr error) error { w := op.controller op.log.Error(restoreErr, "External restore failed") + cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer cancel() // Re-resolve because restore may fail before discovering the placeholder PID. - placeholderHostPID, _, err := w.runtime.ResolveContainer(ctx, op.containerID) + placeholderHostPID, _, err := w.runtime.ResolveContainer(cleanupCtx, op.containerID) if err != nil { return fmt.Errorf("restore failed and placeholder PID could not be resolved: %w", err) } if err := w.sendSignalFn(op.log, placeholderHostPID, syscall.SIGKILL, "restore failed"); err != nil { return fmt.Errorf("restore failed and placeholder could not be killed: %w", err) } - return op.markRestoreFailed(ctx, restoreErr) + return op.markRestoreFailed(cleanupCtx, restoreErr) } func (op *restoreOperation) completeRestore(ctx context.Context, placeholderHostPID int) error { diff --git a/agent/internal/controller/controller_test.go b/agent/internal/controller/controller_test.go index 994d9fd5..baa2654b 100644 --- a/agent/internal/controller/controller_test.go +++ b/agent/internal/controller/controller_test.go @@ -48,12 +48,16 @@ type fakeRuntime struct { containerIDByPod string resolvedContainerIDs []string resolveContainerPID int + requireLiveContext bool } var _ snapshotruntime.Runtime = (*fakeRuntime)(nil) -func (r *fakeRuntime) ResolveContainer(_ context.Context, id string) (int, *specs.Spec, error) { +func (r *fakeRuntime) ResolveContainer(ctx context.Context, id string) (int, *specs.Spec, error) { r.resolvedContainerIDs = append(r.resolvedContainerIDs, id) + if r.requireLiveContext && ctx.Err() != nil { + return 0, nil, ctx.Err() + } if r.resolveContainerPID > 0 { return r.resolveContainerPID, &specs.Spec{}, nil } @@ -1011,7 +1015,7 @@ func TestRunRestoreCleanupFailureStillCompletesRestore(t *testing.T) { assert.Contains(t, string(lastPodStatusApply(t, w).GetPatch()), `"reason":"RestoreSucceeded"`) } -func TestRunRestoreRetriesFullRestoreUntilFailureCleanupSucceeds(t *testing.T) { +func TestRunRestoreDoesNotReplayAfterFailureCleanupWasInterrupted(t *testing.T) { pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) w := makeTestController(t, pod) w.runtime = &fakeRuntime{resolveContainerPID: 4242} @@ -1029,6 +1033,7 @@ func TestRunRestoreRetriesFullRestoreUntilFailureCleanupSucceeds(t *testing.T) { } return nil } + w.controlSentinelExistsFn = func(int, string) (bool, error) { return false, nil } requeue, err := w.runRestore(context.Background(), pod, artifact, "ctr-abc", time.Time{}) require.Error(t, err) @@ -1039,7 +1044,25 @@ func TestRunRestoreRetriesFullRestoreUntilFailureCleanupSucceeds(t *testing.T) { requeue, err = w.runRestore(context.Background(), pod, artifact, "ctr-abc", time.Time{}) require.NoError(t, err) assert.False(t, requeue) - assert.Equal(t, 2, restoreCalls, "CRIU restore should retry when the previous cleanup did not finish") + assert.Equal(t, 1, restoreCalls, "an unknown CRIU/CUDA outcome must not be replayed") + assert.Contains(t, string(lastPodStatusApply(t, w).GetPatch()), `"reason":"RestoreFailed"`) +} + +func TestFailRestoreUsesCleanupContextAfterControllerCancellation(t *testing.T) { + pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + w := makeTestController(t, pod) + w.runtime = &fakeRuntime{resolveContainerPID: 4242, requireLiveContext: true} + signalCalls := 0 + w.sendSignalFn = func(logr.Logger, int, syscall.Signal, string) error { + signalCalls++ + return nil + } + op := w.newRestoreOperation(pod, &restoreArtifact{SnapshotName: "snapshot-a", ContainerName: "main"}, "ctr-abc", time.Time{}) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + require.NoError(t, op.failRestore(ctx, errors.New("restore interrupted"))) + assert.Equal(t, 1, signalCalls) assert.Contains(t, string(lastPodStatusApply(t, w).GetPatch()), `"reason":"RestoreFailed"`) } @@ -1119,3 +1142,42 @@ func TestCheckpointLeaseNameUsesContentAndContainer(t *testing.T) { assert.NotEqual(t, a, b) assert.True(t, strings.HasPrefix(a, "snapshot-capture-")) } + +func TestCheckpointLeasePersistsTargetIdentityAcrossRenewalAndTakeover(t *testing.T) { + w := makeTestController(t, nil) + key := client.ObjectKey{Namespace: "inference", Name: "capture-identity"} + first := captureLeaseTarget{podUID: "pod-a", containerID: "container-a", pid: 101} + + acquired, err := w.acquireLease(context.Background(), key, first) + require.NoError(t, err) + require.True(t, acquired) + lease, err := w.clientset.CoordinationV1().Leases(key.Namespace).Get(context.Background(), key.Name, metav1.GetOptions{}) + require.NoError(t, err) + got, err := captureTargetFromLease(lease) + require.NoError(t, err) + assert.Equal(t, first, got) + + require.NoError(t, w.renewLeaseOnce(context.Background(), key)) + lease, err = w.clientset.CoordinationV1().Leases(key.Namespace).Get(context.Background(), key.Name, metav1.GetOptions{}) + require.NoError(t, err) + got, err = captureTargetFromLease(lease) + require.NoError(t, err) + assert.Equal(t, first, got) + + expired := metav1.NewMicroTime(time.Now().Add(-2 * checkpointLeaseDuration)) + otherHolder := "snapshot-agent/expired" + lease.Spec.HolderIdentity = &otherHolder + lease.Spec.RenewTime = &expired + _, err = w.clientset.CoordinationV1().Leases(key.Namespace).Update(context.Background(), lease, metav1.UpdateOptions{}) + require.NoError(t, err) + + second := captureLeaseTarget{podUID: "pod-b", containerID: "container-b", pid: 202} + acquired, err = w.acquireLease(context.Background(), key, second) + require.NoError(t, err) + require.True(t, acquired) + lease, err = w.clientset.CoordinationV1().Leases(key.Namespace).Get(context.Background(), key.Name, metav1.GetOptions{}) + require.NoError(t, err) + got, err = captureTargetFromLease(lease) + require.NoError(t, err) + assert.Equal(t, second, got) +} diff --git a/agent/internal/controller/podsnapshotcontent.go b/agent/internal/controller/podsnapshotcontent.go index 0617fc41..9b4a21d4 100644 --- a/agent/internal/controller/podsnapshotcontent.go +++ b/agent/internal/controller/podsnapshotcontent.go @@ -14,6 +14,7 @@ import ( "time" "github.com/go-logr/logr" + coordinationv1 "k8s.io/api/coordination/v1" corev1 "k8s.io/api/core/v1" apierrors "k8s.io/apimachinery/pkg/api/errors" "k8s.io/apimachinery/pkg/api/meta" @@ -160,6 +161,18 @@ func (w *NodeController) captureLeaseHeldElsewhere(ctx context.Context, content *lease.Spec.HolderIdentity != w.holderID } +func (w *NodeController) captureLeaseExpired(ctx context.Context, content *snapshotv1alpha1.PodSnapshotContent, contentUID, containerName string) (*coordinationv1.Lease, bool, error) { + key := client.ObjectKey{Namespace: content.Spec.PodSnapshotRef.Namespace, Name: checkpointLeaseName(contentUID, containerName)} + lease, err := w.clientset.CoordinationV1().Leases(key.Namespace).Get(ctx, key.Name, metav1.GetOptions{}) + if err != nil { + if apierrors.IsNotFound(err) { + return nil, false, nil + } + return nil, false, fmt.Errorf("read checkpoint lease %s before capture: %w", key.String(), err) + } + return lease, checkpointLeaseExpired(lease, time.Now()), nil +} + // failContentFromGate records a terminal failure from the pre-bind gate, which has no workqueue // to surface errors to: a failed status write is logged and left to the informer resync. func (w *NodeController) failContentFromGate(ctx context.Context, content *snapshotv1alpha1.PodSnapshotContent, reason string, cause error) { @@ -253,6 +266,34 @@ func (w *NodeController) reconcileSourcePod(ctx context.Context, pod *corev1.Pod if artifactPresent(artifactPath, contentUID, containerName) { return w.markCheckpointReady(ctx, content) } + lease, expiredLease, err := w.captureLeaseExpired(ctx, content, contentUID, containerName) + if err != nil { + return err + } + if expiredLease { + // The abandoned attempt may have parked the CUDA process before the agent + // died. Refuse replay and terminate only the exact source identity persisted + // before mutation. A restarted container is never signaled. + causeMessage := "a prior checkpoint lease expired without a committed artifact; CRIU/CUDA outcome is unknown and will not be replayed" + target, targetErr := captureTargetFromLease(lease) + if targetErr != nil { + causeMessage += fmt.Sprintf("; exact source termination was not attempted because the lease identity is unavailable: %v", targetErr) + } else if target.podUID != string(pod.UID) { + causeMessage += fmt.Sprintf("; exact source termination was not attempted because lease pod UID %q does not match current pod UID %q", target.podUID, pod.UID) + } else if status := containerStatusForName(pod, containerName); status == nil || + snapshotruntime.StripCRIScheme(status.ContainerID) != target.containerID || status.State.Running == nil { + causeMessage += "; the recorded source container is no longer the current running container" + } else if killErr := w.killCheckpointContainer(ctx, logger, target.containerID, target.pid, "unknown checkpoint outcome"); killErr != nil { + var identityChanged *checkpointContainerIdentityChangedError + if errors.As(killErr, &identityChanged) { + causeMessage += "; the recorded source process identity is no longer current" + } else { + return fmt.Errorf("terminate source after unknown checkpoint outcome: %w", killErr) + } + } + w.removeCaptureEligibleLabel(ctx, pod) + return w.setSnapshotContentFailed(ctx, content, "CheckpointOutcomeUnknown", errors.New(causeMessage)) + } // The in-flight guard held above is process-local; overlapping agent instances arbitrate // through the shared capture Lease. A foreign unexpired holder may be between killing the @@ -288,7 +329,11 @@ func (w *NodeController) reconcileSourcePod(ctx context.Context, pod *corev1.Pod return w.setSnapshotContentFailed(ctx, content, "ContainerNotResolved", fmt.Errorf("resolve container %q: %w", containerName, err)) } leaseKey := client.ObjectKey{Namespace: content.Spec.PodSnapshotRef.Namespace, Name: checkpointLeaseName(contentUID, containerName)} - acquired, err := w.acquireLease(ctx, leaseKey) + acquired, err := w.acquireLease(ctx, leaseKey, captureLeaseTarget{ + podUID: string(pod.UID), + containerID: containerID, + pid: containerPID, + }) if err != nil { return fmt.Errorf("acquire checkpoint lease %s: %w", leaseKey.String(), err) } @@ -370,7 +415,7 @@ func (w *NodeController) runCheckpoint( } } } - if killErr := w.killCheckpointProcess(logger, containerPID, "checkpoint lease cancelled"); killErr != nil { + if killErr := w.killCheckpointContainer(ctx, logger, containerID, containerPID, "checkpoint lease cancelled"); killErr != nil { logger.Error(killErr, "Failed to kill target after lease cancellation", "content", content.Name) } return @@ -573,8 +618,9 @@ func (w *NodeController) setSnapshotContentFailed(ctx context.Context, content * // executorCheckpoint is the production checkpointFn. The reconciler has already resolved the // container ID and host PID. It runs executor.Checkpoint to the destination and verifies the -// artifact directory. On dump or verification failure it SIGKILLs the CUDA-locked process before -// returning the error; on success the dump itself has already terminated the source process. +// artifact directory. Proven pre-mutation failures preserve the source; failures that may have +// mutated it and post-capture verification failures re-resolve the container before terminating it. +// On success the dump itself has already terminated the source process. func (w *NodeController) executorCheckpoint(ctx context.Context, params CheckpointParams) error { log := logr.FromContextOrDiscard(ctx) @@ -590,8 +636,10 @@ func (w *NodeController) executorCheckpoint(ctx context.Context, params Checkpoi Clientset: w.clientset, } if err := executor.Checkpoint(ctx, w.runtime, log, req, w.config); err != nil { - if killErr := w.killCheckpointProcess(log, params.ContainerPID, "checkpoint failed"); killErr != nil { - log.Error(killErr, "Failed to kill target after checkpoint failure") + if checkpointFailureRequiresTermination(err) { + if killErr := w.killCheckpointContainer(ctx, log, params.ContainerID, params.ContainerPID, "checkpoint failed"); killErr != nil { + log.Error(killErr, "Failed to kill target after checkpoint failure") + } } return fmt.Errorf("checkpoint: %w", err) } @@ -604,7 +652,7 @@ func (w *NodeController) executorCheckpoint(ctx context.Context, params Checkpoi } else { verifyErr = fmt.Errorf("verify checkpoint path %s: not a directory", params.HostPath) } - if killErr := w.killCheckpointProcess(log, params.ContainerPID, "checkpoint verification failed"); killErr != nil { + if killErr := w.killCheckpointContainer(ctx, log, params.ContainerID, params.ContainerPID, "checkpoint verification failed"); killErr != nil { log.Error(killErr, "Failed to kill target after checkpoint verification failure") } return verifyErr @@ -613,10 +661,30 @@ func (w *NodeController) executorCheckpoint(ctx context.Context, params Checkpoi return nil } -// killCheckpointProcess SIGKILLs the CUDA-locked process so it does not hang after a failed dump. -// ESRCH (already exited) is success. Any other signal error is returned so callers can fail closed. -func (w *NodeController) killCheckpointProcess(log logr.Logger, pid int, reason string) error { - if err := snapshotruntime.SendSignalToPID(log, pid, syscall.SIGKILL, reason); err != nil { +// checkpointFailureRequiresTermination preserves the source only when every +// failing layer proves that no CRIU or CUDA mutation began. +func checkpointFailureRequiresTermination(err error) bool { + return !executor.CheckpointFailedBeforeTargetMutation(err) +} + +// killCheckpointContainer re-resolves the CRI container immediately before +// signaling. It refuses to use a stale PID if the container disappeared or +// restarted while the potentially long checkpoint operation was running. +func (w *NodeController) killCheckpointContainer(ctx context.Context, log logr.Logger, containerID string, expectedPID int, reason string) error { + resolveCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer cancel() + currentPID, _, err := w.runtime.ResolveContainer(resolveCtx, containerID) + if err != nil { + return fmt.Errorf("refusing to terminate checkpoint container %q without current runtime identity: %w", containerID, err) + } + if currentPID != expectedPID { + return &checkpointContainerIdentityChangedError{ + containerID: containerID, + currentPID: currentPID, + expectedPID: expectedPID, + } + } + if err := w.sendSignalFn(log, currentPID, syscall.SIGKILL, reason); err != nil { if errors.Is(err, syscall.ESRCH) { return nil } @@ -626,16 +694,33 @@ func (w *NodeController) killCheckpointProcess(log logr.Logger, pid int, reason return nil } +type checkpointContainerIdentityChangedError struct { + containerID string + currentPID int + expectedPID int +} + +func (e *checkpointContainerIdentityChangedError) Error() string { + return fmt.Sprintf("refusing to terminate checkpoint container %q: current PID %d does not match expected PID %d", e.containerID, e.currentPID, e.expectedPID) +} + // containerIDForName returns the running container's CRI-stripped ID, or "" if absent. func containerIDForName(pod *corev1.Pod, containerName string) string { - for _, cs := range pod.Status.ContainerStatuses { - if cs.Name == containerName { - return snapshotruntime.StripCRIScheme(cs.ContainerID) - } + if status := containerStatusForName(pod, containerName); status != nil { + return snapshotruntime.StripCRIScheme(status.ContainerID) } return "" } +func containerStatusForName(pod *corev1.Pod, containerName string) *corev1.ContainerStatus { + for i := range pod.Status.ContainerStatuses { + if pod.Status.ContainerStatuses[i].Name == containerName { + return &pod.Status.ContainerStatuses[i] + } + } + return nil +} + // isContentTerminal reports whether the work order already has a terminal condition. func isContentTerminal(content *snapshotv1alpha1.PodSnapshotContent) bool { return isContentReady(content) || isContentFailed(content) diff --git a/agent/internal/controller/podsnapshotcontent_coverage_test.go b/agent/internal/controller/podsnapshotcontent_coverage_test.go index 3717f6ae..df486aaa 100644 --- a/agent/internal/controller/podsnapshotcontent_coverage_test.go +++ b/agent/internal/controller/podsnapshotcontent_coverage_test.go @@ -23,6 +23,7 @@ import ( crfake "sigs.k8s.io/controller-runtime/pkg/client/fake" "sigs.k8s.io/controller-runtime/pkg/client/interceptor" + snapshotruntime "github.com/ai-dynamo/snapshot/agent/internal/runtime" snapshottypes "github.com/ai-dynamo/snapshot/agent/internal/types" snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" ) @@ -49,11 +50,51 @@ func makeNodeControllerWithInterceptor(t *testing.T, fc *fakeCheckpointer, funcs holderID: "snapshot-agent/test", inFlight: make(map[string]struct{}), contentIndexer: idx, + sendSignalFn: snapshotruntime.SendSignalToPID, } w.checkpointFn = fc.fn return w } +func TestUnknownCheckpointRecoveryRetriesStatusAfterExactTermination(t *testing.T) { + content := makeWorkOrder("podsnapshotcontent-abc", "node-a", "abc") + pod := makeSourcePod() + failedOnce := false + funcs := interceptor.Funcs{ + SubResourcePatch: func(ctx context.Context, c client.Client, sub string, obj client.Object, patch client.Patch, opts ...client.SubResourcePatchOption) error { + if snapshotContent, ok := obj.(*snapshotv1alpha1.PodSnapshotContent); ok { + condition := meta.FindStatusCondition(snapshotContent.Status.Conditions, snapshotv1alpha1.PodSnapshotConditionFailed) + if condition != nil && condition.Reason == "CheckpointOutcomeUnknown" && !failedOnce { + failedOnce = true + return errors.New("status temporarily unavailable") + } + } + return c.Status().Patch(ctx, obj, patch, opts...) + }, + } + w := makeNodeControllerWithInterceptor(t, &fakeCheckpointer{}, funcs, content, pod) + w.runtime = &fakeRuntime{resolveContainerPID: 123} + signalCount := 0 + w.sendSignalFn = func(logr.Logger, int, syscall.Signal, string) error { + signalCount++ + return nil + } + foreignCaptureLease(t, w, content, true) + + err := w.reconcileSourcePod(context.Background(), pod) + require.ErrorContains(t, err, "status temporarily unavailable") + assert.Equal(t, 1, signalCount) + + pod.Status.ContainerStatuses[0].State = corev1.ContainerState{ + Terminated: &corev1.ContainerStateTerminated{ExitCode: 137}, + } + require.NoError(t, w.reconcileSourcePod(context.Background(), pod)) + assert.Equal(t, 1, signalCount, "a terminated exact target must not be signaled again") + condition := meta.FindStatusCondition(getContent(t, w, content.Name).Status.Conditions, snapshotv1alpha1.PodSnapshotConditionFailed) + require.NotNil(t, condition) + assert.Equal(t, "CheckpointOutcomeUnknown", condition.Reason) +} + func TestReconcilePodSnapshotContent_ContentGetErrorReturns(t *testing.T) { content := makeWorkOrder("podsnapshotcontent-x", "node-a", "x") pod := makeSourcePod() diff --git a/agent/internal/controller/podsnapshotcontent_test.go b/agent/internal/controller/podsnapshotcontent_test.go index 515427fc..05a9c712 100644 --- a/agent/internal/controller/podsnapshotcontent_test.go +++ b/agent/internal/controller/podsnapshotcontent_test.go @@ -9,6 +9,7 @@ import ( "os" "os/exec" "path/filepath" + "strconv" "sync" "syscall" "testing" @@ -30,6 +31,8 @@ import ( crfake "sigs.k8s.io/controller-runtime/pkg/client/fake" "sigs.k8s.io/controller-runtime/pkg/client/interceptor" + "github.com/ai-dynamo/snapshot/agent/internal/executor" + snapshotruntime "github.com/ai-dynamo/snapshot/agent/internal/runtime" snapshottypes "github.com/ai-dynamo/snapshot/agent/internal/types" snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" ) @@ -97,6 +100,7 @@ func makeNodeController(t *testing.T, fc *fakeCheckpointer, objs ...client.Objec holderID: "snapshot-agent/test", inFlight: make(map[string]struct{}), contentIndexer: idx, + sendSignalFn: snapshotruntime.SendSignalToPID, } w.checkpointFn = fc.fn return w @@ -131,7 +135,12 @@ func makeSourcePod() *corev1.Pod { Status: corev1.PodStatus{ Phase: corev1.PodRunning, ContainerStatuses: []corev1.ContainerStatus{ - {Name: "main", Ready: true, ContainerID: "containerd://abc123"}, + { + Name: "main", + Ready: true, + ContainerID: "containerd://abc123", + State: corev1.ContainerState{Running: &corev1.ContainerStateRunning{}}, + }, }, }, } @@ -501,6 +510,7 @@ func TestRunCheckpoint_LeaseCancelledAfterDumpFailsAndKills(t *testing.T) { } } ctx, target := startKillableTarget(t) + w.runtime = &fakeRuntime{resolveContainerPID: target.Process.Pid} pod := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{Name: "worker-0", Namespace: "inference", UID: types.UID("pod-uid")}} leaseKey := client.ObjectKey{Namespace: "inference", Name: "checkpoint-lease-abc"} artifactPath := filepath.Join(w.config.Storage.BasePath, "abc", "versions", "1") @@ -515,6 +525,27 @@ func TestRunCheckpoint_LeaseCancelledAfterDumpFailsAndKills(t *testing.T) { assert.Equal(t, "LeaseCancelled", failed.Reason) } +func TestKillCheckpointContainerRefusesReusedPID(t *testing.T) { + w := makeNodeController(t, &fakeCheckpointer{}) + w.runtime = &fakeRuntime{resolveContainerPID: 99} + + err := w.killCheckpointContainer(context.Background(), logr.Discard(), "abc123", 42, "checkpoint failed") + require.Error(t, err) + assert.Contains(t, err.Error(), "current PID 99 does not match expected PID 42") +} + +func TestCheckpointPreflightFailurePreservesSource(t *testing.T) { + err := executor.Checkpoint( + context.Background(), + &fakeRuntime{}, + logr.Discard(), + executor.CheckpointRequest{}, + &snapshottypes.AgentConfig{}, + ) + require.Error(t, err) + assert.False(t, checkpointFailureRequiresTermination(err)) +} + func TestRunCheckpoint_LeaseCancelledConflictReadyDoesNotKill(t *testing.T) { orig := checkpointLeaseRenewInterval checkpointLeaseRenewInterval = time.Millisecond @@ -814,6 +845,11 @@ func foreignCaptureLease(t *testing.T, w *NodeController, content *snapshotv1alp ObjectMeta: metav1.ObjectMeta{ Name: checkpointLeaseName(string(content.UID), "main"), Namespace: content.Spec.PodSnapshotRef.Namespace, + Annotations: captureLeaseTarget{ + podUID: "pod-uid", + containerID: "abc123", + pid: 123, + }.annotations(), }, Spec: coordinationv1.LeaseSpec{ HolderIdentity: &holder, @@ -881,6 +917,93 @@ func TestReconcilePodSnapshotContent_ExpiredForeignLeaseStillFails(t *testing.T) assert.Equal(t, "SourcePodGone", cond.Reason) } +func TestReconcileSourcePod_ExpiredLeaseDoesNotReplayLiveTarget(t *testing.T) { + content := makeWorkOrder("podsnapshotcontent-abc", "node-a", "abc") + pod := makeSourcePod() + checkpointer := &fakeCheckpointer{} + w := makeNodeController(t, checkpointer, content, pod) + ctx, target := startKillableTarget(t) + w.runtime = &fakeRuntime{resolveContainerPID: target.Process.Pid} + foreignCaptureLease(t, w, content, true) + leaseName := checkpointLeaseName(string(content.UID), "main") + lease, err := w.clientset.CoordinationV1().Leases("inference").Get(context.Background(), leaseName, metav1.GetOptions{}) + require.NoError(t, err) + lease.Annotations[captureLeasePIDAnnotation] = strconv.Itoa(target.Process.Pid) + _, err = w.clientset.CoordinationV1().Leases("inference").Update(context.Background(), lease, metav1.UpdateOptions{}) + require.NoError(t, err) + + require.NoError(t, w.reconcileSourcePod(context.Background(), pod)) + + assert.False(t, checkpointer.wasCalled(), "an unknown CRIU/CUDA outcome must not be replayed") + requireKilledBySIGKILL(t, ctx, target) + cond := meta.FindStatusCondition(getContent(t, w, content.Name).Status.Conditions, snapshotv1alpha1.PodSnapshotConditionFailed) + require.NotNil(t, cond) + assert.Equal(t, "CheckpointOutcomeUnknown", cond.Reason) +} + +func TestReconcileSourcePod_ExpiredLeaseTerminalizesWithoutUnsafeSignal(t *testing.T) { + for _, tc := range []struct { + name string + runtimePID int + mutate func(*corev1.Pod, *coordinationv1.Lease) + }{ + { + name: "legacy lease without target identity", + mutate: func(_ *corev1.Pod, lease *coordinationv1.Lease) { + lease.Annotations = nil + }, + }, + { + name: "replacement container", + mutate: func(pod *corev1.Pod, _ *coordinationv1.Lease) { + pod.Status.ContainerStatuses[0].ContainerID = "containerd://replacement" + }, + }, + { + name: "recorded container already terminated", + mutate: func(pod *corev1.Pod, _ *coordinationv1.Lease) { + pod.Status.ContainerStatuses[0].State = corev1.ContainerState{ + Terminated: &corev1.ContainerStateTerminated{ExitCode: 137}, + } + }, + }, + { + name: "recorded container PID changed", + runtimePID: 456, + mutate: func(*corev1.Pod, *coordinationv1.Lease) {}, + }, + } { + t.Run(tc.name, func(t *testing.T) { + content := makeWorkOrder("podsnapshotcontent-abc", "node-a", "abc") + pod := makeSourcePod() + checkpointer := &fakeCheckpointer{} + w := makeNodeController(t, checkpointer, content, pod) + runtimePID := tc.runtimePID + if runtimePID == 0 { + runtimePID = 123 + } + w.runtime = &fakeRuntime{resolveContainerPID: runtimePID} + w.sendSignalFn = func(logr.Logger, int, syscall.Signal, string) error { + t.Fatal("replacement or unidentifiable process must not be signaled") + return nil + } + foreignCaptureLease(t, w, content, true) + leaseName := checkpointLeaseName(string(content.UID), "main") + lease, err := w.clientset.CoordinationV1().Leases("inference").Get(context.Background(), leaseName, metav1.GetOptions{}) + require.NoError(t, err) + tc.mutate(pod, lease) + _, err = w.clientset.CoordinationV1().Leases("inference").Update(context.Background(), lease, metav1.UpdateOptions{}) + require.NoError(t, err) + + require.NoError(t, w.reconcileSourcePod(context.Background(), pod)) + assert.False(t, checkpointer.wasCalled()) + cond := meta.FindStatusCondition(getContent(t, w, content.Name).Status.Conditions, snapshotv1alpha1.PodSnapshotConditionFailed) + require.NotNil(t, cond) + assert.Equal(t, "CheckpointOutcomeUnknown", cond.Reason) + }) + } +} + // TestReconcileSourcePod_ForeignLeaseDefersContainerExitFailure covers the capture path's window: // B's reconcileSourcePod sees the target already killed by A's in-flight dump (exit 137, no // artifact yet) and must not write CheckpointContainerFailed or SIGKILL anything under A's Lease. diff --git a/agent/internal/controller/util.go b/agent/internal/controller/util.go index 991d7476..47fc89b7 100644 --- a/agent/internal/controller/util.go +++ b/agent/internal/controller/util.go @@ -7,6 +7,7 @@ import ( "context" "crypto/sha256" "fmt" + "strconv" "time" "github.com/go-logr/logr" @@ -21,6 +22,46 @@ import ( const checkpointLeaseDuration = 30 * time.Second +const ( + captureLeasePodUIDAnnotation = "snapshot.nvidia.com/capture-pod-uid" + captureLeaseContainerIDAnnotation = "snapshot.nvidia.com/capture-container-id" + captureLeasePIDAnnotation = "snapshot.nvidia.com/capture-host-pid" +) + +type captureLeaseTarget struct { + podUID string + containerID string + pid int +} + +func (t captureLeaseTarget) annotations() map[string]string { + return map[string]string{ + captureLeasePodUIDAnnotation: t.podUID, + captureLeaseContainerIDAnnotation: t.containerID, + captureLeasePIDAnnotation: strconv.Itoa(t.pid), + } +} + +func captureTargetFromLease(lease *coordinationv1.Lease) (captureLeaseTarget, error) { + if lease == nil { + return captureLeaseTarget{}, fmt.Errorf("capture lease is missing") + } + annotations := lease.GetAnnotations() + pid, err := strconv.Atoi(annotations[captureLeasePIDAnnotation]) + if err != nil || pid <= 0 { + return captureLeaseTarget{}, fmt.Errorf("capture lease has invalid host PID %q", annotations[captureLeasePIDAnnotation]) + } + target := captureLeaseTarget{ + podUID: annotations[captureLeasePodUIDAnnotation], + containerID: annotations[captureLeaseContainerIDAnnotation], + pid: pid, + } + if target.podUID == "" || target.containerID == "" { + return captureLeaseTarget{}, fmt.Errorf("capture lease is missing source identity") + } + return target, nil +} + // checkpointLeaseRenewInterval is a package-level var (not const) so tests can shorten the // renewal loop without a fake clock. Only same-package _test.go files should mutate it. var checkpointLeaseRenewInterval = 10 * time.Second @@ -72,7 +113,7 @@ func checkpointLeaseName(contentUID, containerName string) string { // acquireLease acquires or renews a checkpoint lease at an arbitrary namespace/name key, // returning false when another live holder owns it. -func (w *NodeController) acquireLease(ctx context.Context, key client.ObjectKey) (bool, error) { +func (w *NodeController) acquireLease(ctx context.Context, key client.ObjectKey, target captureLeaseTarget) (bool, error) { now := metav1.NewMicroTime(time.Now()) leaseDurationSeconds := int32(checkpointLeaseDuration.Seconds()) @@ -83,7 +124,7 @@ func (w *NodeController) acquireLease(ctx context.Context, key client.ObjectKey) return false, fmt.Errorf("get checkpoint lease %s: %w", key.String(), err) } lease := &coordinationv1.Lease{ - ObjectMeta: metav1.ObjectMeta{Name: key.Name, Namespace: key.Namespace}, + ObjectMeta: metav1.ObjectMeta{Name: key.Name, Namespace: key.Namespace, Annotations: target.annotations()}, Spec: coordinationv1.LeaseSpec{ HolderIdentity: &w.holderID, LeaseDurationSeconds: &leaseDurationSeconds, @@ -106,6 +147,7 @@ func (w *NodeController) acquireLease(ctx context.Context, key client.ObjectKey) return false, nil } existing.Spec.HolderIdentity = &w.holderID + existing.Annotations = target.annotations() existing.Spec.LeaseDurationSeconds = &leaseDurationSeconds if existing.Spec.AcquireTime == nil || checkpointLeaseExpired(existing, now.Time) { existing.Spec.AcquireTime = &now diff --git a/agent/internal/criu/restore.go b/agent/internal/criu/restore.go index ce1d14f6..5b3db431 100644 --- a/agent/internal/criu/restore.go +++ b/agent/internal/criu/restore.go @@ -40,10 +40,10 @@ func ExecuteRestore( settings := m.CRIUDump.CRIU var prepare, restore time.Duration - // Return the FD closers as cleanup() rather than deferring them here, so the - // caller can run them after cuda unlock instead of between the CRIU restore - // and unlock. That keeps the window where the restored process runs with CUDA - // still locked as short as possible. cleanup is called on the error paths below. + // Return the FD closers as cleanup() rather than deferring them here so the + // caller controls their lifetime. The current nsrestore path releases these + // CRIU-only resources before returning the restored process identities to the + // host agent for deferred CUDA restore. cleanup is called on error paths below. var openFiles, inheritedFiles []*os.File scratchDir, removeScratch, err := restoreScratchDir(settings.WorkDir) if err != nil { diff --git a/agent/internal/cuda/cuda.go b/agent/internal/cuda/cuda.go index 67430b83..50172f83 100644 --- a/agent/internal/cuda/cuda.go +++ b/agent/internal/cuda/cuda.go @@ -9,6 +9,7 @@ import ( "errors" "fmt" "os/exec" + "path/filepath" "regexp" "strconv" "strings" @@ -19,6 +20,9 @@ import ( "google.golang.org/grpc/credentials/insecure" "k8s.io/client-go/kubernetes" podresourcesv1 "k8s.io/kubelet/pkg/apis/podresources/v1" + + snapshotruntime "github.com/ai-dynamo/snapshot/agent/internal/runtime" + "github.com/ai-dynamo/snapshot/agent/internal/types" ) const ( @@ -36,10 +40,123 @@ var podResourcesSocketPath = "/var/lib/kubelet/pod-resources/kubelet.sock" var gpuUUIDPattern = regexp.MustCompile(`^GPU-[a-fA-F0-9]{8}-[a-fA-F0-9]{4}-[a-fA-F0-9]{4}-[a-fA-F0-9]{4}-[a-fA-F0-9]{12}$`) +// Each Snapshot agent process deliberately executes one CUDA operation at a time. +// Hold the slot across each process-tree sequence so checkpoints/restores from +// different workloads cannot interleave lock, transfer, and unlock requests. +var cudaOperationSlot = make(chan struct{}, 1) + +func acquireCUDAOperation(ctx context.Context) error { + select { + case cudaOperationSlot <- struct{}{}: + return nil + case <-ctx.Done(): + return fmt.Errorf("wait for process-local CUDA operation slot: %w", context.Cause(ctx)) + } +} + +func acquireCUDAOperationLogged(ctx context.Context, log logr.Logger) error { + started := time.Now() + err := acquireCUDAOperation(ctx) + log.Info("CUDA operation slot wait finished", "wait_duration", time.Since(started), "acquired", err == nil) + return err +} + +func releaseCUDAOperation() { <-cudaOperationSlot } + type CheckpointPhaseTimings struct { TotalDuration time.Duration } +type checkpointOperationError struct { + err error + targetMayBeMutated bool +} + +func (e *checkpointOperationError) Error() string { return e.err.Error() } +func (e *checkpointOperationError) Unwrap() error { return e.err } + +// FailedBeforeTargetMutation is true only when the helper could not be +// contacted before any CUDA lock succeeded. Unknown RPC outcomes and CUDA +// failures remain conservatively classified as possibly mutating the target. +func FailedBeforeTargetMutation(err error) bool { + var operationErr *checkpointOperationError + return errors.As(err, &operationErr) && !operationErr.targetMayBeMutated +} + +// LockAndCheckpointProcessTreeValidated locks and checkpoints the identities +// captured before any destructive CUDA operation and revalidates each identity +// immediately before every driver call. +func LockAndCheckpointProcessTreeValidated( + ctx context.Context, + processes []snapshotruntime.ProcessDetails, + jobFile, + storageMode, + checkpointDir string, + gpuUUIDs []string, + transferSettings types.CUDATransferSettings, + log logr.Logger, +) (CheckpointPhaseTimings, error) { + if err := validateCUDAOperationBudget(ctx, actionCheckpoint, len(processes)); err != nil { + return CheckpointPhaseTimings{}, &checkpointOperationError{ + err: err, + targetMayBeMutated: false, + } + } + if err := acquireCUDAOperationLogged(ctx, log); err != nil { + return CheckpointPhaseTimings{}, &checkpointOperationError{ + err: err, + targetMayBeMutated: false, + } + } + defer releaseCUDAOperation() + pids, targetIDs, identities, err := validatedProcessIdentities(processes) + if err != nil { + return CheckpointPhaseTimings{}, &checkpointOperationError{ + err: err, + targetMayBeMutated: false, + } + } + return lockAndCheckpointProcessTree( + ctx, + pids, + targetIDs, + jobFile, + storageMode, + checkpointDir, + gpuUUIDs, + transferSettings, + identityValidatingRunner{ + runner: commandHelperActionRunner{}, + procRoot: snapshotruntime.HostProcPath, + identities: identities, + }, + log, + ) +} + +func validatedProcessIdentities(processes []snapshotruntime.ProcessDetails) ([]int, []int, map[int]snapshotruntime.ProcessDetails, error) { + pids := make([]int, 0, len(processes)) + targetIDs := make([]int, 0, len(processes)) + identities := make(map[int]snapshotruntime.ProcessDetails, len(processes)) + seenTargetIDs := make(map[int]struct{}, len(processes)) + for _, process := range processes { + if process.OutermostPID <= 0 || process.InnermostPID <= 0 || process.StartTimeTicks == 0 || process.Cgroup == "" { + return nil, nil, nil, fmt.Errorf("invalid host process identity") + } + if _, exists := identities[process.OutermostPID]; exists { + return nil, nil, nil, fmt.Errorf("duplicate host PID %d", process.OutermostPID) + } + if _, exists := seenTargetIDs[process.InnermostPID]; exists { + return nil, nil, nil, fmt.Errorf("duplicate CUDA target namespace PID %d", process.InnermostPID) + } + pids = append(pids, process.OutermostPID) + targetIDs = append(targetIDs, process.InnermostPID) + identities[process.OutermostPID] = process + seenTargetIDs[process.InnermostPID] = struct{}{} + } + return pids, targetIDs, identities, nil +} + type RestorePhaseTimings struct { TotalDuration time.Duration } @@ -243,7 +360,23 @@ func orderDRAUUIDsByRuntime(allocatedUUIDs, visibleUUIDs []string) ([]string, er // --get-state, because --get-state incorrectly matches coordinator processes like // cuda-checkpoint --launch-job that share a /proc namespace with CUDA processes but // don't hold CUDA contexts themselves. -func FilterProcesses(ctx context.Context, allPIDs []int, log logr.Logger) []int { +func parseRestoreTIDProbeOutput(output []byte) (bool, error) { + value := strings.TrimSpace(string(output)) + if value == "none" { + return false, nil + } + tid, err := strconv.Atoi(value) + if err != nil || tid <= 0 { + return false, fmt.Errorf("invalid CUDA restore-tid probe response %q", value) + } + return true, nil +} + +func FilterProcesses(ctx context.Context, allPIDs []int, log logr.Logger) ([]int, error) { + if err := acquireCUDAOperationLogged(ctx, log); err != nil { + return nil, fmt.Errorf("acquire process-local CUDA operation slot for restore-tid discovery: %w", err) + } + defer releaseCUDAOperation() cudaPIDs := make([]int, 0, len(allPIDs)) for _, pid := range allPIDs { if pid <= 0 { @@ -253,16 +386,22 @@ func FilterProcesses(ctx context.Context, allPIDs []int, log logr.Logger) []int output, err := cmd.CombinedOutput() if err != nil { if ctx.Err() != nil { - break + return nil, fmt.Errorf("discover CUDA process identity: %w", context.Cause(ctx)) } + return nil, fmt.Errorf("CUDA restore-tid probe failed for PID %d: %w (output: %s)", pid, err, strings.TrimSpace(string(output))) + } + ownsCUDAState, err := parseRestoreTIDProbeOutput(output) + if err != nil { + return nil, fmt.Errorf("CUDA restore-tid probe failed for PID %d: %w", pid, err) + } + if !ownsCUDAState { log.V(1).Info("CUDA restore-tid probe negative", "pid", pid) continue } - tid := strings.TrimSpace(string(output)) - log.V(1).Info("CUDA restore-tid probe positive", "pid", pid, "tid", tid) + log.V(1).Info("CUDA restore-tid probe positive", "pid", pid, "tid", strings.TrimSpace(string(output))) cudaPIDs = append(cudaPIDs, pid) } - return cudaPIDs + return cudaPIDs, nil } // BuildDeviceMap creates a cuda-checkpoint-helper --device-map value from source and target GPU UUID lists. @@ -328,57 +467,180 @@ func BuildDeviceMap(sourceUUIDs, targetUUIDs []string, log logr.Logger) (string, return strings.Join(pairs, ","), nil } -// CheckpointProcessTree locks and checkpoints CUDA state for all given PIDs, -// then persists the launch-job state needed to restore them. -// On failure, the caller is expected to fail the operation and terminate the workload. -func CheckpointProcessTree(ctx context.Context, cudaPIDs []int, jobFile, checkpointDir string, log logr.Logger) (CheckpointPhaseTimings, error) { +func validateTransferSettings(transferSettings types.CUDATransferSettings) (types.CUDATransferSettings, error) { + transferSettings = transferSettings.WithDefaults() + if err := transferSettings.Validate(); err != nil { + return types.CUDATransferSettings{}, fmt.Errorf("invalid CUDA transfer settings: %w", err) + } + return transferSettings, nil +} + +func lockAndCheckpointProcessTree( + ctx context.Context, + cudaPIDs []int, + targetIDs []int, + jobFile, + storageMode, + checkpointDir string, + gpuUUIDs []string, + transferSettings types.CUDATransferSettings, + runner helperActionRunner, + log logr.Logger, +) (CheckpointPhaseTimings, error) { + // Once any target has been locked, an error is intentionally not recovered + // in this package. The caller must treat the process tree as unsafe and + // terminate it after revalidating every process identity. Unlocking only a + // subset could resume a workload whose CUDA targets no longer agree. var timings CheckpointPhaseTimings + var err error + transferSettings, err = validateTransferSettings(transferSettings) + if err != nil { + return timings, err + } + if storageMode == types.CUDAStorageModePOSIX && len(targetIDs) != len(cudaPIDs) { + return timings, fmt.Errorf("CUDA target identity count %d does not match PID count %d", len(targetIDs), len(cudaPIDs)) + } start := time.Now() + locked := 0 for _, pid := range cudaPIDs { - if err := lockWithJobFile(ctx, pid, jobFile, log); err != nil { + if err := runner.run(ctx, helperAction{PID: pid, Action: actionLock, StorageMode: storageMode, JobFile: jobFile, Transfer: transferSettings}, log); err != nil { timings.TotalDuration = time.Since(start) - return timings, err + return timings, &checkpointOperationError{ + err: err, + targetMayBeMutated: locked > 0 || + (!errors.Is(err, errDaemonUnavailable) && + !errors.Is(err, errCheckpointLockNotAcquired) && + !errors.Is(err, errProcessIdentityChangedBeforeCUDA)), + } } + locked++ } - for _, pid := range cudaPIDs { - if err := checkpointWithJobFile(ctx, pid, jobFile, log); err != nil { + for index, pid := range cudaPIDs { + processDir := "" + var selectedDevices []string + if storageMode == types.CUDAStorageModePOSIX { + processDir = customStorageProcessDir(checkpointDir, targetIDs[index]) + selectedDevices = gpuUUIDs + } + if err := runner.run(ctx, helperAction{PID: pid, Action: actionCheckpoint, StorageMode: storageMode, StorageDir: processDir, JobFile: jobFile, GPUUUIDs: selectedDevices, Transfer: transferSettings}, log); err != nil { timings.TotalDuration = time.Since(start) - return timings, err + return timings, &checkpointOperationError{err: err, targetMayBeMutated: true} } } if err := refreshJobFileArtifact(jobFile, checkpointDir); err != nil { timings.TotalDuration = time.Since(start) - return timings, err + return timings, &checkpointOperationError{err: err, targetMayBeMutated: true} } timings.TotalDuration = time.Since(start) return timings, nil } -// RestoreAndUnlockProcessTree restores and unlocks CUDA state for the given PIDs. -// helperBinaryPath must be the absolute path to cuda-checkpoint-helper: DefaultHelperBinaryPath -// on the agent, or filepath.Join(bundleDir, HelperBinaryName) inside the placeholder namespace. -func RestoreAndUnlockProcessTree(ctx context.Context, cudaPIDs []int, deviceMap, helperBinaryPath string, log logr.Logger) (RestorePhaseTimings, error) { +func customStorageProcessDir(checkpointDir string, namespacePID int) string { + return filepath.Join(checkpointDir, "cuda-custom-storage", fmt.Sprintf("process-nspid-%d", namespacePID)) +} + +// RestoreAndUnlockProcessTreeValidated restores and unlocks identities resolved +// after CRIU and revalidates them immediately before every CUDA driver call. +func RestoreAndUnlockProcessTreeValidated( + ctx context.Context, + processes []snapshotruntime.ProcessDetails, + deviceMap, + storageMode, + checkpointDir, + jobFile string, + targetGPUUUIDs []string, + transferSettings types.CUDATransferSettings, + log logr.Logger, +) (RestorePhaseTimings, error) { + if err := validateCUDAOperationBudget(ctx, actionRestore, len(processes)); err != nil { + return RestorePhaseTimings{}, err + } + if err := acquireCUDAOperationLogged(ctx, log); err != nil { + return RestorePhaseTimings{}, err + } + defer releaseCUDAOperation() + pids, targetIDs, identities, err := validatedProcessIdentities(processes) + if err != nil { + return RestorePhaseTimings{}, err + } + return restoreAndUnlockProcessTree( + ctx, + pids, + targetIDs, + deviceMap, + storageMode, + checkpointDir, + jobFile, + targetGPUUUIDs, + transferSettings, + identityValidatingRunner{ + runner: commandHelperActionRunner{}, + procRoot: snapshotruntime.HostProcPath, + identities: identities, + }, + log, + ) +} + +func restoreAndUnlockProcessTree( + ctx context.Context, + cudaPIDs []int, + targetIDs []int, + deviceMap, + storageMode, + checkpointDir, + jobFile string, + targetGPUUUIDs []string, + transferSettings types.CUDATransferSettings, + runner helperActionRunner, + log logr.Logger, +) (RestorePhaseTimings, error) { var timings RestorePhaseTimings + var err error + transferSettings, err = validateTransferSettings(transferSettings) + if err != nil { + return timings, err + } + if storageMode == types.CUDAStorageModePOSIX && len(targetIDs) != len(cudaPIDs) { + return timings, fmt.Errorf("CUDA target identity count %d does not match PID count %d", len(targetIDs), len(cudaPIDs)) + } start := time.Now() - for _, pid := range cudaPIDs { - if err := restoreProcess(ctx, pid, deviceMap, helperBinaryPath, log); err != nil { + for index, pid := range cudaPIDs { + processDir := "" + var selectedDevices []string + requestJobFile := jobFile + if jobFile != "" { + requestJobFile, err = HostJobFilePath(pid) + if err != nil { + timings.TotalDuration = time.Since(start) + return timings, err + } + } + if storageMode == types.CUDAStorageModePOSIX { + processDir = customStorageProcessDir(checkpointDir, targetIDs[index]) + selectedDevices = targetGPUUUIDs + } + if err := runner.run(ctx, helperAction{PID: pid, Action: actionRestore, DeviceMap: deviceMap, StorageMode: storageMode, StorageDir: processDir, JobFile: requestJobFile, GPUUUIDs: selectedDevices, Transfer: transferSettings}, log); err != nil { timings.TotalDuration = time.Since(start) return timings, err } } for _, pid := range cudaPIDs { - if err := unlock(ctx, pid, helperBinaryPath, log); err != nil { - timings.TotalDuration = time.Since(start) - state, stateErr := getState(ctx, pid, helperBinaryPath) - if stateErr == nil && state == "running" { - log.Info("cuda-checkpoint-helper unlock returned error but process is already running", "pid", pid) - continue + requestJobFile := jobFile + if jobFile != "" { + requestJobFile, err = HostJobFilePath(pid) + if err != nil { + timings.TotalDuration = time.Since(start) + return timings, err } + } + if err := runner.run(ctx, helperAction{PID: pid, Action: actionUnlock, StorageMode: storageMode, JobFile: requestJobFile, Transfer: transferSettings}, log); err != nil { + timings.TotalDuration = time.Since(start) return timings, err } } diff --git a/agent/internal/cuda/cuda_test.go b/agent/internal/cuda/cuda_test.go index f486dc76..b506595a 100644 --- a/agent/internal/cuda/cuda_test.go +++ b/agent/internal/cuda/cuda_test.go @@ -6,6 +6,7 @@ package cuda import ( "context" "errors" + "fmt" "net" "path/filepath" "strings" @@ -21,8 +22,203 @@ import ( metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/client-go/kubernetes/fake" podresourcesv1 "k8s.io/kubelet/pkg/apis/podresources/v1" + + snapshotruntime "github.com/ai-dynamo/snapshot/agent/internal/runtime" + "github.com/ai-dynamo/snapshot/agent/internal/types" ) +func TestCUDAOperationSlotHonorsCancellation(t *testing.T) { + if err := acquireCUDAOperation(context.Background()); err != nil { + t.Fatalf("acquireCUDAOperation() = %v", err) + } + defer releaseCUDAOperation() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if err := acquireCUDAOperation(ctx); err == nil || !errors.Is(err, context.Canceled) { + t.Fatalf("acquireCUDAOperation(canceled) = %v, want context.Canceled", err) + } + + _, err := LockAndCheckpointProcessTreeValidated( + ctx, nil, "", "", "", nil, types.CUDATransferSettings{}, logr.Discard(), + ) + if err == nil || !errors.Is(err, context.Canceled) { + t.Fatalf("LockAndCheckpointProcessTreeValidated(canceled) = %v, want context.Canceled", err) + } + if !FailedBeforeTargetMutation(err) { + t.Fatalf("slot-acquisition failure = %v, want pre-mutation classification", err) + } +} + +func TestParseRestoreTIDProbeOutput(t *testing.T) { + for _, test := range []struct { + name string + output string + want bool + wantErr bool + }{ + {name: "CUDA owner", output: "42\n", want: true}, + {name: "no CUDA context", output: "none\n", want: false}, + {name: "empty", output: "", wantErr: true}, + {name: "zero TID", output: "0\n", wantErr: true}, + {name: "diagnostic", output: "CUDA_ERROR_UNKNOWN\n", wantErr: true}, + } { + t.Run(test.name, func(t *testing.T) { + got, err := parseRestoreTIDProbeOutput([]byte(test.output)) + if (err != nil) != test.wantErr || got != test.want { + t.Fatalf("parseRestoreTIDProbeOutput(%q) = %v, %v; want %v, error=%v", test.output, got, err, test.want, test.wantErr) + } + }) + } +} + +func TestFailedBeforeTargetMutation(t *testing.T) { + safe := &checkpointOperationError{ + err: errDaemonUnavailable, + targetMayBeMutated: false, + } + if !FailedBeforeTargetMutation(safe) { + t.Fatal("unavailable helper before the first lock should be classified as pre-mutation") + } + unsafe := &checkpointOperationError{ + err: errDaemonUnavailable, + targetMayBeMutated: true, + } + if FailedBeforeTargetMutation(unsafe) || FailedBeforeTargetMutation(errors.New("unknown")) { + t.Fatal("mutated or unclassified failures must not be classified as pre-mutation") + } +} + +func TestCheckpointRejectsShortBudgetBeforeMutation(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + + _, err := LockAndCheckpointProcessTreeValidated( + ctx, + []snapshotruntime.ProcessDetails{{OutermostPID: 42, StartTimeTicks: 1, Cgroup: "0::/test\n"}}, + "", "", "", nil, types.CUDATransferSettings{}, logr.Discard(), + ) + if err == nil || !FailedBeforeTargetMutation(err) { + t.Fatalf("short-budget checkpoint error = %v, want pre-mutation classification", err) + } +} + +type helperActionRunnerFunc func(context.Context, helperAction, logr.Logger) error + +func (f helperActionRunnerFunc) run(ctx context.Context, request helperAction, log logr.Logger) error { + return f(ctx, request, log) +} + +func TestIdentityValidatingRunnerRejectsMissingAndChangedIdentityBeforeDriver(t *testing.T) { + for _, test := range []struct { + name string + identities map[int]snapshotruntime.ProcessDetails + }{ + {name: "missing", identities: map[int]snapshotruntime.ProcessDetails{}}, + {name: "changed", identities: map[int]snapshotruntime.ProcessDetails{42: { + OutermostPID: 42, InnermostPID: 42, StartTimeTicks: 1, Cgroup: "0::/test\n", + }}}, + } { + t.Run(test.name, func(t *testing.T) { + called := false + runner := identityValidatingRunner{ + runner: helperActionRunnerFunc(func(context.Context, helperAction, logr.Logger) error { + called = true + return nil + }), + procRoot: t.TempDir(), + identities: test.identities, + } + err := runner.run(context.Background(), helperAction{PID: 42, Action: actionLock}, logr.Discard()) + if !errors.Is(err, errProcessIdentityChangedBeforeCUDA) { + t.Fatalf("identityValidatingRunner.run() error = %v, want identity-change sentinel", err) + } + if called { + t.Fatal("identityValidatingRunner called the driver-facing runner after identity rejection") + } + }) + } +} + +func TestCheckpointIdentityRaceBeforeFirstLockIsPreMutation(t *testing.T) { + _, err := lockAndCheckpointProcessTree( + context.Background(), []int{41}, nil, "", types.CUDAStorageModeLegacy, "", nil, + types.CUDATransferSettings{}, + helperActionRunnerFunc(func(context.Context, helperAction, logr.Logger) error { + return fmt.Errorf("%w: test race", errProcessIdentityChangedBeforeCUDA) + }), + logr.Discard(), + ) + if err == nil || !FailedBeforeTargetMutation(err) { + t.Fatalf("first-lock identity race = %v, want pre-mutation classification", err) + } +} + +func TestCheckpointIdentityRaceAfterEarlierLockIsMutating(t *testing.T) { + calls := 0 + _, err := lockAndCheckpointProcessTree( + context.Background(), []int{41, 42}, nil, "", types.CUDAStorageModeLegacy, "", nil, + types.CUDATransferSettings{}, + helperActionRunnerFunc(func(context.Context, helperAction, logr.Logger) error { + calls++ + if calls == 2 { + return fmt.Errorf("%w: test race", errProcessIdentityChangedBeforeCUDA) + } + return nil + }), + logr.Discard(), + ) + if err == nil || FailedBeforeTargetMutation(err) { + t.Fatalf("second-lock identity race = %v, want possibly-mutating classification", err) + } +} + +func TestCheckpointRejectsInvalidIdentityBeforeMutation(t *testing.T) { + _, err := LockAndCheckpointProcessTreeValidated( + context.Background(), + []snapshotruntime.ProcessDetails{{OutermostPID: 42}}, + "", "", "", nil, types.CUDATransferSettings{}, logr.Discard(), + ) + if err == nil || !FailedBeforeTargetMutation(err) { + t.Fatalf("invalid-identity checkpoint error = %v, want pre-mutation classification", err) + } +} + +func TestCustomStorageProcessDirectoryUsesStableNamespacePID(t *testing.T) { + got := customStorageProcessDir("/checkpoint", 42) + if want := "/checkpoint/cuda-custom-storage/process-nspid-42"; got != want { + t.Fatalf("customStorageProcessDir() = %q, want %q", got, want) + } +} + +func TestRestoreDerivesJobFilePathForEachTargetPID(t *testing.T) { + var requests []helperAction + runner := helperActionRunnerFunc(func(_ context.Context, request helperAction, _ logr.Logger) error { + requests = append(requests, request) + return nil + }) + _, err := restoreAndUnlockProcessTree( + context.Background(), []int{101, 202}, nil, "", types.CUDAStorageModeLegacy, + "/checkpoint", "job-file-present", nil, types.CUDATransferSettings{}, runner, logr.Discard(), + ) + if err != nil { + t.Fatalf("restoreAndUnlockProcessTree() error = %v", err) + } + if len(requests) != 4 { + t.Fatalf("request count = %d, want 4", len(requests)) + } + for index, request := range requests { + pid := []int{101, 202, 101, 202}[index] + want, err := HostJobFilePath(pid) + if err != nil { + t.Fatal(err) + } + if request.PID != pid || request.JobFile != want { + t.Fatalf("request[%d] = PID %d job %q, want PID %d job %q", index, request.PID, request.JobFile, pid, want) + } + } +} + func TestBuildDeviceMap(t *testing.T) { tests := []struct { name string diff --git a/agent/internal/cuda/daemon_client.go b/agent/internal/cuda/daemon_client.go new file mode 100644 index 00000000..7038b515 --- /dev/null +++ b/agent/internal/cuda/daemon_client.go @@ -0,0 +1,453 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package cuda + +import ( + "context" + "encoding/binary" + "errors" + "fmt" + "net" + "os" + "strings" + "time" + + "github.com/go-logr/logr" + + snapshotruntime "github.com/ai-dynamo/snapshot/agent/internal/runtime" + "github.com/ai-dynamo/snapshot/agent/internal/types" +) + +const ( + daemonProtocolMagic = uint32(0x50484344) + daemonProtocolVersion = uint16(6) + daemonRequestHeader = 56 + daemonResponseHeader = 24 + daemonMaxRequest = 64 * 1024 + daemonMaxResponse = 128 * 1024 + daemonMaxCgroup = 4096 + daemonMaxJobFile = 4096 + + daemonActionHealth = uint16(0) + daemonActionCheckpoint = uint16(1) + daemonActionRestore = uint16(2) + daemonActionLock = uint16(3) + daemonActionUnlock = uint16(4) + + daemonResponseFatal = uint32(1 << 0) + daemonCapabilityDeferredCUDA = uint32(1 << 1) + daemonCapabilityCustomStorage = uint32(1 << 2) + daemonResponseLockNotAcquired = uint32(1 << 3) + // Match the helper's chart startup probe (2 seconds * 150 attempts). The + // helper and agent are regular containers and Kubernetes does not guarantee + // their startup order. + daemonHealthWait = 5 * time.Minute + daemonHealthRetryInterval = 100 * time.Millisecond + // The daemon rejects watchdogs longer than one hour. Keep the client wait + // five minutes longer so the helper, not the caller, resolves an operation + // that has already reached CUDA. + daemonRPCTimeout = time.Hour + 5*time.Minute +) + +var errDaemonUnavailable = errors.New("CUDA helper daemon unavailable") +var errDaemonFatal = errors.New("CUDA helper daemon entered fatal state") +var errCheckpointLockNotAcquired = errors.New("CUDA checkpoint lock was not acquired before mutation") +var daemonSocketPath = types.CUDAHelperSocketPath + +func daemonRequest(request helperAction) ([]byte, error) { + pid := request.PID + action := request.Action + deviceMap := request.DeviceMap + storageMode := request.StorageMode + storageDir := request.StorageDir + jobFile := request.JobFile + selectedDevices := request.GPUUUIDs + transfer := request.Transfer + identity := request.Identity + var daemonAction uint16 + switch action { + case "": + daemonAction = daemonActionHealth + case actionCheckpoint: + daemonAction = daemonActionCheckpoint + case actionRestore: + daemonAction = daemonActionRestore + case actionLock: + daemonAction = daemonActionLock + case actionUnlock: + daemonAction = daemonActionUnlock + default: + return nil, fmt.Errorf("action %q is not supported by CUDA helper daemon", action) + } + health := daemonAction == daemonActionHealth + var backend uint16 + switch storageMode { + case "": + case types.CUDAStorageModeLegacy: + backend = 1 + case types.CUDAStorageModePOSIX: + backend = 2 + default: + return nil, fmt.Errorf("unsupported CUDA storage mode %q", storageMode) + } + if health { + pid = 0 + identity = snapshotruntime.ProcessDetails{} + deviceMap = "" + storageDir = "" + jobFile = "" + selectedDevices = nil + transfer = types.CUDATransferSettings{} + if backend != 0 { + return nil, errors.New("CUDA helper daemon health request has a backend") + } + } else if pid <= 0 || identity.OutermostPID != pid || identity.StartTimeTicks == 0 || + identity.Cgroup == "" || len(identity.Cgroup) > daemonMaxCgroup { + return nil, errors.New("invalid CUDA helper daemon process identity") + } + if daemonAction == daemonActionLock || daemonAction == daemonActionUnlock { + if backend == 0 || deviceMap != "" || storageDir != "" || len(selectedDevices) != 0 { + return nil, errors.New("CUDA helper daemon lock/unlock request has transfer arguments") + } + transfer = types.CUDATransferSettings{} + } else if !health { + if backend == 0 { + return nil, errors.New("CUDA helper daemon checkpoint/restore request has no backend") + } + if daemonAction == daemonActionCheckpoint && deviceMap != "" { + return nil, errors.New("CUDA helper daemon checkpoint request has a device map") + } + if backend == 1 { + if storageDir != "" || len(selectedDevices) != 0 { + return nil, errors.New("regular CUDA helper request has a storage directory") + } + transfer = types.CUDATransferSettings{} + } else { + if storageDir == "" || storageDir[0] != '/' || len(selectedDevices) == 0 { + return nil, errors.New("invalid CUDA helper daemon POSIX storage directory") + } + if err := transfer.Validate(); err != nil { + return nil, fmt.Errorf("invalid CUDA helper daemon transfer settings: %w", err) + } + } + } + seenDevices := make(map[string]struct{}, len(selectedDevices)) + for _, uuid := range selectedDevices { + if !gpuUUIDPattern.MatchString(uuid) { + return nil, fmt.Errorf("invalid selected CUDA device %q", uuid) + } + canonicalUUID := strings.ToLower(uuid) + if _, duplicate := seenDevices[canonicalUUID]; duplicate { + return nil, fmt.Errorf("duplicate selected CUDA device %q", uuid) + } + seenDevices[canonicalUUID] = struct{}{} + } + selectedDeviceList := strings.Join(selectedDevices, ",") + for name, value := range map[string]string{ + "device map": deviceMap, + "storage directory": storageDir, + "process cgroup": identity.Cgroup, + "job file": jobFile, + "selected devices": selectedDeviceList, + } { + if strings.ContainsRune(value, '\x00') { + return nil, fmt.Errorf("CUDA helper daemon %s contains NUL", name) + } + } + if len(jobFile) > daemonMaxJobFile || (jobFile != "" && jobFile[0] != '/') { + return nil, errors.New("invalid CUDA helper daemon job file") + } + if len(deviceMap)+len(storageDir)+len(identity.Cgroup)+len(jobFile)+len(selectedDeviceList) > daemonMaxRequest-daemonRequestHeader { + return nil, errors.New("CUDA helper daemon request is too large") + } + packet := make([]byte, daemonRequestHeader+len(deviceMap)+len(storageDir)+len(identity.Cgroup)+len(jobFile)+len(selectedDeviceList)) + binary.LittleEndian.PutUint32(packet[0:4], daemonProtocolMagic) + binary.LittleEndian.PutUint16(packet[4:6], daemonProtocolVersion) + binary.LittleEndian.PutUint16(packet[6:8], daemonRequestHeader) + binary.LittleEndian.PutUint16(packet[8:10], daemonAction) + binary.LittleEndian.PutUint16(packet[10:12], backend) + binary.LittleEndian.PutUint32(packet[12:16], uint32(pid)) + binary.LittleEndian.PutUint32(packet[16:20], uint32(transfer.BufferCount)) + binary.LittleEndian.PutUint64(packet[20:28], transfer.ChunkBytes) + binary.LittleEndian.PutUint32(packet[28:32], uint32(len(deviceMap))) + binary.LittleEndian.PutUint32(packet[32:36], uint32(len(storageDir))) + binary.LittleEndian.PutUint32(packet[36:40], uint32(len(identity.Cgroup))) + binary.LittleEndian.PutUint64(packet[40:48], identity.StartTimeTicks) + binary.LittleEndian.PutUint32(packet[48:52], uint32(len(jobFile))) + binary.LittleEndian.PutUint32(packet[52:56], uint32(len(selectedDeviceList))) + copy(packet[daemonRequestHeader:], deviceMap) + copy(packet[daemonRequestHeader+len(deviceMap):], storageDir) + copy(packet[daemonRequestHeader+len(deviceMap)+len(storageDir):], identity.Cgroup) + copy(packet[daemonRequestHeader+len(deviceMap)+len(storageDir)+len(identity.Cgroup):], jobFile) + copy(packet[daemonRequestHeader+len(deviceMap)+len(storageDir)+len(identity.Cgroup)+len(jobFile):], selectedDeviceList) + return packet, nil +} + +func parseDaemonResponse(packet []byte) (int32, uint32, string, string, error) { + if len(packet) < daemonResponseHeader || len(packet) > daemonMaxResponse { + return 0, 0, "", "", fmt.Errorf("invalid CUDA helper daemon response size %d", len(packet)) + } + if binary.LittleEndian.Uint32(packet[0:4]) != daemonProtocolMagic || + binary.LittleEndian.Uint16(packet[4:6]) != daemonProtocolVersion || + binary.LittleEndian.Uint16(packet[6:8]) != daemonResponseHeader { + return 0, 0, "", "", errors.New("invalid CUDA helper daemon response header") + } + flags := binary.LittleEndian.Uint32(packet[12:16]) + if flags & ^(daemonResponseFatal|daemonCapabilityDeferredCUDA|daemonCapabilityCustomStorage|daemonResponseLockNotAcquired) != 0 { + return 0, 0, "", "", errors.New("invalid CUDA helper daemon response flags") + } + outputSize := int(binary.LittleEndian.Uint32(packet[16:20])) + errorSize := int(binary.LittleEndian.Uint32(packet[20:24])) + if outputSize < 0 || errorSize < 0 || outputSize+errorSize != len(packet)-daemonResponseHeader { + return 0, 0, "", "", errors.New("invalid CUDA helper daemon response payload lengths") + } + payload := packet[daemonResponseHeader:] + return int32(binary.LittleEndian.Uint32(packet[8:12])), + flags, string(payload[:outputSize]), string(payload[outputSize:]), nil +} + +func daemonRPC( + ctx context.Context, + socket string, + packet []byte, +) (int32, uint32, string, string, time.Duration, error) { + dialer := net.Dialer{} + conn, err := dialer.DialContext(ctx, "unixpacket", socket) + if err != nil { + return 0, 0, "", "", 0, fmt.Errorf("%w at %s: %v", errDaemonUnavailable, socket, err) + } + defer conn.Close() + stop := context.AfterFunc(ctx, func() { _ = conn.Close() }) + defer stop() + deadline := time.Now().Add(daemonRPCTimeout) + if contextDeadline, ok := ctx.Deadline(); ok && contextDeadline.Before(deadline) { + deadline = contextDeadline + } + if err := conn.SetDeadline(deadline); err != nil { + return 0, 0, "", "", 0, fmt.Errorf("set CUDA helper daemon RPC deadline: %w", err) + } + + start := time.Now() + written, err := conn.Write(packet) + if err != nil { + return 0, 0, "", "", time.Since(start), + fmt.Errorf("CUDA helper daemon request write failed; operation outcome is unknown and will not be replayed: %w", err) + } + if written != len(packet) { + return 0, 0, "", "", time.Since(start), + fmt.Errorf("CUDA helper daemon request write was short (%d of %d bytes); operation outcome is unknown and will not be replayed", written, len(packet)) + } + response := make([]byte, daemonMaxResponse+1) + read, err := conn.Read(response) + rpcWall := time.Since(start) + if err != nil { + if errors.Is(err, os.ErrDeadlineExceeded) || ctx.Err() != nil { + cause := ctx.Err() + if cause == nil { + cause = err + } + return 0, 0, "", "", rpcWall, + fmt.Errorf("CUDA helper daemon RPC ended after %s; operation outcome is unknown and will not be replayed: %w", rpcWall, cause) + } + return 0, 0, "", "", rpcWall, + fmt.Errorf("CUDA helper daemon disconnected after request; operation outcome is unknown and will not be replayed: %w", err) + } + if read > daemonMaxResponse { + return 0, 0, "", "", rpcWall, fmt.Errorf("CUDA helper daemon response exceeded %d bytes", daemonMaxResponse) + } + status, flags, stdout, stderr, err := parseDaemonResponse(response[:read]) + if err != nil { + return 0, 0, "", "", rpcWall, err + } + return status, flags, stdout, stderr, rpcWall, nil +} + +func runDaemonAction( + ctx context.Context, + request helperAction, + log logr.Logger, +) error { + packet, err := daemonRequest(request) + if err != nil { + return err + } + pid := request.PID + action := request.Action + storageMode := request.StorageMode + status, flags, stdout, stderr, rpcWall, err := daemonRPC(ctx, daemonSocketPath, packet) + if err != nil { + return err + } + output := stdout + stderr + if status != 0 { + if action == actionLock && flags&daemonResponseLockNotAcquired != 0 { + classification := errCheckpointLockNotAcquired + if flags&daemonResponseFatal != 0 { + classification = errors.Join(classification, errDaemonFatal) + } + return fmt.Errorf("%w: pid %d after %s with CUDA status %d (output: %s)", + classification, pid, rpcWall, status, output) + } + if flags&daemonResponseFatal != 0 { + return fmt.Errorf("%w: CUDA helper daemon %s failed for pid %d after %s with CUDA status %d (output: %s)", + errDaemonFatal, action, pid, rpcWall, status, output) + } + return fmt.Errorf("CUDA helper daemon %s failed for pid %d after %s with CUDA status %d (output: %s)", + action, pid, rpcWall, status, output) + } + if action == actionLock || action == actionUnlock { + log.V(1).Info("CUDA helper daemon action succeeded", + "pid", pid, + "action", action, + "daemon_rpc_wall_duration", rpcWall, + "output", output, + ) + return nil + } + if storageMode == types.CUDAStorageModeLegacy { + log.V(1).Info("CUDA helper daemon action succeeded", + "pid", pid, + "action", action, + "backend", storageMode, + "daemon_rpc_wall_duration", rpcWall, + "output", output, + ) + return nil + } + telemetry := parseCustomStorageTelemetry(output, rpcWall) + values := []any{ + "pid", pid, + "action", action, + "transport", "daemon", + "daemon_rpc_wall_duration", rpcWall, + "helper_telemetry_status", telemetry.status, + } + if telemetry.status == "valid" { + values = append(values, "helper_operation_to_telemetry_duration", telemetry.helperMainDuration) + } else { + values = append(values, "helper_telemetry_error", telemetry.err) + } + log.Info("CUDA custom-storage transfer succeeded", append(values, "output", output)...) + return nil +} + +// validateCUDAOperationBudget rejects a long-running CUDA sequence before its +// first driver call. Per-request validation in runDaemonAction is too late for +// checkpoint because the process has already been locked by then. +func validateCUDAOperationBudget(ctx context.Context, action string, targetCount int) error { + if targetCount < 0 { + return fmt.Errorf("CUDA helper %s target count must not be negative", action) + } + callsPerTarget := 1 + if action == actionCheckpoint { + // Checkpoint locks each target before asking the driver to checkpoint it. + // Both calls may consume the full daemon RPC timeout. Restore unlock is a + // bounded terminal action and is deliberately allowed under a short + // caller deadline, so only the restore call contributes to this preflight. + callsPerTarget = 2 + } + maxTargets := int((time.Duration(1<<63-1) / daemonRPCTimeout) / time.Duration(callsPerTarget)) + if targetCount > maxTargets { + return fmt.Errorf("CUDA helper %s target count %d exceeds operation-budget capacity", action, targetCount) + } + required := time.Duration(targetCount) * time.Duration(callsPerTarget) * daemonRPCTimeout + if deadline, ok := ctx.Deadline(); ok && time.Until(deadline) <= required { + return fmt.Errorf( + "CUDA helper %s for %d target(s) requires more than %s of caller budget before state-changing work", + action, targetCount, required, + ) + } + return nil +} + +func daemonCapabilities(ctx context.Context) (uint32, error) { + packet, err := daemonRequest(helperAction{}) + if err != nil { + return 0, err + } + // Health has a separate listener so it cannot queue behind a long-running + // operation. Both listeners are created and owned by the same daemon process; + // any accept/poll failure on either listener is fatal and terminates the + // daemon, so a surviving health listener cannot mask a failed operation loop. + status, flags, _, stderr, _, err := daemonRPC(ctx, daemonSocketPath+".health", packet) + if err != nil { + return 0, err + } + if status != 0 || flags&daemonCapabilityDeferredCUDA == 0 { + return 0, fmt.Errorf("CUDA helper daemon health/capability check failed: status=%d flags=%#x error=%s", + status, flags, stderr) + } + return flags, nil +} + +// SelectCUDAStorageMode applies the operator's checkpoint-creation policy. +// It chooses one artifact storage mode for the entire checkpoint before any +// CUDA process is locked. Restore selection is separate: it always follows the +// published manifest. +var waitForDaemon = WaitForDaemon + +func SelectCUDAStorageMode(ctx context.Context, configuredMode string) (string, error) { + if err := waitForDaemon(ctx, configuredMode); err != nil { + return "", err + } + switch configuredMode { + case types.CUDAStorageModeLegacy: + return configuredMode, nil + case types.CUDAStorageModePOSIX: + return configuredMode, nil + default: + return "", fmt.Errorf("unsupported configured CUDA storage mode %q", configuredMode) + } +} + +// ValidateCUDAStorageMode rejects unsupported artifacts before rootfs or CRIU +// restore changes the placeholder. +func ValidateCUDAStorageMode(ctx context.Context, storageMode string) error { + flags, err := daemonCapabilities(ctx) + if err != nil { + return err + } + return validateCUDAStorageModeCapabilities(storageMode, flags) +} + +func validateCUDAStorageModeCapabilities(storageMode string, flags uint32) error { + switch storageMode { + case types.CUDAStorageModeLegacy: + return nil + case types.CUDAStorageModePOSIX: + if flags&daemonCapabilityCustomStorage == 0 { + return errors.New("CUDA POSIX CustomStorage artifact requires daemon CustomStorage capability") + } + return nil + default: + return fmt.Errorf("unsupported CUDA storage mode %q", storageMode) + } +} + +// WaitForDaemon waits for protocol-level health and every capability required +// by the configured checkpoint creation mode. +func WaitForDaemon(ctx context.Context, storageMode string) error { + waitCtx, cancel := context.WithTimeout(ctx, daemonHealthWait) + defer cancel() + ticker := time.NewTicker(daemonHealthRetryInterval) + defer ticker.Stop() + var lastErr error + for { + flags, err := daemonCapabilities(waitCtx) + if err == nil { + err = validateCUDAStorageModeCapabilities(storageMode, flags) + } + if err == nil { + return nil + } + lastErr = err + select { + case <-waitCtx.Done(): + return fmt.Errorf("wait for CUDA helper daemon at %s: %w (last error: %v)", daemonSocketPath, waitCtx.Err(), lastErr) + case <-ticker.C: + } + } +} diff --git a/agent/internal/cuda/daemon_client_test.go b/agent/internal/cuda/daemon_client_test.go new file mode 100644 index 00000000..ef25becb --- /dev/null +++ b/agent/internal/cuda/daemon_client_test.go @@ -0,0 +1,577 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package cuda + +import ( + "bytes" + "context" + "encoding/binary" + "encoding/hex" + "errors" + "net" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/go-logr/logr" + + snapshotruntime "github.com/ai-dynamo/snapshot/agent/internal/runtime" + "github.com/ai-dynamo/snapshot/agent/internal/types" +) + +func TestDaemonRequestMatchesSharedGoldenFixture(t *testing.T) { + encodedFixture, err := os.ReadFile(filepath.Join("..", "..", "cmd", "cuda-checkpoint-helper", "testdata", "daemon_request_v6.hex")) + if err != nil { + t.Fatal(err) + } + want, err := hex.DecodeString(strings.TrimSpace(string(encodedFixture))) + if err != nil { + t.Fatal(err) + } + identity := testDaemonIdentity(42) + identity.StartTimeTicks = 12345 + identity.Cgroup = "0::/kubepods/test\n" + got, err := daemonRequest(helperAction{ + PID: 42, Action: actionRestore, + DeviceMap: "GPU-aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee=GPU-11111111-2222-3333-4444-555555555555", + StorageMode: types.CUDAStorageModePOSIX, + StorageDir: "/checkpoints/process-nspid-42", + JobFile: "/host/proc/42/root/tmp/cuda-job", + GPUUUIDs: []string{"GPU-12345678-1234-1234-1234-123456789abc"}, + Transfer: types.CUDATransferSettings{BufferCount: 2, ChunkBytes: 8 * 1024 * 1024}, + Identity: identity, + }) + if err != nil { + t.Fatal(err) + } + if !bytes.Equal(got, want) { + t.Fatalf("daemonRequest() bytes differ from shared v6 fixture\n got: %x\nwant: %x", got, want) + } +} + +func daemonTestResponse(status int32, flags uint32) []byte { + packet := make([]byte, daemonResponseHeader) + binary.LittleEndian.PutUint32(packet[0:4], daemonProtocolMagic) + binary.LittleEndian.PutUint16(packet[4:6], daemonProtocolVersion) + binary.LittleEndian.PutUint16(packet[6:8], daemonResponseHeader) + binary.LittleEndian.PutUint32(packet[8:12], uint32(status)) + binary.LittleEndian.PutUint32(packet[12:16], flags) + return packet +} + +func testDaemonIdentity(pid int) snapshotruntime.ProcessDetails { + return snapshotruntime.ProcessDetails{ + ObservedPID: pid, OutermostPID: pid, InnermostPID: pid, + NamespacePIDs: []int{pid}, StartTimeTicks: 12345, + Cgroup: "0::/kubepods/test\n", + } +} + +func withOperationServer(t *testing.T, handler func(*net.UnixConn)) { + t.Helper() + socket := filepath.Join(t.TempDir(), "helper.sock") + oldSocket := daemonSocketPath + daemonSocketPath = socket + t.Cleanup(func() { daemonSocketPath = oldSocket }) + listener, err := net.ListenUnix("unixpacket", &net.UnixAddr{Name: socket, Net: "unixpacket"}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = listener.Close() }) + go func() { + conn, err := listener.AcceptUnix() + if err != nil { + return + } + defer conn.Close() + handler(conn) + }() +} + +func TestCommandRunnerSendsDaemonOperations(t *testing.T) { + identity := testDaemonIdentity(42) + transfer := types.CUDATransferSettings{BufferCount: 2, ChunkBytes: 8 * 1024 * 1024} + for _, test := range []struct { + action, backend, deviceMap, storageDir string + wantAction, wantBackend uint16 + wantStorageDir string + wantBufferCount uint32 + wantChunkBytes uint64 + }{ + {action: actionLock, backend: types.CUDAStorageModeLegacy, storageDir: "/ignored", wantAction: daemonActionLock, wantBackend: 1}, + {action: actionCheckpoint, backend: types.CUDAStorageModeLegacy, storageDir: "/ignored", wantAction: daemonActionCheckpoint, wantBackend: 1}, + { + action: actionRestore, backend: types.CUDAStorageModePOSIX, deviceMap: "0=1", + storageDir: "/checkpoints/process-0000", wantAction: daemonActionRestore, wantBackend: 2, + wantStorageDir: "/checkpoints/process-0000", wantBufferCount: 2, wantChunkBytes: 8 * 1024 * 1024, + }, + {action: actionUnlock, backend: types.CUDAStorageModePOSIX, storageDir: "/ignored", wantAction: daemonActionUnlock, wantBackend: 2}, + } { + t.Run(test.action, func(t *testing.T) { + var gpuUUIDs []string + if test.backend == types.CUDAStorageModePOSIX && + (test.action == actionCheckpoint || test.action == actionRestore) { + gpuUUIDs = []string{"GPU-12345678-1234-1234-1234-123456789abc"} + } + request := make(chan []byte, 1) + withOperationServer(t, func(conn *net.UnixConn) { + packet := make([]byte, daemonMaxRequest) + n, err := conn.Read(packet) + if err != nil { + t.Error(err) + return + } + request <- packet[:n] + _, _ = conn.Write(daemonTestResponse(0, 0)) + }) + if err := (commandHelperActionRunner{}).run( + context.Background(), helperAction{ + PID: 42, Action: test.action, DeviceMap: test.deviceMap, + StorageMode: test.backend, StorageDir: test.storageDir, + JobFile: "/host/proc/42/root/tmp/cuda-job", Transfer: transfer, + GPUUUIDs: gpuUUIDs, + Identity: identity, + }, logr.Discard(), + ); err != nil { + t.Fatalf("run() error = %v", err) + } + + packet := <-request + if got := binary.LittleEndian.Uint16(packet[8:10]); got != test.wantAction { + t.Errorf("action = %d, want %d", got, test.wantAction) + } + if got := binary.LittleEndian.Uint16(packet[10:12]); got != test.wantBackend { + t.Errorf("backend = %d, want %d", got, test.wantBackend) + } + if got := binary.LittleEndian.Uint32(packet[12:16]); got != 42 { + t.Errorf("pid = %d, want 42", got) + } + if got := binary.LittleEndian.Uint32(packet[16:20]); got != test.wantBufferCount { + t.Errorf("buffer count = %d, want %d", got, test.wantBufferCount) + } + if got := binary.LittleEndian.Uint64(packet[20:28]); got != test.wantChunkBytes { + t.Errorf("chunk bytes = %d, want %d", got, test.wantChunkBytes) + } + deviceMapSize := int(binary.LittleEndian.Uint32(packet[28:32])) + storageDirSize := int(binary.LittleEndian.Uint32(packet[32:36])) + cgroupSize := int(binary.LittleEndian.Uint32(packet[36:40])) + jobFileSize := int(binary.LittleEndian.Uint32(packet[48:52])) + selectedDevicesSize := int(binary.LittleEndian.Uint32(packet[52:56])) + if got := binary.LittleEndian.Uint64(packet[40:48]); got != identity.StartTimeTicks { + t.Errorf("start time = %d, want %d", got, identity.StartTimeTicks) + } + payload := packet[daemonRequestHeader:] + if got := string(payload[:deviceMapSize]); got != test.deviceMap { + t.Errorf("device map = %q, want %q", got, test.deviceMap) + } + payload = payload[deviceMapSize:] + if got := string(payload[:storageDirSize]); got != test.wantStorageDir { + t.Errorf("storage directory = %q, want %q", got, test.wantStorageDir) + } + payload = payload[storageDirSize:] + if got := string(payload[:cgroupSize]); got != identity.Cgroup { + t.Errorf("cgroup = %q, want %q", got, identity.Cgroup) + } + payload = payload[cgroupSize:] + if got := string(payload[:jobFileSize]); got != "/host/proc/42/root/tmp/cuda-job" { + t.Errorf("job file = %q, want host-visible launch-job path", got) + } + payload = payload[jobFileSize:] + wantSelectedDevices := "" + if test.backend == types.CUDAStorageModePOSIX && + (test.action == actionCheckpoint || test.action == actionRestore) { + wantSelectedDevices = "GPU-12345678-1234-1234-1234-123456789abc" + } + if got := string(payload[:selectedDevicesSize]); got != wantSelectedDevices { + t.Errorf("selected devices = %q, want %q", got, wantSelectedDevices) + } + }) + } +} + +func TestDaemonRequestRejectsBackendArgumentMismatch(t *testing.T) { + identity := testDaemonIdentity(42) + transfer := types.CUDATransferSettings{BufferCount: 1, ChunkBytes: types.DefaultCUDATransferChunkBytes} + for _, test := range []struct { + name, backend, storageDir string + }{ + {name: "regular with directory", backend: types.CUDAStorageModeLegacy, storageDir: "/checkpoints/process-0000"}, + {name: "posix without directory", backend: types.CUDAStorageModePOSIX}, + {name: "unknown", backend: "auto"}, + } { + t.Run(test.name, func(t *testing.T) { + if _, err := daemonRequest(helperAction{PID: 42, Action: actionCheckpoint, StorageMode: test.backend, StorageDir: test.storageDir, Transfer: transfer, Identity: identity}); err == nil { + t.Fatal("daemonRequest() accepted mismatched backend arguments") + } + }) + } +} + +func TestDaemonRequestRejectsRelativeJobFile(t *testing.T) { + identity := testDaemonIdentity(42) + _, err := daemonRequest(helperAction{PID: 42, Action: actionCheckpoint, + StorageMode: types.CUDAStorageModeLegacy, JobFile: "tmp/cuda-job", Identity: identity}) + if err == nil || !strings.Contains(err.Error(), "job file") { + t.Fatalf("daemonRequest() error = %v, want invalid job-file rejection", err) + } +} + +func TestDaemonRequestRejectsParserInvariantViolations(t *testing.T) { + identity := testDaemonIdentity(42) + validTransfer := types.CUDATransferSettings{BufferCount: 1, ChunkBytes: types.DefaultCUDATransferChunkBytes} + validUUID := "GPU-12345678-1234-1234-1234-123456789abc" + for _, test := range []struct { + name string + action string + deviceMap string + storageDir string + jobFile string + selectedDevices []string + transfer types.CUDATransferSettings + identity snapshotruntime.ProcessDetails + }{ + {name: "checkpoint device map", action: actionCheckpoint, deviceMap: "GPU-a=GPU-b", storageDir: "/checkpoints/process", selectedDevices: []string{validUUID}, transfer: validTransfer, identity: identity}, + {name: "zero transfer settings", action: actionRestore, storageDir: "/checkpoints/process", selectedDevices: []string{validUUID}, identity: identity}, + {name: "case-insensitive duplicate UUID", action: actionRestore, storageDir: "/checkpoints/process", selectedDevices: []string{validUUID, strings.ToUpper(validUUID)}, transfer: validTransfer, identity: identity}, + {name: "NUL in storage path", action: actionRestore, storageDir: "/checkpoints/process\x00other", selectedDevices: []string{validUUID}, transfer: validTransfer, identity: identity}, + {name: "NUL in job file", action: actionRestore, storageDir: "/checkpoints/process", jobFile: "/tmp/job\x00other", selectedDevices: []string{validUUID}, transfer: validTransfer, identity: identity}, + } { + t.Run(test.name, func(t *testing.T) { + _, err := daemonRequest(helperAction{PID: 42, Action: test.action, DeviceMap: test.deviceMap, + StorageMode: types.CUDAStorageModePOSIX, StorageDir: test.storageDir, JobFile: test.jobFile, + GPUUUIDs: test.selectedDevices, Transfer: test.transfer, Identity: test.identity}) + if err == nil { + t.Fatal("daemonRequest() accepted a request rejected by the daemon parser") + } + }) + } +} + +func withHealthServer(t *testing.T, flags uint32) { + t.Helper() + socket := filepath.Join(t.TempDir(), "helper.sock") + oldSocket := daemonSocketPath + daemonSocketPath = socket + t.Cleanup(func() { daemonSocketPath = oldSocket }) + listener, err := net.ListenUnix("unixpacket", &net.UnixAddr{Name: socket + ".health", Net: "unixpacket"}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = listener.Close() }) + go func() { + conn, err := listener.AcceptUnix() + if err != nil { + return + } + defer conn.Close() + request := make([]byte, daemonMaxRequest) + _, _ = conn.Read(request) + _, _ = conn.Write(daemonTestResponse(0, flags|daemonCapabilityDeferredCUDA)) + }() +} + +func TestWaitForDaemonRetriesUntilHealthListenerIsReady(t *testing.T) { + socket := filepath.Join(t.TempDir(), "helper.sock") + oldSocket := daemonSocketPath + daemonSocketPath = socket + t.Cleanup(func() { daemonSocketPath = oldSocket }) + serverErr := make(chan error, 1) + go func() { + time.Sleep(2 * daemonHealthRetryInterval) + listener, err := net.ListenUnix("unixpacket", &net.UnixAddr{Name: socket + ".health", Net: "unixpacket"}) + if err != nil { + serverErr <- err + return + } + defer listener.Close() + conn, err := listener.AcceptUnix() + if err == nil { + defer conn.Close() + request := make([]byte, daemonMaxRequest) + _, _ = conn.Read(request) + _, err = conn.Write(daemonTestResponse(0, daemonCapabilityDeferredCUDA)) + } + serverErr <- err + }() + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + if err := WaitForDaemon(ctx, types.CUDAStorageModeLegacy); err != nil { + t.Fatalf("WaitForDaemon() error = %v, want retry success", err) + } + if err := <-serverErr; err != nil { + t.Fatalf("health server error = %v", err) + } +} + +func TestWaitForDaemonHonorsCallerCancellation(t *testing.T) { + oldSocket := daemonSocketPath + daemonSocketPath = filepath.Join(t.TempDir(), "missing.sock") + t.Cleanup(func() { daemonSocketPath = oldSocket }) + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + err := WaitForDaemon(ctx, types.CUDAStorageModeLegacy) + if err == nil || !strings.Contains(err.Error(), "wait for CUDA helper daemon") { + t.Fatalf("WaitForDaemon() error = %v, want bounded readiness failure", err) + } +} + +func TestValidateCUDAStorageModeCapabilitiesRequiresConfiguredBackend(t *testing.T) { + if err := validateCUDAStorageModeCapabilities(types.CUDAStorageModeLegacy, daemonCapabilityDeferredCUDA); err != nil { + t.Fatalf("legacy capability validation failed: %v", err) + } + err := validateCUDAStorageModeCapabilities(types.CUDAStorageModePOSIX, daemonCapabilityDeferredCUDA) + if err == nil || !strings.Contains(err.Error(), "requires daemon CustomStorage capability") { + t.Fatalf("POSIX capability validation error = %v, want missing CustomStorage capability", err) + } + if err := validateCUDAStorageModeCapabilities(types.CUDAStorageModePOSIX, daemonCapabilityDeferredCUDA|daemonCapabilityCustomStorage); err != nil { + t.Fatalf("POSIX capability validation failed: %v", err) + } +} + +func TestSelectCUDAStorageModeRequiresExplicitOptIn(t *testing.T) { + for _, test := range []struct { + name string + configured string + flags uint32 + want string + wantError string + }{ + { + name: "legacy does not auto-enable CustomStorage", + configured: types.CUDAStorageModeLegacy, + flags: daemonCapabilityCustomStorage, + want: types.CUDAStorageModeLegacy, + }, + { + name: "posix requires capability", + configured: types.CUDAStorageModePOSIX, + wantError: "requires daemon CustomStorage capability", + }, + { + name: "posix selects CustomStorage when capable", + configured: types.CUDAStorageModePOSIX, + flags: daemonCapabilityCustomStorage, + want: types.CUDAStorageModePOSIX, + }, + } { + t.Run(test.name, func(t *testing.T) { + oldWaitForDaemon := waitForDaemon + waitForDaemon = func(_ context.Context, storageMode string) error { + return validateCUDAStorageModeCapabilities( + storageMode, + test.flags|daemonCapabilityDeferredCUDA, + ) + } + t.Cleanup(func() { waitForDaemon = oldWaitForDaemon }) + got, err := SelectCUDAStorageMode(context.Background(), test.configured) + if test.wantError != "" { + if err == nil || !strings.Contains(err.Error(), test.wantError) { + t.Fatalf("SelectCUDAStorageMode() error = %v, want %q", err, test.wantError) + } + return + } + if err != nil || got != test.want { + t.Fatalf("SelectCUDAStorageMode() = %q, %v; want %q", got, err, test.want) + } + manifest := types.NewCUDAManifest([]int{42}, []string{"GPU-aaa"}, got) + if manifest.StorageMode != test.want { + t.Fatalf("checkpoint manifest mode = %q, want %q", manifest.StorageMode, test.want) + } + }) + } +} + +func TestValidateCUDAStorageModeObeysManifest(t *testing.T) { + withHealthServer(t, daemonCapabilityCustomStorage) + if err := ValidateCUDAStorageMode(context.Background(), types.CUDAStorageModeLegacy); err != nil { + t.Fatalf("legacy restore on capable daemon failed: %v", err) + } + + withHealthServer(t, 0) + err := ValidateCUDAStorageMode(context.Background(), types.CUDAStorageModePOSIX) + if err == nil || !strings.Contains(err.Error(), "requires daemon CustomStorage capability") { + t.Fatalf("POSIX restore error = %v, want capability rejection", err) + } +} + +func TestRunDaemonActionMapsFatalResponse(t *testing.T) { + withOperationServer(t, func(conn *net.UnixConn) { + request := make([]byte, daemonMaxRequest) + _, _ = conn.Read(request) + _, _ = conn.Write(daemonTestResponse(2, daemonResponseFatal)) + }) + err := runDaemonAction(context.Background(), helperAction{PID: 42, Action: actionCheckpoint, + StorageMode: types.CUDAStorageModeLegacy, Identity: testDaemonIdentity(42)}, logr.Discard()) + if !errors.Is(err, errDaemonFatal) { + t.Fatalf("runDaemonAction() error = %v, want errDaemonFatal", err) + } +} + +func TestRunDaemonActionMapsLockTimeoutBeforeMutation(t *testing.T) { + withOperationServer(t, func(conn *net.UnixConn) { + request := make([]byte, daemonMaxRequest) + _, _ = conn.Read(request) + _, _ = conn.Write(daemonTestResponse(600, daemonResponseLockNotAcquired)) + }) + err := runDaemonAction(context.Background(), helperAction{PID: 42, Action: actionLock, + StorageMode: types.CUDAStorageModeLegacy, Identity: testDaemonIdentity(42)}, logr.Discard()) + if !errors.Is(err, errCheckpointLockNotAcquired) { + t.Fatalf("runDaemonAction() error = %v, want pre-mutation lock timeout", err) + } +} + +func TestRunDaemonActionPreservesFatalLockRejectionClassification(t *testing.T) { + withOperationServer(t, func(conn *net.UnixConn) { + request := make([]byte, daemonMaxRequest) + _, _ = conn.Read(request) + _, _ = conn.Write(daemonTestResponse(600, daemonResponseFatal|daemonResponseLockNotAcquired)) + }) + err := runDaemonAction(context.Background(), helperAction{PID: 42, Action: actionLock, + StorageMode: types.CUDAStorageModeLegacy, Identity: testDaemonIdentity(42)}, logr.Discard()) + if !errors.Is(err, errCheckpointLockNotAcquired) || !errors.Is(err, errDaemonFatal) { + t.Fatalf("runDaemonAction() error = %v, want both pre-mutation and fatal classifications", err) + } +} + +func TestValidateCUDAOperationBudgetRejectsShortCallerDeadline(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + err := validateCUDAOperationBudget(ctx, actionCheckpoint, 1) + if err == nil || !strings.Contains(err.Error(), "before state-changing work") { + t.Fatalf("validateCUDAOperationBudget() error = %v, want caller-budget rejection", err) + } +} + +func TestValidateCUDAOperationBudgetCoversWholeMultiTargetSequence(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 3*daemonRPCTimeout) + defer cancel() + err := validateCUDAOperationBudget(ctx, actionCheckpoint, 2) + if err == nil || !strings.Contains(err.Error(), "for 2 target(s)") { + t.Fatalf("validateCUDAOperationBudget() error = %v, want whole-sequence rejection", err) + } +} + +func TestValidateCUDAOperationBudgetAllowsOneTargetRestoreAtDefaultDeadline(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Hour) + defer cancel() + if err := validateCUDAOperationBudget(ctx, actionRestore, 1); err != nil { + t.Fatalf("validateCUDAOperationBudget() error = %v, want default one-target restore to fit", err) + } +} + +func TestValidateCUDAOperationBudgetAllowsQualifiedTwoTargetRestoreAtChartDefault(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 135*time.Minute) + defer cancel() + if err := validateCUDAOperationBudget(ctx, actionRestore, 2); err != nil { + t.Fatalf("validateCUDAOperationBudget() error = %v, want qualified two-PID restore to fit chart default", err) + } +} + +func TestValidateCUDAOperationBudgetCoversAllRestoreTargets(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 3*daemonRPCTimeout) + defer cancel() + err := validateCUDAOperationBudget(ctx, actionRestore, 4) + if err == nil || !strings.Contains(err.Error(), "for 4 target(s)") { + t.Fatalf("validateCUDAOperationBudget() error = %v, want whole restore-sequence rejection", err) + } +} + +func TestRunDaemonActionAllowsUnlockWithShortCallerDeadline(t *testing.T) { + withOperationServer(t, func(conn *net.UnixConn) { + request := make([]byte, daemonMaxRequest) + _, _ = conn.Read(request) + _, _ = conn.Write(daemonTestResponse(0, 0)) + }) + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + if err := runDaemonAction(ctx, helperAction{PID: 42, Action: actionUnlock, + StorageMode: types.CUDAStorageModeLegacy, Identity: testDaemonIdentity(42)}, logr.Discard()); err != nil { + t.Fatalf("runDaemonAction(unlock) error = %v", err) + } +} + +func TestRunDaemonActionRejectsInvalidResponses(t *testing.T) { + for _, test := range []struct { + name string + response func() []byte + want string + }{ + { + name: "malformed payload lengths", + response: func() []byte { + packet := daemonTestResponse(0, 0) + binary.LittleEndian.PutUint32(packet[16:20], 1) + return packet + }, + want: "payload lengths", + }, + { + name: "oversized packet", + response: func() []byte { + return make([]byte, daemonMaxResponse+1) + }, + want: "exceeded", + }, + } { + t.Run(test.name, func(t *testing.T) { + withOperationServer(t, func(conn *net.UnixConn) { + request := make([]byte, daemonMaxRequest) + _, _ = conn.Read(request) + _, _ = conn.Write(test.response()) + }) + err := runDaemonAction(context.Background(), helperAction{PID: 42, Action: actionCheckpoint, + StorageMode: types.CUDAStorageModeLegacy, Identity: testDaemonIdentity(42)}, logr.Discard()) + if err == nil || !strings.Contains(err.Error(), test.want) { + t.Fatalf("runDaemonAction() error = %v, want error containing %q", err, test.want) + } + }) + } +} + +func TestRunDaemonActionDisconnectAfterSendIsNotReplayed(t *testing.T) { + socket := filepath.Join(t.TempDir(), "helper.sock") + oldSocket := daemonSocketPath + daemonSocketPath = socket + t.Cleanup(func() { daemonSocketPath = oldSocket }) + listener, err := net.ListenUnix("unixpacket", &net.UnixAddr{Name: socket, Net: "unixpacket"}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = listener.Close() }) + accepted := make(chan int, 1) + go func() { + count := 0 + conn, acceptErr := listener.AcceptUnix() + if acceptErr != nil { + accepted <- count + return + } + count++ + request := make([]byte, daemonMaxRequest) + _, _ = conn.Read(request) + _ = conn.Close() + _ = listener.SetDeadline(time.Now().Add(100 * time.Millisecond)) + conn, acceptErr = listener.AcceptUnix() + if acceptErr == nil { + count++ + _ = conn.Close() + } + accepted <- count + }() + err = runDaemonAction(context.Background(), helperAction{PID: 42, Action: actionCheckpoint, + StorageMode: types.CUDAStorageModeLegacy, Identity: testDaemonIdentity(42)}, logr.Discard()) + if err == nil || !strings.Contains(err.Error(), "outcome is unknown and will not be replayed") { + t.Fatalf("runDaemonAction() error = %v, want ambiguous non-replayable operation error", err) + } + if errors.Is(err, errDaemonUnavailable) { + t.Fatalf("disconnect after request must not be treated as pre-send unavailability: %v", err) + } + if count := <-accepted; count != 1 { + t.Fatalf("daemon accepted %d operation requests after an ambiguous disconnect, want 1", count) + } +} diff --git a/agent/internal/cuda/job.go b/agent/internal/cuda/job.go index 25ef350a..3ecc3c66 100644 --- a/agent/internal/cuda/job.go +++ b/agent/internal/cuda/job.go @@ -8,8 +8,10 @@ import ( "io" "os" "path/filepath" + "strconv" "strings" + snapshotruntime "github.com/ai-dynamo/snapshot/agent/internal/runtime" snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" "golang.org/x/sys/unix" ) @@ -17,6 +19,20 @@ import ( // JobFileEnv is the CUDA launch-job environment variable consumed by the driver. const JobFileEnv = "CUDA_CHECKPOINT_JOB_FILE" +// HostJobFilePath returns the host-visible path to the fixed launch-job file +// inside a restored CUDA process's mount namespace. +func HostJobFilePath(hostPID int) (string, error) { + if hostPID <= 0 { + return "", fmt.Errorf("invalid host PID %d", hostPID) + } + return filepath.Join( + snapshotruntime.HostProcPath, + strconv.Itoa(hostPID), + "root", + strings.TrimPrefix(snapshotv1alpha1.CUDAJobFilePath, string(os.PathSeparator)), + ), nil +} + // StageJobFile copies a launch-job file into the checkpoint artifact and // returns the host-visible path to the source pod's live job file. Capture // helpers must use that live file so they join the same CUDA job as the target diff --git a/agent/internal/cuda/job_test.go b/agent/internal/cuda/job_test.go index 9c1b210d..2543c8bf 100644 --- a/agent/internal/cuda/job_test.go +++ b/agent/internal/cuda/job_test.go @@ -5,6 +5,7 @@ package cuda import ( "context" + "fmt" "os" "path/filepath" "strings" @@ -13,9 +14,33 @@ import ( "github.com/go-logr/logr" "golang.org/x/sys/unix" + "github.com/ai-dynamo/snapshot/agent/internal/types" snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" ) +type jobFileRunner struct { + jobFile string + trace []string +} + +func (r *jobFileRunner) run( + _ context.Context, + request helperAction, + _ logr.Logger, +) error { + r.trace = append(r.trace, fmt.Sprintf("%s %d", request.Action, request.PID)) + if request.Action != actionCheckpoint { + return nil + } + file, err := os.OpenFile(r.jobFile, os.O_APPEND|os.O_WRONLY, 0600) + if err != nil { + return err + } + defer file.Close() + _, err = fmt.Fprintf(file, "|%d", request.PID) + return err +} + func TestStageJobFile(t *testing.T) { sourceRoot := t.TempDir() checkpointDir := t.TempDir() @@ -45,6 +70,20 @@ func TestStageJobFile(t *testing.T) { } } +func TestHostJobFilePath(t *testing.T) { + got, err := HostJobFilePath(42) + if err != nil { + t.Fatalf("HostJobFilePath() error = %v", err) + } + want := filepath.Join("/host/proc/42/root", snapshotv1alpha1.CUDAJobFilePath) + if got != want { + t.Fatalf("HostJobFilePath() = %q, want %q", got, want) + } + if _, err := HostJobFilePath(0); err == nil { + t.Fatal("HostJobFilePath() accepted an invalid PID") + } +} + func TestStageJobFileRejectsSymlink(t *testing.T) { sourceRoot := t.TempDir() checkpointDir := t.TempDir() @@ -109,36 +148,8 @@ func TestStageJobFileRequiresLaunchJobStateForMultiGPU(t *testing.T) { } } -func TestCheckpointProcessTreePersistsStateAfterEveryProcessCheckpoint(t *testing.T) { +func TestCheckpointProcessTreePersistsPostCheckpointJobState(t *testing.T) { tempDir := t.TempDir() - trace := filepath.Join(tempDir, "trace") - helper := filepath.Join(tempDir, "cuda-checkpoint-helper") - script := `#!/bin/sh -action="" -pid="" -job_file="" -while [ "$#" -gt 0 ]; do - case "$1" in - --action) action="$2"; shift 2 ;; - --pid) pid="$2"; shift 2 ;; - --job-file) job_file="$2"; shift 2 ;; - *) shift ;; - esac -done -if [ "$job_file" != "$DYNAMO_TEST_JOB_FILE" ]; then - printf 'job file = %s, want %s\n' "$job_file" "$DYNAMO_TEST_JOB_FILE" >&2 - exit 1 -fi -printf '%s %s\n' "$action" "$pid" >> "$DYNAMO_TEST_TRACE" -if [ "$action" = checkpoint ]; then printf '|%s' "$pid" >> "$job_file"; fi -` - if err := os.WriteFile(helper, []byte(script), 0700); err != nil { - t.Fatal(err) - } - originalHelper := cudaCheckpointHelperBinary - cudaCheckpointHelperBinary = helper - t.Cleanup(func() { cudaCheckpointHelperBinary = originalHelper }) - liveJobFile := filepath.Join(tempDir, "live-job") checkpointDir := filepath.Join(tempDir, "checkpoint") if err := os.Mkdir(checkpointDir, 0700); err != nil { @@ -150,17 +161,22 @@ if [ "$action" = checkpoint ]; then printf '|%s' "$pid" >> "$job_file"; fi if err := os.WriteFile(filepath.Join(checkpointDir, snapshotv1alpha1.CUDAJobFileName), []byte("validation-copy"), 0600); err != nil { t.Fatal(err) } - t.Setenv("DYNAMO_TEST_TRACE", trace) - t.Setenv("DYNAMO_TEST_JOB_FILE", liveJobFile) - - if _, err := CheckpointProcessTree(context.Background(), []int{101, 202}, liveJobFile, checkpointDir, logr.Discard()); err != nil { - t.Fatalf("CheckpointProcessTree() error = %v", err) - } - traceContent, err := os.ReadFile(trace) - if err != nil { - t.Fatal(err) - } - if got, want := string(traceContent), "lock 101\nlock 202\ncheckpoint 101\ncheckpoint 202\n"; got != want { + runner := &jobFileRunner{jobFile: liveJobFile} + if _, err := lockAndCheckpointProcessTree( + context.Background(), + []int{101, 202}, + nil, + liveJobFile, + types.CUDAStorageModeLegacy, + checkpointDir, + nil, + types.CUDATransferSettings{}.WithDefaults(), + runner, + logr.Discard(), + ); err != nil { + t.Fatalf("lockAndCheckpointProcessTree() error = %v", err) + } + if got, want := strings.Join(runner.trace, "\n")+"\n", "lock 101\nlock 202\ncheckpoint 101\ncheckpoint 202\n"; got != want { t.Fatalf("helper call order = %q, want %q", got, want) } artifact, err := os.ReadFile(filepath.Join(checkpointDir, snapshotv1alpha1.CUDAJobFileName)) diff --git a/agent/internal/cuda/prefetch.go b/agent/internal/cuda/prefetch.go new file mode 100644 index 00000000..e00eadf4 --- /dev/null +++ b/agent/internal/cuda/prefetch.go @@ -0,0 +1,129 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package cuda + +import ( + "context" + "errors" + "fmt" + "io/fs" + "os" + "path/filepath" + "regexp" + "sort" + "time" + + "golang.org/x/sys/unix" +) + +const customStoragePrefetchBufferBytes = 8 * 1024 * 1024 + +var customStorageExtentFilePattern = regexp.MustCompile(`^device-[0-9]{4}\.bin(?:\.part-[0-9]{4})?$`) + +// CustomStoragePrefetchResult describes a completed best-effort page-cache +// preload. Duration is service time and may overlap CRIU restore. +type CustomStoragePrefetchResult struct { + Files int + Bytes int64 + Duration time.Duration +} + +// PrefetchCustomStorageArtifacts validates and reads every CUDA CustomStorage +// extent into the node page cache. Snapshot starts this before CRIU restore so +// durable storage I/O overlaps process restore; the CUDA helper still performs +// the authoritative read into registered host buffers after target PIDs exist. +func PrefetchCustomStorageArtifacts(ctx context.Context, checkpointDir string) (CustomStoragePrefetchResult, error) { + start := time.Now() + root := filepath.Join(checkpointDir, "cuda-custom-storage") + rootInfo, err := os.Lstat(root) + if err != nil { + return CustomStoragePrefetchResult{}, fmt.Errorf("inspect CUDA CustomStorage artifact directory: %w", err) + } + if !rootInfo.IsDir() { + return CustomStoragePrefetchResult{}, fmt.Errorf("CUDA CustomStorage artifact path is not a directory") + } + + var paths []string + err = filepath.WalkDir(root, func(path string, entry fs.DirEntry, walkErr error) error { + if walkErr != nil { + return walkErr + } + if entry.Type()&os.ModeSymlink != 0 { + if customStorageExtentFilePattern.MatchString(entry.Name()) { + return fmt.Errorf("CUDA CustomStorage extent %s is a symlink", path) + } + return nil + } + if entry.IsDir() || !customStorageExtentFilePattern.MatchString(entry.Name()) { + return nil + } + if !entry.Type().IsRegular() { + return fmt.Errorf("CUDA CustomStorage extent %s is not a regular file", path) + } + paths = append(paths, path) + return nil + }) + if err != nil { + return CustomStoragePrefetchResult{}, fmt.Errorf("discover CUDA CustomStorage extents: %w", err) + } + if len(paths) == 0 { + return CustomStoragePrefetchResult{}, fmt.Errorf("CUDA CustomStorage artifact contains no extent files") + } + sort.Strings(paths) + + buffer := make([]byte, customStoragePrefetchBufferBytes) + result := CustomStoragePrefetchResult{Files: len(paths)} + for _, path := range paths { + bytesRead, err := prefetchCustomStorageFile(ctx, path, buffer) + if err != nil { + return CustomStoragePrefetchResult{}, err + } + result.Bytes += bytesRead + } + result.Duration = time.Since(start) + return result, nil +} + +func prefetchCustomStorageFile(ctx context.Context, path string, buffer []byte) (int64, error) { + fd, err := unix.Open(path, unix.O_RDONLY|unix.O_CLOEXEC|unix.O_NOFOLLOW, 0) + if err != nil { + return 0, fmt.Errorf("open CUDA CustomStorage extent %s: %w", path, err) + } + defer unix.Close(fd) + + var stat unix.Stat_t + if err := unix.Fstat(fd, &stat); err != nil { + return 0, fmt.Errorf("stat CUDA CustomStorage extent %s: %w", path, err) + } + if stat.Mode&unix.S_IFMT != unix.S_IFREG || stat.Size <= 0 { + return 0, fmt.Errorf("CUDA CustomStorage extent %s is not a nonempty regular file", path) + } + _ = unix.Fadvise(fd, 0, stat.Size, unix.FADV_WILLNEED) + + // FADV_WILLNEED is only an asynchronous hint and may return before the + // extent reaches the page cache. Reading the complete file makes this + // best-effort prefetch observable and ensures the later CUDA restore can + // consume cached pages when the filesystem honors normal buffered I/O. + var total int64 + for { + if err := ctx.Err(); err != nil { + return 0, fmt.Errorf("prefetch CUDA CustomStorage extent %s: %w", path, err) + } + read, err := unix.Read(fd, buffer) + if errors.Is(err, unix.EINTR) { + continue + } + if err != nil { + return 0, fmt.Errorf("read CUDA CustomStorage extent %s: %w", path, err) + } + total += int64(read) + if read == 0 { + break + } + } + if total != stat.Size { + return 0, fmt.Errorf("CUDA CustomStorage extent %s changed size while prefetching: read %d, expected %d", path, total, stat.Size) + } + return total, nil +} diff --git a/agent/internal/cuda/prefetch_test.go b/agent/internal/cuda/prefetch_test.go new file mode 100644 index 00000000..2448face --- /dev/null +++ b/agent/internal/cuda/prefetch_test.go @@ -0,0 +1,125 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package cuda + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" +) + +func TestPrefetchCustomStorageArtifacts(t *testing.T) { + checkpointDir := t.TempDir() + processDir := filepath.Join(checkpointDir, "cuda-custom-storage", "process-nspid-42") + if err := os.MkdirAll(processDir, 0o700); err != nil { + t.Fatal(err) + } + files := map[string][]byte{ + "device-0000.bin.part-0000": []byte("first"), + "device-0000.bin.part-0001": []byte("second"), + "manifest.txt": []byte("ignored"), + } + for name, contents := range files { + if err := os.WriteFile(filepath.Join(processDir, name), contents, 0o600); err != nil { + t.Fatal(err) + } + } + + result, err := PrefetchCustomStorageArtifacts(context.Background(), checkpointDir) + if err != nil { + t.Fatal(err) + } + if result.Files != 2 || result.Bytes != int64(len("first")+len("second")) { + t.Fatalf("unexpected prefetch result: %+v", result) + } +} + +func TestPrefetchCustomStorageArtifactsRejectsExtentSymlink(t *testing.T) { + checkpointDir := t.TempDir() + processDir := filepath.Join(checkpointDir, "cuda-custom-storage", "process-nspid-42") + if err := os.MkdirAll(processDir, 0o700); err != nil { + t.Fatal(err) + } + target := filepath.Join(checkpointDir, "target") + if err := os.WriteFile(target, []byte("data"), 0o600); err != nil { + t.Fatal(err) + } + if err := os.Symlink(target, filepath.Join(processDir, "device-0000.bin")); err != nil { + t.Fatal(err) + } + + _, err := PrefetchCustomStorageArtifacts(context.Background(), checkpointDir) + if err == nil || !strings.Contains(err.Error(), "symlink") { + t.Fatalf("expected symlink rejection, got %v", err) + } +} + +func TestPrefetchCustomStorageArtifactsHonorsCancellation(t *testing.T) { + checkpointDir := t.TempDir() + processDir := filepath.Join(checkpointDir, "cuda-custom-storage", "process-nspid-42") + if err := os.MkdirAll(processDir, 0o700); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(processDir, "device-0000.bin"), []byte("data"), 0o600); err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err := PrefetchCustomStorageArtifacts(ctx, checkpointDir) + if err == nil || !strings.Contains(err.Error(), "context canceled") { + t.Fatalf("expected cancellation, got %v", err) + } +} + +func TestPrefetchCustomStorageArtifactsRejectsIncompleteArtifacts(t *testing.T) { + tests := []struct { + name string + prepare func(t *testing.T, checkpointDir string) + wantError string + }{ + { + name: "missing artifact directory", + prepare: func(*testing.T, string) {}, + wantError: "inspect CUDA CustomStorage artifact directory", + }, + { + name: "no extent files", + prepare: func(t *testing.T, checkpointDir string) { + t.Helper() + if err := os.MkdirAll(filepath.Join(checkpointDir, "cuda-custom-storage"), 0o700); err != nil { + t.Fatal(err) + } + }, + wantError: "contains no extent files", + }, + { + name: "empty extent", + prepare: func(t *testing.T, checkpointDir string) { + t.Helper() + processDir := filepath.Join(checkpointDir, "cuda-custom-storage", "process-nspid-42") + if err := os.MkdirAll(processDir, 0o700); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(processDir, "device-0000.bin"), nil, 0o600); err != nil { + t.Fatal(err) + } + }, + wantError: "not a nonempty regular file", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + checkpointDir := t.TempDir() + test.prepare(t, checkpointDir) + _, err := PrefetchCustomStorageArtifacts(context.Background(), checkpointDir) + if err == nil || !strings.Contains(err.Error(), test.wantError) { + t.Fatalf("expected error containing %q, got %v", test.wantError, err) + } + }) + } +} diff --git a/agent/internal/cuda/shim.go b/agent/internal/cuda/shim.go index fb60092e..08219cbb 100644 --- a/agent/internal/cuda/shim.go +++ b/agent/internal/cuda/shim.go @@ -5,114 +5,134 @@ package cuda import ( "context" + "encoding/json" "errors" "fmt" - "os" - "os/exec" + "math" "strconv" "strings" - "syscall" "time" "github.com/go-logr/logr" snapshotruntime "github.com/ai-dynamo/snapshot/agent/internal/runtime" + "github.com/ai-dynamo/snapshot/agent/internal/types" ) const ( - helperWaitDelay = 2 * time.Second - actionLock = "lock" actionCheckpoint = "checkpoint" actionRestore = "restore" actionUnlock = "unlock" ) -var cudaCheckpointHelperBinary = DefaultHelperBinaryPath +var errProcessIdentityChangedBeforeCUDA = errors.New("process identity changed before CUDA driver call") -func lock(ctx context.Context, pid int, log logr.Logger) error { - return runAction(ctx, pid, actionLock, "", DefaultHelperBinaryPath, log) +type helperActionRunner interface { + run(context.Context, helperAction, logr.Logger) error } -func checkpoint(ctx context.Context, pid int, log logr.Logger) error { - return runAction(ctx, pid, actionCheckpoint, "", DefaultHelperBinaryPath, log) +type helperAction struct { + PID int + Action string + DeviceMap string + StorageMode string + StorageDir string + JobFile string + GPUUUIDs []string + Transfer types.CUDATransferSettings + Identity snapshotruntime.ProcessDetails } -func restoreProcess(ctx context.Context, pid int, deviceMap, helperBinaryPath string, log logr.Logger) error { - return runAction(ctx, pid, actionRestore, deviceMap, helperBinaryPath, log) +type commandHelperActionRunner struct{} + +type identityValidatingRunner struct { + runner helperActionRunner + procRoot string + identities map[int]snapshotruntime.ProcessDetails } -func unlock(ctx context.Context, pid int, helperBinaryPath string, log logr.Logger) error { - return runAction(ctx, pid, actionUnlock, "", helperBinaryPath, log) +type customStorageTelemetry struct { + Event string `json:"event"` + HelperMainToTelemetrySeconds json.RawMessage `json:"helper_main_to_telemetry_seconds"` } -func getState(ctx context.Context, pid int, helperBinaryPath string) (string, error) { - cmd := exec.CommandContext(ctx, helperBinaryPath, "--get-state", "--pid", strconv.Itoa(pid)) - output, err := cmd.CombinedOutput() - state := strings.TrimSpace(string(output)) - if err != nil { - return "", fmt.Errorf("cuda-checkpoint-helper --get-state failed for pid %d: %w (output: %s)", pid, err, state) - } - if state == "" { - return "", fmt.Errorf("cuda-checkpoint-helper --get-state returned empty state for pid %d", pid) - } - return state, nil +type customStorageTelemetryParse struct { + status string + err string + helperMainDuration time.Duration } -func runAction(ctx context.Context, pid int, action, deviceMap, helperBinaryPath string, log logr.Logger) error { - args := []string{"--action", action, "--pid", strconv.Itoa(pid)} - if action == actionRestore && deviceMap != "" { - args = append(args, "--device-map", deviceMap) - } - cmd := exec.CommandContext(ctx, helperBinaryPath, args...) - cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} - cmd.Cancel = func() error { - return normalizeProcessGroupKillError(syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)) +func parseCustomStorageTelemetry(output string, processWall time.Duration) customStorageTelemetryParse { + sawMalformedJSON := false + for _, line := range strings.Split(output, "\n") { + line = strings.TrimSpace(line) + if !strings.HasPrefix(line, "{") { + continue + } + var telemetry customStorageTelemetry + if err := json.Unmarshal([]byte(line), &telemetry); err != nil { + sawMalformedJSON = true + continue + } + if telemetry.Event != "cuda_custom_storage_transfer" { + continue + } + if len(telemetry.HelperMainToTelemetrySeconds) == 0 || string(telemetry.HelperMainToTelemetrySeconds) == "null" { + return customStorageTelemetryParse{status: "missing-duration", err: "expected helper_main_to_telemetry_seconds"} + } + var seconds json.Number + if err := json.Unmarshal(telemetry.HelperMainToTelemetrySeconds, &seconds); err != nil { + return customStorageTelemetryParse{status: "invalid-duration", err: "helper_main_to_telemetry_seconds is not a number"} + } + value, err := strconv.ParseFloat(seconds.String(), 64) + if err != nil || math.IsNaN(value) || math.IsInf(value, 0) || value < 0 { + return customStorageTelemetryParse{status: "invalid-duration", err: "helper_main_to_telemetry_seconds is not a finite non-negative number"} + } + const roundingToleranceSeconds = 1e-6 + processWallSeconds := processWall.Seconds() + if value > processWallSeconds+roundingToleranceSeconds { + return customStorageTelemetryParse{status: "duration-exceeds-process-wall", err: "helper_main_to_telemetry_seconds exceeds process wall duration"} + } + if value >= processWallSeconds || value*float64(time.Second) >= float64(math.MaxInt64) { + return customStorageTelemetryParse{status: "valid", helperMainDuration: processWall} + } + return customStorageTelemetryParse{status: "valid", helperMainDuration: time.Duration(value * float64(time.Second))} } - cmd.WaitDelay = helperWaitDelay - details := snapshotruntime.ProcessDetails{ - ObservedPID: pid, - OutermostPID: pid, - InnermostPID: pid, - NamespacePIDs: []int{pid}, + if sawMalformedJSON { + return customStorageTelemetryParse{status: "malformed-json", err: "malformed JSON telemetry output"} } - if process, err := snapshotruntime.ReadProcessDetails("/proc", pid); err == nil { - details = process + return customStorageTelemetryParse{status: "event-absent", err: "cuda_custom_storage_transfer event not found"} +} + +func (commandHelperActionRunner) run( + ctx context.Context, + request helperAction, + log logr.Logger, +) error { + if request.Identity.OutermostPID != request.PID || + request.Identity.StartTimeTicks == 0 || request.Identity.Cgroup == "" { + return fmt.Errorf("incomplete process identity for host PID %d", request.PID) } - start := time.Now() - output, err := cmd.CombinedOutput() - duration := time.Since(start) - out := strings.TrimSpace(string(output)) - if err != nil { - if ctx.Err() != nil { - err = ctx.Err() - } - log.Error(err, "cuda-checkpoint-helper command failed", - "pid", pid, - "outermost_pid", details.OutermostPID, - "innermost_pid", details.InnermostPID, - "cmdline", details.Cmdline, - "action", action, - "duration", duration, - "output", out, - ) - return fmt.Errorf("cuda-checkpoint-helper %v failed for pid %d after %s: %w (output: %s)", args, pid, duration, err, out) + if request.Action == actionLock || request.Action == actionUnlock || + request.StorageMode == types.CUDAStorageModeLegacy { + request.StorageDir = "" } - log.V(1).Info("cuda-checkpoint-helper command succeeded", - "pid", pid, - "outermost_pid", details.OutermostPID, - "innermost_pid", details.InnermostPID, - "cmdline", details.Cmdline, - "action", action, - "duration", duration, - "output", out, - ) - return nil + return runDaemonAction(ctx, request, log) } -func normalizeProcessGroupKillError(err error) error { - if errors.Is(err, syscall.ESRCH) { - return os.ErrProcessDone +func (r identityValidatingRunner) run( + ctx context.Context, + request helperAction, + log logr.Logger, +) error { + expected, ok := r.identities[request.PID] + if !ok { + return fmt.Errorf("%w: missing expected process identity for host PID %d", errProcessIdentityChangedBeforeCUDA, request.PID) + } + if err := snapshotruntime.ValidateProcessIdentity(r.procRoot, expected); err != nil { + return fmt.Errorf("%w: validate host PID %d immediately before CUDA %s: %v", errProcessIdentityChangedBeforeCUDA, request.PID, request.Action, err) } - return err + request.Identity = expected + return r.runner.run(ctx, request, log) } diff --git a/agent/internal/cuda/shim_job_file.go b/agent/internal/cuda/shim_job_file.go deleted file mode 100644 index 88213e84..00000000 --- a/agent/internal/cuda/shim_job_file.go +++ /dev/null @@ -1,80 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -package cuda - -import ( - "context" - "fmt" - "os/exec" - "strconv" - "strings" - "syscall" - "time" - - "github.com/go-logr/logr" - - snapshotruntime "github.com/ai-dynamo/snapshot/agent/internal/runtime" -) - -func lockWithJobFile(ctx context.Context, pid int, jobFile string, log logr.Logger) error { - if jobFile == "" { - return lock(ctx, pid, log) - } - return runActionWithJobFile(ctx, pid, actionLock, jobFile, log) -} - -func checkpointWithJobFile(ctx context.Context, pid int, jobFile string, log logr.Logger) error { - if jobFile == "" { - return checkpoint(ctx, pid, log) - } - return runActionWithJobFile(ctx, pid, actionCheckpoint, jobFile, log) -} - -func runActionWithJobFile(ctx context.Context, pid int, action, jobFile string, log logr.Logger) error { - args := []string{"--action", action, "--pid", strconv.Itoa(pid), "--job-file", jobFile} - cmd := exec.CommandContext(ctx, cudaCheckpointHelperBinary, args...) - cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} - cmd.Cancel = func() error { - return normalizeProcessGroupKillError(syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)) - } - cmd.WaitDelay = helperWaitDelay - details := snapshotruntime.ProcessDetails{ - ObservedPID: pid, - OutermostPID: pid, - InnermostPID: pid, - NamespacePIDs: []int{pid}, - } - if process, err := snapshotruntime.ReadProcessDetails("/proc", pid); err == nil { - details = process - } - start := time.Now() - output, err := cmd.CombinedOutput() - duration := time.Since(start) - out := strings.TrimSpace(string(output)) - if err != nil { - if ctx.Err() != nil { - err = ctx.Err() - } - log.Error(err, "cuda-checkpoint-helper command failed", - "pid", pid, - "outermost_pid", details.OutermostPID, - "innermost_pid", details.InnermostPID, - "cmdline", details.Cmdline, - "action", action, - "duration", duration, - "output", out, - ) - return fmt.Errorf("cuda-checkpoint-helper %v failed for pid %d after %s: %w (output: %s)", args, pid, duration, err, out) - } - log.V(1).Info("cuda-checkpoint-helper command succeeded", - "pid", pid, - "outermost_pid", details.OutermostPID, - "innermost_pid", details.InnermostPID, - "cmdline", details.Cmdline, - "action", action, - "duration", duration, - "output", out, - ) - return nil -} diff --git a/agent/internal/cuda/shim_restore_job_file_test.go b/agent/internal/cuda/shim_restore_job_file_test.go deleted file mode 100644 index e7d0d469..00000000 --- a/agent/internal/cuda/shim_restore_job_file_test.go +++ /dev/null @@ -1,32 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -package cuda - -import ( - "context" - "os" - "path/filepath" - "testing" - - "github.com/go-logr/logr" - - snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" -) - -func TestRunActionInheritsJobFileEnvironment(t *testing.T) { - trace := filepath.Join(t.TempDir(), "trace") - installFakeCUDAHelper(t, "printf '%s' \"$CUDA_CHECKPOINT_JOB_FILE\" > \""+trace+"\"\n") - t.Setenv(JobFileEnv, snapshotv1alpha1.CUDAJobFilePath) - - if err := runAction(context.Background(), 11, actionRestore, "", cudaCheckpointHelperBinary, logr.Discard()); err != nil { - t.Fatalf("runAction() error = %v", err) - } - content, err := os.ReadFile(trace) - if err != nil { - t.Fatal(err) - } - if got, want := string(content), snapshotv1alpha1.CUDAJobFilePath; got != want { - t.Fatalf("helper environment = %q, want %q", got, want) - } -} diff --git a/agent/internal/cuda/shim_test.go b/agent/internal/cuda/shim_test.go deleted file mode 100644 index 6dc9bb17..00000000 --- a/agent/internal/cuda/shim_test.go +++ /dev/null @@ -1,50 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -package cuda - -import ( - "context" - "errors" - "os" - "path/filepath" - "strings" - "syscall" - "testing" - "time" - - "github.com/go-logr/logr" -) - -func installFakeCUDAHelper(t *testing.T, script string) { - t.Helper() - helper := filepath.Join(t.TempDir(), "cuda-checkpoint-helper") - if err := os.WriteFile(helper, []byte("#!/bin/sh\n"+script), 0700); err != nil { - t.Fatal(err) - } - originalHelper := cudaCheckpointHelperBinary - cudaCheckpointHelperBinary = helper - t.Cleanup(func() { cudaCheckpointHelperBinary = originalHelper }) -} - -func TestRunActionCancellationIsBounded(t *testing.T) { - installFakeCUDAHelper(t, "sleep 300 &\nwait\n") - ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) - defer cancel() - - started := time.Now() - err := runAction(ctx, 11, actionRestore, "", cudaCheckpointHelperBinary, logr.Discard()) - duration := time.Since(started) - if err == nil || !strings.Contains(err.Error(), context.DeadlineExceeded.Error()) { - t.Fatalf("runAction() error = %v", err) - } - if duration > helperWaitDelay+time.Second { - t.Fatalf("runAction() took %s after cancellation", duration) - } -} - -func TestNormalizeProcessGroupKillErrorReportsFinishedProcess(t *testing.T) { - if err := normalizeProcessGroupKillError(syscall.ESRCH); !errors.Is(err, os.ErrProcessDone) { - t.Fatalf("normalizeProcessGroupKillError() error = %v, want %v", err, os.ErrProcessDone) - } -} diff --git a/agent/internal/executor/checkpoint.go b/agent/internal/executor/checkpoint.go index 1f965012..856ee19f 100644 --- a/agent/internal/executor/checkpoint.go +++ b/agent/internal/executor/checkpoint.go @@ -7,9 +7,11 @@ package executor import ( "context" + "errors" "fmt" "os" "path/filepath" + "syscall" "time" criurpc "github.com/checkpoint-restore/go-criu/v8/rpc" @@ -43,6 +45,26 @@ type checkpointPhaseTimings struct { OverlayCaptureDuration time.Duration } +type checkpointMutationError struct { + err error + targetMayBeMutated bool +} + +func (e *checkpointMutationError) Error() string { return e.err.Error() } +func (e *checkpointMutationError) Unwrap() error { return e.err } + +func checkpointPreMutationError(err error) error { + return &checkpointMutationError{err: err, targetMayBeMutated: false} +} + +// CheckpointFailedBeforeTargetMutation is true only when checkpoint +// preflight failed before CRIU or any CUDA driver operation could mutate the +// source workload. Unclassified failures remain fail-closed. +func CheckpointFailedBeforeTargetMutation(err error) bool { + var checkpointErr *checkpointMutationError + return errors.As(err, &checkpointErr) && !checkpointErr.targetMayBeMutated +} + // Checkpoint performs a CRIU dump of a container. // // The checkpoint directory is staged under the content-owned .tmp directory. @@ -54,42 +76,73 @@ func Checkpoint(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger finalDir, err := nsmount.ResolveArtifactPath(cfg.Storage.BasePath, req.ContentUID, req.ContainerName) if err != nil { - return fmt.Errorf("resolve checkpoint artifact path: %w", err) + return checkpointPreMutationError(fmt.Errorf("resolve checkpoint artifact path: %w", err)) } tmpRoot, err := nsmount.ResolveArtifactStagingRoot(cfg.Storage.BasePath, req.ContentUID) if err != nil { - return fmt.Errorf("resolve checkpoint staging root: %w", err) + return checkpointPreMutationError(fmt.Errorf("resolve checkpoint staging root: %w", err)) } if err := os.MkdirAll(tmpRoot, 0700); err != nil { - return fmt.Errorf("failed to create checkpoint staging root: %w", err) + return checkpointPreMutationError(fmt.Errorf("failed to create checkpoint staging root: %w", err)) } if err := os.MkdirAll(filepath.Dir(finalDir), 0700); err != nil { - return fmt.Errorf("failed to create checkpoint container root: %w", err) + return checkpointPreMutationError(fmt.Errorf("failed to create checkpoint container root: %w", err)) } tmpDir := filepath.Join(tmpRoot, uuid.NewString()) if err := os.Mkdir(tmpDir, 0700); err != nil { - return fmt.Errorf("failed to create checkpoint staging directory: %w", err) + return checkpointPreMutationError(fmt.Errorf("failed to create checkpoint staging directory: %w", err)) } defer os.RemoveAll(tmpDir) state, gpuDeviceMapDuration, err := inspectContainer(ctx, rt, log, req) if err != nil { - return err + return checkpointPreMutationError(err) } cudaJobFile := "" + cudaStorageMode := types.CUDAStorageModeLegacy if len(state.CUDAHostPIDs) > 0 { cudaJobFile, err = cuda.StageJobFile(state.RootFS, tmpDir, len(state.GPUUUIDs)) if err != nil { - return err + return checkpointPreMutationError(err) + } + if cfg.CUDACheckpoint.StorageMode == types.CUDAStorageModePOSIX { + if err := validatePOSIXCustomStorageTopology(len(state.CUDAHostPIDs), len(state.GPUUUIDs)); err != nil { + return checkpointPreMutationError(err) + } + } + cudaStorageMode, err = cuda.SelectCUDAStorageMode( + ctx, + cfg.CUDACheckpoint.StorageMode, + ) + if err != nil { + return checkpointPreMutationError(fmt.Errorf("select CUDA storage mode before locking target: %w", err)) + } + if cudaStorageMode == types.CUDAStorageModePOSIX { + log.Info("CUDA CustomStorage explicitly enabled and available; using the Snapshot-local NIXL POSIX path", + "cuda_storage_mode", cudaStorageMode) + } else { + log.Info("CUDA CustomStorage disabled for new checkpoints; using legacy CUDA checkpoint storage", + "cuda_storage_mode", cudaStorageMode) } } - criuOpts, data, err := configureCheckpoint(log, state, req, cfg, tmpDir) + criuOpts, data, err := configureCheckpoint(log, state, req, cfg, tmpDir, cudaStorageMode) if err != nil { - return err + return checkpointPreMutationError(err) } - captureTimings, err := captureCheckpoint(ctx, criuOpts, &cfg.CRIU, data, state, tmpDir, cudaJobFile, log) + captureTimings, err := captureCheckpoint( + ctx, + criuOpts, + &cfg.CRIU, + cfg.CUDACheckpoint.TransferSettings(), + data, + state, + tmpDir, + cudaJobFile, + cudaStorageMode, + log, + ) if err != nil { return err } @@ -132,6 +185,14 @@ func Checkpoint(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger return nil } +func validatePOSIXCustomStorageTopology(processCount, gpuCount int) error { + if processCount < 1 || gpuCount != 1 { + return fmt.Errorf("CUDA POSIX CustomStorage is qualified only for one or more CUDA processes on one GPU; found processes=%d GPUs=%d", + processCount, gpuCount) + } + return nil +} + func inspectContainer(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, req CheckpointRequest) (*types.CheckpointContainerSnapshot, time.Duration, error) { containerID := req.ContainerID pid, ociSpec, err := rt.ResolveContainer(ctx, containerID) @@ -178,7 +239,10 @@ func inspectContainer(ctx context.Context, rt snapshotruntime.Runtime, log logr. // Discover CUDA processes and GPU UUIDs allPIDs := snapshotruntime.ProcessTreePIDs(pid) - cudaHostPIDs := cuda.FilterProcesses(ctx, allPIDs, log) + cudaHostPIDs, err := cuda.FilterProcesses(ctx, allPIDs, log) + if err != nil { + return nil, 0, fmt.Errorf("discover CUDA processes: %w", err) + } cudaNamespacePIDs := make([]int, 0, len(cudaHostPIDs)) for _, cudaHostPID := range cudaHostPIDs { process, err := snapshotruntime.ReadProcessDetails(snapshotruntime.HostProcPath, cudaHostPID) @@ -234,6 +298,7 @@ func configureCheckpoint( req CheckpointRequest, cfg *types.AgentConfig, checkpointDir string, + cudaStorageMode string, ) (*criurpc.CriuOpts, *types.CheckpointManifest, error) { criuOpts, err := criu.BuildDumpOptions(state, &cfg.CRIU, checkpointDir, log) if err != nil { @@ -248,7 +313,7 @@ func configureCheckpoint( types.NewOverlayManifest(cfg.Overlay, state.UpperDir, state.OCISpec), ) if len(state.CUDANSPIDs) > 0 { - m.CUDA = types.NewCUDAManifest(state.CUDANSPIDs, state.GPUUUIDs) + m.CUDA = types.NewCUDAManifest(state.CUDANSPIDs, state.GPUUUIDs, cudaStorageMode) } if err := types.WriteManifest(checkpointDir, m); err != nil { @@ -258,14 +323,53 @@ func configureCheckpoint( return criuOpts, m, nil } -func captureCheckpoint(ctx context.Context, criuOpts *criurpc.CriuOpts, criuSettings *types.CRIUSettings, data *types.CheckpointManifest, state *types.CheckpointContainerSnapshot, checkpointDir, cudaJobFile string, log logr.Logger) (*checkpointPhaseTimings, error) { +func captureCheckpoint( + ctx context.Context, + criuOpts *criurpc.CriuOpts, + criuSettings *types.CRIUSettings, + cudaTransfer types.CUDATransferSettings, + data *types.CheckpointManifest, + state *types.CheckpointContainerSnapshot, + checkpointDir, + cudaJobFile, + cudaStorageMode string, + log logr.Logger, +) (*checkpointPhaseTimings, error) { timings := &checkpointPhaseTimings{} // CUDA lock+checkpoint must happen before CRIU dump if len(state.CUDAHostPIDs) > 0 { - cudaTimings, err := cuda.CheckpointProcessTree(ctx, state.CUDAHostPIDs, cudaJobFile, checkpointDir, log) + processes, err := readCUDAProcessDetailsForCheckpoint( + snapshotruntime.HostProcPath, + state.CUDAHostPIDs, + ) + if err != nil { + return nil, err + } + cudaTimings, err := cuda.LockAndCheckpointProcessTreeValidated( + ctx, + processes, + cudaJobFile, + cudaStorageMode, + checkpointDir, + state.GPUUUIDs, + cudaTransfer, + log, + ) if err != nil { - return nil, fmt.Errorf("CUDA checkpoint failed: %w", err) + checkpointErr := fmt.Errorf("CUDA checkpoint failed: %w", err) + if cuda.FailedBeforeTargetMutation(err) { + return nil, checkpointPreMutationError(checkpointErr) + } + cleanupErr := terminateCUDAProcessesAfterOperationFailure( + processes, + snapshotruntime.HostProcPath, + "checkpoint", + log, + snapshotruntime.ValidateProcessIdentity, + snapshotruntime.SendSignalToPID, + ) + return nil, errors.Join(checkpointErr, cleanupErr) } timings.CUDACheckpointDuration = cudaTimings.TotalDuration } @@ -292,3 +396,57 @@ func captureCheckpoint(ctx context.Context, criuOpts *criurpc.CriuOpts, criuSett return timings, nil } + +func readCUDAProcessDetailsForCheckpoint( + procRoot string, + pids []int, +) ([]snapshotruntime.ProcessDetails, error) { + processes := make([]snapshotruntime.ProcessDetails, 0, len(pids)) + for _, pid := range pids { + process, err := snapshotruntime.ReadProcessDetails(procRoot, pid) + if err != nil { + return nil, checkpointPreMutationError( + fmt.Errorf("capture CUDA process identity for PID %d: %w", pid, err), + ) + } + processes = append(processes, process) + } + return processes, nil +} + +type signalProcessFunc func(logr.Logger, int, syscall.Signal, string) error + +type validateProcessIdentityFunc func(string, snapshotruntime.ProcessDetails) error + +func terminateCUDAProcessesAfterOperationFailure( + processes []snapshotruntime.ProcessDetails, + procRoot string, + operation string, + log logr.Logger, + validateProcessIdentity validateProcessIdentityFunc, + signalProcess signalProcessFunc, +) error { + var cleanupErr error + for _, process := range processes { + pid := process.OutermostPID + if err := validateProcessIdentity(procRoot, process); err != nil { + cleanupErr = errors.Join( + cleanupErr, + fmt.Errorf( + "refusing to terminate CUDA PID %d after %s identity validation failed: %w", + pid, + operation, + err, + ), + ) + continue + } + if err := signalProcess(log, pid, syscall.SIGKILL, "CUDA "+operation+" failed"); err != nil { + cleanupErr = errors.Join(cleanupErr, err) + } + } + if cleanupErr != nil { + return fmt.Errorf("failed to terminate one or more CUDA processes after %s failure: %w", operation, cleanupErr) + } + return nil +} diff --git a/agent/internal/executor/checkpoint_test.go b/agent/internal/executor/checkpoint_test.go index 6674987a..7cd7ef02 100644 --- a/agent/internal/executor/checkpoint_test.go +++ b/agent/internal/executor/checkpoint_test.go @@ -7,15 +7,17 @@ import ( "context" "errors" "path/filepath" + "strings" + "syscall" "testing" + "github.com/ai-dynamo/snapshot/agent/internal/nsmount" + snapshotruntime "github.com/ai-dynamo/snapshot/agent/internal/runtime" + "github.com/ai-dynamo/snapshot/agent/internal/types" "github.com/go-logr/logr" - specs "github.com/opencontainers/runtime-spec/specs-go" + "github.com/opencontainers/runtime-spec/specs-go" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" - - "github.com/ai-dynamo/snapshot/agent/internal/nsmount" - "github.com/ai-dynamo/snapshot/agent/internal/types" ) type checkpointPathRuntime struct{} @@ -47,3 +49,122 @@ func TestCheckpointPreparesContentArtifactParents(t *testing.T) { assert.DirExists(t, filepath.Dir(finalDir)) assert.DirExists(t, filepath.Join(cfg.Storage.BasePath, "artifacts", "content-uid", ".tmp")) } + +func TestValidatePOSIXCustomStorageTopology(t *testing.T) { + for _, processCount := range []int{1, 2, 8} { + if err := validatePOSIXCustomStorageTopology(processCount, 1); err != nil { + t.Fatalf("validatePOSIXCustomStorageTopology(%d, 1) = %v", processCount, err) + } + } + for _, topology := range []struct { + processes int + gpus int + }{{0, 1}, {1, 2}, {2, 2}} { + err := validatePOSIXCustomStorageTopology(topology.processes, topology.gpus) + if err == nil || !strings.Contains(err.Error(), "qualified only for one or more CUDA processes on one GPU") { + t.Fatalf("validatePOSIXCustomStorageTopology(%d, %d) = %v, want qualification error", + topology.processes, topology.gpus, err) + } + } +} + +func TestReadCUDAProcessDetailsFailureIsPreMutation(t *testing.T) { + _, err := readCUDAProcessDetailsForCheckpoint(t.TempDir(), []int{424242}) + if err == nil { + t.Fatal("readCUDAProcessDetailsForCheckpoint() error = nil, want missing process error") + } + if !CheckpointFailedBeforeTargetMutation(err) { + t.Fatalf("CheckpointFailedBeforeTargetMutation(%v) = false, want true", err) + } +} + +func TestTerminateCUDAProcessesAfterOperationFailurePropagatesCleanupError(t *testing.T) { + var attempted []int + err := terminateCUDAProcessesAfterOperationFailure( + []snapshotruntime.ProcessDetails{ + {OutermostPID: 41, StartTimeTicks: 100, Cgroup: "first"}, + {OutermostPID: 42, StartTimeTicks: 200, Cgroup: "second"}, + }, + "/test/proc", + "checkpoint", + logr.Discard(), + func(procRoot string, process snapshotruntime.ProcessDetails) error { + if procRoot != "/test/proc" { + t.Fatalf("proc root = %q, want /test/proc", procRoot) + } + return nil + }, + func(_ logr.Logger, pid int, signal syscall.Signal, reason string) error { + attempted = append(attempted, pid) + if signal != syscall.SIGKILL || reason != "CUDA checkpoint failed" { + t.Fatalf("signal call = (%d, %q), want SIGKILL and checkpoint reason", signal, reason) + } + if pid == 41 { + return errors.New("permission denied") + } + return nil + }, + ) + if len(attempted) != 2 || attempted[0] != 41 || attempted[1] != 42 { + t.Fatalf("attempted PIDs = %v, want [41 42]", attempted) + } + if err == nil || !strings.Contains(err.Error(), "permission denied") { + t.Fatalf("cleanup error = %v, want propagated signal failure", err) + } +} + +func TestTerminateCUDAProcessesAfterOperationFailureRejectsChangedIdentity(t *testing.T) { + var attempted []int + err := terminateCUDAProcessesAfterOperationFailure( + []snapshotruntime.ProcessDetails{ + {OutermostPID: 41, StartTimeTicks: 100, Cgroup: "first"}, + {OutermostPID: 42, StartTimeTicks: 200, Cgroup: "second"}, + }, + "/test/proc", + "checkpoint", + logr.Discard(), + func(_ string, process snapshotruntime.ProcessDetails) error { + if process.OutermostPID == 41 { + return errors.New("process identity changed") + } + return nil + }, + func(_ logr.Logger, pid int, _ syscall.Signal, _ string) error { + attempted = append(attempted, pid) + return nil + }, + ) + if len(attempted) != 1 || attempted[0] != 42 { + t.Fatalf("attempted PIDs = %v, want [42]", attempted) + } + if err == nil || !strings.Contains(err.Error(), "process identity changed") { + t.Fatalf("cleanup error = %v, want identity validation failure", err) + } +} + +func TestTerminateCUDAProcessesAfterRestoreFailureKillsEveryValidatedTarget(t *testing.T) { + var attempted []int + err := terminateCUDAProcessesAfterOperationFailure( + []snapshotruntime.ProcessDetails{ + {OutermostPID: 51, StartTimeTicks: 300, Cgroup: "first"}, + {OutermostPID: 52, StartTimeTicks: 400, Cgroup: "second"}, + }, + "/test/proc", + "restore", + logr.Discard(), + func(string, snapshotruntime.ProcessDetails) error { return nil }, + func(_ logr.Logger, pid int, signal syscall.Signal, reason string) error { + attempted = append(attempted, pid) + if signal != syscall.SIGKILL || reason != "CUDA restore failed" { + t.Fatalf("signal call = (%d, %q), want SIGKILL and restore reason", signal, reason) + } + return nil + }, + ) + if err != nil { + t.Fatalf("restore cleanup error = %v", err) + } + if len(attempted) != 2 || attempted[0] != 51 || attempted[1] != 52 { + t.Fatalf("attempted PIDs = %v, want [51 52]", attempted) + } +} diff --git a/agent/internal/executor/nsrestore.go b/agent/internal/executor/nsrestore.go index 7ab89d3f..e99bec02 100644 --- a/agent/internal/executor/nsrestore.go +++ b/agent/internal/executor/nsrestore.go @@ -7,7 +7,6 @@ import ( "context" "fmt" "os" - "path/filepath" "syscall" "time" @@ -32,12 +31,13 @@ type RestoreOptions struct { } type RestoreInNamespaceResult struct { - RestoredPID int `json:"restoredPID"` - CleanupError *CleanupError `json:"cleanupError,omitempty"` - OverlayCaptureDuration time.Duration `json:"overlayCaptureDuration"` - CRIUPrepareDuration time.Duration `json:"criuPrepareDuration"` - CRIURestoreDuration time.Duration `json:"criuRestoreDuration"` - CUDARestoreDuration time.Duration `json:"cudaRestoreDuration"` + RestoredPID int `json:"restoredPID"` + CleanupError *CleanupError `json:"cleanupError,omitempty"` + OverlayCaptureDuration time.Duration `json:"overlayCaptureDuration"` + CRIUPrepareDuration time.Duration `json:"criuPrepareDuration"` + CRIURestoreDuration time.Duration `json:"criuRestoreDuration"` + CUDARestoreDuration time.Duration `json:"cudaRestoreDuration"` + DeferredCUDAProcesses []snapshotruntime.ProcessDetails `json:"deferredCUDAProcesses,omitempty"` } // CleanupError is the wire representation of a successful restore whose @@ -99,6 +99,7 @@ func RestoreInNamespace(ctx context.Context, opts RestoreOptions, log logr.Logge CRIUPrepareDuration: executeTimings.criuPrepareDuration, CRIURestoreDuration: executeTimings.criuRestoreDuration, CUDARestoreDuration: executeTimings.cudaRestoreDuration, + DeferredCUDAProcesses: executeTimings.deferredCUDAProcesses, } if cleanupErr != nil { result.CleanupError = &CleanupError{ @@ -114,6 +115,7 @@ type nsrestorePhaseTimings struct { criuPrepareDuration time.Duration criuRestoreDuration time.Duration cudaRestoreDuration time.Duration + deferredCUDAProcesses []snapshotruntime.ProcessDetails } func executeRestore( @@ -160,24 +162,6 @@ func executeRestore( } }() - // Open the cuda-checkpoint-helper fd BEFORE CRIU runs. CRIU restores the - // original mount namespace of the checkpointed process, which did not include - // the bundle mount at /tmp/snapshot-binaries. The C helper's umount code - // tolerates ENOENT from umount2 with the comment "Already gone (CRIU removed - // it during namespace restore)" — confirming this is observed behaviour. By - // opening the binary now and exec'ing via /proc/self/fd/N after CRIU returns, - // the fd remains valid even if the mount is gone. - var cudaHelperFdPath string - if !m.CUDA.IsEmpty() { - helperPath := filepath.Join(opts.BundleDir, cuda.HelperBinaryName) - f, err := os.Open(helperPath) - if err != nil { - return nil, 0, nil, fmt.Errorf("failed to open cuda-checkpoint-helper before CRIU restore: %w", err) - } - defer f.Close() - cudaHelperFdPath = fmt.Sprintf("/proc/self/fd/%d", f.Fd()) - } - // The restore-complete sentinel lives on the pod emptyDir mounted at // SnapshotControlMountPath. Clear it here, in that mount namespace, so a // leftover from an earlier incarnation cannot release the restored process @@ -192,7 +176,9 @@ func executeRestore( return nil, 0, nil, err } restoredPID = int(criuPID) - // Cleanup runs after CUDA unlock. A cleanup-only failure is returned + // Cleanup releases only CRIU-owned files and scratch paths. It runs when + // nsrestore returns, before the host agent asks the CUDA daemon to restore + // and unlock the parked processes. A cleanup-only failure is returned // separately so the host controller can warn without killing the workload. defer func() { if err := cleanup(); err != nil { @@ -235,8 +221,9 @@ func executeRestore( ) } - // CUDA restore — remap checkpoint-time innermost namespace PIDs onto the - // current visible restored PIDs before invoking cuda-checkpoint. + // Resolve the restored CUDA processes while still inside the restored PID + // namespace. The host agent maps these identities to host PIDs and delegates + // restore to the always-on CUDA helper daemon after nsrestore returns. if !m.CUDA.IsEmpty() { restorePIDs, err := snapshotruntime.ResolveManifestPIDsToObservedPIDs(processes, restoredPID, m.CUDA.PIDs) if err != nil { @@ -247,11 +234,12 @@ func executeRestore( "restored_cuda_pids", restorePIDs, "criu_callback_pid", restoredPID, ) - cudaStart := time.Now() - _, err = cuda.RestoreAndUnlockProcessTree(ctx, restorePIDs, opts.CUDADeviceMap, cudaHelperFdPath, log) - timings.cudaRestoreDuration = time.Since(cudaStart) - if err != nil { - return nil, 0, nil, fmt.Errorf("CUDA restore failed: %w", err) + for _, pid := range restorePIDs { + process, err := snapshotruntime.ReadProcessDetails("/proc", pid) + if err != nil { + return nil, 0, nil, fmt.Errorf("capture restored CUDA process identity for PID %d: %w", pid, err) + } + timings.deferredCUDAProcesses = append(timings.deferredCUDAProcesses, process) } } diff --git a/agent/internal/executor/restore.go b/agent/internal/executor/restore.go index 6847a3fc..f9b78f95 100644 --- a/agent/internal/executor/restore.go +++ b/agent/internal/executor/restore.go @@ -51,6 +51,34 @@ type restoreMount struct { point nsmount.MountPoint } +type customStoragePrefetchOutcome struct { + result cuda.CustomStoragePrefetchResult + err error +} + +func waitForCustomStoragePrefetch( + ctx context.Context, + outcomes <-chan customStoragePrefetchOutcome, + cancel context.CancelFunc, + discard bool, +) (cuda.CustomStoragePrefetchResult, error) { + if outcomes == nil { + return cuda.CustomStoragePrefetchResult{}, nil + } + if discard { + cancel() + return cuda.CustomStoragePrefetchResult{}, nil + } + select { + case outcome := <-outcomes: + cancel() + return outcome.result, outcome.err + case <-ctx.Done(): + cancel() + return cuda.CustomStoragePrefetchResult{}, fmt.Errorf("wait for CUDA CustomStorage artifact prefetch: %w", ctx.Err()) + } +} + func cleanupRestoreMounts(ctx context.Context, mounts []restoreMount) error { var cleanupErr error cleanupCtx := context.WithoutCancel(ctx) @@ -73,8 +101,16 @@ type RestoreRequest struct { TargetPodIP string ContainerName string Clientset kubernetes.Interface + CUDATransfer types.CUDATransferSettings } +var ( + waitForCUDAStorageMode = cuda.WaitForDaemon + readRestoredHostProcessTable = snapshotruntime.ReadProcessTable + validateRestoredProcessIdentity = snapshotruntime.ValidateProcessIdentity + restoreAndUnlockCUDAProcessTree = cuda.RestoreAndUnlockProcessTreeValidated +) + // Restore performs external restore for the given request. // Returns the namespace-relative PID of the restored process. // The DaemonSet side inspects the placeholder and launches nsrestore, @@ -148,13 +184,71 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r point: artifactMount, }) + var prefetch <-chan customStoragePrefetchOutcome + var cancelPrefetch context.CancelFunc + if snap.CUDAStorageMode == types.CUDAStorageModePOSIX { + prefetchCtx, cancel := context.WithCancel(ctx) + cancelPrefetch = cancel + outcomes := make(chan customStoragePrefetchOutcome, 1) + prefetch = outcomes + go func() { + result, err := cuda.PrefetchCustomStorageArtifacts(prefetchCtx, artifactPath) + outcomes <- customStoragePrefetchOutcome{result: result, err: err} + }() + } + awaitPrefetch := func(cancel bool) (cuda.CustomStoragePrefetchResult, error) { + if prefetch == nil { + return cuda.CustomStoragePrefetchResult{}, nil + } + result, err := waitForCustomStoragePrefetch(ctx, prefetch, cancelPrefetch, cancel) + prefetch = nil + cancelPrefetch = nil + return result, err + } + defer func() { + if prefetch != nil { + _, _ = awaitPrefetch(true) + } + }() + + // NodeController.failRestore owns placeholder-wide termination for every + // non-cleanup error returned after execution begins. Keeping cleanup in the + // controller guarantees that RestoreFailed is not persisted until the + // runtime-owned placeholder has actually been resolved and terminated. result, err := execNSRestore(ctx, log, req, snap, bundleMount, nsmount.CheckpointDst) if err != nil { + _, _ = awaitPrefetch(true) return 0, fmt.Errorf("nsrestore failed: %w", err) } + prefetchResult, err := awaitPrefetch(false) + if err != nil { + if ctx.Err() != nil { + return 0, fmt.Errorf("CUDA CustomStorage artifact prefetch interrupted after CRIU restore: %w", err) + } + // Prefetch only overlaps durable-storage reads with CRIU. The CUDA + // helper performs the authoritative read and validation, so an + // optimization failure must not strand a process after CRIU restore. + log.Error(err, "CUDA CustomStorage artifact prefetch failed; continuing with authoritative restore") + } else if prefetchResult.Files > 0 { + log.Info("CUDA CustomStorage artifact prefetch completed", + "files", prefetchResult.Files, + "bytes", prefetchResult.Bytes, + "service_duration", prefetchResult.Duration, + "overlapped_with_criu", true, + ) + } if result.CleanupError != nil { cleanupErr = errors.Join(cleanupErr, result.CleanupError) } + if len(result.DeferredCUDAProcesses) > 0 { + cudaTimings, err := restoreDeferredCUDAProcesses( + ctx, result.DeferredCUDAProcesses, snap, artifactPath, req.CUDATransfer, log, + ) + if err != nil { + return 0, err + } + result.CUDARestoreDuration += cudaTimings + } if err := validateRestoredProcess(snap.TargetRoot, result.RestoredPID, log); err != nil { return 0, err } @@ -191,6 +285,54 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r return snap.PlaceholderPID, nil } +func restoreDeferredCUDAProcesses( + ctx context.Context, + namespaceProcesses []snapshotruntime.ProcessDetails, + snap *types.RestoreContainerSnapshot, + artifactPath string, + transferSettings types.CUDATransferSettings, + log logr.Logger, +) (time.Duration, error) { + processTable, err := readRestoredHostProcessTable(snapshotruntime.HostProcPath) + if err != nil { + return 0, fmt.Errorf("snapshot restored host process table: %w", err) + } + hostProcesses := make([]snapshotruntime.ProcessDetails, 0, len(namespaceProcesses)) + for _, namespaceProcess := range namespaceProcesses { + process, err := snapshotruntime.ResolveHostProcessIdentityFromTable(processTable, namespaceProcess) + if err != nil { + return 0, fmt.Errorf("resolve restored CUDA host process identity: %w", err) + } + if err := validateRestoredProcessIdentity(snapshotruntime.HostProcPath, process); err != nil { + return 0, fmt.Errorf("validate restored CUDA process identity: %w", err) + } + hostProcesses = append(hostProcesses, process) + } + cudaJobFile := "" + if stagedJobFile, err := cuda.JobFileFromCheckpoint(artifactPath); err != nil { + return 0, err + } else if stagedJobFile != "" { + // The CUDA layer uses this only as a presence signal and derives a + // host-visible path from each identity-validated target PID. + cudaJobFile = stagedJobFile + } + cudaTimings, err := restoreAndUnlockCUDAProcessTree( + ctx, + hostProcesses, + snap.CUDADeviceMap, + snap.CUDAStorageMode, + artifactPath, + cudaJobFile, + snap.TargetGPUUUIDs, + transferSettings, + log, + ) + if err != nil { + return 0, fmt.Errorf("host CUDA restore failed: %w", err) + } + return cudaTimings.TotalDuration, nil +} + func remainingDuration(wall time.Duration, parts ...time.Duration) time.Duration { var sum time.Duration for _, part := range parts { @@ -222,6 +364,17 @@ func validateRestoreManifest(req RestoreRequest, manifest *types.CheckpointManif req.ContainerName, ) } + mode, err := manifest.CUDA.EffectiveStorageMode() + if err != nil { + return fmt.Errorf("validate CUDA artifact storage mode: %w", err) + } + if mode == types.CUDAStorageModePOSIX { + if err := validatePOSIXCustomStorageTopology( + len(manifest.CUDA.PIDs), len(manifest.CUDA.SourceGPUUUIDs), + ); err != nil { + return err + } + } return nil } @@ -232,6 +385,17 @@ func inspectRestore( req RestoreRequest, manifest *types.CheckpointManifest, ) (*types.RestoreContainerSnapshot, time.Duration, error) { + cudaStorageMode := types.CUDAStorageModeLegacy + if !manifest.CUDA.IsEmpty() { + var err error + cudaStorageMode, err = manifest.CUDA.EffectiveStorageMode() + if err != nil { + return nil, 0, fmt.Errorf("invalid CUDA artifact metadata: %w", err) + } + if err := waitForCUDAStorageMode(ctx, cudaStorageMode); err != nil { + return nil, 0, fmt.Errorf("CUDA storage mode %q is unavailable before restore: %w", cudaStorageMode, err) + } + } var ( placeholderPID int err error @@ -253,13 +417,14 @@ func inspectRestore( } cudaDeviceMap := "" + var targetGPUUUIDs []string var gpuDeviceMapDuration time.Duration if !manifest.CUDA.IsEmpty() { if len(manifest.CUDA.SourceGPUUUIDs) == 0 { return nil, 0, fmt.Errorf("missing source GPU UUIDs in checkpoint manifest") } gpuStart := time.Now() - targetGPUUUIDs, err := cuda.DiscoverGPUUUIDs( + targetGPUUUIDs, err = cuda.DiscoverGPUUUIDs( ctx, req.Clientset, req.PodName, @@ -288,10 +453,12 @@ func inspectRestore( } return &types.RestoreContainerSnapshot{ - PlaceholderPID: placeholderPID, - TargetRoot: fmt.Sprintf("%s/%d/root", snapshotruntime.HostProcPath, placeholderPID), - CgroupRoot: cgroupRoot, - CUDADeviceMap: cudaDeviceMap, + PlaceholderPID: placeholderPID, + TargetRoot: fmt.Sprintf("%s/%d/root", snapshotruntime.HostProcPath, placeholderPID), + CgroupRoot: cgroupRoot, + CUDADeviceMap: cudaDeviceMap, + TargetGPUUUIDs: append([]string(nil), targetGPUUUIDs...), + CUDAStorageMode: cudaStorageMode, }, gpuDeviceMapDuration, nil } diff --git a/agent/internal/executor/restore_test.go b/agent/internal/executor/restore_test.go index 811f9d13..fe3105cc 100644 --- a/agent/internal/executor/restore_test.go +++ b/agent/internal/executor/restore_test.go @@ -12,10 +12,13 @@ import ( "testing" "time" + "github.com/go-logr/logr" "github.com/go-logr/logr/testr" specs "github.com/opencontainers/runtime-spec/specs-go" + "github.com/ai-dynamo/snapshot/agent/internal/cuda" "github.com/ai-dynamo/snapshot/agent/internal/nsmount" + snapshotruntime "github.com/ai-dynamo/snapshot/agent/internal/runtime" "github.com/ai-dynamo/snapshot/agent/internal/types" ) @@ -129,6 +132,19 @@ func TestValidateRestoreManifest(t *testing.T) { } }) } + + manifest.CUDA = types.NewCUDAManifest( + []int{42, 43}, + []string{"GPU-aaa", "GPU-bbb"}, + types.CUDAStorageModePOSIX, + ) + err := validateRestoreManifest( + RestoreRequest{ContentUID: "content-uid-123", ContainerName: "main", PodNamespace: "team-a"}, + manifest, + ) + if err == nil || !strings.Contains(err.Error(), "qualified only for one or more CUDA processes on one GPU") { + t.Fatalf("validateRestoreManifest(multi-GPU POSIX) = %v, want qualification error", err) + } } func TestRestoreInNamespaceRejectsMultiGPUCheckpointWithoutLaunchJobState(t *testing.T) { @@ -140,7 +156,7 @@ func TestRestoreInNamespaceRejectsMultiGPUCheckpointWithoutLaunchJobState(t *tes types.NewSourcePodManifest("source-id", 456, "node-1", "source-pod", "default", "10.0.0.11", nil), types.OverlayManifest{}, ) - manifest.CUDA = types.NewCUDAManifest([]int{42, 43}, []string{"GPU-aaa", "GPU-bbb"}) + manifest.CUDA = types.NewCUDAManifest([]int{42, 43}, []string{"GPU-aaa", "GPU-bbb"}, types.CUDAStorageModeLegacy) if err := types.WriteManifest(checkpointDir, manifest); err != nil { t.Fatalf("WriteManifest: %v", err) } @@ -160,3 +176,128 @@ func TestRemainingDuration(t *testing.T) { t.Fatal("remainingDuration should not go negative") } } + +func TestWaitForCustomStoragePrefetchDiscardDoesNotWait(t *testing.T) { + outcomes := make(chan customStoragePrefetchOutcome, 1) + ctx, cancel := context.WithCancel(context.Background()) + + done := make(chan struct{}) + go func() { + defer close(done) + if _, err := waitForCustomStoragePrefetch(ctx, outcomes, cancel, true); err != nil { + t.Errorf("waitForCustomStoragePrefetch(discard=true): %v", err) + } + }() + + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("discarded prefetch waited for an outcome") + } + select { + case <-ctx.Done(): + case <-time.After(time.Second): + t.Fatal("discarded prefetch did not cancel its context") + } +} + +func TestWaitForCustomStoragePrefetchHonorsContextCancellation(t *testing.T) { + outcomes := make(chan customStoragePrefetchOutcome, 1) + ctx, cancelContext := context.WithCancel(context.Background()) + cancelContext() + prefetchCtx, cancelPrefetch := context.WithCancel(context.Background()) + + _, err := waitForCustomStoragePrefetch(ctx, outcomes, cancelPrefetch, false) + if !errors.Is(err, context.Canceled) { + t.Fatalf("waitForCustomStoragePrefetch() error = %v, want context.Canceled", err) + } + select { + case <-prefetchCtx.Done(): + case <-time.After(time.Second): + t.Fatal("canceled wait did not cancel prefetch") + } +} + +func TestRestoreDeferredCUDAProcessesResolvesAndValidatesHostIdentity(t *testing.T) { + namespaceProcess := snapshotruntime.ProcessDetails{ + InnermostPID: 7, + StartTimeTicks: 101, + Cgroup: "0::/restored\n", + } + hostProcess := snapshotruntime.ProcessDetails{ + ObservedPID: 9007, + OutermostPID: 9007, + InnermostPID: 7, + NamespacePIDs: []int{9007, 7}, + StartTimeTicks: 101, + Cgroup: "0::/restored\n", + } + + originalRead := readRestoredHostProcessTable + originalValidate := validateRestoredProcessIdentity + originalRestore := restoreAndUnlockCUDAProcessTree + t.Cleanup(func() { + readRestoredHostProcessTable = originalRead + validateRestoredProcessIdentity = originalValidate + restoreAndUnlockCUDAProcessTree = originalRestore + }) + + readRestoredHostProcessTable = func(procRoot string) ([]snapshotruntime.ProcessDetails, error) { + if procRoot != snapshotruntime.HostProcPath { + t.Fatalf("ReadProcessTable proc root = %q, want %q", procRoot, snapshotruntime.HostProcPath) + } + return []snapshotruntime.ProcessDetails{hostProcess}, nil + } + validated := false + validateRestoredProcessIdentity = func(procRoot string, process snapshotruntime.ProcessDetails) error { + if procRoot != snapshotruntime.HostProcPath || process.OutermostPID != hostProcess.OutermostPID { + t.Fatalf("ValidateProcessIdentity(%q, PID %d), want (%q, PID %d)", + procRoot, process.OutermostPID, snapshotruntime.HostProcPath, hostProcess.OutermostPID) + } + validated = true + return nil + } + restored := false + restoreAndUnlockCUDAProcessTree = func( + ctx context.Context, + processes []snapshotruntime.ProcessDetails, + deviceMap, storageMode, checkpointDir, jobFile string, + targetGPUUUIDs []string, + transferSettings types.CUDATransferSettings, + _ logr.Logger, + ) (cuda.RestorePhaseTimings, error) { + if len(processes) != 1 || processes[0].OutermostPID != hostProcess.OutermostPID { + t.Fatalf("CUDA restore processes = %+v, want host PID %d", processes, hostProcess.OutermostPID) + } + if deviceMap != "0=1" || storageMode != types.CUDAStorageModePOSIX || checkpointDir == "" { + t.Fatalf("CUDA restore args = device map %q, storage mode %q, checkpoint dir %q", deviceMap, storageMode, checkpointDir) + } + if len(targetGPUUUIDs) != 1 || targetGPUUUIDs[0] != "GPU-target" { + t.Fatalf("target GPU UUIDs = %v", targetGPUUUIDs) + } + restored = true + return cuda.RestorePhaseTimings{TotalDuration: 250 * time.Millisecond}, nil + } + + duration, err := restoreDeferredCUDAProcesses( + context.Background(), + []snapshotruntime.ProcessDetails{namespaceProcess}, + &types.RestoreContainerSnapshot{ + CUDADeviceMap: "0=1", + CUDAStorageMode: types.CUDAStorageModePOSIX, + TargetGPUUUIDs: []string{"GPU-target"}, + }, + t.TempDir(), + types.CUDATransferSettings{}, + testr.New(t), + ) + if err != nil { + t.Fatalf("restoreDeferredCUDAProcesses: %v", err) + } + if !validated || !restored { + t.Fatalf("validated = %t, restored = %t; want both true", validated, restored) + } + if duration != 250*time.Millisecond { + t.Fatalf("duration = %s, want 250ms", duration) + } +} diff --git a/agent/internal/runtime/process.go b/agent/internal/runtime/process.go index ad7fb7ac..0103faf1 100644 --- a/agent/internal/runtime/process.go +++ b/agent/internal/runtime/process.go @@ -22,12 +22,53 @@ const HostProcPath = "/host/proc" // ProcessDetails captures the parent link plus the observed, outermost, and innermost // PID views for one proc entry. ObservedPID is relative to the proc root being read. type ProcessDetails struct { - ObservedPID int - ParentPID int - OutermostPID int - InnermostPID int - NamespacePIDs []int - Cmdline string + ObservedPID int + ParentPID int + OutermostPID int + InnermostPID int + NamespacePIDs []int + Cmdline string + StartTimeTicks uint64 + Cgroup string +} + +// ResolveHostProcessIdentity maps an identity captured through a container's +// proc mount to the unique host /proc entry with the same namespace PID, +// process start time, and cgroup. +func ResolveHostProcessIdentity(procRoot string, expected ProcessDetails) (ProcessDetails, error) { + if expected.InnermostPID <= 0 || expected.StartTimeTicks == 0 || expected.Cgroup == "" { + return ProcessDetails{}, fmt.Errorf("incomplete restored process identity") + } + processes, err := ReadProcessTable(procRoot) + if err != nil { + return ProcessDetails{}, err + } + return ResolveHostProcessIdentityFromTable(processes, expected) +} + +// ResolveHostProcessIdentityFromTable maps a restored identity against one +// host-process snapshot. Callers resolving multiple targets should reuse the +// same table and then revalidate each live PID before a destructive operation. +func ResolveHostProcessIdentityFromTable(processes []ProcessDetails, expected ProcessDetails) (ProcessDetails, error) { + if expected.InnermostPID <= 0 || expected.StartTimeTicks == 0 || expected.Cgroup == "" { + return ProcessDetails{}, fmt.Errorf("incomplete restored process identity") + } + var match ProcessDetails + for _, process := range processes { + if process.InnermostPID != expected.InnermostPID || + process.StartTimeTicks != expected.StartTimeTicks || + process.Cgroup != expected.Cgroup { + continue + } + if match.OutermostPID != 0 { + return ProcessDetails{}, fmt.Errorf("restored process identity for namespace PID %d is not unique", expected.InnermostPID) + } + match = process + } + if match.OutermostPID == 0 { + return ProcessDetails{}, fmt.Errorf("restored process identity for namespace PID %d not found in host proc", expected.InnermostPID) + } + return match, nil } // ReadProcessFilesystemIDs returns the filesystem UID and GID from a proc status entry. @@ -141,17 +182,66 @@ func ReadProcessDetails(procRoot string, pid int) (ProcessDetails, error) { cmdline = strings.TrimSpace(string(comm)) } } + statPath := filepath.Join(procRoot, strconv.Itoa(pid), "stat") + statBytes, err := os.ReadFile(statPath) + if err != nil { + return ProcessDetails{}, fmt.Errorf("failed to read %s: %w", statPath, err) + } + startTimeTicks, err := ParseProcStartTime(string(statBytes)) + if err != nil { + return ProcessDetails{}, fmt.Errorf("failed to parse process start time: %w", err) + } + cgroupPath := filepath.Join(procRoot, strconv.Itoa(pid), "cgroup") + cgroupBytes, err := os.ReadFile(cgroupPath) + if err != nil { + return ProcessDetails{}, fmt.Errorf("failed to read %s: %w", cgroupPath, err) + } + cgroup := string(cgroupBytes) return ProcessDetails{ - ObservedPID: pid, - ParentPID: parentPID, - OutermostPID: nspids[0], - InnermostPID: nspids[len(nspids)-1], - NamespacePIDs: nspids, - Cmdline: cmdline, + ObservedPID: pid, + ParentPID: parentPID, + OutermostPID: nspids[0], + InnermostPID: nspids[len(nspids)-1], + NamespacePIDs: nspids, + Cmdline: cmdline, + StartTimeTicks: startTimeTicks, + Cgroup: cgroup, }, nil } +// ParseProcStartTime extracts field 22 from /proc//stat. +func ParseProcStartTime(statLine string) (uint64, error) { + statLine = strings.TrimSpace(statLine) + paren := strings.LastIndex(statLine, ")") + if paren < 0 || paren+2 > len(statLine) { + return 0, fmt.Errorf("malformed stat line") + } + fields := strings.Fields(statLine[paren+2:]) + if len(fields) < 20 { + return 0, fmt.Errorf("malformed stat fields") + } + return strconv.ParseUint(fields[19], 10, 64) +} + +// ValidateProcessIdentity rejects PID reuse or namespace/cgroup changes. +func ValidateProcessIdentity(procRoot string, expected ProcessDetails) error { + if expected.StartTimeTicks == 0 || expected.Cgroup == "" { + return fmt.Errorf("incomplete process identity for host PID %d", expected.OutermostPID) + } + current, err := ReadProcessDetails(procRoot, expected.OutermostPID) + if err != nil { + return err + } + if current.OutermostPID != expected.OutermostPID || + current.InnermostPID != expected.InnermostPID || + current.StartTimeTicks != expected.StartTimeTicks || + current.Cgroup != expected.Cgroup { + return fmt.Errorf("process identity changed for host PID %d", expected.OutermostPID) + } + return nil +} + // ReadProcessDetailsOrDefault preserves pid-scoped logging even when proc parsing fails. func ReadProcessDetailsOrDefault(procRoot string, pid int) ProcessDetails { details := ProcessDetails{ diff --git a/agent/internal/runtime/process_test.go b/agent/internal/runtime/process_test.go index 259478b3..a3161b8d 100644 --- a/agent/internal/runtime/process_test.go +++ b/agent/internal/runtime/process_test.go @@ -7,6 +7,7 @@ import ( "os" "path/filepath" "strconv" + "strings" "testing" ) @@ -99,6 +100,79 @@ func TestParseProcExitCode(t *testing.T) { } } +func TestResolveHostProcessIdentity(t *testing.T) { + procRoot := t.TempDir() + procDir := filepath.Join(procRoot, "900") + if err := os.MkdirAll(procDir, 0o755); err != nil { + t.Fatal(err) + } + files := map[string]string{ + "status": "Name:\tworker\nPPid:\t1\nNSpid:\t900 77\n", + "stat": "900 (worker with spaces) S 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 424242 20\n", + "cgroup": "0::/kubepods/pod/container\n", + "comm": "worker\n", + } + for name, data := range files { + if err := os.WriteFile(filepath.Join(procDir, name), []byte(data), 0o644); err != nil { + t.Fatal(err) + } + } + + got, err := ResolveHostProcessIdentity(procRoot, ProcessDetails{ + InnermostPID: 77, + StartTimeTicks: 424242, + Cgroup: "0::/kubepods/pod/container\n", + }) + if err != nil { + t.Fatalf("ResolveHostProcessIdentity() error = %v", err) + } + if got.OutermostPID != 900 { + t.Fatalf("OutermostPID = %d, want 900", got.OutermostPID) + } + if err := ValidateProcessIdentity(procRoot, got); err != nil { + t.Fatalf("ValidateProcessIdentity() error = %v", err) + } + changedStart := got + changedStart.StartTimeTicks++ + if err := ValidateProcessIdentity(procRoot, changedStart); err == nil { + t.Fatal("ValidateProcessIdentity() accepted changed start time") + } + changedCgroup := got + changedCgroup.Cgroup = "0::/different\n" + if err := ValidateProcessIdentity(procRoot, changedCgroup); err == nil { + t.Fatal("ValidateProcessIdentity() accepted changed cgroup") + } + + duplicateDir := filepath.Join(procRoot, "901") + if err := os.MkdirAll(duplicateDir, 0o755); err != nil { + t.Fatal(err) + } + for name, data := range files { + if name == "status" { + data = "Name:\tworker\nPPid:\t1\nNSpid:\t901 77\n" + } else if name == "stat" { + data = strings.Replace(data, "900 (", "901 (", 1) + } + if err := os.WriteFile(filepath.Join(duplicateDir, name), []byte(data), 0o644); err != nil { + t.Fatal(err) + } + } + _, err = ResolveHostProcessIdentity(procRoot, ProcessDetails{ + InnermostPID: 77, StartTimeTicks: 424242, + Cgroup: "0::/kubepods/pod/container\n", + }) + if err == nil || !strings.Contains(err.Error(), "not unique") { + t.Fatalf("expected non-unique identity error, got %v", err) + } + _, err = ResolveHostProcessIdentity(procRoot, ProcessDetails{ + InnermostPID: 999, StartTimeTicks: 424242, + Cgroup: "0::/kubepods/pod/container\n", + }) + if err == nil || !strings.Contains(err.Error(), "not found") { + t.Fatalf("expected not-found identity error, got %v", err) + } +} + func TestReadProcessDetails(t *testing.T) { procRoot := t.TempDir() pid := 1018 @@ -112,6 +186,12 @@ func TestReadProcessDetails(t *testing.T) { if err := os.WriteFile(filepath.Join(procDir, "cmdline"), []byte("python3\x00-m\x00dynamo.vllm\x00"), 0644); err != nil { t.Fatalf("WriteFile(cmdline): %v", err) } + if err := os.WriteFile(filepath.Join(procDir, "stat"), []byte("2402711 (python3) S 0 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 424242 20\n"), 0644); err != nil { + t.Fatalf("WriteFile(stat): %v", err) + } + if err := os.WriteFile(filepath.Join(procDir, "cgroup"), []byte("0::/kubepods/test\n"), 0644); err != nil { + t.Fatalf("WriteFile(cgroup): %v", err) + } details, err := ReadProcessDetails(procRoot, pid) if err != nil { @@ -167,6 +247,13 @@ func TestReadProcessTable(t *testing.T) { if err := os.WriteFile(filepath.Join(procDir, "cmdline"), []byte(cmdline), 0644); err != nil { t.Fatalf("WriteFile(cmdline): %v", err) } + stat := strconv.Itoa(pid) + " (process) S 0 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 " + strconv.Itoa(pid+1000) + " 20\n" + if err := os.WriteFile(filepath.Join(procDir, "stat"), []byte(stat), 0644); err != nil { + t.Fatalf("WriteFile(stat): %v", err) + } + if err := os.WriteFile(filepath.Join(procDir, "cgroup"), []byte("0::/kubepods/test\n"), 0644); err != nil { + t.Fatalf("WriteFile(cgroup): %v", err) + } } writeEntry(768, "Name:\tworker\nPPid:\t1\nNSpid:\t2444000 768\n", "VLLM::Worker_TP0\x00") diff --git a/agent/internal/types/config.go b/agent/internal/types/config.go index 2d3a29ea..e3fdb578 100644 --- a/agent/internal/types/config.go +++ b/agent/internal/types/config.go @@ -18,11 +18,87 @@ const CheckpointBasePath = "/checkpoints" // AgentConfig holds the full agent configuration: static checkpoint settings // from the ConfigMap YAML, plus runtime fields from environment variables. type AgentConfig struct { - NodeName string `yaml:"-"` - Storage StorageSpec `yaml:"storage"` - Overlay OverlaySettings `yaml:"overlay"` - Restore RestoreSpec `yaml:"restore"` - CRIU CRIUSettings `yaml:"criu"` + NodeName string `yaml:"-"` + Storage StorageSpec `yaml:"storage"` + CUDACheckpoint CUDACheckpointSettings `yaml:"cudaCheckpoint"` + Overlay OverlaySettings `yaml:"overlay"` + Restore RestoreSpec `yaml:"restore"` + CRIU CRIUSettings `yaml:"criu"` +} + +// CUDACheckpointSettings holds CUDA CustomStorage transfer settings. +type CUDACheckpointSettings struct { + StorageMode string `yaml:"storageMode"` + TransferBufferCount *int `yaml:"transferBufferCount"` + TransferChunkBytes *uint64 `yaml:"transferChunkBytes"` +} + +// CUDATransferSettings is the validated, concrete transfer configuration used +// by the CUDA helper daemon. +type CUDATransferSettings struct { + BufferCount int + ChunkBytes uint64 +} + +const ( + // CUDAStorageModeLegacy uses the CUDA driver's original in-memory storage. + CUDAStorageModeLegacy = "legacy" + // CUDAStorageModePOSIX writes CUDA CustomStorage extents into checkpoint files. + CUDAStorageModePOSIX = "posix" + + DefaultCUDATransferBufferCount = 1 + DefaultCUDATransferChunkBytes = 64 * 1024 * 1024 + maxCUDATransferBufferCount = 8 + minCUDATransferChunkBytes = 1 * 1024 * 1024 + maxCUDATransferChunkBytes = 256 * 1024 * 1024 + maxCUDAPinnedBytesPerDevice = 1 * 1024 * 1024 * 1024 + cudaTransferBufferAlignment = 4096 + + CUDAHelperSocketDirectory = "/run/cuda-checkpoint-helper" + CUDAHelperSocketPath = CUDAHelperSocketDirectory + "/helper.sock" +) + +func (c CUDACheckpointSettings) TransferSettings() CUDATransferSettings { + settings := CUDATransferSettings{ + BufferCount: DefaultCUDATransferBufferCount, + ChunkBytes: DefaultCUDATransferChunkBytes, + } + if c.TransferBufferCount != nil { + settings.BufferCount = *c.TransferBufferCount + } + if c.TransferChunkBytes != nil { + settings.ChunkBytes = *c.TransferChunkBytes + } + return settings +} + +func (c CUDATransferSettings) WithDefaults() CUDATransferSettings { + settings := c + if settings.BufferCount == 0 { + settings.BufferCount = DefaultCUDATransferBufferCount + } + if settings.ChunkBytes == 0 { + settings.ChunkBytes = DefaultCUDATransferChunkBytes + } + return settings +} + +func (c CUDATransferSettings) Validate() error { + if c.BufferCount < 1 || c.BufferCount > maxCUDATransferBufferCount { + return fmt.Errorf("buffer count must be between 1 and %d", maxCUDATransferBufferCount) + } + if c.ChunkBytes < minCUDATransferChunkBytes || c.ChunkBytes > maxCUDATransferChunkBytes || c.ChunkBytes%cudaTransferBufferAlignment != 0 { + return fmt.Errorf( + "chunk bytes must be a %d-byte multiple between %d and %d", + cudaTransferBufferAlignment, + minCUDATransferChunkBytes, + maxCUDATransferChunkBytes, + ) + } + if uint64(c.BufferCount) > maxCUDAPinnedBytesPerDevice/c.ChunkBytes { + return fmt.Errorf("buffers exceed the 1 GiB per-device pinned-memory limit") + } + return nil } func (c *AgentConfig) LoadEnvOverrides() { @@ -58,6 +134,35 @@ func (c *AgentConfig) Validate() error { Message: fmt.Sprintf("unsupported imageIoMode %q; expected %q, %q, or empty", c.CRIU.ImageIoMode, "writeback", "direct"), } } + if c.CUDACheckpoint.TransferBufferCount == nil { + value := DefaultCUDATransferBufferCount + c.CUDACheckpoint.TransferBufferCount = &value + } + if c.CUDACheckpoint.TransferChunkBytes == nil { + value := uint64(DefaultCUDATransferChunkBytes) + c.CUDACheckpoint.TransferChunkBytes = &value + } + if err := c.CUDACheckpoint.TransferSettings().Validate(); err != nil { + return &ConfigError{Field: "cudaCheckpoint", Message: err.Error()} + } + storageMode := strings.ToLower(strings.TrimSpace(c.CUDACheckpoint.StorageMode)) + if storageMode == "" { + storageMode = CUDAStorageModeLegacy + } + switch storageMode { + case CUDAStorageModeLegacy, CUDAStorageModePOSIX: + c.CUDACheckpoint.StorageMode = storageMode + default: + return &ConfigError{ + Field: "cudaCheckpoint.storageMode", + Message: fmt.Sprintf( + "unsupported CUDA storage mode %q; expected %q or %q", + c.CUDACheckpoint.StorageMode, + CUDAStorageModeLegacy, + CUDAStorageModePOSIX, + ), + } + } return c.Restore.Validate() } diff --git a/agent/internal/types/config_test.go b/agent/internal/types/config_test.go index 27e1af81..1ac9e8e5 100644 --- a/agent/internal/types/config_test.go +++ b/agent/internal/types/config_test.go @@ -26,3 +26,100 @@ func TestAgentConfigValidateRequiresFixedStorageBasePath(t *testing.T) { } } } + +func TestCUDATransferSettingsWithDefaults(t *testing.T) { + got := (CUDATransferSettings{}).WithDefaults() + if got.BufferCount != DefaultCUDATransferBufferCount || got.ChunkBytes != DefaultCUDATransferChunkBytes { + t.Fatalf("WithDefaults() = %+v, want 1 slot and 64 MiB", got) + } +} + +func TestAgentConfigValidateCUDATransferSettings(t *testing.T) { + cfg := validAgentConfig() + bufferCount := 4 + chunkBytes := uint64(32 * 1024 * 1024) + cfg.CUDACheckpoint.TransferBufferCount = &bufferCount + cfg.CUDACheckpoint.TransferChunkBytes = &chunkBytes + if err := cfg.Validate(); err != nil { + t.Fatalf("Validate() error = %v", err) + } + if got := cfg.CUDACheckpoint.TransferSettings(); got.BufferCount != bufferCount || got.ChunkBytes != chunkBytes { + t.Fatalf("CUDA transfer settings = %+v, want count=%d chunk=%d", got, bufferCount, chunkBytes) + } + + tooManyBuffers := 8 + tooLargeChunk := uint64(256 * 1024 * 1024) + cfg.CUDACheckpoint.TransferBufferCount = &tooManyBuffers + cfg.CUDACheckpoint.TransferChunkBytes = &tooLargeChunk + if err := cfg.Validate(); err == nil { + t.Fatal("expected excessive per-device pinned memory to be rejected") + } +} + +func TestAgentConfigValidateDefaultsCUDATransferSettings(t *testing.T) { + cfg := validAgentConfig() + if err := cfg.Validate(); err != nil { + t.Fatalf("Validate() error = %v", err) + } + if cfg.CUDACheckpoint.TransferBufferCount == nil || cfg.CUDACheckpoint.TransferChunkBytes == nil { + t.Fatal("Validate() did not populate unset CUDA transfer fields") + } + settings := cfg.CUDACheckpoint.TransferSettings() + if settings.BufferCount != DefaultCUDATransferBufferCount || settings.ChunkBytes != DefaultCUDATransferChunkBytes { + t.Fatalf("CUDA transfer settings = %+v, want defaults", settings) + } + if cfg.CUDACheckpoint.StorageMode != CUDAStorageModeLegacy { + t.Fatalf("CUDA storage mode = %q, want default %q", cfg.CUDACheckpoint.StorageMode, CUDAStorageModeLegacy) + } +} + +func TestAgentConfigValidateCUDAStorageMode(t *testing.T) { + for _, test := range []struct { + name string + mode string + want string + wantError bool + }{ + {name: "unset defaults legacy", want: CUDAStorageModeLegacy}, + {name: "legacy", mode: CUDAStorageModeLegacy, want: CUDAStorageModeLegacy}, + {name: "posix", mode: CUDAStorageModePOSIX, want: CUDAStorageModePOSIX}, + {name: "normalizes", mode: " POSIX ", want: CUDAStorageModePOSIX}, + {name: "rejects auto", mode: "auto", wantError: true}, + } { + t.Run(test.name, func(t *testing.T) { + cfg := validAgentConfig() + cfg.CUDACheckpoint.StorageMode = test.mode + err := cfg.Validate() + if test.wantError { + if err == nil { + t.Fatal("Validate() accepted unsupported CUDA storage mode") + } + return + } + if err != nil { + t.Fatalf("Validate() error = %v", err) + } + if cfg.CUDACheckpoint.StorageMode != test.want { + t.Fatalf("CUDA storage mode = %q, want %q", cfg.CUDACheckpoint.StorageMode, test.want) + } + }) + } +} + +func TestCUDATransferSettingsValidateRejectsBadChunkBytes(t *testing.T) { + tests := []struct { + name string + chunk uint64 + }{ + {name: "below minimum", chunk: minCUDATransferChunkBytes - cudaTransferBufferAlignment}, + {name: "misaligned", chunk: minCUDATransferChunkBytes + 1}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + err := (CUDATransferSettings{BufferCount: 1, ChunkBytes: test.chunk}).Validate() + if err == nil { + t.Fatalf("Validate accepted chunk size %d", test.chunk) + } + }) + } +} diff --git a/agent/internal/types/inspect.go b/agent/internal/types/inspect.go index 740f3001..3473c59b 100644 --- a/agent/internal/types/inspect.go +++ b/agent/internal/types/inspect.go @@ -35,8 +35,10 @@ type CheckpointContainerSnapshot struct { // RestoreContainerSnapshot holds inspected state for the restore target. type RestoreContainerSnapshot struct { - PlaceholderPID int - TargetRoot string - CgroupRoot string - CUDADeviceMap string + PlaceholderPID int + TargetRoot string + CgroupRoot string + CUDADeviceMap string + TargetGPUUUIDs []string + CUDAStorageMode string } diff --git a/agent/internal/types/manifest.go b/agent/internal/types/manifest.go index 75e35606..6defa8c9 100644 --- a/agent/internal/types/manifest.go +++ b/agent/internal/types/manifest.go @@ -142,12 +142,14 @@ func NewOverlayManifest(exclusions OverlaySettings, upperDir string, ociSpec *sp type CUDAManifest struct { PIDs []int `yaml:"pids"` SourceGPUUUIDs []string `yaml:"sourceGpuUuids"` + StorageMode string `yaml:"storageMode,omitempty"` } -func NewCUDAManifest(pids []int, sourceGPUUUIDs []string) CUDAManifest { +func NewCUDAManifest(pids []int, sourceGPUUUIDs []string, storageMode string) CUDAManifest { return CUDAManifest{ PIDs: append([]int(nil), pids...), SourceGPUUUIDs: append([]string(nil), sourceGPUUUIDs...), + StorageMode: storageMode, } } @@ -155,6 +157,21 @@ func (m CUDAManifest) IsEmpty() bool { return len(m.PIDs) == 0 } +// EffectiveStorageMode preserves compatibility with manifests written before +// CustomStorage was introduced. +func (m CUDAManifest) EffectiveStorageMode() (string, error) { + mode := strings.ToLower(strings.TrimSpace(m.StorageMode)) + if mode == "" { + return CUDAStorageModeLegacy, nil + } + switch mode { + case CUDAStorageModeLegacy, CUDAStorageModePOSIX: + return mode, nil + default: + return "", fmt.Errorf("unsupported CUDA artifact storage mode %q", m.StorageMode) + } +} + // WriteManifest writes a checkpoint manifest file in the checkpoint directory. func WriteManifest(checkpointDir string, data *CheckpointManifest) error { if data == nil { @@ -163,6 +180,14 @@ func WriteManifest(checkpointDir string, data *CheckpointManifest) error { if err := validateArtifactManifest(data.Artifact); err != nil { return err } + if !data.CUDA.IsEmpty() && strings.TrimSpace(data.CUDA.StorageMode) == "" { + return fmt.Errorf("checkpoint manifest CUDA section is missing storageMode") + } + if !data.CUDA.IsEmpty() { + if _, err := data.CUDA.EffectiveStorageMode(); err != nil { + return err + } + } content, err := yaml.Marshal(data) if err != nil { diff --git a/agent/internal/types/manifest_test.go b/agent/internal/types/manifest_test.go index ac46a85b..6126a765 100644 --- a/agent/internal/types/manifest_test.go +++ b/agent/internal/types/manifest_test.go @@ -6,6 +6,7 @@ package types import ( "os" "path/filepath" + "strings" "testing" criurpc "github.com/checkpoint-restore/go-criu/v8/rpc" @@ -36,7 +37,7 @@ func TestManifestRoundTrip(t *testing.T) { BindMountDests: []string{"/data"}, }, ) - original.CUDA = NewCUDAManifest([]int{42, 43}, []string{"GPU-aaa", "GPU-bbb"}) + original.CUDA = NewCUDAManifest([]int{42, 43}, []string{"GPU-aaa", "GPU-bbb"}, CUDAStorageModePOSIX) if err := WriteManifest(dir, original); err != nil { t.Fatalf("WriteManifest: %v", err) @@ -90,6 +91,75 @@ func TestManifestRoundTrip(t *testing.T) { if len(loaded.CUDA.SourceGPUUUIDs) != 2 || loaded.CUDA.SourceGPUUUIDs[0] != "GPU-aaa" { t.Errorf("CUDA.SourceGPUUUIDs = %v", loaded.CUDA.SourceGPUUUIDs) } + if loaded.CUDA.StorageMode != CUDAStorageModePOSIX { + t.Errorf("CUDA.StorageMode = %q, want %q", loaded.CUDA.StorageMode, CUDAStorageModePOSIX) + } +} + +func TestCUDAManifestEffectiveStorageMode(t *testing.T) { + tests := []struct { + name string + recorded string + want string + wantErr bool + }{ + {name: "old manifest defaults legacy", want: CUDAStorageModeLegacy}, + {name: "explicit legacy", recorded: CUDAStorageModeLegacy, want: CUDAStorageModeLegacy}, + {name: "posix", recorded: CUDAStorageModePOSIX, want: CUDAStorageModePOSIX}, + {name: "unknown", recorded: "s3", wantErr: true}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + got, err := (CUDAManifest{StorageMode: test.recorded}).EffectiveStorageMode() + if test.wantErr { + if err == nil { + t.Fatalf("EffectiveStorageMode() = %q, want error", got) + } + return + } + if err != nil || got != test.want { + t.Fatalf("EffectiveStorageMode() = %q, %v, want %q, nil", got, err, test.want) + } + }) + } +} + +func TestRegularCUDAManifestRoundTripPersistsLegacy(t *testing.T) { + dir := t.TempDir() + original := NewCheckpointManifest("content-uid", "main", CRIUDumpManifest{}, SourcePodManifest{}, OverlayManifest{}) + original.CUDA = NewCUDAManifest([]int{42}, []string{"GPU-aaa"}, CUDAStorageModeLegacy) + if err := WriteManifest(dir, original); err != nil { + t.Fatalf("WriteManifest: %v", err) + } + content, err := os.ReadFile(filepath.Join(dir, manifestFilename)) + if err != nil { + t.Fatalf("ReadFile: %v", err) + } + if !strings.Contains(string(content), "storageMode: legacy\n") { + t.Fatalf("new regular manifest does not explicitly persist legacy storage mode:\n%s", content) + } +} + +func TestReadLegacyManifestWithoutStorageModeDefaultsLegacy(t *testing.T) { + dir := t.TempDir() + content := []byte("artifact:\n contentUID: legacy-content\n containerName: main\ncudaRestore:\n pids:\n - 42\n sourceGpuUuids:\n - GPU-aaa\n") + if err := os.WriteFile(filepath.Join(dir, manifestFilename), content, 0o600); err != nil { + t.Fatalf("WriteFile: %v", err) + } + manifest, err := ReadManifest(dir) + if err != nil { + t.Fatalf("ReadManifest: %v", err) + } + mode, err := manifest.CUDA.EffectiveStorageMode() + if err != nil || mode != CUDAStorageModeLegacy { + t.Fatalf("EffectiveStorageMode() = %q, %v, want %q, nil", mode, err, CUDAStorageModeLegacy) + } + if len(manifest.CUDA.PIDs) != 1 || manifest.CUDA.PIDs[0] != 42 { + t.Fatalf("CUDA.PIDs = %v, want [42]; legacy CUDA section was not parsed", manifest.CUDA.PIDs) + } + if manifest.CUDA.StorageMode != "" { + t.Fatalf("CUDA.StorageMode = %q, want empty", manifest.CUDA.StorageMode) + } } func TestNewCRIUDumpManifest(t *testing.T) { @@ -157,6 +227,17 @@ func TestWriteManifestRejectsMissingArtifactIdentity(t *testing.T) { } } +func TestWriteManifestRequiresCUDAStorageMode(t *testing.T) { + dir := t.TempDir() + manifest := NewCheckpointManifest("content-uid", "main", CRIUDumpManifest{}, SourcePodManifest{}, OverlayManifest{}) + manifest.CUDA = NewCUDAManifest([]int{42}, []string{"GPU-aaa"}, "") + + err := WriteManifest(dir, manifest) + if err == nil || !strings.Contains(err.Error(), "missing storageMode") { + t.Fatalf("expected missing CUDA storageMode error, got %v", err) + } +} + func TestReadManifestRejectsMissingArtifactIdentity(t *testing.T) { dir := t.TempDir() diff --git a/charts/snapshot/README.md b/charts/snapshot/README.md index 24a5e973..669c1fa1 100644 --- a/charts/snapshot/README.md +++ b/charts/snapshot/README.md @@ -137,6 +137,11 @@ kubectl get pods -n ${NAMESPACE} -l app.kubernetes.io/name=snapshot -o wide | `storage.pvc.size` | Requested PVC size | `1Ti` | | `storage.pvc.storageClass` | Storage class name | `""` | | `storage.pvc.basePath` | Fixed checkpoint mount path enforced by the privileged helper | `/checkpoints` | +| `config.cudaCheckpoint.storageMode` | Storage mode for newly created CUDA checkpoints: `legacy` or explicitly enabled `posix` CustomStorage | `legacy` | +| `config.cudaCheckpoint.transferBufferCount` | Pinned CustomStorage pipeline slots per CUDA device (1-8) | `4` | +| `config.cudaCheckpoint.transferChunkBytes` | Bytes per pinned slot (1-256 MiB, 4096-byte aligned) | `67108864` | +| `config.cudaCheckpoint.daemon.maxOperationSeconds` | Cooperative extent-transfer/health watchdog (maximum one hour; CUDA driver calls are not forcibly interruptible) | `3600` | +| `config.restore.restoreTimeoutSeconds` | Overall restore deadline; default covers the qualified two-CUDA-PID workload within one target container and must scale by 65 minutes per additional CUDA-owning process | `8100` | | `seccomp.deploy` | Deploy the CRIU seccomp profile ConfigMap and init container. Use this field name; `seccomp.enabled` is not a chart value | `true` | | `runtime.type` | CRI backend: `containerd` or `crio` | `containerd` | | `runtime.socketPath` | CRI socket (empty = default for `runtime.type`) | `""` | @@ -148,6 +153,50 @@ kubectl get pods -n ${NAMESPACE} -l app.kubernetes.io/name=snapshot -o wide Reserved `s3` and `oci` values remain chart-owned placeholders for future snapshot backends, but only `pvc` is implemented today. +CustomStorage is opt-in for new checkpoints. Set +`config.cudaCheckpoint.storageMode=posix` only on nodes whose helper advertises +the CUDA 13.4 CustomStorage completion API and the Snapshot-local POSIX adapter. +Snapshot rejects the checkpoint before locking the target when the requested +capability is unavailable; it does not silently produce a legacy artifact. +The first rollout is limited to one GPU (TP1). A container may have multiple +CUDA-owning processes in that GPU's process tree. Checkpoint creation and +restore reject larger POSIX topologies before CUDA or +CRIU mutation. Four 64 MiB transfer slots are the qualified TP1 setting. +Changing the value back to `legacy` affects new checkpoints only. Restore uses +the storage mode recorded in each checkpoint manifest so already published +POSIX checkpoints remain restorable. + +The CUDA helper sidecar is mandatory for CUDA-bearing `legacy` and `posix` +operations in this release, but the agent controller starts independently so +CPU-only checkpoint/restore remains available if the helper is unhealthy. +Before a CUDA-bearing operation mutates its target, Snapshot waits for helper +health and the capabilities required by the selected storage mode. This +prevents a systemic helper or driver failure from being mistaken for a +negative CUDA process probe and producing a CRIU-only checkpoint for a GPU +workload. Deploy the agent and sidecar together; changing the ConfigMap rolls +the DaemonSet pods. +Before a planned agent upgrade or restart, drain every live target that has +completed a CustomStorage checkpoint or restore on that helper. Unexpected +helper restart while such a target remains live is not qualified in V1. + +V1 serializes CUDA checkpoint and restore sequences within one agent pod. A +sequence may cover multiple CUDA-owning PIDs from one workload, but another +workload handled by that agent waits until the active sequence completes. Run +at most one Snapshot agent installation on a node: separate DaemonSets are not +coordinated and can issue overlapping CUDA operations. Host-scoped coordination +and per-GPU concurrency are follow-ups. + +POSIX manifests require a reader that understands `cudaRestore.storageMode`. +Do not roll the agent back to a release predating that field while any POSIX +artifacts remain eligible for restore. Disable new POSIX creation, retire or +migrate those artifacts according to their retention policy, and only then +roll back. Legacy manifests from older releases remain readable. + +`transferBufferCount * transferChunkBytes` must not exceed 1 GiB +(1073741824 bytes) of pinned memory per CUDA device. Increase the CUDA helper's +memory limit when increasing either transfer setting or the number of GPUs used +by one operation. + See [values.yaml](./values.yaml) for the full configuration surface. ## Uninstall diff --git a/charts/snapshot/templates/_helpers.tpl b/charts/snapshot/templates/_helpers.tpl index 7a3927ec..d3213915 100644 --- a/charts/snapshot/templates/_helpers.tpl +++ b/charts/snapshot/templates/_helpers.tpl @@ -101,3 +101,11 @@ reads for rootfs-diff capture, and CRI-O config.json fallback). {{- define "snapshot.runtimeStorageDir" -}} {{- if eq .Values.runtime.type "crio" -}}/var/lib/containers{{- else -}}/var/lib/containerd{{- end -}} {{- end }} + +{{/* Require an integer-valued Helm number and reject strings, booleans, and null. */}} +{{- define "snapshot.requireIntegral" -}} +{{- $value := .value -}} +{{- if not (or (kindIs "int" $value) (kindIs "int64" $value) (and (kindIs "float64" $value) (eq $value (floor $value)))) -}} +{{- fail (printf "snapshot.%s must be an integral numeric value; fractional numbers, booleans, strings, and null are not accepted" .path) -}} +{{- end -}} +{{- end }} diff --git a/charts/snapshot/templates/configmap.yaml b/charts/snapshot/templates/configmap.yaml index 5f5f9983..da2d8f56 100644 --- a/charts/snapshot/templates/configmap.yaml +++ b/charts/snapshot/templates/configmap.yaml @@ -9,6 +9,35 @@ kind: ConfigMap {{- if ne .Values.storage.pvc.basePath "/checkpoints" }} {{- fail "snapshot.storage.pvc.basePath is fixed at /checkpoints so the privileged mount helper can enforce its source allowlist" }} {{- end }} +{{- $maxOperationSecondsValue := .Values.config.cudaCheckpoint.daemon.maxOperationSeconds }} +{{- include "snapshot.requireIntegral" (dict "value" $maxOperationSecondsValue "path" "config.cudaCheckpoint.daemon.maxOperationSeconds") }} +{{- if or (lt (int64 $maxOperationSecondsValue) 1) (gt (int64 $maxOperationSecondsValue) 3600) }} +{{- fail "snapshot.config.cudaCheckpoint.daemon.maxOperationSeconds must be between 1 and 3600" }} +{{- end }} +{{- $restoreTimeoutSecondsValue := .Values.config.restore.restoreTimeoutSeconds }} +{{- include "snapshot.requireIntegral" (dict "value" $restoreTimeoutSecondsValue "path" "config.restore.restoreTimeoutSeconds") }} +{{- if lt (int64 $restoreTimeoutSecondsValue) 1 }} +{{- fail "snapshot.config.restore.restoreTimeoutSeconds must be greater than zero" }} +{{- end }} +{{- $transferBufferCountValue := .Values.config.cudaCheckpoint.transferBufferCount }} +{{- include "snapshot.requireIntegral" (dict "value" $transferBufferCountValue "path" "config.cudaCheckpoint.transferBufferCount") }} +{{- $transferChunkBytesValue := .Values.config.cudaCheckpoint.transferChunkBytes }} +{{- include "snapshot.requireIntegral" (dict "value" $transferChunkBytesValue "path" "config.cudaCheckpoint.transferChunkBytes") }} +{{- $transferBufferCount := int $transferBufferCountValue }} +{{- $transferChunkBytes := int64 $transferChunkBytesValue }} +{{- if or (lt $transferBufferCount 1) (gt $transferBufferCount 8) }} +{{- fail "snapshot.config.cudaCheckpoint.transferBufferCount must be between 1 and 8" }} +{{- end }} +{{- if or (lt $transferChunkBytes 1048576) (gt $transferChunkBytes 268435456) (ne (mod $transferChunkBytes 4096) 0) }} +{{- fail "snapshot.config.cudaCheckpoint.transferChunkBytes must be a 4096-byte multiple between 1048576 and 268435456" }} +{{- end }} +{{- if gt (mul $transferBufferCount $transferChunkBytes) 1073741824 }} +{{- fail "snapshot.config.cudaCheckpoint transfer buffers exceed the 1 GiB per-device pinned-memory limit" }} +{{- end }} +{{- $cudaStorageMode := lower (trim .Values.config.cudaCheckpoint.storageMode) }} +{{- if not (has $cudaStorageMode (list "legacy" "posix")) }} +{{- fail "snapshot.config.cudaCheckpoint.storageMode must be either legacy or posix" }} +{{- end }} metadata: name: {{ include "snapshot.fullname" . }}-config namespace: {{ .Release.Namespace }} @@ -22,6 +51,11 @@ data: basePath: {{ .Values.storage.pvc.basePath | quote }} {{- end }} + cudaCheckpoint: + storageMode: {{ $cudaStorageMode | quote }} + transferBufferCount: {{ $transferBufferCount }} + transferChunkBytes: {{ $transferChunkBytes }} + overlay: exclusions: {{ toYaml .Values.config.overlay.exclusions | nindent 8 }} diff --git a/charts/snapshot/templates/daemonset.yaml b/charts/snapshot/templates/daemonset.yaml index 1cbb0642..d56564f6 100644 --- a/charts/snapshot/templates/daemonset.yaml +++ b/charts/snapshot/templates/daemonset.yaml @@ -26,8 +26,8 @@ spec: {{- with .Values.daemonset.podLabels }} {{- toYaml . | nindent 8 }} {{- end }} - {{- if or .Values.openshift.enabled .Values.daemonset.podAnnotations }} annotations: + checksum/config: {{ include (print $.Template.BasePath "/configmap.yaml") . | sha256sum }} {{- if .Values.openshift.enabled }} # Pin this DS to the built-in "privileged" SCC even if a higher-priority # custom SCC is added later. See openshift/enhancements custom-scc-preemption-prevention. @@ -41,7 +41,6 @@ spec: {{- toYaml . | nindent 8 }} {{- end }} {{- end }} - {{- end }} spec: serviceAccountName: {{ include "snapshot.serviceAccountName" . }} hostPID: true @@ -127,6 +126,8 @@ spec: - name: config mountPath: /etc/snapshot readOnly: true + - name: cuda-helper-socket + mountPath: /run/cuda-checkpoint-helper {{- if eq .Values.storage.type "pvc" }} # Every agent owns direct access to the shared checkpoint store. - name: checkpoints @@ -168,11 +169,66 @@ spec: {{- end }} resources: {{- toYaml .Values.daemonset.resources | nindent 12 }} + - name: cuda-checkpoint-helper + image: "{{ .Values.image.agent.repository }}:{{ .Values.image.agent.tag | default .Chart.AppVersion }}" + imagePullPolicy: {{ .Values.image.agent.pullPolicy }} + command: + - /usr/local/bin/cuda-checkpoint-helper + args: + - --daemon + - --socket + - /run/cuda-checkpoint-helper/helper.sock + - --max-operation-seconds + - {{ .Values.config.cudaCheckpoint.daemon.maxOperationSeconds | quote }} + securityContext: + privileged: true + startupProbe: + exec: + command: + - /usr/local/bin/cuda-checkpoint-helper + - --health + - --socket + - /run/cuda-checkpoint-helper/helper.sock + periodSeconds: 2 + failureThreshold: 150 + readinessProbe: + exec: + command: + - /usr/local/bin/cuda-checkpoint-helper + - --health + - --socket + - /run/cuda-checkpoint-helper/helper.sock + periodSeconds: 5 + timeoutSeconds: 2 + failureThreshold: 2 + livenessProbe: + exec: + command: + - /usr/local/bin/cuda-checkpoint-helper + - --health + - --socket + - /run/cuda-checkpoint-helper/helper.sock + periodSeconds: 30 + timeoutSeconds: 2 + failureThreshold: 3 + volumeMounts: + - name: cuda-helper-socket + mountPath: /run/cuda-checkpoint-helper + - name: host-proc + mountPath: /host/proc + {{- if eq .Values.storage.type "pvc" }} + - name: checkpoints + mountPath: {{ .Values.storage.pvc.basePath }} + {{- end }} + resources: + {{- toYaml .Values.config.cudaCheckpoint.daemon.resources | nindent 12 }} volumes: # Configuration ConfigMap - name: config configMap: name: {{ include "snapshot.fullname" . }}-config + - name: cuda-helper-socket + emptyDir: {} {{- if .Values.seccomp.deploy }} # Seccomp profile ConfigMap (used by initContainer) - name: seccomp-profiles diff --git a/charts/snapshot/tests/config_test.yaml b/charts/snapshot/tests/config_test.yaml index 49e2603d..b51afb55 100644 --- a/charts/snapshot/tests/config_test.yaml +++ b/charts/snapshot/tests/config_test.yaml @@ -11,3 +11,21 @@ tests: asserts: - failedTemplate: errorPattern: "storage.pvc.basePath is fixed at /checkpoints" + - it: rejects an implicit or unknown CUDA storage mode + set: + config.cudaCheckpoint.storageMode: auto + asserts: + - failedTemplate: + errorPattern: "config.cudaCheckpoint.storageMode must be either legacy or posix" + - it: rejects a CUDA helper watchdog beyond the supported operation cap + set: + config.cudaCheckpoint.daemon.maxOperationSeconds: 3601 + asserts: + - failedTemplate: + errorPattern: "config.cudaCheckpoint.daemon.maxOperationSeconds must be between 1 and 3600" + - it: rejects a non-positive restore timeout + set: + config.restore.restoreTimeoutSeconds: 0 + asserts: + - failedTemplate: + errorPattern: "config.restore.restoreTimeoutSeconds must be greater than zero" diff --git a/charts/snapshot/values.yaml b/charts/snapshot/values.yaml index c1d89ca7..bef4348d 100644 --- a/charts/snapshot/values.yaml +++ b/charts/snapshot/values.yaml @@ -162,6 +162,35 @@ rbac: # Static agent configuration (loaded from ConfigMap) # Dynamic values such as NODE_NAME come from environment variables. config: + cudaCheckpoint: + # Storage mode for newly created CUDA checkpoints. "legacy" keeps the + # existing driver-managed path. "posix" explicitly requires the CUDA 13.4 + # CustomStorage capability and the Snapshot-local NIXL POSIX adapter. + # Restore always follows the mode recorded in the checkpoint manifest. + storageMode: legacy + # Pinned transfer slots per CUDA device. More slots add pipeline depth, + # not worker threads. Four 64 MiB slots are the qualified TP1 benchmark + # configuration; POSIX creation is rejected outside that topology for V1. + transferBufferCount: 4 + # Bytes in each pinned transfer slot (1-256 MiB, 4096-byte aligned). + transferChunkBytes: 67108864 + daemon: + # Cooperative watchdog for one CUDA operation's extent transfers and + # health state. CUDA driver calls are not forcibly interruptible. The + # client reserves a five-minute response margin before long-running work. + maxOperationSeconds: 3600 + # Pinned transfer memory is bufferCount * chunkBytes per active CUDA + # device. Increase this limit when increasing those settings or the + # number of GPUs used by one operation; the helper also needs room for + # NIXL, CUDA, and process overhead beyond the pinned buffers. + resources: + limits: + cpu: 2 + memory: 1Gi + requests: + cpu: 250m + memory: 256Mi + overlay: # Rootfs diff tar exclusions. Absolute-looking paths are normalized # relative to the tar root, and patterns starting with * are passed @@ -176,7 +205,11 @@ config: restore: # Maximum seconds to allow a restore attempt before the agent marks it failed - restoreTimeoutSeconds: 7200 + # Each CUDA target can require one long-running restore call with a + # 65-minute caller budget. Terminal unlock remains permitted under a short + # deadline. The default covers the qualified two-CUDA-PID TP1 workload; + # scale it for workloads with more CUDA-owning processes. + restoreTimeoutSeconds: 8100 criu: # Path to the criu binary