From 2afb549ee418401f661e588c360566ab3614ed57 Mon Sep 17 00:00:00 2001 From: Balaji Ganesan Date: Thu, 24 Sep 2026 08:14:47 -0700 Subject: [PATCH 1/3] refactor(nvsnap): remove the LD_PRELOAD interception stack criu-v2 replaced criu-v2 handles the workload at the OS level and needs nothing injected into it. The webhook still carried the machinery to inject the old stack, and the tree still carried the stack itself. It was not merely unused, it was non-functional. The injection is opt-in per pod via nvsnap.io/auto-inject, and no manifest here sets it. Worse, the DaemonSet passed the four required image refs as positional arguments while the agent bound them to named flags, and nothing reads flag.Args(), so AutoInject.Valid() was always false. The branch could not fire from any chart deployment. On main both isMultiGPU and useCUDAInterposition are hardcoded false, so every interposition block in checkpoint.go was already unreachable by construction. Removed: lib/nvsnap_intercept (13732 lines of C and headers), lib/sitecustomize, internal/webhook/auto_inject.go, the uvloop/libuv/libzmq/pyzmq builder images and the combined init image, 9 dedicated scripts, the four webhook-image flags, the builder wiring in the chart and both DaemonSets, and the orphaned version variables and fork refs in versions.sh. Dockerfile.app.criuv2 goes too: it was an unused injection-free variant, and keeping it now would be misleading. The legacy multi-GPU D2H path goes with it rather than separately. It requires the interposer preloaded into the workload, so it cannot outlive the stack. restore_entrypoint.go stays. It is called from l2_mount.go on the live cachedir restore path, so it keeps the two nvsnap-lib constants it is now the only user of. Whether that mount is still needed is a separate question. Orphans the balajinvda forks of libzmq, libuv, uvloop and pyzmq. The libuv fork existed for the lost SQ-array identity map, which was fixed on the CRIU side, and versions.sh already recorded that rings survive checkpoint and restore with UV_USE_IO_URING=1. Those repositories need their own decision. The implementation stays reachable at tag archive/nvsnap-injection-stack. Verified on dev1 with a rebuilt agent: vllm-small on criu-v2 passes (checkpoint 1m36s/27G, restore 46s) and nim-qwen3-32b on cachedir passes (capture 61G, restore 2m34s). Building the image caught three references that compilation and the unit tests could not: the build context still copied the deleted library, an untracked-source guard pointed at it, and the image self-check ls-ed a file that no longer exists. Deploying caught a fourth, a log line still calling cachedir patches "auto-inject only". Relates to #2094 Co-Authored-By: Balaji Ganesan --- .../nvsnap/cmd/agent/main.go | 13 - .../deploy/helm/nvsnap/templates/_helpers.tpl | 11 - .../nvsnap/templates/agent-daemonset.yaml | 9 - .../nvsnap/deploy/helm/nvsnap/values.yaml | 6 - .../nvsnap/deploy/k8s/agent-daemonset.yaml | 10 - .../nvsnap/docker/agent/Dockerfile.app | 68 +- .../nvsnap/docker/agent/Dockerfile.app.criuv2 | 55 - .../nvsnap/docker/init/Dockerfile | 54 - .../nvsnap/docker/init/init.sh | 27 - .../nvsnap/docker/libuv/Dockerfile | 52 - .../nvsnap/docker/libzmq/Dockerfile | 61 - .../nvsnap/docker/pyzmq/Dockerfile | 69 - .../nvsnap/docker/uvloop/Dockerfile | 64 - .../internal/agent/webhook_integration.go | 7 - .../nvsnap/internal/webhook/BUILD.bazel | 1 - .../nvsnap/internal/webhook/auto_inject.go | 200 -- .../nvsnap/internal/webhook/mutate.go | 20 +- .../internal/webhook/restore_entrypoint.go | 9 + .../nvsnap/lib/nvsnap_intercept/.gitignore | 1 - .../nvsnap/lib/nvsnap_intercept/BUILD.bazel | 79 - .../nvsnap/lib/nvsnap_intercept/Makefile | 197 -- .../nvsnap/lib/nvsnap_intercept/README.md | 117 - .../include/nvsnap/gpu/checkpoint.h | 149 - .../include/nvsnap/gpu/config.h | 39 - .../include/nvsnap/gpu/cuda_types.h | 722 ----- .../include/nvsnap/gpu/fault.h | 53 - .../include/nvsnap/gpu/interpose.h | 119 - .../include/nvsnap/gpu/metrics.h | 83 - .../include/nvsnap/gpu/tracker.h | 180 -- .../include/nvsnap_intercept.h | 239 -- .../nvsnap/lib/nvsnap_intercept/libnvsnap.map | 5 - .../nvsnap_intercept/src/abort_intercept.c | 100 - .../lib/nvsnap_intercept/src/cuda_intercept.c | 18 - .../nvsnap_intercept/src/gpu/checkpoint.cpp | 1234 -------- .../lib/nvsnap_intercept/src/gpu/config.c | 102 - .../lib/nvsnap_intercept/src/gpu/init.c | 229 -- .../src/gpu/interpose_cudart.c | 268 -- .../src/gpu/interpose_cudrv.c | 212 -- .../nvsnap_intercept/src/gpu/interpose_nccl.c | 104 - .../lib/nvsnap_intercept/src/gpu/metrics.c | 110 - .../nvsnap_intercept/src/gpu/symbol_table.c | 122 - .../lib/nvsnap_intercept/src/gpu/tracker.cpp | 363 --- .../lib/nvsnap_intercept/src/gpu_checkpoint.c | 145 - .../nvsnap/lib/nvsnap_intercept/src/init.c | 347 --- .../nvsnap_intercept/src/io_uring_intercept.c | 1284 -------- .../nvsnap_intercept/src/libuv_intercept.c | 1043 ------- .../lib/nvsnap_intercept/src/nccl_intercept.c | 714 ----- .../nvsnap/lib/nvsnap_intercept/src/quiesce.c | 878 ------ .../nvsnap_intercept/src/seccomp_intercept.c | 767 ----- .../lib/nvsnap_intercept/src/self_disable.c | 73 - .../lib/nvsnap_intercept/src/zmq_intercept.c | 2618 ----------------- .../nvsnap_intercept/tests/run_cuda_tests.sh | 54 - .../tests/run_local_criu_test.sh | 137 - .../tests/test_criu_integration.sh | 149 - .../tests/test_cuda_intercept.c | 637 ---- .../tests/test_dlsym_recursion.c | 144 - .../tests/test_full_checkpoint_restore.sh | 303 -- .../nvsnap_intercept/tests/test_in_docker.sh | 124 - .../tests/test_library_safety.c | 375 --- .../tests/test_libuv_intercept.c | 94 - .../tests/test_libuv_simple.c | 60 - .../tests/test_local_checkpoint.sh | 242 -- .../tests/test_local_criu_uvloop.py | 273 -- .../nvsnap_intercept/tests/test_local_vllm.sh | 110 - .../nvsnap_intercept/tests/test_pytorch.py | 287 -- .../tests/test_pytorch_distributed.py | 104 - .../nvsnap_intercept/tests/test_quiesce.py | 197 -- .../nvsnap_intercept/tests/test_seccomp.py | 82 - .../tests/test_seccomp_direct.c | 107 - .../tests/test_seccomp_uvloop.py | 136 - .../tests/test_sglang_local.py | 102 - .../tests/test_sglang_simple.py | 48 - .../tests/test_uring_simple.py | 24 - .../tests/test_uvloop_simple.py | 31 - .../lib/nvsnap_intercept/tests/test_vllm.py | 217 -- .../tests/test_vllm_checkpoint.sh | 194 -- .../nvsnap_intercept/tests/test_vllm_local.py | 93 - .../nvsnap/lib/sitecustomize/sitecustomize.py | 44 - .../nvsnap/scripts/auto-inject-init.sh | 53 - .../nvsnap/scripts/build-agent-app.sh | 88 - .../nvsnap/scripts/build-agent.sh | 49 +- .../nvsnap/scripts/build-deps.sh | 264 -- .../scripts/build-intercept-lib-local.sh | 28 - .../nvsnap/scripts/build-libzmq-image.sh | 116 - .../nvsnap/scripts/build-pyzmq-wheel.sh | 161 - .../nvsnap/scripts/build-uvloop-wheel.sh | 102 - .../nvsnap/scripts/local-uvloop-test.sh | 235 -- .../nvsnap/scripts/restore-bundle-init.sh | 2 +- .../nvsnap/scripts/sync-versions.sh | 5 - .../nvsnap/scripts/test-gpu-checkpoint.sh | 253 -- .../nvsnap/scripts/test-vllm-zmq.sh | 294 -- .../nvsnap/scripts/validate-libzmq-fork.sh | 230 -- .../nvsnap/scripts/validate-test-env.sh | 16 +- .../nvsnap/scripts/versions.sh | 18 - 94 files changed, 29 insertions(+), 19703 deletions(-) delete mode 100644 src/compute-plane-services/nvsnap/docker/agent/Dockerfile.app.criuv2 delete mode 100644 src/compute-plane-services/nvsnap/docker/init/Dockerfile delete mode 100644 src/compute-plane-services/nvsnap/docker/init/init.sh delete mode 100644 src/compute-plane-services/nvsnap/docker/libuv/Dockerfile delete mode 100644 src/compute-plane-services/nvsnap/docker/libzmq/Dockerfile delete mode 100644 src/compute-plane-services/nvsnap/docker/pyzmq/Dockerfile delete mode 100644 src/compute-plane-services/nvsnap/docker/uvloop/Dockerfile delete mode 100644 src/compute-plane-services/nvsnap/internal/webhook/auto_inject.go delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/.gitignore delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/BUILD.bazel delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/Makefile delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/README.md delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/checkpoint.h delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/config.h delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/cuda_types.h delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/fault.h delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/interpose.h delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/metrics.h delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/tracker.h delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap_intercept.h delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/libnvsnap.map delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/abort_intercept.c delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/cuda_intercept.c delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/checkpoint.cpp delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/config.c delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/init.c delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/interpose_cudart.c delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/interpose_cudrv.c delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/interpose_nccl.c delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/metrics.c delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/symbol_table.c delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/tracker.cpp delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu_checkpoint.c delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/init.c delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/io_uring_intercept.c delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/libuv_intercept.c delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/nccl_intercept.c delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/quiesce.c delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/seccomp_intercept.c delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/self_disable.c delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/zmq_intercept.c delete mode 100755 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/run_cuda_tests.sh delete mode 100755 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/run_local_criu_test.sh delete mode 100755 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_criu_integration.sh delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_cuda_intercept.c delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_dlsym_recursion.c delete mode 100755 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_full_checkpoint_restore.sh delete mode 100755 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_in_docker.sh delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_library_safety.c delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_libuv_intercept.c delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_libuv_simple.c delete mode 100755 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_local_checkpoint.sh delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_local_criu_uvloop.py delete mode 100755 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_local_vllm.sh delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_pytorch.py delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_pytorch_distributed.py delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_quiesce.py delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_seccomp.py delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_seccomp_direct.c delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_seccomp_uvloop.py delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_sglang_local.py delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_sglang_simple.py delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_uring_simple.py delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_uvloop_simple.py delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_vllm.py delete mode 100755 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_vllm_checkpoint.sh delete mode 100644 src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_vllm_local.py delete mode 100644 src/compute-plane-services/nvsnap/lib/sitecustomize/sitecustomize.py delete mode 100644 src/compute-plane-services/nvsnap/scripts/auto-inject-init.sh delete mode 100755 src/compute-plane-services/nvsnap/scripts/build-agent-app.sh delete mode 100755 src/compute-plane-services/nvsnap/scripts/build-deps.sh delete mode 100755 src/compute-plane-services/nvsnap/scripts/build-intercept-lib-local.sh delete mode 100755 src/compute-plane-services/nvsnap/scripts/build-libzmq-image.sh delete mode 100755 src/compute-plane-services/nvsnap/scripts/build-pyzmq-wheel.sh delete mode 100755 src/compute-plane-services/nvsnap/scripts/build-uvloop-wheel.sh delete mode 100755 src/compute-plane-services/nvsnap/scripts/local-uvloop-test.sh delete mode 100755 src/compute-plane-services/nvsnap/scripts/test-gpu-checkpoint.sh delete mode 100755 src/compute-plane-services/nvsnap/scripts/test-vllm-zmq.sh delete mode 100755 src/compute-plane-services/nvsnap/scripts/validate-libzmq-fork.sh diff --git a/src/compute-plane-services/nvsnap/cmd/agent/main.go b/src/compute-plane-services/nvsnap/cmd/agent/main.go index 55b85a5872..c0d8405f7b 100644 --- a/src/compute-plane-services/nvsnap/cmd/agent/main.go +++ b/src/compute-plane-services/nvsnap/cmd/agent/main.go @@ -146,19 +146,6 @@ func main() { flag.StringVar(&config.Webhook.Path, "webhook-path", "/mutate", "HTTP path the in-agent webhook listens on") - // BYOC auto-inject: image refs the webhook stamps into the four - // init containers when a pod has nvsnap.io/auto-inject: "true". All - // four must be set for the auto-inject branch to fire; otherwise - // the webhook fails open (admits the pod unchanged). - flag.StringVar(&config.Webhook.AutoInject.Uvloop, "webhook-image-uvloop", "", //nolint:staticcheck // deprecated field intentionally bound for flag back-compat - "Image ref for the auto-inject get-uvloop init container (multi-python uvloop wheels)") - flag.StringVar(&config.Webhook.AutoInject.LibUV, "webhook-image-libuv", "", - "Image ref for the auto-inject get-libuv init container") - flag.StringVar(&config.Webhook.AutoInject.LibZMQ, "webhook-image-libzmq", "", - "Image ref for the auto-inject get-libzmq init container") - flag.StringVar(&config.Webhook.AutoInject.Agent, "webhook-image-agent", "", - "Image ref for the auto-inject get-nvsnap init container (nvsnap-agent — must match running agent)") - // nvsnap#147: nvsnap-l2-wait init container ref. When set, the // webhook prepends a nvsnap-l2-wait init container on restore pods // that polls nvsnap-server until the L2 PVC promote is ready. diff --git a/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/templates/_helpers.tpl b/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/templates/_helpers.tpl index 5fc465d8cb..50af0fbeef 100644 --- a/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/templates/_helpers.tpl +++ b/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/templates/_helpers.tpl @@ -155,17 +155,6 @@ registry. Keeping our registry under our own namespace prevents that. hostPath mounts from there onto function pods. No cross-registry mirror required. */ -}} -{{- define "nvsnap.builder.image" -}} -{{- /* args: ctx, name (e.g. "uvloop-builder"). Same mandatory-tag rule - as nvsnap.image — if nvsnap.builderTag is empty the user has to - set it explicitly. */ -}} -{{- $reg := .ctx.Values.nvsnap.imageRegistry -}} -{{- if not .ctx.Values.nvsnap.builderTag -}} -{{- fail (printf "nvsnap.builder.image: nvsnap.builderTag is required (rendering %q)" .name) -}} -{{- end -}} -{{- printf "%s/%s:%s" $reg .name .ctx.Values.nvsnap.builderTag -}} -{{- end }} - {{/* Image-pull-secrets block, rendered if any are configured. Used at the PodSpec level (not Deployment-level — Kubernetes requires it on pods). diff --git a/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/templates/agent-daemonset.yaml b/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/templates/agent-daemonset.yaml index a9fc51e89b..9bd41e9ee9 100644 --- a/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/templates/agent-daemonset.yaml +++ b/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/templates/agent-daemonset.yaml @@ -154,15 +154,6 @@ spec: - --webhook-mount-prep-init-image={{ .Values.webhook.mountPrepInitImage | default (include "nvsnap.agent.image" .) }} - --webhook-agent-host-port={{ .Values.webhook.agentHostPort | default 8081 }} {{- end }} - # BYOC auto-inject: when a pod carries nvsnap.io/auto-inject: "true" - # the webhook stamps in the sitecustomize plumbing using these - # images. All four must be set; the agent image MUST match - # this DaemonSet's image so libnvsnap_intercept.so build-IDs - # line up at restore time. - - {{ include "nvsnap.builder.image" (dict "ctx" . "name" "uvloop-builder") }} - - {{ include "nvsnap.builder.image" (dict "ctx" . "name" "libuv-builder") }} - - {{ include "nvsnap.builder.image" (dict "ctx" . "name" "libzmq-builder") }} - - {{ include "nvsnap.agent.image" . }} {{- end }} securityContext: privileged: true diff --git a/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/values.yaml b/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/values.yaml index bc5a67fcaf..136aa8dae3 100644 --- a/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/values.yaml +++ b/src/compute-plane-services/nvsnap/deploy/helm/nvsnap/values.yaml @@ -80,12 +80,6 @@ nvsnap: imagePullSecrets: - nvsnap-pull-secret - # Tag used for dependency builder images (uvloop-builder, libuv-builder, - # libzmq-builder). The agent passes these to its in-process webhook so - # it can auto-inject the right init container into BYOC workload pods. - # Must match what's actually been built+pushed to imageRegistry — - # there's no AppVersion fallback (see _helpers.tpl). - builderTag: "v0.0.1" # ─── NvSnap Agent (DaemonSet on GPU nodes) ─────────────────────────────── diff --git a/src/compute-plane-services/nvsnap/deploy/k8s/agent-daemonset.yaml b/src/compute-plane-services/nvsnap/deploy/k8s/agent-daemonset.yaml index 4ae17a3f6d..440c7a956f 100644 --- a/src/compute-plane-services/nvsnap/deploy/k8s/agent-daemonset.yaml +++ b/src/compute-plane-services/nvsnap/deploy/k8s/agent-daemonset.yaml @@ -131,16 +131,6 @@ spec: - --webhook-listen=:8443 - --webhook-cert=/etc/nvsnap/webhook/tls.crt - --webhook-key=/etc/nvsnap/webhook/tls.key - # BYOC auto-inject — when a pod carries - # nvsnap.io/auto-inject: "true", the webhook stamps in - # the sitecustomize plumbing using these images. All four - # must be set or the webhook fails open. Agent image MUST - # match this DaemonSet's container image so the - # libnvsnap_intercept.so build-ID lines up at restore time. - - nvcr.io/0651155215864979/ncp-dev/uvloop-builder:v0.0.1 - - nvcr.io/0651155215864979/ncp-dev/libuv-builder:v0.0.1 - - nvcr.io/0651155215864979/ncp-dev/libzmq-builder:v0.0.1 - - nvcr.io/0651155215864979/ncp-dev/nvsnap-agent:v0.2.42 securityContext: privileged: true # Required for CRIU and containerd access env: diff --git a/src/compute-plane-services/nvsnap/docker/agent/Dockerfile.app b/src/compute-plane-services/nvsnap/docker/agent/Dockerfile.app index a0363e37e0..29067b0912 100644 --- a/src/compute-plane-services/nvsnap/docker/agent/Dockerfile.app +++ b/src/compute-plane-services/nvsnap/docker/agent/Dockerfile.app @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 # NVSNAP Agent Application Image -# Builds on top of nvsnap-agent-base with Go binaries and intercept library +# Builds on top of nvsnap-agent-base with the Go binaries and C helpers. # This is the image that gets rebuilt frequently during development # # Prerequisites: Build base image first with Dockerfile.base @@ -10,15 +10,6 @@ ARG BASE_IMAGE=nvcr.io/0651155215864979/ncp-dev/nvsnap-agent-base:v0.0.19 -# Builder images we fold INTO the agent. Lets the BYOC auto-inject -# webhook use ONE init container (this image) instead of four — -# saves init container startup overhead and image-pull bandwidth on -# cold nodes. Trade-off is a ~30MB increase in the agent image, -# which is fine because the agent already carries the CRIU bundle. -ARG UVLOOP_IMAGE=nvcr.io/0651155215864979/ncp-dev/uvloop-builder:v0.0.1 -ARG LIBUV_IMAGE=nvcr.io/0651155215864979/ncp-dev/libuv-builder:v0.0.1 -ARG LIBZMQ_IMAGE=nvcr.io/0651155215864979/ncp-dev/libzmq-builder:v0.0.1 - # ============================================================================ # Stage 1: Build Go binaries # ============================================================================ @@ -35,9 +26,9 @@ RUN CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -ldflags="-s -w" -o /bin/nvsn RUN CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -ldflags="-s -w" -o /bin/nvsnap-rootfs-restore ./cmd/nvsnap-rootfs-restore # ============================================================================ -# Stage 2: Build intercept library +# Stage 2: Build the standalone C helpers # ============================================================================ -FROM ubuntu:22.04 AS intercept-builder +FROM ubuntu:22.04 AS c-builder RUN apt-get update && apt-get install -y --no-install-recommends \ build-essential \ @@ -46,13 +37,6 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ cmake \ && rm -rf /var/lib/apt/lists/* -WORKDIR /app -COPY lib/nvsnap_intercept/ . - -ARG ENABLE_LIBUV_INTERCEPT=0 -RUN make clean && make ENABLE_LIBUV_INTERCEPT=${ENABLE_LIBUV_INTERCEPT} && \ - ls -la libnvsnap_intercept.so - # Build nvsnap-gpu-restore (standalone C binary, links libcuda at runtime via dlopen) COPY cmd/nvsnap-gpu-restore/main.c /tmp/nvsnap-gpu-restore.c RUN gcc -O2 -o /tmp/nvsnap-gpu-restore /tmp/nvsnap-gpu-restore.c -ldl && \ @@ -67,13 +51,6 @@ RUN cd /tmp/nvsnap_restore_helper && make && \ # Build NvSnap from source (ensures glibc compatibility with Ubuntu 22.04) -# ============================================================================ -# Builder-payload stages (just pulled for COPY --from in final stage) -# ============================================================================ -FROM ${UVLOOP_IMAGE} AS uvloop-payload -FROM ${LIBUV_IMAGE} AS libuv-payload -FROM ${LIBZMQ_IMAGE} AS libzmq-payload - # ============================================================================ # Stage 3: Final image (fast - just copies binaries into base) # ============================================================================ @@ -96,45 +73,21 @@ COPY --from=go-builder /bin/restore-entrypoint /criu-bundle/restore-entrypoint COPY --from=go-builder /bin/nvsnap-mount-prep /criu-bundle/nvsnap-mount-prep COPY --from=go-builder /bin/nvsnap-rootfs-restore /criu-bundle/nvsnap-rootfs-restore -# Copy intercept library (includes uv_loop_fork call via Python C API) -COPY --from=intercept-builder /app/libnvsnap_intercept.so /criu-bundle/lib/libnvsnap_intercept.so - -# Bundle the runtime sitecustomize. Init container get-criu (or any -# init that copies /criu-bundle into /nvsnap-lib) delivers this onto the -# emptyDir volume; workload Pods then set PYTHONPATH=/nvsnap-lib/sitecustomize -# and our patched site-packages-cpXY wins on import. See -# docs/GENERIC-PYTHON-INJECTION-DESIGN.md. -COPY lib/sitecustomize/sitecustomize.py /criu-bundle/sitecustomize/sitecustomize.py - -# Bundle the BYOC auto-inject payload: uvloop wheels (multi-python), -# patched libuv.so, patched libzmq.so. The webhook injects ONE init -# container running this agent image with auto-inject-init.sh, which -# copies these payloads into the workload pod's /nvsnap-lib emptyDir. -# Replaces the previous 4-init-container fan-out and removes the -# separate nvsnap-init image (one less artifact to keep version-matched). -COPY --from=uvloop-payload /wheels/ /criu-bundle/payload/wheels/ -COPY --from=libuv-payload /usr/local/lib/libuv.so /criu-bundle/payload/lib/libuv.so -COPY --from=libuv-payload /usr/local/lib/libuv.so.1 /criu-bundle/payload/lib/libuv.so.1 -COPY --from=libzmq-payload /usr/local/lib/libzmq.so /criu-bundle/payload/lib/libzmq.so -COPY --from=libzmq-payload /usr/local/lib/libzmq.so.5 /criu-bundle/payload/lib/libzmq.so.5 -COPY scripts/auto-inject-init.sh /criu-bundle/auto-inject-init.sh -RUN chmod +x /criu-bundle/auto-inject-init.sh - # nvsnap#147: Restore-bundle init payload. The mutating webhook injects # one init container running this same agent image with restore-bundle-init.sh # as its command; the script copies /criu-bundle/. → /nvsnap onto a shared # emptyDir so the rewritten workload command /nvsnap/restore-entrypoint can -# exec the CRIU restore. Companion to auto-inject-init.sh. +# exec the CRIU restore. COPY scripts/restore-bundle-init.sh /criu-bundle/restore-bundle-init.sh RUN chmod +x /criu-bundle/restore-bundle-init.sh -# Copy NvSnap GPU interposition library (built from source in intercept-builder) +# Standalone C helpers built in the c-builder stage above. # Copy nvsnap-gpu-restore binary (restores GPU memory via CUDA VMM APIs after CRIU restore) -COPY --from=intercept-builder /tmp/nvsnap-gpu-restore /criu-bundle/nvsnap-gpu-restore -COPY --from=intercept-builder /tmp/nvsnap_restore_helper/nvsnap-restore-helper /criu-bundle/nvsnap-restore-helper +COPY --from=c-builder /tmp/nvsnap-gpu-restore /criu-bundle/nvsnap-gpu-restore +COPY --from=c-builder /tmp/nvsnap_restore_helper/nvsnap-restore-helper /criu-bundle/nvsnap-restore-helper -# Override cuda-checkpoint wrapper: isolate from LD_PRELOAD intercept library. -# Without this, the intercept library's log messages corrupt cuda-checkpoint's +# Override cuda-checkpoint wrapper: keep stray stdout out of the tool's output. +# Historically this isolated it from the LD_PRELOAD interceptor, whose logs # output, causing the CRIU CUDA plugin to get tid=0 and GPU resume to fail. COPY cuda-checkpoint-wrapper.sh /criu-bundle/cuda-checkpoint RUN chmod +x /criu-bundle/cuda-checkpoint @@ -146,8 +99,7 @@ RUN ln -sf /criu-bundle/nvsnap-agent /usr/local/bin/nvsnap-agent && \ # Verify everything works RUN echo "=== Agent ===" && /criu-bundle/nvsnap-agent --help 2>&1 | head -3 || true && \ - echo "=== Bundle contents ===" && ls -la /criu-bundle/ && \ - echo "=== Intercept library ===" && ls -la /criu-bundle/lib/libnvsnap_intercept.so + echo "=== Bundle contents ===" && ls -la /criu-bundle/ # Set ENTRYPOINT (not CMD) so K8s args append properly ENTRYPOINT ["/criu-bundle/nvsnap-agent"] diff --git a/src/compute-plane-services/nvsnap/docker/agent/Dockerfile.app.criuv2 b/src/compute-plane-services/nvsnap/docker/agent/Dockerfile.app.criuv2 deleted file mode 100644 index 638c7d64b3..0000000000 --- a/src/compute-plane-services/nvsnap/docker/agent/Dockerfile.app.criuv2 +++ /dev/null @@ -1,55 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# NvSnap agent — criu-v2-only image, buildable entirely from PUBLIC inputs. -# -# Unlike Dockerfile.app, this pulls no private builder images. criu-v2 dropped -# the legacy injection stack (patched uvloop/libuv/libzmq + libnvsnap_intercept), -# so the agent image needs only: -# - the CRIU bundle from nvsnap-agent-base (built by Dockerfile.base from the -# PUBLIC fork github.com/balajinvda/criu — criu + cuda_plugin.so) -# - cuda-checkpoint (shipped in-repo; the public NVIDIA release is x86-64 only, -# so we carry the binary — arm64 needs the arm64 binary committed too) -# - the Go agent (this repo) -# -# glibc: Go binaries are CGO_ENABLED=0 -> fully static, arch/libc-agnostic. The -# base runtime is ubuntu:22.04 (glibc 2.35, the portability floor) and criu -# ships its own ld-linux + libc in /criu-bundle/lib. Do NOT bump the base off -# 22.04 and do NOT enable CGO — that is where the glibc-mismatch failures come -# from (a binary linked against a newer glibc aborts inside older workload -# containers). - -ARG BASE_IMAGE=nvsnap-agent-base:v0.0.7 - -# --------------------------------------------------------------------------- -# Stage 1: build the Go agent (static, no glibc dependency) -# --------------------------------------------------------------------------- -FROM golang:1.25-bookworm AS go-builder -ARG TARGETARCH=amd64 -WORKDIR /app -COPY go.mod go.sum ./ -RUN go mod download -COPY . . -RUN CGO_ENABLED=0 GOOS=linux GOARCH=${TARGETARCH} go build -ldflags="-s -w" \ - -o /bin/nvsnap-agent ./cmd/agent - -# --------------------------------------------------------------------------- -# Stage 2: assemble the agent image on top of the CRIU base bundle -# --------------------------------------------------------------------------- -FROM ${BASE_IMAGE} - -COPY --from=go-builder /bin/nvsnap-agent /criu-bundle/nvsnap-agent - -# cuda-checkpoint wrapper: unsets LD_PRELOAD / disables /etc/ld.so.preload -# before exec'ing cuda-checkpoint.real, so the CUDA plugin can parse its stdout -# for the restore tid (a stray LD_PRELOAD log line corrupts it -> tid=0 -> -# GPU resume fails). cuda-checkpoint.real rides in the base bundle. -COPY docker/agent/cuda-checkpoint-wrapper.sh /criu-bundle/cuda-checkpoint -RUN chmod +x /criu-bundle/cuda-checkpoint && \ - ln -sf /criu-bundle/nvsnap-agent /usr/local/bin/nvsnap-agent && \ - ln -sf /criu-bundle/criu /usr/local/sbin/criu - -RUN echo "=== bundle ===" && ls -la /criu-bundle/ && \ - echo "=== criu ===" && /criu-bundle/criu --version - -ENTRYPOINT ["/criu-bundle/nvsnap-agent"] diff --git a/src/compute-plane-services/nvsnap/docker/init/Dockerfile b/src/compute-plane-services/nvsnap/docker/init/Dockerfile deleted file mode 100644 index f234adfaa6..0000000000 --- a/src/compute-plane-services/nvsnap/docker/init/Dockerfile +++ /dev/null @@ -1,54 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# Combined init container for NvSnap restore pods. -# Replaces 5 separate init containers (uvloop, libuv, libzmq, pyzmq, criu) -# with a single container that copies all dependencies in one step. -# -# Build: docker build -t nvsnap-init:VERSION --build-arg AGENT_IMAGE=... -f docker/init/Dockerfile . - -ARG UVLOOP_IMAGE=nvcr.io/0651155215864979/ncp-dev/uvloop-builder:v0.0.1 -ARG LIBUV_IMAGE=nvcr.io/0651155215864979/ncp-dev/libuv-builder:v0.0.1 -ARG LIBZMQ_IMAGE=nvcr.io/0651155215864979/ncp-dev/libzmq-builder:v0.0.1 -ARG PYZMQ_IMAGE=nvcr.io/0651155215864979/ncp-dev/pyzmq-builder:v0.0.1 -ARG AGENT_IMAGE=nvcr.io/0651155215864979/ncp-dev/nvsnap-agent:v0.0.1 - -# Stage 1-4: Pull artifacts from builder images -FROM ${UVLOOP_IMAGE} AS uvloop -FROM ${LIBUV_IMAGE} AS libuv -FROM ${LIBZMQ_IMAGE} AS libzmq -FROM ${PYZMQ_IMAGE} AS pyzmq -FROM ${AGENT_IMAGE} AS agent - -# Stage 5: Combine everything into a minimal image -FROM ubuntu:22.04 - -# uvloop wheel -> site-packages -COPY --from=uvloop /wheels/ /staging/wheels/uvloop/ -# libuv shared libs -COPY --from=libuv /usr/local/lib/libuv.so* /staging/lib/ -# libzmq shared libs -COPY --from=libzmq /usr/local/lib/libzmq.so* /staging/lib/ -# pyzmq wheel -> site-packages -COPY --from=pyzmq /wheels/ /staging/wheels/pyzmq/ -# criu bundle (agent binaries, CRIU, intercept lib, cuda-checkpoint) -COPY --from=agent /criu-bundle/ /staging/criu-bundle/ -# py-spy (optional debug tool) -COPY --from=agent /usr/local/bin/py-spy /staging/criu-bundle/py-spy - -# Pre-extract wheels so init just copies dirs. Loop because the multi- -# python uvloop builder produces multiple wheels (cp310/311/312/313) -# and `python3 -m zipfile -e` only accepts a single input file. Same -# target dir for all wheels is fine — Python picks the right -# cpython--*.so at runtime by interpreter tag. -RUN apt-get update && apt-get install -y --no-install-recommends python3 && rm -rf /var/lib/apt/lists/* && \ - mkdir -p /staging/site-packages && \ - for w in /staging/wheels/uvloop/uvloop-*.whl; do python3 -m zipfile -e "$w" /staging/site-packages/; done && \ - for w in /staging/wheels/pyzmq/pyzmq-*.whl; do python3 -m zipfile -e "$w" /staging/site-packages/; done && \ - rm -rf /staging/wheels - -# The init entrypoint copies everything to the shared volumes -COPY docker/init/init.sh /init.sh -RUN chmod +x /init.sh - -ENTRYPOINT ["/init.sh"] diff --git a/src/compute-plane-services/nvsnap/docker/init/init.sh b/src/compute-plane-services/nvsnap/docker/init/init.sh deleted file mode 100644 index d7b73f59b5..0000000000 --- a/src/compute-plane-services/nvsnap/docker/init/init.sh +++ /dev/null @@ -1,27 +0,0 @@ -#!/bin/sh -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# Combined init script for NvSnap restore pods. -# Copies all dependencies to shared volumes in one step. -set -e - -echo "=== NvSnap init: copying dependencies ===" - -# Copy Python packages (uvloop, pyzmq) to nvsnap-lib volume -mkdir -p /nvsnap-lib/site-packages -cp -r /staging/site-packages/* /nvsnap-lib/site-packages/ - -# Copy shared libraries (libuv, libzmq, libnvsnap_intercept) to nvsnap-lib volume -cp /staging/lib/libuv.so* /nvsnap-lib/ 2>/dev/null || true -cp /staging/lib/libzmq.so* /nvsnap-lib/ -cp /staging/criu-bundle/lib/libnvsnap_intercept.so /nvsnap-lib/ - -# Copy CRIU bundle (criu, restore-entrypoint, agent, cuda-checkpoint) to nvsnap-tools volume -cp -r /staging/criu-bundle/. /nvsnap/ -# Copy optional debug tools -if [ -f /staging/criu-bundle/py-spy ]; then - cp /staging/criu-bundle/py-spy /nvsnap/ -fi - -echo "=== NvSnap init complete ===" diff --git a/src/compute-plane-services/nvsnap/docker/libuv/Dockerfile b/src/compute-plane-services/nvsnap/docker/libuv/Dockerfile deleted file mode 100644 index 51087455f0..0000000000 --- a/src/compute-plane-services/nvsnap/docker/libuv/Dockerfile +++ /dev/null @@ -1,52 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# libuv builder image for NVSNAP -# -# Builds our forked libuv with CRIU checkpoint/restore support: -# - Detect CRIU restore in uv__io_poll via /run/criu-restored marker -# - Call uv_loop_fork() to reinitialize epoll backend after restore -# - Fixes segfault in c10d::LibUVStoreDaemon (PyTorch TCPStore) -# -# Build context: the libuv fork source directory -# docker build -t libuv-builder:TAG -f docker/libuv/Dockerfile /path/to/libuv - -FROM ubuntu:22.04 - -RUN apt-get update && apt-get install -y --no-install-recommends \ - build-essential \ - cmake \ - binutils \ - && rm -rf /var/lib/apt/lists/* - -COPY . /libuv-src -WORKDIR /libuv-src - -# Clean any stale build artifacts from host -RUN rm -rf build/ - -RUN mkdir -p build && cd build && \ - cmake -DCMAKE_INSTALL_PREFIX=/usr/local \ - -DCMAKE_BUILD_TYPE=Release \ - -DBUILD_TESTING=OFF \ - .. && \ - make -j$(nproc) && \ - make install && \ - ldconfig - -# Verify CRIU restore detection code is present -RUN nm -D /usr/local/lib/libuv.so.1 | grep -q "uv_loop_fork" || \ - (echo "ERROR: uv_loop_fork symbol not found"; exit 1) - -# Verify GLIBC compatibility with Ubuntu 24.04 (GLIBC 2.39) -RUN echo "=== libuv build info ===" && \ - echo "Version:" && strings /usr/local/lib/libuv.so.1 | grep -E "^1\.[0-9]+\.[0-9]" | head -1 && \ - echo "CRIU restore: enabled (uv__io_poll marker check)" && \ - echo "GLIBC deps:" && \ - objdump -p /usr/local/lib/libuv.so.1.* | grep GLIBC | sort -u && \ - echo "File size:" && ls -la /usr/local/lib/libuv.so.1.* - -# Smoke test: verify the .so loads -RUN ldconfig -p | grep libuv && ldd /usr/local/lib/libuv.so.1 - -CMD ["/bin/bash"] diff --git a/src/compute-plane-services/nvsnap/docker/libzmq/Dockerfile b/src/compute-plane-services/nvsnap/docker/libzmq/Dockerfile deleted file mode 100644 index 61a0dbfbaa..0000000000 --- a/src/compute-plane-services/nvsnap/docker/libzmq/Dockerfile +++ /dev/null @@ -1,61 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# libzmq builder image for NVSNAP -# -# Builds our forked libzmq with CRIU checkpoint/restore support: -# - Epoll reinit after CRIU restore (epoll.cpp) -# - EINTR handling for blocked recv/send (socket_base.cpp) -# - Checkpoint/restore draft API (zmq_ctx_checkpoint, zmq_ctx_restore) -# -# Build context: the libzmq fork source directory -# docker build -t libzmq-builder:TAG -f docker/libzmq/Dockerfile /path/to/libzmq -# -# IMPORTANT: ENABLE_DRAFTS must be ON — the checkpoint API is a draft API. -# Without it, pyzmq can't find zmq_ctx_checkpoint and fails to load. - -FROM ubuntu:22.04 - -RUN apt-get update && apt-get install -y --no-install-recommends \ - build-essential \ - cmake \ - pkg-config \ - libsodium-dev \ - binutils \ - && rm -rf /var/lib/apt/lists/* - -COPY . /libzmq-src -WORKDIR /libzmq-src - -# Clean any stale build artifacts from host -RUN rm -rf build/ - -RUN mkdir -p build && cd build && \ - cmake -DCMAKE_INSTALL_PREFIX=/usr/local \ - -DCMAKE_BUILD_TYPE=Release \ - -DBUILD_TESTS=OFF \ - -DENABLE_DRAFTS=ON \ - .. && \ - make -j$(nproc) && \ - make install && \ - ldconfig - -# Verify checkpoint API symbols exist (draft API) -RUN nm -D /usr/local/lib/libzmq.so | grep -E "zmq_ctx_checkpoint|zmq_ctx_restore|zmq_get_all_contexts" || \ - (echo "ERROR: Checkpoint API symbols not found. Was -DENABLE_DRAFTS=ON set?"; exit 1) - -# Verify no GLIBC dependency newer than 2.35 (Ubuntu 22.04) -RUN ! objdump -p /usr/local/lib/libzmq.so.5.* | grep -E "GLIBC_2\.(3[6-9]|[4-9][0-9])" || \ - (echo "ERROR: libzmq requires GLIBC newer than 2.35 — host build artifacts leaked"; exit 1) - -# Print build info -RUN echo "=== libzmq build info ===" && \ - echo "Version:" && strings /usr/local/lib/libzmq.so | grep -E "^4\.[0-9]\.[0-9]" | head -1 && \ - echo "Draft API: enabled" && \ - echo "Checkpoint symbols:" && \ - nm -D /usr/local/lib/libzmq.so | grep checkpoint && \ - echo "GLIBC deps:" && \ - objdump -p /usr/local/lib/libzmq.so.5.* | grep GLIBC | sort -u && \ - echo "File size:" && ls -la /usr/local/lib/libzmq.so.5.* - -CMD ["/bin/bash"] diff --git a/src/compute-plane-services/nvsnap/docker/pyzmq/Dockerfile b/src/compute-plane-services/nvsnap/docker/pyzmq/Dockerfile deleted file mode 100644 index b758e4501c..0000000000 --- a/src/compute-plane-services/nvsnap/docker/pyzmq/Dockerfile +++ /dev/null @@ -1,69 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# pyzmq wheel builder for NVSNAP -# -# Builds pyzmq linked against our patched libzmq (not bundled stock libzmq). -# The patched libzmq has CRIU checkpoint/restore support. -# -# Build context must contain: -# pyzmq-src/ - the pyzmq fork source -# libzmq-src/ - the libzmq fork source (built first, then pyzmq links against it) -# -# Build: -# BUILD_CTX=$(mktemp -d) -# rsync -a --exclude='.git' --exclude='build/' /path/to/pyzmq/ $BUILD_CTX/pyzmq-src/ -# rsync -a --exclude='.git' --exclude='build/' /path/to/libzmq/ $BUILD_CTX/libzmq-src/ -# docker build -t pyzmq-builder:TAG -f docker/pyzmq/Dockerfile $BUILD_CTX -# -# Output: wheel at /wheels/pyzmq-*.whl -# IMPORTANT: PYZMQ_NO_BUNDLE=1 ensures pyzmq links against system libzmq -# (our patched version), not a bundled copy. - -# Build stage -FROM vllm/vllm-openai:v0.11.2 AS builder - -RUN apt-get update && apt-get install -y --no-install-recommends \ - build-essential \ - cmake \ - pkg-config \ - libsodium-dev \ - && rm -rf /var/lib/apt/lists/* - -RUN pip install --no-cache-dir "cython>=3.0.0" "packaging" "scikit-build-core>=0.10" - -# Build patched libzmq first (pyzmq links against it) -COPY libzmq-src/ /libzmq-src/ -WORKDIR /libzmq-src -RUN rm -rf build/ && mkdir -p build && cd build && \ - cmake -DCMAKE_INSTALL_PREFIX=/usr/local \ - -DCMAKE_BUILD_TYPE=Release \ - -DBUILD_TESTS=OFF \ - -DENABLE_DRAFTS=ON \ - .. && \ - make -j$(nproc) && \ - make install && \ - ldconfig - -# Build pyzmq wheel linked against our libzmq -COPY pyzmq-src/ /pyzmq-src/ -WORKDIR /pyzmq-src -RUN rm -rf build/ dist/ *.egg-info - -# PYZMQ_NO_BUNDLE=1: don't bundle stock libzmq, link against system (ours) -# ZMQ_PREFIX: tell build where to find our libzmq -RUN ZMQ_PREFIX=/usr/local PYZMQ_NO_BUNDLE=1 \ - pip wheel . -w /wheels/ --no-build-isolation - -RUN ls -la /wheels/pyzmq-*.whl || (echo "ERROR: No pyzmq wheel built"; exit 1) - -# Final image: just the wheel -FROM python:3.12-slim -COPY --from=builder /wheels/ /wheels/ - -# Smoke test: install and import -RUN pip install /wheels/pyzmq-*.whl && \ - python3 -c "import zmq; print(f'pyzmq {zmq.__version__} zmq {zmq.zmq_version()} OK')" && \ - pip uninstall -y pyzmq - -CMD ["/bin/bash"] diff --git a/src/compute-plane-services/nvsnap/docker/uvloop/Dockerfile b/src/compute-plane-services/nvsnap/docker/uvloop/Dockerfile deleted file mode 100644 index a1fa6bec27..0000000000 --- a/src/compute-plane-services/nvsnap/docker/uvloop/Dockerfile +++ /dev/null @@ -1,64 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# uvloop wheel builder for NvSnap — multi-Python (cp310/311/312/313) -# -# Builds our forked uvloop with CRIU checkpoint/restore support -# (uv_loop_fork() inside C-level uv_run, reinit libuv after restore). -# -# Builder base: quay.io/pypa/manylinux_2_28_x86_64. -# - AlmaLinux 8 → GLIBC 2.28 -# - ships cp310 / cp311 / cp312 / cp313 under /opt/python/cp3XX-cp3XX/ -# - gcc-toolset-14, autotools, libtool already installed -# -# Output: 4 wheels at /wheels/ named -# uvloop--cpXY-cpXY-linux_x86_64.whl for XY in 310, 311, 312, 313 -# -# Wheels built against GLIBC 2.28 run forward on the GLIBC 2.35 hosts -# we deploy on. arm64 is deferred (see docs/GENERIC-PYTHON-INJECTION-DESIGN.md). -# -# Build context: the uvloop fork source directory -# docker build -t uvloop-builder:TAG -f docker/uvloop/Dockerfile /path/to/uvloop - -FROM quay.io/pypa/manylinux_2_28_x86_64 AS builder - -COPY . /uvloop-src -WORKDIR /uvloop-src - -# Loop over the four interpreters. For each: -# - wipe Cython output so loop.c is regenerated for that Python -# - install matching setuptools / wheel / Cython (no build isolation) -# - pip wheel writes uvloop-*-cpXY-cpXY-linux_x86_64.whl into /wheels/ -RUN set -eux; \ - for py in /opt/python/cp310-cp310 \ - /opt/python/cp311-cp311 \ - /opt/python/cp312-cp312 \ - /opt/python/cp313-cp313; do \ - echo "=== uvloop build: $($py/bin/python --version) ==="; \ - rm -rf build/ dist/ *.egg-info .eggs/ uvloop/loop.c; \ - "$py/bin/pip" install --no-cache-dir "setuptools>=60,<80" wheel "Cython>=3.1,<3.2"; \ - "$py/bin/pip" wheel . -w /wheels/ --no-build-isolation; \ - done - -# Smoke test: each wheel must import under its matching interpreter. -# `cd /tmp` to avoid importing the source tree rather than the installed wheel. -RUN set -eux; \ - for py in /opt/python/cp310-cp310 \ - /opt/python/cp311-cp311 \ - /opt/python/cp312-cp312 \ - /opt/python/cp313-cp313; do \ - tag=$(basename "$py" | cut -d- -f1); \ - whl=$(ls /wheels/uvloop-*-${tag}-${tag}-*.whl); \ - echo "=== smoke test $whl under $tag ==="; \ - "$py/bin/pip" install --force-reinstall --no-deps "$whl"; \ - (cd /tmp && "$py/bin/python" -c "import uvloop, sys; print(f' {sys.version_info.major}.{sys.version_info.minor}: uvloop', uvloop.__version__, 'OK')"); \ - done - -# Wheel count gate — fail loud if any interpreter silently produced nothing -RUN test "$(ls /wheels/uvloop-*.whl | wc -l)" -eq 4 || (ls -la /wheels/; echo 'expected 4 uvloop wheels'; exit 1) - -# Output stage — wheels only, small image -FROM python:3.12-slim -COPY --from=builder /wheels/ /wheels/ -RUN ls -la /wheels/uvloop-*.whl -CMD ["/bin/bash"] diff --git a/src/compute-plane-services/nvsnap/internal/agent/webhook_integration.go b/src/compute-plane-services/nvsnap/internal/agent/webhook_integration.go index 273ae84f73..ac58cdd7c2 100644 --- a/src/compute-plane-services/nvsnap/internal/agent/webhook_integration.go +++ b/src/compute-plane-services/nvsnap/internal/agent/webhook_integration.go @@ -42,12 +42,6 @@ type WebhookConfig struct { KeyFile string // PEM key Path string // default "/mutate" - // AutoInject configures the image refs the webhook stamps into - // auto-injected init containers when a pod carries - // nvsnap.io/auto-inject: "true". Empty fields disable that branch - // (the webhook fails open and admits the pod unchanged). - AutoInject webhook.AutoInjectImages - // L2WaitImage is the nvsnap-l2-wait init-container image ref // (nvsnap#147). When set, restore pods admitted with // nvsnap.io/restore-from get a nvsnap-l2-wait init container that @@ -160,7 +154,6 @@ func (a *Agent) startWebhook(ctx context.Context, cfg WebhookConfig, backend che // behind the agent's HTTP API for future "any GPU node" flows // where an init container fetches the bytes asynchronously. Log: a.log.WithField("subsys", "webhook.mutate"), - AutoInject: cfg.AutoInject, // nvsnap#147: L2 restore gating. When L2WaitImage is set, // the webhook prepends a nvsnap-l2-wait init container that // blocks the main container on pvc_promote_state == "ready". diff --git a/src/compute-plane-services/nvsnap/internal/webhook/BUILD.bazel b/src/compute-plane-services/nvsnap/internal/webhook/BUILD.bazel index bd2d0ef45d..bc0256dca2 100644 --- a/src/compute-plane-services/nvsnap/internal/webhook/BUILD.bazel +++ b/src/compute-plane-services/nvsnap/internal/webhook/BUILD.bazel @@ -4,7 +4,6 @@ go_library( name = "webhook", srcs = [ "admission.go", - "auto_inject.go", "cachedir.go", "cert.go", "extract_coalesce.go", diff --git a/src/compute-plane-services/nvsnap/internal/webhook/auto_inject.go b/src/compute-plane-services/nvsnap/internal/webhook/auto_inject.go deleted file mode 100644 index 4090f49afa..0000000000 --- a/src/compute-plane-services/nvsnap/internal/webhook/auto_inject.go +++ /dev/null @@ -1,200 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -*/ - -package webhook - -import ( - corev1 "k8s.io/api/core/v1" -) - -// AutoInjectAnnotation opts a pod into auto-injection of the NvSnap -// sitecustomize plumbing — emptyDir volume, one init container -// running nvsnap-agent's auto-inject-init.sh (which lays down the -// uvloop wheels + patched libuv/libzmq + libnvsnap_intercept.so + -// sitecustomize.py), and three env vars on the main container -// (PYTHONPATH, LD_LIBRARY_PATH, LD_PRELOAD). -// -// Set "nvsnap.io/auto-inject": "true" on a customer pod and the -// webhook stamps everything in. Lets BYOC pods (anyone's inference -// pod, NVCA tenants, third-party customers) participate in NvSnap -// checkpoint/restore without touching their pod yaml. -// -// Idempotent: if the pod already has a "nvsnap-lib" volume, the -// webhook assumes the operator wired things manually and skips -// auto-injection. -const AutoInjectAnnotation = "nvsnap.io/auto-inject" - -// AutoInjectImages was the four-image config used when the webhook -// fanned out across four init containers. The unified -// auto-inject-init.sh in the nvsnap-agent image now does everything in -// one container — only Agent is consulted. -// -// Kept as a struct (rather than a single string) so future additions -// (e.g., a debug-tools image override) can land without breaking the -// flag surface again. Existing flag plumbing on the agent still -// fills Uvloop/LibUV/LibZMQ but they're ignored. -type AutoInjectImages struct { - // Agent is the nvsnap-agent image ref the webhook injects as the - // single init container. Must match the agent runtime image so - // the libnvsnap_intercept.so build-ID lines up at restore time. - Agent string - - // Deprecated: kept for backward-compat with existing flag plumbing. - // The unified init container in nvsnap-agent already carries these - // payloads, so these fields are not read. - Uvloop string - LibUV string - LibZMQ string -} - -// Valid reports whether the Agent image is set. The other fields are -// deprecated and ignored — only Agent matters. Webhook fails open -// (skips auto-injection) when Agent is empty. -func (i AutoInjectImages) Valid() bool { - return i.Agent != "" -} - -// autoInjectPatches returns the JSON Patch ops needed to inject the -// NvSnap sitecustomize plumbing into pod. Returns nil if auto-inject -// is not requested, already done, or the agent image is unconfigured. -func (m *Mutator) autoInjectPatches(pod *corev1.Pod) []PatchOp { - if pod.Annotations[AutoInjectAnnotation] != "true" { - return nil - } - if !m.AutoInject.Valid() { - m.logger().Warn("auto-inject requested but Agent image unset; admitting pod unchanged") - return nil - } - // If the pod already has a nvsnap-lib volume, the operator wired - // it manually. Don't double-inject. - for i := range pod.Spec.Volumes { - if pod.Spec.Volumes[i].Name == nvsnapLibVolumeName { - return nil - } - } - if m.MainContainer < 0 || m.MainContainer >= len(pod.Spec.Containers) { - m.logger().Warn("auto-inject: MainContainer index out of range; admitting pod unchanged") - return nil - } - - var patches []PatchOp - - // 1. emptyDir volume. - nvsnapLibVol := map[string]any{ - "name": nvsnapLibVolumeName, - "emptyDir": map[string]any{}, - } - if len(pod.Spec.Volumes) == 0 { - patches = append(patches, PatchOp{Op: "add", Path: "/spec/volumes", Value: []any{nvsnapLibVol}}) - } else { - patches = append(patches, PatchOp{Op: "add", Path: "/spec/volumes/-", Value: nvsnapLibVol}) - } - - // 2. One init container running nvsnap-agent's auto-inject-init.sh. - // Replaces the previous four-init-container fan-out (get-uvloop, - // get-libuv, get-libzmq, get-nvsnap) — same effect, fewer container - // startups + one image pull on cold nodes. - initContainer := autoInjectInitContainer(m.AutoInject.Agent) - if len(pod.Spec.InitContainers) == 0 { - patches = append(patches, PatchOp{Op: "add", Path: "/spec/initContainers", Value: []any{initContainer}}) - } else { - patches = append(patches, PatchOp{Op: "add", Path: "/spec/initContainers/-", Value: initContainer}) - } - - // 3. Mount nvsnap-lib on the main container. - mount := map[string]any{ - "name": nvsnapLibVolumeName, - "mountPath": nvsnapLibMountPath, - } - mainPath := func(field string) string { - return "/spec/containers/" + intToStr(m.MainContainer) + "/" + field - } - if len(pod.Spec.Containers[m.MainContainer].VolumeMounts) == 0 { - patches = append(patches, PatchOp{Op: "add", Path: mainPath("volumeMounts"), Value: []any{mount}}) - } else { - patches = append(patches, PatchOp{Op: "add", Path: mainPath("volumeMounts/-"), Value: mount}) - } - - // 4. PYTHONPATH / LD_LIBRARY_PATH / LD_PRELOAD on the main - // container. We skip any env var the customer already set so an - // explicit override (e.g. PATH-style LD_LIBRARY_PATH with extra - // dirs) wins. Customer can always opt out per-var. - existingEnv := map[string]bool{} - for _, e := range pod.Spec.Containers[m.MainContainer].Env { - existingEnv[e.Name] = true - } - envAdds := []map[string]any{ - {"name": "PYTHONPATH", "value": nvsnapLibMountPath + "/sitecustomize"}, - {"name": "LD_LIBRARY_PATH", "value": nvsnapLibMountPath + ":/usr/local/nvidia/lib64:/usr/local/cuda/lib64"}, - {"name": "LD_PRELOAD", "value": nvsnapLibMountPath + "/libnvsnap_intercept.so"}, - } - for i, e := range envAdds { - if existingEnv[e["name"].(string)] { - continue - } - if len(pod.Spec.Containers[m.MainContainer].Env) == 0 && i == 0 { - patches = append(patches, PatchOp{Op: "add", Path: mainPath("env"), Value: []any{e}}) - } else { - patches = append(patches, PatchOp{Op: "add", Path: mainPath("env/-"), Value: e}) - } - } - - return patches -} - -// autoInjectInitContainer returns the single init container spec that -// the webhook stamps into pods carrying nvsnap.io/auto-inject. The -// container runs the agent image's bundled auto-inject-init.sh which -// lays the four payloads under /nvsnap-lib. -func autoInjectInitContainer(agentImage string) map[string]any { - return map[string]any{ - "name": "nvsnap-init", - "image": agentImage, - "imagePullPolicy": "IfNotPresent", - // Override the agent ENTRYPOINT — this image's default is to - // start the nvsnap-agent process; here we just need the - // payload-install script to run and exit. - "command": []any{"/criu-bundle/auto-inject-init.sh"}, - "volumeMounts": []any{ - map[string]any{ - "name": nvsnapLibVolumeName, - "mountPath": nvsnapLibMountPath, - }, - }, - } -} - -const ( - nvsnapLibVolumeName = "nvsnap-lib" - nvsnapLibMountPath = "/nvsnap-lib" -) - -// intToStr is a tiny strconv.Itoa alias so we don't pull strconv into -// every callsite for the one-time MainContainer index format. -func intToStr(n int) string { - if n == 0 { - return "0" - } - // Single-digit cases cover us — MainContainer is almost always 0, - // rarely above 9. Format defensively for two-digit indices. - digits := []byte{} - for n > 0 { - digits = append([]byte{byte('0' + n%10)}, digits...) - n /= 10 - } - return string(digits) -} diff --git a/src/compute-plane-services/nvsnap/internal/webhook/mutate.go b/src/compute-plane-services/nvsnap/internal/webhook/mutate.go index 889254fc32..ace75f07f9 100644 --- a/src/compute-plane-services/nvsnap/internal/webhook/mutate.go +++ b/src/compute-plane-services/nvsnap/internal/webhook/mutate.go @@ -301,12 +301,6 @@ type Mutator struct { // Log is the structured logger; nil disables logging. Log logrus.FieldLogger - // AutoInject configures the image refs used when the webhook - // auto-injects sitecustomize plumbing for pods carrying the - // nvsnap.io/auto-inject: "true" annotation. Empty/zero = the - // auto-inject branch is a no-op (failing open). - AutoInject AutoInjectImages - // OverlayPreparer hands the webhook a per-restore-pod writable // OverlayFS union layered on top of any captured volume — both // rootfs-extract subpaths (nvsnap#194) AND hostPath/emptyDir @@ -412,11 +406,11 @@ func (m *Mutator) Mutate(ctx context.Context, pod *corev1.Pod) ([]PatchOp, error span.SetAttributes(attribute.String("nvsnap.pod", pod.Namespace+"/"+pod.Name)) } - // Auto-inject sitecustomize plumbing first so the restore-from - // branch sees a pod that already has /nvsnap-lib volume + mounts - // + env vars. Both can run on the same pod (auto-injected - // boilerplate + restore mounts). - injectPatches := m.autoInjectPatches(pod) + // Patches that must land before the restore-from branch runs, so + // that branch sees the pod they produce. Only the cachedir capture + // patches below use this now; the LD_PRELOAD auto-injection that + // used to seed it was removed with the interception stack. + var injectPatches []PatchOp raw, ok := pod.Annotations[RestoreFromAnnotation] if !ok || raw == "" { @@ -431,7 +425,7 @@ func (m *Mutator) Mutate(ctx context.Context, pod *corev1.Pod) ([]PatchOp, error if len(injectPatches) > 0 { m.logger().WithField("pod", pod.Namespace+"/"+pod.Name). WithField("patches", len(injectPatches)). - Info("auto-inject only") + Info("capture pod: cache-dir patches only") } return injectPatches, nil } @@ -563,7 +557,7 @@ func (m *Mutator) Mutate(ctx context.Context, pod *corev1.Pod) ([]PatchOp, error "hash": checkpointstore.ShortHash(hash), "pod": pod.Namespace + "/" + pod.Name, "patches": len(patches), - "auto_inject": len(injectPatches), + "pre_patches": len(injectPatches), }).Info("rootfs-only mutation applied") } return patches, nil diff --git a/src/compute-plane-services/nvsnap/internal/webhook/restore_entrypoint.go b/src/compute-plane-services/nvsnap/internal/webhook/restore_entrypoint.go index 80479e4a36..a52c81ec88 100644 --- a/src/compute-plane-services/nvsnap/internal/webhook/restore_entrypoint.go +++ b/src/compute-plane-services/nvsnap/internal/webhook/restore_entrypoint.go @@ -142,6 +142,15 @@ const ( // pre-wired or webhook disabled). Never returns an error today — // the signature carries one so future validation has a place to // land. +// The restore bundle is staged on the host and mounted into the restore pod: +// nvsnap-tools carries the criu bundle, nvsnap-lib the shared libraries. These +// constants lived in the interception stack until it was removed; this is now +// their only user. +const ( + nvsnapLibVolumeName = "nvsnap-lib" + nvsnapLibMountPath = "/nvsnap-lib" +) + func (m *Mutator) restoreBundleInjectPatches(pod *corev1.Pod) ([]PatchOp, error) { if m.MainContainer < 0 || m.MainContainer >= len(pod.Spec.Containers) { return nil, fmt.Errorf("restoreBundleInjectPatches: MainContainer index %d out of range (have %d containers)", diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/.gitignore b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/.gitignore deleted file mode 100644 index 5761abcfdf..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/.gitignore +++ /dev/null @@ -1 +0,0 @@ -*.o diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/BUILD.bazel b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/BUILD.bazel deleted file mode 100644 index 4b27109ee1..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/BUILD.bazel +++ /dev/null @@ -1,79 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -load("@rules_cc//cc:defs.bzl", "cc_binary") - -# libnvsnap_intercept.so: LD_PRELOAD interposition library loaded into -# checkpointed workloads. Source breakdown matches the legacy -# Makefile: -# - OS-level intercepts (C): src/*.c -# - GPU sub-module helpers (C): src/gpu/*.c -# - GPU sub-module core (C++17): src/gpu/*.cpp -# The library uses its own cuda_types.h header so we do not need an -# actual CUDA SDK at build time. libuv_intercept.c is gated by -# ENABLE_LIBUV_INTERCEPT in the Makefile; included unconditionally -# here because rules_cc has no equivalent toggle and the symbol set -# is always shipped. -# -# Single cc_binary(linkshared=True) instead of cc_library + -# cc_shared_library, because rules_cc's cc_shared_library forces -# its dep cc_libraries to build their own intermediate `.so` outputs. -# Those intermediates have versioned glibc symbol refs -# (sigaction@@GLIBC_2.2.5 etc.) that fail to link cleanly without -# the cc_binary's full link recipe (proper -lc, version script, -# etc.). cc_binary linkshared=True is the long-standing Bazel -# pattern for shared libraries and produces exactly the same .so -# layout cc_shared_library does. - -cc_binary( - name = "libnvsnap_intercept.so", - srcs = [ - "src/abort_intercept.c", - "src/cuda_intercept.c", - "src/gpu_checkpoint.c", - "src/init.c", - "src/io_uring_intercept.c", - "src/libuv_intercept.c", - "src/nccl_intercept.c", - "src/quiesce.c", - "src/seccomp_intercept.c", - "src/zmq_intercept.c", - "src/gpu/checkpoint.cpp", - "src/gpu/config.c", - "src/gpu/init.c", - "src/gpu/interpose_cudart.c", - "src/gpu/interpose_cudrv.c", - "src/gpu/interpose_nccl.c", - "src/gpu/metrics.c", - "src/gpu/symbol_table.c", - "src/gpu/tracker.cpp", - ] + glob(["include/**/*.h"]), - additional_linker_inputs = ["libnvsnap.map"], - copts = [ - "-Wall", - "-Wextra", - "-fPIC", - "-O2", - "-D_GNU_SOURCE", - # C++17 only takes effect on .cpp files; gcc silently ignores - # it on .c. rules_cc does not split conlyopts/cxxopts on - # cc_binary so the joint copts list is the simplest option. - "-std=gnu++17", - ], - includes = ["include"], - linkopts = [ - # ld.bfd is the gcc default the upstream Makefile uses; - # Bazel's auto-detected C++ toolchain on ubuntu prefers - # ld.gold which is too strict on versioned glibc symbol refs - # (sigaction@@GLIBC_2.2.5, etc.). Force bfd to match the - # Makefile's link recipe. - "-fuse-ld=bfd", - "-ldl", - "-lpthread", - "-rdynamic", - "-Wl,--version-script,$(location libnvsnap.map)", - ], - linkshared = True, - linkstatic = True, - visibility = ["//visibility:public"], -) diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/Makefile b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/Makefile deleted file mode 100644 index 4c0db83149..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/Makefile +++ /dev/null @@ -1,197 +0,0 @@ -# NVSNAP Interception Library + NvSnap GPU Hypervisor — Merged Build -# -# Single library: libnvsnap_intercept.so -# Contains: io_uring/libuv/ZMQ interception (NvSnap) + GPU allocation tracking, -# NCCL lifecycle, D2H checkpoint save (NvSnap). - -CC = gcc -CXX = g++ - -# Compiler flags -CFLAGS = -Wall -Wextra -fPIC -O2 -g -D_GNU_SOURCE -Iinclude -CXXFLAGS = -Wall -Wextra -fPIC -O2 -g -std=c++17 -D_GNU_SOURCE -Iinclude - -# Linker flags — use g++ for C++ object linking -LDFLAGS = -shared -ldl -lpthread -rdynamic -static-libstdc++ -static-libgcc -Wl,--version-script=libnvsnap.map - -# ── NvSnap source files (C) ────────────────────────────────────────────── -NVSNAP_SRCS = src/init.c \ - src/self_disable.c \ - src/quiesce.c \ - src/io_uring_intercept.c \ - src/seccomp_intercept.c \ - src/zmq_intercept.c \ - src/nccl_intercept.c \ - src/abort_intercept.c \ - src/cuda_intercept.c \ - src/gpu_checkpoint.c - -# ── NvSnap GPU-interposition source files (C + C++17) ───────────────── -# Sub-module under src/gpu/. Logically distinct from the OS-level -# intercept group above (cuda_intercept, io_uring_intercept, -# libuv_intercept, nccl_intercept, quiesce, seccomp_intercept, -# zmq_intercept) — the GPU sub-module hooks CUDA driver/runtime/NCCL -# at the call level for allocation tracking + checkpoint save/restore. -NVSNAP_GPU_C_SRCS = src/gpu/init.c \ - src/gpu/config.c \ - src/gpu/metrics.c \ - src/gpu/interpose_cudart.c \ - src/gpu/interpose_cudrv.c \ - src/gpu/interpose_nccl.c \ - src/gpu/symbol_table.c - -NVSNAP_GPU_CXX_SRCS = src/gpu/tracker.cpp \ - src/gpu/checkpoint.cpp - -ENABLE_LIBUV_INTERCEPT ?= 0 -ifeq ($(ENABLE_LIBUV_INTERCEPT),1) -NVSNAP_SRCS += src/libuv_intercept.c -endif - -# Object files -C_OBJS = $(NVSNAP_SRCS:.c=.o) $(NVSNAP_GPU_C_SRCS:.c=.o) -CXX_OBJS = $(NVSNAP_GPU_CXX_SRCS:.cpp=.o) -OBJS = $(C_OBJS) $(CXX_OBJS) - -# Output library -LIB = libnvsnap_intercept.so - -.PHONY: all clean test test-cuda test-quick install debug help - -all: $(LIB) - -$(LIB): $(OBJS) - $(CXX) $(CFLAGS) $(LDFLAGS) -o $@ $^ - @echo "Built $(LIB) (OS-level intercepts + GPU sub-module merged)" - @$(MAKE) --no-print-directory verify-deps - @$(MAKE) --no-print-directory verify-self-disable - -# This library is force-loaded into EVERY process in the target mount namespace -# via /etc/ld.so.preload -- including CRIU, which runs against the bundled glibc -# rather than the container's. Any dependency beyond libc therefore gets resolved -# from the *container*, and pairing a newer container libstdc++/libm with the -# bundle's older libc fails with "GLIBC_2.xx not found" before CRIU can start. -# The C++ runtime is statically linked (-static-libstdc++ -static-libgcc) to keep -# the DT_NEEDED set at libc alone; libc itself is safe because glibc is backward -# compatible, so a library built here runs against any newer container glibc. -# Guard it so a new dependency cannot reintroduce the failure silently. -ALLOWED_NEEDED = libc.so.6 libdl.so.2 libpthread.so.0 ld-linux-x86-64.so.2 - -.PHONY: verify-deps -verify-deps: $(LIB) - @bad=""; \ - for n in $$(readelf -d $(LIB) 2>/dev/null | awk '/NEEDED/{gsub(/[][]/,"",$$5); print $$5}'); do \ - case " $(ALLOWED_NEEDED) " in *" $$n "*) ;; *) bad="$$bad $$n";; esac; \ - done; \ - if [ -n "$$bad" ]; then \ - echo "ERROR: $(LIB) has disallowed dependencies:$$bad"; \ - echo " It is preloaded into CRIU, which uses the bundled glibc;"; \ - echo " anything outside [$(ALLOWED_NEEDED)] resolves from the target"; \ - echo " container and breaks on any newer-glibc image."; \ - echo " Static-link it (see -static-libstdc++) or drop the dependency."; \ - exit 1; \ - fi; \ - echo "verify-deps: OK (only $(ALLOWED_NEEDED) permitted)" - -# Workloads enable this library through /etc/ld.so.preload, which the loader -# applies to every process in the mount namespace -- including the CRIU that -# nsenters in to dump the container. Loaded into CRIU it wedges the dump (CRIU -# blocks in wait4() before finishing seize), so every constructor must call -# nvsnap_self_disabled() first and bail out. No environment variable can undo -# it later: constructors run before any NVSNAP_* gate is read. Adding a -# constructor without the guard silently reintroduces the hang, so check it -# here rather than discovering it in a customer's checkpoint. -CTOR_SRCS = $(NVSNAP_SRCS) $(NVSNAP_GPU_C_SRCS) src/libuv_intercept.c - -.PHONY: verify-self-disable -verify-self-disable: - @bad=0; \ - for f in $(CTOR_SRCS); do \ - [ -f "$$f" ] || continue; \ - case "$$f" in */self_disable.c) continue;; esac; \ - awk -v F="$$f" ' \ - /^[[:space:]]*[*]/ { next } \ - /^[[:space:]]*\/\// { next } \ - /__attribute__\(\(constructor/ { want=1; L=NR; next } \ - want && /^[[:space:]]*(static[[:space:]]+)?void[[:space:]]+[A-Za-z_].*\(/ { next } \ - want && /^[[:space:]]*\{[[:space:]]*$$/ { next } \ - want && /^[[:space:]]*$$/ { next } \ - want { \ - if ($$0 ~ /if[[:space:]]*\([[:space:]]*nvsnap_self_disabled\(\)[[:space:]]*\)/) { want=0; next } \ - print F ":" L; bad=1; want=0 \ - } \ - END { exit bad?1:0 }' "$$f" || bad=1; \ - done; \ - if [ "$$bad" != "0" ]; then \ - echo "ERROR: the constructors listed above do not call nvsnap_self_disabled()"; \ - echo " as their FIRST statement. They will run inside CRIU (via"; \ - echo " /etc/ld.so.preload) and hang the dump."; \ - echo " Required first line of the body:"; \ - echo " if (nvsnap_self_disabled()) return;"; \ - exit 1; \ - fi; \ - echo "verify-self-disable: OK (guard is the first statement in every constructor)" - -# C compilation -src/%.o: src/%.c - $(CC) $(CFLAGS) -c -o $@ $< - -# C++ compilation -src/%.o: src/%.cpp - $(CXX) $(CXXFLAGS) -c -o $@ $< - -# ── Tests ─────────────────────────────────────────────────────────────── - -test: $(LIB) - @echo "=== Quick load test ===" - NVSNAP_LOG_LEVEL=3 LD_PRELOAD=./$(LIB) /bin/true - @echo "Library loads successfully" - @echo "" - @echo "=== io_uring test ===" - NVSNAP_LOG_LEVEL=4 LD_PRELOAD=./$(LIB) python3 tests/test_uring_simple.py 2>&1 || echo "io_uring test skipped" - @echo "" - @echo "=== uvloop test ===" - NVSNAP_LOG_LEVEL=4 LD_PRELOAD=./$(LIB) python3 tests/test_uvloop_simple.py 2>&1 || echo "uvloop test skipped" - @echo "" - @echo "=== All tests completed ===" - -test-cuda: $(LIB) - @bash tests/run_cuda_tests.sh - -test-quick: $(LIB) tests/test_dlsym_recursion tests/test_library_safety - @echo "=== dlsym recursion test ===" - @NVSNAP_LOG_LEVEL=1 LD_PRELOAD=./$(LIB) ./tests/test_dlsym_recursion 2>/dev/null - @echo "" - @echo "=== Library safety test ===" - @NVSNAP_LOG_LEVEL=1 NVSNAP_QUIESCE_SIGNALS=1 NVSNAP_NCCL_INTERCEPT=1 NVSNAP_CUDA_INTERCEPT=1 \ - LD_PRELOAD=./$(LIB) ./tests/test_library_safety 2>/dev/null - -tests/test_dlsym_recursion: tests/test_dlsym_recursion.c $(LIB) - $(CC) -g -o $@ $< -ldl - -tests/test_library_safety: tests/test_library_safety.c $(LIB) - $(CC) -g -o $@ $< -ldl -lpthread - -clean: - rm -f $(OBJS) $(LIB) - rm -f tests/test_libuv_simple tests/test_libuv_intercept tests/test_cuda_intercept - -install: $(LIB) - install -d $(DESTDIR)/usr/lib - install -m 755 $(LIB) $(DESTDIR)/usr/lib/ - -debug: CFLAGS += -DDEBUG -O0 -debug: CXXFLAGS += -DDEBUG -O0 -debug: $(LIB) - -help: - @echo "NVSNAP + NvSnap Merged Interception Library" - @echo "" - @echo "Targets:" - @echo " all Build libnvsnap_intercept.so" - @echo " test Run all tests" - @echo " test-cuda Run GPU tests (requires GPU)" - @echo " test-quick Quick load test" - @echo " clean Remove built files" - @echo " install Install to system" - @echo " debug Build with debug flags" diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/README.md b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/README.md deleted file mode 100644 index 2c43b87a20..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/README.md +++ /dev/null @@ -1,117 +0,0 @@ -# NVSNAP Interception Library - -An LD_PRELOAD library for intercepting io_uring and libuv to enable CRIU checkpoint/restore of modern async applications. - -## Purpose - -CRIU (Checkpoint/Restore In Userspace) has limited support for: - -- **io_uring**: Kernel-side state, SQPOLL threads, registered buffers -- **libuv/uvloop**: Internal C pointers that become stale after restore - -This library intercepts these subsystems to: - -1. Track io_uring instances and drain them before checkpoint -2. Track libuv loops and reinitialize them after restore -3. Enable checkpoint/restore of applications using uvloop (vLLM, SGLang, etc.) - -**Note**: GPU/CUDA state is handled separately by `cuda-checkpoint` (NVIDIA's tool). This library only handles io_uring and libuv. - -## Building - -```bash -make -``` - -## Usage - -```bash -LD_PRELOAD=/path/to/libnvsnap_intercept.so your_application -``` - -### Environment Variables - -| Variable | Values | Description | -|----------|--------|-------------| -| `NVSNAP_LOG_LEVEL` | 0-5 | 0=off, 1=error, 2=warn, 3=info, 4=debug, 5=trace | -| `NVSNAP_LOG_FILE` | path or "stderr" | Where to write logs | -| `NVSNAP_ENABLED` | 0 or 1 | Disable interception entirely | - -## How It Works - -### io_uring Interception - -```text -Application → io_uring_setup() syscall - ↓ - Our intercept (via syscall hook) - ↓ - Track: fd, sq_entries, cq_entries, flags - ↓ - Before checkpoint: drain all pending I/O - After restore: recreate rings with same params -``` - -### libuv Interception - -```text -Application (uvloop) → uv_loop_init() - ↓ - Our intercept (via dlsym) - ↓ - Track loop pointer - ↓ - After restore detected: call uv_loop_fork() - to reinitialize kernel-side handles -``` - -### Restore Detection - -The library detects restore via a marker file: - -- `restore-entrypoint` creates `/var/run/nvsnap/.restored` -- Library checks for this file and triggers reinitialization - -## Testing - -```bash -# Quick test - verify library loads -make test-quick - -# Test io_uring interception -make test-uring - -# Test uvloop interception -make test-uvloop - -# Test quiescence signal handling -make test-quiesce -``` - -## Integration with NVSNAP - -This library is bundled in the NVSNAP agent image and injected into containers via: - -1. Init container copies `libnvsnap_intercept.so` to a shared volume -2. Source pod sets `LD_PRELOAD` to load the library -3. On checkpoint: CRIU dumps the process, io_uring is drained -4. On restore: Library detects restore marker and reinits libuv loops - -## Files - -```text -lib/nvsnap_intercept/ -├── include/ -│ └── nvsnap_intercept.h # Public API -├── src/ -│ ├── init.c # Initialization, logging -│ ├── quiesce.c # io_uring + libuv tracking -│ ├── io_uring_intercept.c # io_uring syscall hooks -│ └── libuv_intercept.c # libuv function hooks -├── tests/ -│ ├── test_uring_simple.py # io_uring test -│ ├── test_uvloop_simple.py# uvloop test -│ └── ... -├── Makefile -└── README.md -``` diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/checkpoint.h b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/checkpoint.h deleted file mode 100644 index ec43404c6b..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/checkpoint.h +++ /dev/null @@ -1,149 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * NvSnap -- GPU checkpoint/restore API. - * - * Saves and restores GPU state (allocations, streams, events, metadata) - * to/from a directory on disk. Foundation for live GPU migration. - */ -#ifndef NVSNAP_CHECKPOINT_H -#define NVSNAP_CHECKPOINT_H - -#include -#include - -#ifdef __cplusplus -extern "C" { -#endif - -/* ── Checkpoint file format structures ───────────────────────────────── */ - -#define NVSNAP_CHECKPOINT_MAGIC 0x57524150 /* "WRAP" */ -#define NVSNAP_CHECKPOINT_VERSION 1 - -typedef struct { - uint32_t magic; /* 0x57524150 ("WRAP") */ - uint32_t version; /* 1 */ - uint32_t num_allocations; - uint32_t num_streams; - uint32_t num_events; - uint32_t num_nccl_comms; - uint32_t num_vmm_mappings; - uint32_t source_device; /* original GPU device ordinal */ - uint64_t total_gpu_bytes; /* total GPU memory saved */ - uint64_t timestamp; /* checkpoint creation time (ns since epoch) */ -} NvSnapCheckpointHeader; - -typedef struct { - uint64_t ptr; /* original GPU virtual address */ - uint64_t size; /* allocation size */ - int32_t device; /* device ordinal */ - uint32_t alloc_type; /* AllocType enum value */ - uint64_t data_offset; /* offset into gpu_data.bin where content is stored */ - uint32_t seq_num; /* allocation sequence number (for replay ordering) */ - uint32_t pad; /* maintain 8-byte alignment */ -} NvSnapCheckpointAlloc; - -typedef struct { - uint64_t handle; /* stream handle (opaque) */ - int32_t device; - uint32_t flags; -} NvSnapCheckpointStream; - -typedef struct { - uint64_t handle; /* event handle (opaque) */ - int32_t device; - uint32_t flags; -} NvSnapCheckpointEvent; - -typedef struct { - uint64_t comm; /* ncclComm_t handle (opaque) */ - int32_t nranks; - int32_t rank; - uint8_t unique_id[128]; - int32_t device; - uint32_t _pad; -} NvSnapCheckpointNcclComm; - -/* ── API ─────────────────────────────────────────────────────────────── */ - -/* - * Checkpoint GPU state to a directory. - * 1. Quiesces all GPU operations (cudaDeviceSynchronize) - * 2. Saves all tracked GPU allocations (D2H copy) - * 3. Saves metadata (allocation map, streams, events, NCCL comms) - * - * Files created: - * /gpu-/meta.bin -- header + allocation/stream/event records - * /gpu-/gpu_data.bin -- raw GPU memory contents - * - * Returns 0 on success, -1 on error. - */ -int nvsnap_checkpoint_save(const char *checkpoint_dir); - -/* - * Restore GPU state from a checkpoint directory. - * 1. Initializes CUDA driver (cuInit) and acquires fresh GPU context - * 2. Reads metadata from /gpu-/meta.bin - * 3. Reserves original GPU virtual addresses (cuMemAddressReserve) - * - Falls back to any VA if original is unavailable - * 4. Creates physical allocations (cuMemCreate) - * 5. Maps physical to VA (cuMemMap + cuMemSetAccess) - * 6. Copies data from host to GPU (H2D) - * - * Returns 0 on success (all VAs preserved), 1 if some VAs were not preserved, - * -1 on error. - */ -int nvsnap_checkpoint_restore(const char *checkpoint_dir); - -/* - * Restore GPU state for the CURRENT process only. - * Looks for /gpu-/meta.bin. - * - * Use this when called from INSIDE a restored process (e.g., from - * libnvsnap_intercept.so after CRIU restore). Each TP worker restores - * its own GPU — CRIU preserves the original PID so gpu-/ matches. - * - * Returns 0 on success, 1 if VA not preserved, -1 on error. - */ -int nvsnap_checkpoint_restore_self(const char *checkpoint_dir); - -/* - * Pre-checkpoint quiesce: prepare GPU state for checkpoint. - * - * 1. Destroys all tracked NCCL communicators - * 2. Disables P2P access between all GPU pairs - * 3. Synchronizes all GPU devices - * - * Call this BEFORE cuCheckpointProcessLock/Checkpoint. - * Without this, multi-GPU checkpoint hangs on NVLink P2P state. - * - * Returns 0 on success, -1 on error. - */ -int nvsnap_pre_checkpoint_quiesce(void); - -/* - * Post-restore resume: re-enable GPU state after restore. - * - * 1. Re-enables P2P access between all GPU pairs - * - * Call this AFTER cuCheckpointProcessRestore/Unlock. - * - * Returns 0 on success, -1 on error. - */ -int nvsnap_post_restore_resume(void); - -/* - * Query tracked state (for logging/diagnostics). - */ -int nvsnap_get_alloc_count(void); -uint64_t nvsnap_get_total_bytes(void); - -#ifdef __cplusplus -} -#endif - -#endif /* NVSNAP_CHECKPOINT_H */ diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/config.h b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/config.h deleted file mode 100644 index e1390ca202..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/config.h +++ /dev/null @@ -1,39 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * NvSnap — Runtime configuration loaded from environment variables. - */ -#ifndef NVSNAP_CONFIG_H -#define NVSNAP_CONFIG_H - -#include -#include - -#ifdef __cplusplus -extern "C" { -#endif - -typedef struct { - int log_level; /* 0=off .. 4=debug */ - int metrics_enabled; /* 0/1 */ - int fault_injection_enabled; /* 0/1 */ - size_t host_pool_size; /* bytes */ - char agent_socket_path[256]; - double oversubscription_ratio; - int detailed_tracing; /* 0/1 */ -} NvSnapConfig; - -void nvsnap_config_init(void); -const NvSnapConfig *nvsnap_config_get(void); - -/* Utility: parse a human-readable size string ("4G", "512M", "1024"). */ -size_t nvsnap_parse_size(const char *str); - -#ifdef __cplusplus -} -#endif - -#endif /* NVSNAP_CONFIG_H */ diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/cuda_types.h b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/cuda_types.h deleted file mode 100644 index cd440c7dce..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/cuda_types.h +++ /dev/null @@ -1,722 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * NvSnap — Vendored CUDA/NCCL type definitions. - * Allows compilation without the CUDA toolkit installed. - * Guarded so these don't conflict if real CUDA headers are also included. - */ -#ifndef NVSNAP_CUDA_TYPES_H -#define NVSNAP_CUDA_TYPES_H - -#include -#include - -#ifdef __cplusplus -extern "C" { -#endif - -/* ── CUDA Runtime types ──────────────────────────────────────────────── */ - -#ifndef __DRIVER_TYPES_H__ -#ifndef NVSNAP_CUDA_RUNTIME_TYPES_DEFINED -#define NVSNAP_CUDA_RUNTIME_TYPES_DEFINED - -typedef enum { - cudaSuccess = 0, - cudaErrorInvalidValue = 1, - cudaErrorMemoryAllocation = 2, - cudaErrorInitializationError = 3, - cudaErrorLaunchFailure = 719, - cudaErrorECCUncorrectable = 72, - cudaErrorNotReady = 600, - cudaErrorInvalidDevice = 101, - cudaErrorUnknown = 999 -} cudaError_t; - -typedef enum { - cudaMemcpyHostToHost = 0, - cudaMemcpyHostToDevice = 1, - cudaMemcpyDeviceToHost = 2, - cudaMemcpyDeviceToDevice = 3, - cudaMemcpyDefault = 4 -} cudaMemcpyKind; - -typedef void *cudaStream_t; -typedef void *cudaEvent_t; - -typedef struct { - unsigned int x, y, z; -} dim3; - -/* Minimal cudaDeviceProp — enough fields to not crash callers. */ -typedef struct { - char name[256]; - size_t totalGlobalMem; - size_t sharedMemPerBlock; - int regsPerBlock; - int warpSize; - int maxThreadsPerBlock; - int maxThreadsDim[3]; - int maxGridSize[3]; - int clockRate; - int major; - int minor; - size_t totalConstMem; - int multiProcessorCount; - int l2CacheSize; - int maxThreadsPerMultiProcessor; - int computeMode; - /* Pad to a reasonable size to avoid ABI mismatches with real struct. */ - char _pad[4096 - 256 - 7 * sizeof(size_t)]; -} cudaDeviceProp; - -/* ── Stream capture mode ──────────────────────────────────────────────── */ - -typedef enum { - cudaStreamCaptureModeGlobal = 0, - cudaStreamCaptureModeThreadLocal = 1, - cudaStreamCaptureModeRelaxed = 2 -} cudaStreamCaptureMode; - -typedef enum { - cudaStreamCaptureStatusNone = 0, - cudaStreamCaptureStatusActive = 1, - cudaStreamCaptureStatusInvalidated = 2 -} cudaStreamCaptureStatus; - -/* ── Memory pool types ───────────────────────────────────────────────── */ - -typedef void *cudaMemPool_t; - -typedef enum { - cudaMemPoolAttrReuseFollowEventDependencies = 0x1, - cudaMemPoolAttrReuseAllowOpportunistic = 0x2, - cudaMemPoolAttrReuseAllowInternalDependencies = 0x3, - cudaMemPoolAttrReleaseThreshold = 0x4, - cudaMemPoolAttrReservedMemCurrent = 0x5, - cudaMemPoolAttrReservedMemHigh = 0x6, - cudaMemPoolAttrUsedMemCurrent = 0x7, - cudaMemPoolAttrUsedMemHigh = 0x8 -} cudaMemPoolAttr; - -/* Minimal cudaMemPoolProps. */ -typedef struct { - unsigned char _opaque[256]; -} cudaMemPoolProps; - -/* ── CUDA Array types ────────────────────────────────────────────────── */ - -typedef void *cudaArray_t; -typedef const void *cudaArray_const_t; -typedef void *cudaMipmappedArray_t; -typedef const void *cudaMipmappedArray_const_t; - -/* Minimal channel format desc. */ -typedef struct { - int x, y, z, w; - int f; /* cudaChannelFormatKind */ -} cudaChannelFormatDesc; - -/* Minimal cudaExtent. */ -typedef struct { - size_t width, height, depth; -} cudaExtent; - -/* Minimal cudaPitchedPtr. */ -typedef struct { - void *ptr; - size_t pitch; - size_t xsize; - size_t ysize; -} cudaPitchedPtr; - -/* Minimal cudaMemcpy3DParms. */ -typedef struct { - unsigned char _opaque[256]; -} cudaMemcpy3DParms; - -/* Minimal cudaMemcpy3DPeerParms. */ -typedef struct { - unsigned char _opaque[256]; -} cudaMemcpy3DPeerParms; - -/* Minimal cudaPos. */ -typedef struct { - size_t x, y, z; -} cudaPos; - -/* ── Pointer attributes ──────────────────────────────────────────────── */ - -typedef struct { - int type; /* cudaMemoryType */ - int device; - void *devicePointer; - void *hostPointer; - int isManaged; -} cudaPointerAttributes; - -/* ── Memory range attribute ──────────────────────────────────────────── */ - -typedef enum { - cudaMemRangeAttributeReadMostly = 1, - cudaMemRangeAttributePreferredLocation = 2, - cudaMemRangeAttributeAccessedBy = 3, - cudaMemRangeAttributeLastPrefetchLocation = 4 -} cudaMemRangeAttribute; - -/* ── Memory advise ───────────────────────────────────────────────────── */ - -typedef enum { - cudaMemAdviseSetReadMostly = 1, - cudaMemAdviseUnsetReadMostly = 2, - cudaMemAdviseSetPreferredLocation = 3, - cudaMemAdviseUnsetPreferredLocation = 4, - cudaMemAdviseSetAccessedBy = 5, - cudaMemAdviseUnsetAccessedBy = 6 -} cudaMemoryAdvise; - -/* ── Device limit enum ───────────────────────────────────────────────── */ - -typedef enum { - cudaLimitStackSize = 0x00, - cudaLimitPrintfFifoSize = 0x01, - cudaLimitMallocHeapSize = 0x02, - cudaLimitDevRuntimeSyncDepth = 0x03, - cudaLimitDevRuntimePendingLaunchCount = 0x04, - cudaLimitMaxL2FetchGranularity = 0x05 -} cudaLimit; - -/* ── Cache / shared mem config enums ─────────────────────────────────── */ - -typedef enum { - cudaFuncCachePreferNone = 0, - cudaFuncCachePreferShared = 1, - cudaFuncCachePreferL1 = 2, - cudaFuncCachePreferEqual = 3 -} cudaFuncCache; - -typedef enum { - cudaSharedMemBankSizeDefault = 0, - cudaSharedMemBankSizeFourByte = 1, - cudaSharedMemBankSizeEightByte = 2 -} cudaSharedMemConfig; - -/* ── Device P2P attribute ────────────────────────────────────────────── */ - -typedef enum { - cudaDevP2PAttrPerformanceRank = 1, - cudaDevP2PAttrAccessSupported = 2, - cudaDevP2PAttrNativeAtomicSupported = 3, - cudaDevP2PAttrCudaArrayAccessSupported = 4 -} cudaDeviceP2PAttr; - -/* ── IPC types ───────────────────────────────────────────────────────── */ - -typedef struct { - char reserved[64]; -} cudaIpcMemHandle_t; - -typedef struct { - char reserved[64]; -} cudaIpcEventHandle_t; - -/* ── CUDA Graph types ────────────────────────────────────────────────── */ - -typedef void *cudaGraph_t; -typedef void *cudaGraphExec_t; -typedef void *cudaGraphNode_t; - -typedef enum { - cudaGraphNodeTypeKernel = 0, - cudaGraphNodeTypeMemcpy = 1, - cudaGraphNodeTypeMemset = 2, - cudaGraphNodeTypeHost = 3, - cudaGraphNodeTypeGraph = 4, - cudaGraphNodeTypeEmpty = 5, - cudaGraphNodeTypeWaitEvent = 6, - cudaGraphNodeTypeEventRecord = 7, - cudaGraphNodeTypeMemAlloc = 10, - cudaGraphNodeTypeMemFree = 11, - cudaGraphNodeTypeCount = 12 -} cudaGraphNodeType; - -typedef enum { - cudaGraphExecUpdateSuccess = 0, - cudaGraphExecUpdateError = 1, - cudaGraphExecUpdateErrorTopologyChanged = 2, - cudaGraphExecUpdateErrorNodeTypeChanged = 3, - cudaGraphExecUpdateErrorFunctionChanged = 4, - cudaGraphExecUpdateErrorParametersChanged = 5, - cudaGraphExecUpdateErrorNotSupported = 6, - cudaGraphExecUpdateErrorUnsupportedFunctionChange = 7 -} cudaGraphExecUpdateResult; - -/* Minimal graph instantiation params. */ -typedef struct { - unsigned long long flags; - unsigned char _opaque[64]; -} cudaGraphInstantiateParams; - -/* ── Kernel node params ──────────────────────────────────────────────── */ - -typedef struct { - const void *func; - dim3 gridDim; - dim3 blockDim; - void **kernelParams; - size_t sharedMemBytes; - unsigned char _extra[64]; -} cudaKernelNodeParams; - -typedef struct { - void *dst; - size_t pitch; - int value; - cudaExtent extent; - unsigned char _extra[64]; -} cudaMemsetParams; - -/* ── Host function callback ──────────────────────────────────────────── */ - -typedef void (*cudaHostFn_t)(void *userData); - -/* ── Stream callback ─────────────────────────────────────────────────── */ - -typedef void (*cudaStreamCallback_t)(cudaStream_t stream, cudaError_t status, - void *userData); - -/* ── Texture/Surface types ───────────────────────────────────────────── */ - -typedef unsigned long long cudaTextureObject_t; -typedef unsigned long long cudaSurfaceObject_t; - -/* Minimal texture/surface descriptors. */ -typedef struct { unsigned char _opaque[256]; } cudaResourceDesc; -typedef struct { unsigned char _opaque[128]; } cudaTextureDesc; -typedef struct { unsigned char _opaque[128]; } cudaResourceViewDesc; - -/* ── Function attributes ─────────────────────────────────────────────── */ - -typedef struct { - size_t sharedSizeBytes; - size_t constSizeBytes; - size_t localSizeBytes; - int maxThreadsPerBlock; - int numRegs; - int ptxVersion; - int binaryVersion; - int cacheModeCA; - int maxDynamicSharedSizeBytes; - int preferredShmemCarveout; -} cudaFuncAttributes; - -/* ── Function attribute enum ─────────────────────────────────────────── */ - -typedef enum { - cudaFuncAttributeMaxDynamicSharedMemorySize = 8, - cudaFuncAttributePreferredSharedMemoryCarveout = 9, - cudaFuncAttributeMax = 10 -} cudaFuncAttribute; - -/* ── External resource types ─────────────────────────────────────────── */ - -typedef void *cudaExternalMemory_t; -typedef void *cudaExternalSemaphore_t; - -typedef struct { unsigned char _opaque[256]; } cudaExternalMemoryHandleDesc; -typedef struct { unsigned char _opaque[128]; } cudaExternalMemoryBufferDesc; -typedef struct { unsigned char _opaque[128]; } cudaExternalMemoryMipmappedArrayDesc; -typedef struct { unsigned char _opaque[256]; } cudaExternalSemaphoreHandleDesc; -typedef struct { unsigned char _opaque[128]; } cudaExternalSemaphoreSignalParams; -typedef struct { unsigned char _opaque[128]; } cudaExternalSemaphoreWaitParams; - -/* ── Array sparse properties ─────────────────────────────────────────── */ - -typedef struct { unsigned char _opaque[64]; } cudaArraySparseProperties; - -/* ── Device attribute enum ───────────────────────────────────────────── */ - -typedef enum { - cudaDevAttrMaxThreadsPerBlock = 1, - cudaDevAttrMaxBlockDimX = 2, - cudaDevAttrMaxBlockDimY = 3, - cudaDevAttrMaxBlockDimZ = 4, - cudaDevAttrMaxGridDimX = 5, - cudaDevAttrMaxGridDimY = 6, - cudaDevAttrMaxGridDimZ = 7, - cudaDevAttrMaxSharedMemoryPerBlock = 8, - cudaDevAttrWarpSize = 10, - cudaDevAttrMultiProcessorCount = 16, - cudaDevAttrComputeCapabilityMajor = 75, - cudaDevAttrComputeCapabilityMinor = 76 -} cudaDeviceAttr; - -/* ── Graph dependency add/remove ─────────────────────────────────────── */ - -typedef enum { - cudaStreamAddCaptureDependencies = 0, - cudaStreamSetCaptureDependencies = 1 -} cudaStreamUpdateCaptureDependenciesFlags; - -#endif /* NVSNAP_CUDA_RUNTIME_TYPES_DEFINED */ -#endif /* __DRIVER_TYPES_H__ */ - -/* ── CUDA Driver types ───────────────────────────────────────────────── */ - -#ifndef CUDA_H_ -#ifndef NVSNAP_CUDA_DRIVER_TYPES_DEFINED -#define NVSNAP_CUDA_DRIVER_TYPES_DEFINED - -typedef enum { - CUDA_SUCCESS = 0, - CUDA_ERROR_INVALID_VALUE = 1, - CUDA_ERROR_OUT_OF_MEMORY = 2, - CUDA_ERROR_NOT_INITIALIZED = 3, - CUDA_ERROR_DEINITIALIZED = 4, - CUDA_ERROR_INVALID_CONTEXT = 201, - CUDA_ERROR_INVALID_HANDLE = 400, - CUDA_ERROR_NOT_READY = 600, - CUDA_ERROR_ECC_UNCORRECTABLE = 214, - CUDA_ERROR_UNKNOWN = 999 -} CUresult; - -typedef unsigned long long CUdeviceptr; -typedef int CUdevice; -typedef void *CUcontext; -typedef void *CUfunction; -typedef void *CUmodule; -typedef void *CUstream; -typedef void *CUevent; - -/* ── VMM (Virtual Memory Management) types ───────────────────────────── */ - -typedef unsigned long long CUmemGenericAllocationHandle; - -typedef enum { - CU_MEM_HANDLE_TYPE_NONE = 0, - CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR = 1, - CU_MEM_HANDLE_TYPE_WIN32 = 2, - CU_MEM_HANDLE_TYPE_WIN32_KMT = 4 -} CUmemAllocationHandleType; - -typedef enum { - CU_MEM_ALLOCATION_TYPE_INVALID = 0, - CU_MEM_ALLOCATION_TYPE_PINNED = 1 -} CUmemAllocationType; - -typedef enum { - CU_MEM_LOCATION_TYPE_INVALID = 0, - CU_MEM_LOCATION_TYPE_DEVICE = 1 -} CUmemLocationType; - -typedef struct { - CUmemLocationType type; - int id; -} CUmemLocation; - -typedef struct { - CUmemAllocationType type; - CUmemAllocationHandleType requestedHandleTypes; - CUmemLocation location; - void *win32HandleMetaData; - struct { - unsigned char compressionType; - unsigned char gpuDirectRDMACapable; - unsigned short usage; - unsigned char reserved[4]; - } allocFlags; -} CUmemAllocationProp; - -typedef enum { - CU_MEM_ACCESS_FLAGS_PROT_NONE = 0, - CU_MEM_ACCESS_FLAGS_PROT_READ = 1, - CU_MEM_ACCESS_FLAGS_PROT_READWRITE = 3 -} CUmemAccess_flags; - -typedef struct { - CUmemLocation location; - CUmemAccess_flags flags; -} CUmemAccessDesc; - -typedef enum { - CU_MEM_ALLOC_GRANULARITY_MINIMUM = 0, - CU_MEM_ALLOC_GRANULARITY_RECOMMENDED = 1 -} CUmemAllocationGranularity_flags; - -/* ── Additional Driver API opaque types ──────────────────────────────── */ - -typedef void *CUarray; -typedef void *CUmipmappedArray; -typedef void *CUtexref; -typedef void *CUsurfref; -typedef void *CUgraph; -typedef void *CUgraphExec; -typedef void *CUgraphNode; -typedef void *CUexternalMemory; -typedef void *CUexternalSemaphore; -typedef void *CUmemoryPool; -typedef void *CUlinkState; - -/* ── IPC handles ─────────────────────────────────────────────────────── */ - -typedef struct { char reserved[64]; } CUipcMemHandle; -typedef struct { char reserved[64]; } CUipcEventHandle; - -/* ── UUID ────────────────────────────────────────────────────────────── */ - -typedef struct { char bytes[16]; } CUuuid; - -/* ── Device properties (deprecated struct) ───────────────────────────── */ - -typedef struct { - int maxThreadsPerBlock; - int maxThreadsDim[3]; - int maxGridSize[3]; - int sharedMemPerBlock; - int totalConstantMemory; - int SIMDWidth; - int memPitch; - int regsPerBlock; - int clockRate; - int textureAlign; -} CUdevprop; - -/* ── 2D/3D memcpy descriptors ────────────────────────────────────────── */ - -typedef struct { - size_t srcXInBytes, srcY; - CUdeviceptr srcDevice; - const void *srcHost; - CUarray srcArray; - size_t srcPitch; - size_t dstXInBytes, dstY; - CUdeviceptr dstDevice; - void *dstHost; - CUarray dstArray; - size_t dstPitch; - size_t WidthInBytes; - size_t Height; -} CUDA_MEMCPY2D; - -typedef struct { - size_t srcXInBytes, srcY, srcZ; - size_t srcLOD; - CUdeviceptr srcDevice; - const void *srcHost; - CUarray srcArray; - void *reserved0; - size_t srcPitch; - size_t srcHeight; - size_t dstXInBytes, dstY, dstZ; - size_t dstLOD; - CUdeviceptr dstDevice; - void *dstHost; - CUarray dstArray; - void *reserved1; - size_t dstPitch; - size_t dstHeight; - size_t WidthInBytes; - size_t Height; - size_t Depth; -} CUDA_MEMCPY3D; - -typedef struct { - size_t srcXInBytes, srcY, srcZ; - size_t srcLOD; - CUdeviceptr srcDevice; - const void *srcHost; - CUarray srcArray; - CUcontext srcContext; - size_t srcPitch; - size_t srcHeight; - size_t dstXInBytes, dstY, dstZ; - size_t dstLOD; - CUdeviceptr dstDevice; - void *dstHost; - CUarray dstArray; - CUcontext dstContext; - size_t dstPitch; - size_t dstHeight; - size_t WidthInBytes; - size_t Height; - size_t Depth; -} CUDA_MEMCPY3D_PEER; - -/* ── Array descriptors ───────────────────────────────────────────────── */ - -typedef struct { - size_t Width; - size_t Height; - unsigned int Format; - unsigned int NumChannels; -} CUDA_ARRAY_DESCRIPTOR; - -typedef struct { - size_t Width; - size_t Height; - size_t Depth; - unsigned int Format; - unsigned int NumChannels; - unsigned int Flags; -} CUDA_ARRAY3D_DESCRIPTOR; - -/* ── External resource descriptors (opaque) ──────────────────────────── */ - -typedef struct { char _opaque[256]; } CUDA_EXTERNAL_MEMORY_HANDLE_DESC; -typedef struct { char _opaque[128]; } CUDA_EXTERNAL_MEMORY_BUFFER_DESC; -typedef struct { char _opaque[128]; } CUDA_EXTERNAL_MEMORY_MIPMAPPED_ARRAY_DESC; -typedef struct { char _opaque[256]; } CUDA_EXTERNAL_SEMAPHORE_HANDLE_DESC; -typedef struct { char _opaque[128]; } CUDA_EXTERNAL_SEMAPHORE_SIGNAL_PARAMS; -typedef struct { char _opaque[128]; } CUDA_EXTERNAL_SEMAPHORE_WAIT_PARAMS; - -/* ── Graph node params (opaque) ──────────────────────────────────────── */ - -typedef struct { char _opaque[256]; } CUDA_KERNEL_NODE_PARAMS; -typedef struct { char _opaque[128]; } CUDA_MEMSET_NODE_PARAMS; -typedef struct { char _opaque[64]; } CUDA_HOST_NODE_PARAMS; -typedef struct { char _opaque[128]; } CUgraphInstantiateParams; -typedef struct { char _opaque[128]; } CUDA_MEM_ALLOC_NODE_PARAMS; - -/* ── Memory pool ─────────────────────────────────────────────────────── */ - -typedef struct { char _opaque[64]; } CUmemPoolProps; - -/* ── Context creation v3 exec affinity ───────────────────────────────── */ - -typedef struct { char _opaque[16]; } CUexecAffinityParam; - -/* ── Enums used as int in the driver API ─────────────────────────────── */ - -typedef int CUdevice_attribute; -typedef int CUlimit; -typedef int CUfunc_cache; -typedef int CUsharedconfig; -typedef int CUmem_advise; -typedef int CUmem_range_attribute; - -/* ── Pointer attribute enum ──────────────────────────────────────────── */ - -typedef enum { - CU_POINTER_ATTRIBUTE_CONTEXT = 1, - CU_POINTER_ATTRIBUTE_MEMORY_TYPE = 2, - CU_POINTER_ATTRIBUTE_DEVICE_POINTER = 3, - CU_POINTER_ATTRIBUTE_HOST_POINTER = 4, - CU_POINTER_ATTRIBUTE_P2P_TOKENS = 5, - CU_POINTER_ATTRIBUTE_SYNC_MEMOPS = 6, - CU_POINTER_ATTRIBUTE_BUFFER_ID = 7, - CU_POINTER_ATTRIBUTE_IS_MANAGED = 8, - CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL = 9 -} CUpointer_attribute; - -/* ── P2P attribute enum ──────────────────────────────────────────────── */ - -typedef enum { - CU_DEVICE_P2P_ATTRIBUTE_PERFORMANCE_RANK = 0x01, - CU_DEVICE_P2P_ATTRIBUTE_ACCESS_SUPPORTED = 0x02, - CU_DEVICE_P2P_ATTRIBUTE_NATIVE_ATOMIC_SUPPORTED = 0x03, - CU_DEVICE_P2P_ATTRIBUTE_CUDA_ARRAY_ACCESS_SUPPORTED = 0x04 -} CUdevice_P2PAttribute; - -/* ── Stream capture ──────────────────────────────────────────────────── */ - -typedef enum { - CU_STREAM_CAPTURE_MODE_GLOBAL = 0, - CU_STREAM_CAPTURE_MODE_THREAD_LOCAL = 1, - CU_STREAM_CAPTURE_MODE_RELAXED = 2 -} CUstreamCaptureMode; - -typedef enum { - CU_STREAM_CAPTURE_STATUS_NONE = 0, - CU_STREAM_CAPTURE_STATUS_ACTIVE = 1, - CU_STREAM_CAPTURE_STATUS_INVALIDATED = 2 -} CUstreamCaptureStatus; - -/* ── Graph exec update result ────────────────────────────────────────── */ - -typedef enum { - CU_GRAPH_EXEC_UPDATE_SUCCESS = 0x0, - CU_GRAPH_EXEC_UPDATE_ERROR = 0x1, - CU_GRAPH_EXEC_UPDATE_ERROR_TOPOLOGY_CHANGED = 0x2, - CU_GRAPH_EXEC_UPDATE_ERROR_NODE_TYPE_CHANGED = 0x3, - CU_GRAPH_EXEC_UPDATE_ERROR_FUNCTION_CHANGED = 0x4, - CU_GRAPH_EXEC_UPDATE_ERROR_PARAMETERS_CHANGED = 0x5, - CU_GRAPH_EXEC_UPDATE_ERROR_NOT_SUPPORTED = 0x6 -} CUgraphExecUpdateResult; - -/* ── Stream callback / host function ─────────────────────────────────── */ - -typedef void (*CUstreamCallback)(CUstream stream, CUresult status, void *userData); -typedef void (*CUhostFn)(void *userData); - -/* ── Occupancy callback ──────────────────────────────────────────────── */ - -typedef size_t (*CUoccupancyB2DSize)(int blockSize); - -/* ── Function attribute enum ─────────────────────────────────────────── */ - -typedef int CUfunction_attribute; - -#endif /* NVSNAP_CUDA_DRIVER_TYPES_DEFINED */ -#endif /* CUDA_H_ */ - -/* ── NCCL types ──────────────────────────────────────────────────────── */ - -#ifndef NCCL_H_ -#ifndef NVSNAP_NCCL_TYPES_DEFINED -#define NVSNAP_NCCL_TYPES_DEFINED - -typedef enum { - ncclSuccess = 0, - ncclUnhandledCudaError = 1, - ncclSystemError = 2, - ncclInternalError = 3, - ncclInvalidArgument = 4, - ncclInvalidUsage = 5, - ncclNumResults = 6 -} ncclResult_t; - -typedef void *ncclComm_t; - -typedef struct { - char internal[128]; -} ncclUniqueId; - -typedef enum { - ncclInt8 = 0, - ncclChar = 0, - ncclUint8 = 1, - ncclInt32 = 2, - ncclInt = 2, - ncclUint32 = 3, - ncclInt64 = 4, - ncclUint64 = 5, - ncclFloat16 = 6, - ncclHalf = 6, - ncclFloat32 = 7, - ncclFloat = 7, - ncclFloat64 = 8, - ncclDouble = 8, - ncclBfloat16 = 9, - ncclNumTypes = 10 -} ncclDataType_t; - -typedef enum { - ncclSum = 0, - ncclProd = 1, - ncclMax = 2, - ncclMin = 3, - ncclAvg = 4, - ncclNumOps = 5 -} ncclRedOp_t; - -#endif /* NVSNAP_NCCL_TYPES_DEFINED */ -#endif /* NCCL_H_ */ - -#ifdef __cplusplus -} -#endif - -#endif /* NVSNAP_CUDA_TYPES_H */ diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/fault.h b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/fault.h deleted file mode 100644 index 70a2095a6e..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/fault.h +++ /dev/null @@ -1,53 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * NvSnap — Fault injection. - */ -#ifndef NVSNAP_FAULT_H -#define NVSNAP_FAULT_H - -#include - -#ifdef __cplusplus -/* C++ doesn't support C11 _Atomic. The struct fields are only - * accessed atomically from fault.c (pure C). C++ code only sees - * the struct layout for sizeof/offsetof compatibility. */ -#define _Atomic volatile -extern "C" { -#else -#include -#endif - -typedef enum { - NVSNAP_FAULT_OOM = 0, - NVSNAP_FAULT_KERNEL_DROP, - NVSNAP_FAULT_ECC_ERROR, - NVSNAP_FAULT_MEMCPY_LATENCY, - NVSNAP_FAULT_NCCL_FAIL, - NVSNAP_FAULT_NCCL_TIMEOUT, - NVSNAP_FAULT_TYPE_COUNT -} NvSnapFaultType; - -typedef struct { - NvSnapFaultType type; - double probability; /* 0.0 .. 1.0 */ - uint64_t after_count; /* start injecting after N calls */ - _Atomic uint64_t current_count; - _Atomic int enabled; -} NvSnapFaultRule; - -void nvsnap_fault_init(void); -int nvsnap_fault_check(NvSnapFaultType type); /* returns 1 if fault should fire */ -void nvsnap_fault_add_rule(NvSnapFaultType type, double probability, - uint64_t after_count); -void nvsnap_fault_clear(void); - -#ifdef __cplusplus -} -#undef _Atomic -#endif - -#endif /* NVSNAP_FAULT_H */ diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/interpose.h b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/interpose.h deleted file mode 100644 index e8568d6f9c..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/interpose.h +++ /dev/null @@ -1,119 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * NvSnap — Interposition helper macros. - */ -#ifndef NVSNAP_INTERPOSE_H -#define NVSNAP_INTERPOSE_H - -#ifndef _GNU_SOURCE -#define _GNU_SOURCE -#endif - -#include -#include -#include - -/* C11 conflicts with C++ on GCC 11. - * interpose.h doesn't use atomics directly — the .c files that - * need them include themselves. */ - -#ifdef __cplusplus -extern "C" { -#endif - -/* ── Log levels ──────────────────────────────────────────────────────── */ - -enum { - NVSNAP_GPU_LOG_OFF = 0, - NVSNAP_GPU_LOG_ERROR = 1, - NVSNAP_GPU_LOG_WARN = 2, - NVSNAP_GPU_LOG_INFO = 3, - NVSNAP_GPU_LOG_DEBUG = 4 -}; - -extern int nvsnap_gpu_log_level; - -/* - * Thread-safe logging using write() instead of fprintf(). - * fprintf is NOT thread-safe — concurrent writes from 4+ TP workers - * can corrupt FILE* internal buffer → segfault. - * write() is atomic for small writes (= (level), 0)) { \ - nvsnap_log_write((level), fmt, ##__VA_ARGS__); \ - } \ - } while (0) - -#define NVSNAP_GPU_LOG_ERROR(fmt, ...) NVSNAP_GPU_LOG(NVSNAP_GPU_LOG_ERROR, fmt, ##__VA_ARGS__) -#define NVSNAP_GPU_LOG_WARN(fmt, ...) NVSNAP_GPU_LOG(NVSNAP_GPU_LOG_WARN, fmt, ##__VA_ARGS__) -#define NVSNAP_GPU_LOG_INFO(fmt, ...) NVSNAP_GPU_LOG(NVSNAP_GPU_LOG_INFO, fmt, ##__VA_ARGS__) -#define NVSNAP_GPU_LOG_DEBUG(fmt, ...) NVSNAP_GPU_LOG(NVSNAP_GPU_LOG_DEBUG, fmt, ##__VA_ARGS__) - -/* ── Thread-local state ──────────────────────────────────────────────── */ - -#ifdef __cplusplus -extern thread_local int nvsnap_current_device; -extern thread_local int nvsnap_in_graph_capture; -#else -extern _Thread_local int nvsnap_current_device; -extern _Thread_local int nvsnap_in_graph_capture; -#endif - -/* ── Lazy-loading real function pointers ─────────────────────────────── */ - -/* - * Resolve the real function pointer for an intercepted symbol. - * - * First tries RTLD_NEXT (works when the real library is in the link chain). - * If that fails (e.g., NCCL loaded lazily via dlopen by PyTorch), tries - * opening the library explicitly and resolving from that handle. - */ -void *nvsnap_resolve_real(const char *func_name, const char *lib_hint); - -/* In the merged single-library build, dlsym(RTLD_NEXT) goes through our - * dlsym override which returns our OWN hooks → infinite recursion. - * Use nvsnap_resolve_real() which opens libraries explicitly and checks - * dladdr() to ensure the result is NOT from our own library. */ -#define NVSNAP_LOAD_REAL(func_name) \ - do { \ - if (__builtin_expect(real_##func_name == NULL, 0)) { \ - real_##func_name = nvsnap_resolve_real(#func_name, NULL); \ - if (!real_##func_name) { \ - NVSNAP_GPU_LOG_ERROR("resolve failed for " #func_name); \ - } \ - } \ - } while (0) - -/* - * NVSNAP_INTERPOSE — Declares a real-function pointer, provides a wrapper - * function body preamble that lazy-loads the real function. - * - * Usage: - * static cudaError_t (*real_cudaMalloc)(void**, size_t) = NULL; - * cudaError_t cudaMalloc(void **devPtr, size_t size) { - * NVSNAP_LOAD_REAL(cudaMalloc); - * // ... pre-call logic ... - * cudaError_t err = real_cudaMalloc(devPtr, size); - * // ... post-call logic ... - * return err; - * } - * - * For simple pass-through wrappers, use this convenience macro: - */ -#define NVSNAP_DECLARE_REAL(ret_type, func_name, ...) \ - static ret_type (*real_##func_name)(__VA_ARGS__) = NULL - -#ifdef __cplusplus -} -#endif - -#endif /* NVSNAP_INTERPOSE_H */ diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/metrics.h b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/metrics.h deleted file mode 100644 index bfe1d476e6..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/metrics.h +++ /dev/null @@ -1,83 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * NvSnap — Shared-memory ring buffer for zero-syscall metrics. - */ -#ifndef NVSNAP_METRICS_H -#define NVSNAP_METRICS_H - -#include -#include - -#ifdef __cplusplus -#define _Atomic volatile -extern "C" { -#else -#include -#endif - -/* ── Metric types ────────────────────────────────────────────────────── */ - -typedef enum { - NVSNAP_METRIC_ALLOC = 0, - NVSNAP_METRIC_FREE, - NVSNAP_METRIC_KERNEL_LAUNCH, - NVSNAP_METRIC_MEMCPY, - NVSNAP_METRIC_NCCL_COLLECTIVE, - NVSNAP_METRIC_STREAM_CREATE, - NVSNAP_METRIC_STREAM_DESTROY, - NVSNAP_METRIC_EVENT_CREATE, - NVSNAP_METRIC_EVENT_DESTROY, - NVSNAP_METRIC_DEVICE_SYNC, - NVSNAP_METRIC_VMM_MAP, - NVSNAP_METRIC_VMM_UNMAP, - NVSNAP_METRIC_TYPE_COUNT -} NvSnapMetricType; - -/* ── Metrics entry — cache-line aligned (64 bytes) ───────────────────── */ - -typedef struct { - uint32_t type; /* NvSnapMetricType */ - uint32_t pid; - int32_t device; - uint32_t _pad0; - uint64_t timestamp; /* nanoseconds since epoch */ - uint64_t value; /* size in bytes, latency, count, etc. */ - uint64_t ptr; /* relevant pointer / handle */ - uint8_t _pad1[24]; /* pad to 64 bytes total */ -} __attribute__((aligned(64))) NvSnapMetricsEntry; - -/* ── Ring buffer ─────────────────────────────────────────────────────── */ - -#define NVSNAP_METRICS_RING_SIZE 8192 /* must be power of 2 */ - -typedef struct { - _Atomic uint64_t write_pos; - uint8_t _pad_w[56]; /* separate cache lines */ - _Atomic uint64_t read_pos; - uint8_t _pad_r[56]; - _Atomic uint32_t overflow_count; - uint8_t _pad_o[60]; - NvSnapMetricsEntry entries[NVSNAP_METRICS_RING_SIZE]; -} NvSnapMetricsRingBuffer; - -/* ── C API ───────────────────────────────────────────────────────────── */ - -int nvsnap_metrics_init(void); -void nvsnap_metrics_write(NvSnapMetricType type, int device, - uint64_t value, uint64_t ptr); -int nvsnap_metrics_read(NvSnapMetricsEntry *out); -void nvsnap_metrics_destroy(void); - -/* Access the ring buffer directly (e.g., from an agent process). */ -NvSnapMetricsRingBuffer *nvsnap_metrics_get_buffer(void); - -#ifdef __cplusplus -} -#undef _Atomic -#endif - -#endif /* NVSNAP_METRICS_H */ diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/tracker.h b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/tracker.h deleted file mode 100644 index dea42a825f..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap/gpu/tracker.h +++ /dev/null @@ -1,180 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * NvSnap — GPU resource tracker (C++ header). - */ -#ifndef NVSNAP_TRACKER_H -#define NVSNAP_TRACKER_H - -#include -#include - -#ifdef __cplusplus - -#include -#include -#include -#include -#include -#include - -namespace nvsnap { - -enum class AllocType { DEVICE, MANAGED, HOST_PINNED, VMM }; - -struct AllocInfo { - uintptr_t ptr; - size_t size; - int device; - AllocType type; - uint64_t timestamp; /* nanoseconds since epoch */ - uint32_t seq_num; /* monotonic allocation sequence number */ -}; - -struct StreamInfo { - void *handle; - int device; - unsigned flags; - uint64_t kernel_count; -}; - -struct EventInfo { - void *handle; - int device; - unsigned flags; -}; - -struct NcclCommInfo { - void *comm; - int nranks; - int rank; - uint8_t unique_id[128]; - int device; - uint64_t collective_count; - uint64_t bytes_transferred; -}; - -struct VmmMappingInfo { - uintptr_t va; - size_t size; - uint64_t handle; /* CUmemGenericAllocationHandle */ - int device; -}; - -struct GpuStats { - uint64_t total_allocated_bytes; - uint64_t total_kernel_launches; - uint64_t total_memcpy_bytes; - size_t live_alloc_count; - size_t live_stream_count; - size_t live_event_count; - size_t live_nccl_comm_count; - size_t live_vmm_mapping_count; -}; - -class GpuTracker { -public: - static GpuTracker &instance(); - - /* ── Allocations ─────────────────────────────────────────────────── */ - void track_alloc(uintptr_t ptr, size_t size, int device, AllocType type); - bool untrack_alloc(uintptr_t ptr); - bool lookup_alloc(uintptr_t ptr, AllocInfo *out) const; - std::vector snapshot_allocs() const; - - /* ── Streams ─────────────────────────────────────────────────────── */ - void track_stream(void *handle, int device, unsigned flags); - bool untrack_stream(void *handle); - std::vector snapshot_streams() const; - - /* ── Events ──────────────────────────────────────────────────────── */ - void track_event(void *handle, int device, unsigned flags); - bool untrack_event(void *handle); - - /* ── NCCL communicators ──────────────────────────────────────────── */ - void track_nccl_comm(void *comm, int nranks, int rank, - const uint8_t unique_id[128], int device); - bool untrack_nccl_comm(void *comm); - void record_nccl_collective(void *comm, uint64_t bytes); - std::vector snapshot_nccl_comms() const; - - /* ── VMM mappings ────────────────────────────────────────────────── */ - void track_vmm_mapping(uintptr_t va, size_t size, uint64_t handle, int device); - bool untrack_vmm_mapping(uintptr_t va); - std::vector snapshot_vmm_mappings() const; - - /* ── Atomic counters (safe to call from hot path) ────────────────── */ - void add_allocated_bytes(int64_t delta); - void inc_kernel_launches(); - void add_memcpy_bytes(uint64_t bytes); - void inc_stream_kernel_count(void *stream); - - GpuStats get_stats() const; - - /* ── Fork safety ─────────────────────────────────────────────────── */ - /* Reinitialize after fork(). Clears all tracked state and - * reconstructs mutexes (which are in undefined state post-fork). */ - void reset_after_fork(); - -private: - GpuTracker() = default; - GpuTracker(const GpuTracker &) = delete; - GpuTracker &operator=(const GpuTracker &) = delete; - - mutable std::shared_mutex alloc_mu_; - std::unordered_map allocs_; - - mutable std::shared_mutex stream_mu_; - std::unordered_map streams_; - - mutable std::shared_mutex event_mu_; - std::unordered_map events_; - - mutable std::shared_mutex nccl_mu_; - std::unordered_map nccl_comms_; - - mutable std::shared_mutex vmm_mu_; - std::unordered_map vmm_mappings_; - - std::atomic next_seq_num_{0}; - std::atomic total_allocated_bytes_{0}; - std::atomic total_kernel_launches_{0}; - std::atomic total_memcpy_bytes_{0}; -}; - -} /* namespace nvsnap */ - -/* ── C API for use from .c files ─────────────────────────────────────── */ -extern "C" { -#endif /* __cplusplus */ - -void nvsnap_tracker_track_alloc(uintptr_t ptr, size_t size, int device, int type); -void nvsnap_tracker_untrack_alloc(uintptr_t ptr); -void nvsnap_tracker_add_allocated_bytes(int64_t delta); -void nvsnap_tracker_inc_kernel_launches(void); -void nvsnap_tracker_add_memcpy_bytes(uint64_t bytes); - -void nvsnap_tracker_track_stream(void *handle, int device, unsigned flags); -void nvsnap_tracker_untrack_stream(void *handle); -void nvsnap_tracker_inc_stream_kernel_count(void *stream); - -void nvsnap_tracker_track_event(void *handle, int device, unsigned flags); -void nvsnap_tracker_untrack_event(void *handle); - -void nvsnap_tracker_track_nccl_comm(void *comm, int nranks, int rank, - const unsigned char unique_id[128], int device); -void nvsnap_tracker_untrack_nccl_comm(void *comm); -void nvsnap_tracker_record_nccl_collective(void *comm, uint64_t bytes); - -void nvsnap_tracker_track_vmm_mapping(uintptr_t va, size_t size, uint64_t handle, int device); -void nvsnap_tracker_untrack_vmm_mapping(uintptr_t va); -void nvsnap_tracker_reset_after_fork(void); - -#ifdef __cplusplus -} -#endif - -#endif /* NVSNAP_TRACKER_H */ diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap_intercept.h b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap_intercept.h deleted file mode 100644 index 8f0829e631..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/include/nvsnap_intercept.h +++ /dev/null @@ -1,239 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * NVSNAP Interception Library - * - * This library intercepts io_uring and libuv for checkpoint/restore support. - * - * Key features: - * - io_uring: Tracks and drains rings before checkpoint, recreates after restore - * - libuv: Tracks loops and reinitializes handles after restore - * - * Note: GPU/CUDA state is handled externally by cuda-checkpoint (NVIDIA's tool). - * This library focuses on process-level I/O subsystems that CRIU can't handle natively. - */ - -#ifndef NVSNAP_INTERCEPT_H -#define NVSNAP_INTERCEPT_H - -#include -#include -#include -#include -#include - -#ifdef __cplusplus -extern "C" { -#endif - -/* - * ============================================================================= - * CONFIGURATION - * ============================================================================= - */ - -/* Environment variables */ -#define NVSNAP_ENV_LOG_LEVEL "NVSNAP_LOG_LEVEL" /* 0=off, 1=error, 2=warn, 3=info, 4=debug, 5=trace */ -#define NVSNAP_ENV_LOG_FILE "NVSNAP_LOG_FILE" /* Path to log file, or "stderr" */ -#define NVSNAP_ENV_ENABLED "NVSNAP_ENABLED" /* 0 to disable interception */ - -/* Restore detection marker files */ -#define CRIU_RESTORE_MARKER "/run/criu-restored" /* Generic - preferred */ -#define NVSNAP_RESTORE_MARKER "/var/run/nvsnap/.restored" /* Legacy */ - -/* - * ============================================================================= - * LOGGING - * ============================================================================= - */ - -typedef enum { - NVSNAP_LOG_OFF = 0, - NVSNAP_LOG_ERROR = 1, - NVSNAP_LOG_WARN = 2, - NVSNAP_LOG_INFO = 3, - NVSNAP_LOG_DEBUG = 4, - NVSNAP_LOG_TRACE = 5, -} nvsnap_log_level_t; - -void nvsnap_log(nvsnap_log_level_t level, const char* func, const char* fmt, ...); - -#define NVSNAP_ERROR(fmt, ...) nvsnap_log(NVSNAP_LOG_ERROR, __func__, fmt, ##__VA_ARGS__) -#define NVSNAP_WARN(fmt, ...) nvsnap_log(NVSNAP_LOG_WARN, __func__, fmt, ##__VA_ARGS__) -#define NVSNAP_INFO(fmt, ...) nvsnap_log(NVSNAP_LOG_INFO, __func__, fmt, ##__VA_ARGS__) -#define NVSNAP_DEBUG(fmt, ...) nvsnap_log(NVSNAP_LOG_DEBUG, __func__, fmt, ##__VA_ARGS__) -#define NVSNAP_TRACE(fmt, ...) nvsnap_log(NVSNAP_LOG_TRACE, __func__, fmt, ##__VA_ARGS__) - -/* - * ============================================================================= - * GLOBAL STATE - * ============================================================================= - */ - -typedef struct nvsnap_state { - /* Initialization */ - bool initialized; - bool enabled; - pthread_mutex_t init_mutex; - - /* Logging */ - nvsnap_log_level_t log_level; - FILE* log_file; - pthread_mutex_t log_mutex; - -} nvsnap_state_t; - -/* Global state accessor */ -nvsnap_state_t* nvsnap_get_state(void); - -/* - * ============================================================================= - * INITIALIZATION - * ============================================================================= - */ - -/* Called automatically via __attribute__((constructor)) */ -void nvsnap_init(void); -void nvsnap_fini(void); - -/* Manual initialization (for testing) */ -int nvsnap_init_explicit(void); - -/* - * ============================================================================= - * QUIESCENCE API (io_uring and libuv) - * ============================================================================= - * - * These functions are used to track and manage io_uring and libuv instances - * for checkpoint/restore. They are called automatically by our intercepted - * functions, but can also be called manually for debugging. - */ - -/* Track an io_uring instance */ -int nvsnap_track_io_uring(int fd, uint32_t sq_entries, uint32_t cq_entries, - uint32_t flags); -int nvsnap_untrack_io_uring(int fd); - -/* Update mmap addresses for a tracked io_uring (call after mmap) */ -int nvsnap_update_io_uring_addrs(int fd); - -/* Check if io_uring needs reinit after restore */ -int nvsnap_io_uring_needs_reinit(int fd); - -/* Mark io_uring as validated after restore */ -void nvsnap_mark_io_uring_validated(int fd); - -/* Track a libuv loop */ -int nvsnap_track_libuv_loop(void* loop); -int nvsnap_ensure_libuv_loop_ready(void* loop); - -/* Perform quiescence (drain io_uring, prepare libuv) - called before checkpoint */ -int nvsnap_perform_quiescence(void); - -/* Perform post-restore reinitialization */ -void nvsnap_perform_restore_reinit(void); - -/* Dump quiesce state for debugging */ -void nvsnap_dump_quiesce_state(FILE* out); - -/* No-op stubs: uvloop handles uv_loop_fork() natively now (checkpoint-restore-v1 patch) */ -void nvsnap_dump_uvloop_metadata(void); -void nvsnap_install_uvloop_hook_async(void); - -/* - * ============================================================================= - * ZMQ INTERCEPTION - * ============================================================================= - */ - -/* Reinitialize ZMQ contexts if a restore is detected */ -void nvsnap_zmq_reinit_all_if_restored(void); - -/* - * ============================================================================= - * LIBUV INTERCEPTION - * ============================================================================= - * - * Enable libuv interception after restore is detected. - * This allows us to call uv_loop_fork() on restored loops. - */ - -void nvsnap_libuv_enable_interception(void); - -/* - * ============================================================================= - * SECCOMP-BPF INTERCEPTION - * ============================================================================= - * - * Use seccomp-bpf to intercept io_uring syscalls at the kernel boundary. - * This works regardless of static/dynamic linking (critical for uvloop). - * - * Environment variables: - * NVSNAP_SECCOMP_ENABLED=1 - Enable seccomp interception (default: 0) - * NVSNAP_POST_RESTORE=1 - Mark as post-restore for healing - */ - -#define NVSNAP_ENV_SECCOMP_ENABLED "NVSNAP_SECCOMP_ENABLED" -#define NVSNAP_ENV_POST_RESTORE "NVSNAP_POST_RESTORE" - -/* Install seccomp-bpf filter to trap io_uring syscalls */ -int nvsnap_seccomp_install_filter(void); - -/* Enable post-restore mode - io_uring calls may need healing */ -void nvsnap_seccomp_set_post_restore(bool post_restore); - -/* Check if seccomp is installed */ -bool nvsnap_seccomp_is_installed(void); - -/* Get statistics */ -void nvsnap_seccomp_get_stats(int* enter_count, int* heal_count); - -/* - * ============================================================================= - * NCCL INTERCEPTION - * ============================================================================= - */ - -void nvsnap_nccl_quiesce(void); -void nvsnap_nccl_restore(void); -void nvsnap_nccl_atfork_child(void); -void *nvsnap_nccl_symbol_override(const char *symbol); - -/* - * ============================================================================= - * CUDA MEMORY INTERCEPTION - * ============================================================================= - * - * Tracks GPU allocations for multi-GPU checkpoint/restore without cuda-checkpoint. - * Enable: NVSNAP_CUDA_INTERCEPT=1 - */ - -/* Save all live GPU allocations to files in dir (D2H + manifest JSON) */ -int nvsnap_cuda_save(const char *dir); - -/* dlsym override for cuMemAlloc_v2/cuMemFree_v2 interception */ -void *nvsnap_cuda_symbol_override(const char *symbol); - -#ifdef __cplusplus -} -#endif - -/* ZMQ checkpoint/restore support */ -void nvsnap_zmq_handle_checkpoint(void); -void nvsnap_zmq_handle_restore(void); - -/* Non-zero when this library must stay completely inert in the current - * process, because /etc/ld.so.preload force-loaded it into one of our own - * bundle binaries (CRIU and friends). Every constructor must check this - * first. See src/self_disable.c. */ -int nvsnap_self_disabled(void); - -/* Creates the quiesce worker a fork() child could not create in its atfork - * handler (pthread_create is not async-signal-safe and deadlocks there). - * No-op unless a fork left this process without one. */ -void nvsnap_quiesce_worker_restart_if_needed(void); - -#endif /* NVSNAP_INTERCEPT_H */ diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/libnvsnap.map b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/libnvsnap.map deleted file mode 100644 index f95af30a26..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/libnvsnap.map +++ /dev/null @@ -1,5 +0,0 @@ -GLIBC_2.2.5 { - global: - sigaction; - signal; -}; diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/abort_intercept.c b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/abort_intercept.c deleted file mode 100644 index c09e573c8a..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/abort_intercept.c +++ /dev/null @@ -1,100 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * abort() Interception for NVSNAP - * - * Interposes abort() to: - * 1. Print a C backtrace before crashing (diagnostic) - * 2. After CRIU restore, suppress the first abort for a grace period. - * vLLM's monitor_engine_cores calls abort() when the engine core - * doesn't respond within ~1s. After restore, the engine core is - * alive but needs time to resume its heartbeat. Suppressing the - * abort gives it that time. - */ - -#define _GNU_SOURCE -#include -#include -#include -#include -#include -#include - -static void __attribute__((noreturn)) (*real_abort)(void) = NULL; - -/* From quiesce.c */ -extern int nvsnap_is_restored(void); - -/* Signal-safe unsigned-to-string */ -static int uint_to_str(unsigned val, char *buf) -{ - char tmp[16]; - int len = 0; - do { tmp[len++] = '0' + (val % 10); val /= 10; } while (val); - for (int i = 0; i < len; i++) - buf[i] = tmp[len - 1 - i]; - return len; -} - -void __attribute__((noreturn)) abort(void) -{ - if (!real_abort) { - real_abort = dlsym(RTLD_NEXT, "abort"); - if (!real_abort) - _exit(134); - } - - static const char banner[] = "\n=== NVSNAP ABORT INTERCEPTED ===\n"; - (void)write(STDERR_FILENO, banner, sizeof(banner) - 1); - - /* PID (signal-safe) */ - { - char buf[32] = "PID: "; - int pos = 5; - pos += uint_to_str((unsigned)getpid(), buf + pos); - buf[pos++] = '\n'; - (void)write(STDERR_FILENO, buf, pos); - } - - /* C backtrace */ - { - static const char hdr[] = "Backtrace:\n"; - (void)write(STDERR_FILENO, hdr, sizeof(hdr) - 1); - void *frames[64]; - int n = backtrace(frames, 64); - backtrace_symbols_fd(frames, n, STDERR_FILENO); - } - - static const char end[] = "=== END NVSNAP ABORT ===\n\n"; - (void)write(STDERR_FILENO, end, sizeof(end) - 1); - - /* - * Post-restore grace period: suppress the abort and return (longjmp - * style is unsafe here). Instead, just sleep and _exit with a - * distinct code so we can distinguish "suppressed abort" from - * "real abort". The sleep gives the engine core time to resume. - * - * NOTE: abort() is __noreturn. Suppressing it is undefined behavior - * per the C standard. But in practice, vLLM's monitor_engine_cores - * calls abort() from a daemon thread — if we _exit() instead, the - * process dies cleanly without the SIGABRT cascade that kills all - * child processes. - * - * TODO: A better approach would be to intercept the monitor's - * poll() timeout or inject a heartbeat. This is a stopgap. - */ - if (nvsnap_is_restored()) { - static const char msg[] = "[NVSNAP] Post-restore abort suppressed — sleeping 5s for engine core resume\n"; - (void)write(STDERR_FILENO, msg, sizeof(msg) - 1); - struct timespec ts = { .tv_sec = 5, .tv_nsec = 0 }; - nanosleep(&ts, NULL); - static const char msg2[] = "[NVSNAP] Grace period elapsed, proceeding with abort\n"; - (void)write(STDERR_FILENO, msg2, sizeof(msg2) - 1); - } - - real_abort(); - _exit(134); /* unreachable */ -} diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/cuda_intercept.c b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/cuda_intercept.c deleted file mode 100644 index 8053b45d9b..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/cuda_intercept.c +++ /dev/null @@ -1,18 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * Stub — CUDA allocation tracking is now handled by NvSnap's - * nvsnap_interpose_cudart.c + nvsnap_interpose_cudrv.c + nvsnap_tracker.cpp. - * - * This file only exists for build compatibility (nvsnap_cuda_symbol_override - * was previously called from zmq_intercept.c but is no longer used). - */ -#include "nvsnap_intercept.h" - -void *nvsnap_cuda_symbol_override(const char *symbol) { - (void)symbol; - return NULL; -} diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/checkpoint.cpp b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/checkpoint.cpp deleted file mode 100644 index e293d26e33..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/checkpoint.cpp +++ /dev/null @@ -1,1234 +0,0 @@ -/* - * NvSnap -- GPU checkpoint/restore implementation. - * - * Uses REAL CUDA function pointers (not our wrappers) for all GPU - * operations during save/restore so we don't re-enter interposition. - */ -#ifndef _GNU_SOURCE -#define _GNU_SOURCE -#endif - -/* Paths are bounded by PATH_MAX in practice; GCC's static analysis - * can't prove snprintf(PATH_MAX, "%s/meta.bin", PATH_MAX_input) won't - * truncate, but real filesystem paths never approach PATH_MAX. */ -#pragma GCC diagnostic ignored "-Wformat-truncation" - -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "nvsnap/gpu/checkpoint.h" -#include "nvsnap/gpu/cuda_types.h" -#include "nvsnap/gpu/interpose.h" -#include "nvsnap/gpu/tracker.h" - -/* Defined in interpose_cudart.c */ -extern "C" int nvsnap_get_current_device(void); - -/* ═══════════════════════════════════════════════════════════════════════ - * Real CUDA function pointers — resolved lazily via dlsym(RTLD_NEXT). - * These bypass our interposition layer. - * ═══════════════════════════════════════════════════════════════════════ */ - -static cudaError_t (*real_ckpt_cudaDeviceSynchronize)(void) = nullptr; -static cudaError_t (*real_ckpt_cudaMemcpy)(void *, const void *, size_t, - cudaMemcpyKind) = nullptr; -static cudaError_t (*real_ckpt_cudaStreamCreate)(cudaStream_t *) = nullptr; -static cudaError_t (*real_ckpt_cudaEventCreateWithFlags)(cudaEvent_t *, - unsigned) = nullptr; -static CUresult (*real_ckpt_cuInit)(unsigned) = nullptr; -static CUresult (*real_ckpt_cuDevicePrimaryCtxRetain)(CUcontext *, CUdevice) = nullptr; -static CUresult (*real_ckpt_cuMemAddressReserve)(CUdeviceptr *, size_t, - size_t, CUdeviceptr, - unsigned long long) = nullptr; -static CUresult (*real_ckpt_cuMemCreate)(CUmemGenericAllocationHandle *, - size_t, const CUmemAllocationProp *, - unsigned long long) = nullptr; -static CUresult (*real_ckpt_cuMemMap)(CUdeviceptr, size_t, size_t, - CUmemGenericAllocationHandle, - unsigned long long) = nullptr; -static CUresult (*real_ckpt_cuMemSetAccess)(CUdeviceptr, size_t, - const CUmemAccessDesc *, - size_t) = nullptr; -static CUresult (*real_ckpt_cuMemGetAllocationGranularity)( - size_t *, const CUmemAllocationProp *, - CUmemAllocationGranularity_flags) = nullptr; -static cudaError_t (*real_ckpt_cudaMalloc)(void **, size_t) = nullptr; -static cudaError_t (*real_ckpt_cudaFree)(void *) = nullptr; -static cudaError_t (*real_ckpt_cudaSetDevice)(int) = nullptr; -static const char *(*real_ckpt_cudaGetErrorString)(cudaError_t) = nullptr; - -static void force_resolve_real_functions(void); - -static void ensure_real_functions(void) -{ - if (real_ckpt_cudaDeviceSynchronize) - return; - - /* Resolve ALL function pointers with nvsnap_resolve_real fallback. - * dlsym(RTLD_NEXT) fails for dynamically-loaded libraries (libcuda.so - * loaded via dlopen by PyTorch). nvsnap_resolve_real searches all - * loaded libraries and verifies the result isn't our own wrapper. */ -/* - * Resolve REAL CUDA function pointers, bypassing our own wrappers. - * - * After CRIU restore, dlsym(RTLD_NEXT) can return our own interposed - * functions (broken state). nvsnap_resolve_real() explicitly opens - * the target library and uses dladdr() to verify the result isn't us. - * - * Try the explicit library first (most reliable), RTLD_NEXT as fallback. - */ -#define RESOLVE_REAL(var, type, name) do { \ - var = (type)nvsnap_resolve_real(name, "libcudart.so"); \ - if (!var) var = (type)nvsnap_resolve_real(name, "libcuda.so.1"); \ - if (!var) var = (type)dlsym(RTLD_NEXT, name); \ -} while (0) - - RESOLVE_REAL(real_ckpt_cudaDeviceSynchronize, - cudaError_t (*)(void), "cudaDeviceSynchronize"); - RESOLVE_REAL(real_ckpt_cudaMemcpy, - cudaError_t (*)(void *, const void *, size_t, cudaMemcpyKind), "cudaMemcpy"); - RESOLVE_REAL(real_ckpt_cudaStreamCreate, - cudaError_t (*)(cudaStream_t *), "cudaStreamCreate"); - RESOLVE_REAL(real_ckpt_cudaEventCreateWithFlags, - cudaError_t (*)(cudaEvent_t *, unsigned), "cudaEventCreateWithFlags"); - RESOLVE_REAL(real_ckpt_cuInit, - CUresult (*)(unsigned), "cuInit"); - RESOLVE_REAL(real_ckpt_cuDevicePrimaryCtxRetain, - CUresult (*)(CUcontext *, CUdevice), "cuDevicePrimaryCtxRetain"); - - RESOLVE_REAL(real_ckpt_cuMemAddressReserve, - CUresult (*)(CUdeviceptr *, size_t, size_t, CUdeviceptr, unsigned long long), - "cuMemAddressReserve"); - RESOLVE_REAL(real_ckpt_cuMemCreate, - CUresult (*)(CUmemGenericAllocationHandle *, size_t, const CUmemAllocationProp *, unsigned long long), - "cuMemCreate"); - RESOLVE_REAL(real_ckpt_cuMemMap, - CUresult (*)(CUdeviceptr, size_t, size_t, CUmemGenericAllocationHandle, unsigned long long), - "cuMemMap"); - RESOLVE_REAL(real_ckpt_cuMemSetAccess, - CUresult (*)(CUdeviceptr, size_t, const CUmemAccessDesc *, size_t), - "cuMemSetAccess"); - RESOLVE_REAL(real_ckpt_cuMemGetAllocationGranularity, - CUresult (*)(size_t *, const CUmemAllocationProp *, CUmemAllocationGranularity_flags), - "cuMemGetAllocationGranularity"); - RESOLVE_REAL(real_ckpt_cudaMalloc, - cudaError_t (*)(void **, size_t), "cudaMalloc"); - RESOLVE_REAL(real_ckpt_cudaFree, - cudaError_t (*)(void *), "cudaFree"); - RESOLVE_REAL(real_ckpt_cudaSetDevice, - cudaError_t (*)(int), "cudaSetDevice"); - RESOLVE_REAL(real_ckpt_cudaGetErrorString, - const char *(*)(cudaError_t), "cudaGetErrorString"); - -#undef RESOLVE_REAL -} - -/* - * Force re-resolution of all CUDA function pointers. - * - * After CRIU restore, the cached pointers from the original run are - * stale — CUDA libraries may be at different addresses (ASLR). - * This clears all cached pointers and re-resolves via dlsym. - */ -static void force_resolve_real_functions(void) -{ - real_ckpt_cudaDeviceSynchronize = nullptr; - real_ckpt_cudaMemcpy = nullptr; - real_ckpt_cudaStreamCreate = nullptr; - real_ckpt_cudaEventCreateWithFlags = nullptr; - real_ckpt_cuInit = nullptr; - real_ckpt_cuDevicePrimaryCtxRetain = nullptr; - real_ckpt_cuMemAddressReserve = nullptr; - real_ckpt_cuMemCreate = nullptr; - real_ckpt_cuMemMap = nullptr; - real_ckpt_cuMemSetAccess = nullptr; - real_ckpt_cuMemGetAllocationGranularity = nullptr; - real_ckpt_cudaMalloc = nullptr; - real_ckpt_cudaGetErrorString = nullptr; - real_ckpt_cudaFree = nullptr; - real_ckpt_cudaSetDevice = nullptr; - - ensure_real_functions(); - - int resolved = 0; - if (real_ckpt_cudaDeviceSynchronize) resolved++; - if (real_ckpt_cudaMemcpy) resolved++; - if (real_ckpt_cuInit) resolved++; - if (real_ckpt_cuDevicePrimaryCtxRetain) resolved++; - if (real_ckpt_cuMemAddressReserve) resolved++; - if (real_ckpt_cuMemCreate) resolved++; - if (real_ckpt_cuMemMap) resolved++; - if (real_ckpt_cuMemSetAccess) resolved++; - if (real_ckpt_cuMemGetAllocationGranularity) resolved++; - if (real_ckpt_cudaMalloc) resolved++; - if (real_ckpt_cudaFree) resolved++; - if (real_ckpt_cudaSetDevice) resolved++; - - /* Log where cudaMalloc resolved from — critical for debugging self-resolution */ - if (real_ckpt_cudaMalloc) { - Dl_info minfo = {}; - if (dladdr((void *)real_ckpt_cudaMalloc, &minfo)) { - NVSNAP_GPU_LOG_INFO("checkpoint: cudaMalloc resolved from %s", - minfo.dli_fname ? minfo.dli_fname : "unknown"); - } - } - - NVSNAP_GPU_LOG_INFO("checkpoint: re-resolved %d/12 CUDA function pointers " - "(cuInit=%p, cudaMalloc=%p, cuMemAddressReserve=%p)", - resolved, - (void *)real_ckpt_cuInit, - (void *)real_ckpt_cudaMalloc, - (void *)real_ckpt_cuMemAddressReserve); -} - -/* ═══════════════════════════════════════════════════════════════════════ - * Helpers - * ═══════════════════════════════════════════════════════════════════════ */ - -static uint64_t now_ns(void) -{ - struct timespec ts; - clock_gettime(CLOCK_REALTIME, &ts); - return (uint64_t)ts.tv_sec * 1000000000ULL + (uint64_t)ts.tv_nsec; -} - -static int write_all(const void *buf, size_t size, FILE *fp) -{ - return fwrite(buf, 1, size, fp) == size ? 0 : -1; -} - -static int mkdirs(const char *path) -{ - /* Try mkdir; if it fails because parent doesn't exist, create parents. */ - if (mkdir(path, 0755) == 0 || errno == EEXIST) - return 0; - - /* Simple recursive mkdir. */ - char tmp[PATH_MAX]; - snprintf(tmp, sizeof(tmp), "%s", path); - for (char *p = tmp + 1; *p; p++) { - if (*p == '/') { - *p = '\0'; - mkdir(tmp, 0755); - *p = '/'; - } - } - return mkdir(tmp, 0755) == 0 || errno == EEXIST ? 0 : -1; -} - -/* ═══════════════════════════════════════════════════════════════════════ - * Checkpoint Save - * ═══════════════════════════════════════════════════════════════════════ */ - -/* - * Reset the deferred restore state so that after CRIU checkpoint+restore, - * the restored process will re-check for the marker file. - * Defined in interpose_cudart.c. - */ -extern void nvsnap_reset_restore_state(void); - -extern "C" int nvsnap_checkpoint_save(const char *checkpoint_dir) -{ - if (!checkpoint_dir) { - NVSNAP_GPU_LOG_ERROR("checkpoint_save: NULL directory"); - return -1; - } - - ensure_real_functions(); - - /* Build per-process save directory: /gpu-/ */ - char save_dir[PATH_MAX]; - snprintf(save_dir, sizeof(save_dir), "%s/gpu-%d", - checkpoint_dir, (int)getpid()); - - /* Early exit: if this process has no GPU allocations, skip save entirely. - * Non-GPU processes (uvicorn workers, Python threads) call this but have - * nothing to save. Calling cudaDeviceSynchronize in a non-GPU process - * either hangs (CUDA not initialized) or wastes time. */ - { - auto &tracker = nvsnap::GpuTracker::instance(); - if (tracker.snapshot_allocs().empty()) { - NVSNAP_GPU_LOG_INFO("checkpoint: 0 allocations tracked, skipping save (pid=%d)", - (int)getpid()); - return 0; - } - } - - NVSNAP_GPU_LOG_INFO("checkpoint: saving to %s", save_dir); - - /* Create output directory. */ - if (mkdirs(save_dir) != 0) { - NVSNAP_GPU_LOG_ERROR("checkpoint: cannot create directory %s: %s", - save_dir, strerror(errno)); - return -1; - } - - /* 1. Quiesce all GPU operations. */ - if (real_ckpt_cudaDeviceSynchronize) { - cudaError_t err = real_ckpt_cudaDeviceSynchronize(); - if (err != cudaSuccess) { - NVSNAP_GPU_LOG_ERROR("checkpoint: cudaDeviceSynchronize failed: %d", - (int)err); - return -1; - } - } - - /* 2. Snapshot tracker state. */ - auto &tracker = nvsnap::GpuTracker::instance(); - auto allocs = tracker.snapshot_allocs(); - auto streams = tracker.snapshot_streams(); - auto nccl_comms = tracker.snapshot_nccl_comms(); - auto stats = tracker.get_stats(); - - /* 3. Open output files. */ - char meta_path[PATH_MAX], data_path[PATH_MAX]; - snprintf(meta_path, sizeof(meta_path), "%s/meta.bin", save_dir); - snprintf(data_path, sizeof(data_path), "%s/gpu_data.bin", save_dir); - - FILE *meta_fp = fopen(meta_path, "wb"); - if (!meta_fp) { - NVSNAP_GPU_LOG_ERROR("checkpoint: cannot open %s: %s", meta_path, - strerror(errno)); - return -1; - } - - FILE *data_fp = fopen(data_path, "wb"); - if (!data_fp) { - NVSNAP_GPU_LOG_ERROR("checkpoint: cannot open %s: %s", data_path, - strerror(errno)); - fclose(meta_fp); - return -1; - } - - /* 4. Prepare and write header. */ - NvSnapCheckpointHeader header = {}; - header.magic = NVSNAP_CHECKPOINT_MAGIC; - header.version = NVSNAP_CHECKPOINT_VERSION; - header.num_allocations = (uint32_t)allocs.size(); - header.num_streams = (uint32_t)streams.size(); - header.num_events = (uint32_t)stats.live_event_count; - header.num_nccl_comms = (uint32_t)nccl_comms.size(); - header.num_vmm_mappings = (uint32_t)stats.live_vmm_mapping_count; - /* Use the device from the first allocation, not nvsnap_current_device - * which may be wrong if save is called from a non-CUDA thread. */ - header.source_device = allocs.empty() ? (uint32_t)nvsnap_current_device - : (uint32_t)allocs[0].device; - header.timestamp = now_ns(); - - /* Calculate total GPU bytes. */ - uint64_t total_bytes = 0; - for (auto &a : allocs) { - /* Only save DEVICE and MANAGED allocations (skip HOST_PINNED). */ - if (a.type == nvsnap::AllocType::DEVICE || - a.type == nvsnap::AllocType::MANAGED) { - total_bytes += a.size; - } - } - header.total_gpu_bytes = total_bytes; - - if (write_all(&header, sizeof(header), meta_fp) != 0) { - NVSNAP_GPU_LOG_ERROR("checkpoint: failed to write header"); - fclose(meta_fp); - fclose(data_fp); - return -1; - } - - /* 5. For each allocation, save GPU data and build metadata records. */ - std::vector alloc_recs(allocs.size()); - uint64_t data_offset = 0; - for (size_t i = 0; i < allocs.size(); i++) { - auto &a = allocs[i]; - NvSnapCheckpointAlloc &rec = alloc_recs[i]; - memset(&rec, 0, sizeof(rec)); - rec.ptr = (uint64_t)a.ptr; - rec.size = (uint64_t)a.size; - rec.device = (int32_t)a.device; - rec.alloc_type = (uint32_t)a.type; - rec.seq_num = a.seq_num; - - bool save_data = (a.type == nvsnap::AllocType::DEVICE || - a.type == nvsnap::AllocType::MANAGED); - - if (save_data) { - rec.data_offset = data_offset; - - /* Allocate host staging buffer and copy D2H. */ - void *host_buf = malloc(a.size); - if (!host_buf) { - NVSNAP_GPU_LOG_ERROR("checkpoint: malloc(%zu) failed for ptr 0x%lx", - a.size, (unsigned long)a.ptr); - fclose(meta_fp); - fclose(data_fp); - return -1; - } - - if (real_ckpt_cudaMemcpy) { - cudaError_t err = real_ckpt_cudaMemcpy( - host_buf, (const void *)a.ptr, a.size, - cudaMemcpyDeviceToHost); - if (err != cudaSuccess) { - NVSNAP_GPU_LOG_ERROR( - "checkpoint: cudaMemcpy D2H failed for ptr 0x%lx: %d", - (unsigned long)a.ptr, (int)err); - free(host_buf); - fclose(meta_fp); - fclose(data_fp); - return -1; - } - } - - if (write_all(host_buf, a.size, data_fp) != 0) { - NVSNAP_GPU_LOG_ERROR("checkpoint: failed to write gpu data"); - free(host_buf); - fclose(meta_fp); - fclose(data_fp); - return -1; - } - free(host_buf); - - data_offset += a.size; - } else { - rec.data_offset = UINT64_MAX; /* sentinel: no GPU data */ - } - } - - /* Sort allocation records by seq_num for deterministic replay order. */ - std::sort(alloc_recs.begin(), alloc_recs.end(), - [](const NvSnapCheckpointAlloc &a, - const NvSnapCheckpointAlloc &b) { - return a.seq_num < b.seq_num; - }); - - for (auto &rec : alloc_recs) - (void)write_all(&rec, sizeof(rec), meta_fp); - - /* 6. Write stream metadata. */ - for (auto &s : streams) { - NvSnapCheckpointStream srec = {}; - srec.handle = (uint64_t)(uintptr_t)s.handle; - srec.device = (int32_t)s.device; - srec.flags = s.flags; - (void)write_all(&srec, sizeof(srec), meta_fp); - } - - /* 7. Write NCCL comm metadata. */ - for (auto &c : nccl_comms) { - NvSnapCheckpointNcclComm crec = {}; - crec.comm = (uint64_t)(uintptr_t)c.comm; - crec.nranks = (int32_t)c.nranks; - crec.rank = (int32_t)c.rank; - memcpy(crec.unique_id, c.unique_id, 128); - crec.device = (int32_t)c.device; - (void)write_all(&crec, sizeof(crec), meta_fp); - } - - fclose(meta_fp); - fclose(data_fp); - - NVSNAP_GPU_LOG_INFO("checkpoint: saved %u allocs (%lu bytes), " - "%u streams, %u NCCL comms to %s", - header.num_allocations, - (unsigned long)header.total_gpu_bytes, - header.num_streams, header.num_nccl_comms, - save_dir); - - return 0; -} - -/* ═══════════════════════════════════════════════════════════════════════ - * Restore helpers - * ═══════════════════════════════════════════════════════════════════════ */ - -/* - * VMM-based restore (fallback). Uses cuMemAddressReserve + cuMemCreate + cuMemMap. - * Returns: - * 0 — restored at original VA - * 1 — restored at fallback VA (VA not preserved) - * -1 — failed to restore - */ -__attribute__((unused)) -static int restore_one_alloc_vmm(NvSnapCheckpointAlloc *rec, FILE *data_fp, - nvsnap::GpuTracker &tracker) -{ - CUmemAllocationProp prop = {}; - prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; - prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; - prop.location.id = rec->device; - - NVSNAP_GPU_LOG_DEBUG("restore_one_alloc: ptr=0x%lx size=%lu device=%d", - (unsigned long)rec->ptr, (unsigned long)rec->size, - rec->device); - - size_t granularity = 0; - CUresult gres = real_ckpt_cuMemGetAllocationGranularity( - &granularity, &prop, CU_MEM_ALLOC_GRANULARITY_MINIMUM); - if (gres != CUDA_SUCCESS || granularity == 0) - granularity = 2 * 1024 * 1024; /* 2 MiB default */ - - /* Round size up to granularity. */ - size_t alloc_size = - ((rec->size + granularity - 1) / granularity) * granularity; - - /* Try to reserve at the original VA first. */ - CUdeviceptr reserved_va = 0; - int va_preserved = 1; - CUresult res = real_ckpt_cuMemAddressReserve( - &reserved_va, alloc_size, granularity, (CUdeviceptr)rec->ptr, 0); - if (res != CUDA_SUCCESS) { - NVSNAP_GPU_LOG_WARN( - "checkpoint: cuMemAddressReserve at original VA 0x%lx (size %lu) " - "failed: %d, trying fallback VA", - (unsigned long)rec->ptr, (unsigned long)rec->size, (int)res); - - /* Fallback: reserve at any available VA. */ - reserved_va = 0; - res = real_ckpt_cuMemAddressReserve( - &reserved_va, alloc_size, granularity, 0, 0); - if (res != CUDA_SUCCESS) { - NVSNAP_GPU_LOG_ERROR( - "checkpoint: cuMemAddressReserve fallback also failed for " - "0x%lx (size %lu): %d", - (unsigned long)rec->ptr, (unsigned long)rec->size, (int)res); - return -1; - } - NVSNAP_GPU_LOG_WARN( - "checkpoint: alloc 0x%lx restored at fallback VA 0x%lx " - "(VA NOT preserved)", - (unsigned long)rec->ptr, (unsigned long)reserved_va); - va_preserved = 0; - } - - CUmemGenericAllocationHandle handle = 0; - res = real_ckpt_cuMemCreate(&handle, alloc_size, &prop, 0); - if (res != CUDA_SUCCESS) { - NVSNAP_GPU_LOG_ERROR("checkpoint: cuMemCreate failed for 0x%lx " - "(size=%lu, device=%d): %d", - (unsigned long)rec->ptr, (unsigned long)alloc_size, - rec->device, (int)res); - return -1; - } - NVSNAP_GPU_LOG_DEBUG("checkpoint: cuMemCreate OK handle=0x%lx size=%lu", - (unsigned long)handle, (unsigned long)alloc_size); - - res = real_ckpt_cuMemMap(reserved_va, alloc_size, 0, handle, 0); - if (res != CUDA_SUCCESS) { - NVSNAP_GPU_LOG_ERROR("checkpoint: cuMemMap(va=0x%lx, size=%lu, handle=0x%lx) " - "failed: %d", - (unsigned long)reserved_va, (unsigned long)alloc_size, - (unsigned long)handle, (int)res); - return -1; - } - NVSNAP_GPU_LOG_DEBUG("checkpoint: cuMemMap OK va=0x%lx", (unsigned long)reserved_va); - - CUmemAccessDesc access = {}; - access.location.type = CU_MEM_LOCATION_TYPE_DEVICE; - access.location.id = rec->device; - access.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; - res = real_ckpt_cuMemSetAccess(reserved_va, alloc_size, &access, 1); - if (res != CUDA_SUCCESS) { - NVSNAP_GPU_LOG_ERROR( - "checkpoint: cuMemSetAccess(va=0x%lx, size=%lu, device=%d) failed: %d", - (unsigned long)reserved_va, (unsigned long)alloc_size, - rec->device, (int)res); - return -1; - } - NVSNAP_GPU_LOG_DEBUG("checkpoint: cuMemSetAccess OK va=0x%lx", (unsigned long)reserved_va); - - /* Copy data from file to GPU. */ - void *host_buf = malloc(rec->size); - if (!host_buf) { - NVSNAP_GPU_LOG_ERROR("checkpoint: malloc(%lu) failed", - (unsigned long)rec->size); - return -1; - } - - if (fseek(data_fp, (long)rec->data_offset, SEEK_SET) != 0 || - fread(host_buf, 1, rec->size, data_fp) != rec->size) { - NVSNAP_GPU_LOG_ERROR("checkpoint: short read on gpu_data.bin"); - free(host_buf); - return -1; - } - - if (real_ckpt_cudaMemcpy) { - cudaError_t merr = real_ckpt_cudaMemcpy( - (void *)reserved_va, host_buf, rec->size, - cudaMemcpyHostToDevice); - if (merr != cudaSuccess) { - const char *errstr = (real_ckpt_cudaGetErrorString) - ? real_ckpt_cudaGetErrorString(merr) - : ""; - NVSNAP_GPU_LOG_ERROR( - "checkpoint: H2D memcpy failed for 0x%lx: %d (%s)", - (unsigned long)rec->ptr, (int)merr, errstr ? errstr : ""); - free(host_buf); - return -1; - } - } - free(host_buf); - - /* Re-register in tracker. */ - tracker.track_alloc((uintptr_t)reserved_va, rec->size, rec->device, - static_cast(rec->alloc_type)); - - return va_preserved ? 0 : 1; -} - -/* - * cudaMalloc-replay restore (primary path, post-CRIU). - * After CRIU restore, VMM APIs fail (error 304) but cudaMalloc works. - * Replaying allocations in seq_num order produces deterministic VAs. - * - * Returns: - * 0 — restored, VA matches original - * -2 — VA mismatch (cudaMalloc succeeded but at wrong address) - * -1 — hard error - */ -static int restore_one_alloc(NvSnapCheckpointAlloc *rec, FILE *data_fp, - nvsnap::GpuTracker &tracker) -{ - /* - * After cuda-checkpoint resume, the original GPU allocations already - * exist at their original VAs. The RM state was restored — the memory - * is allocated and mapped. We just need to copy the data back. - * - * NO cudaMalloc needed. The VAs are already valid. - */ - - /* 1. Skip allocations with no saved data. */ - if (rec->data_offset == UINT64_MAX) { - tracker.track_alloc((uintptr_t)rec->ptr, rec->size, rec->device, - static_cast(rec->alloc_type)); - return 0; - } - - /* 2. Read data from checkpoint file into host buffer. */ - void *host_buf = malloc((size_t)rec->size); - if (!host_buf) { - NVSNAP_GPU_LOG_ERROR("restore: host malloc(%lu) failed", (unsigned long)rec->size); - return -1; - } - - if (fseek(data_fp, (long)rec->data_offset, SEEK_SET) != 0 || - fread(host_buf, 1, (size_t)rec->size, data_fp) != (size_t)rec->size) { - NVSNAP_GPU_LOG_ERROR("restore: short read for alloc seq=%u", rec->seq_num); - free(host_buf); - return -1; - } - - /* 3. H2D copy to the ORIGINAL VA (already allocated by cuda-checkpoint resume). */ - cudaError_t merr = real_ckpt_cudaMemcpy((void *)(uintptr_t)rec->ptr, - host_buf, (size_t)rec->size, - cudaMemcpyHostToDevice); - free(host_buf); - if (merr != 0) { - const char *errstr = (real_ckpt_cudaGetErrorString) - ? real_ckpt_cudaGetErrorString(merr) - : ""; - NVSNAP_GPU_LOG_ERROR("restore: H2D to 0x%lx size=%lu failed: %d (%s)", - (unsigned long)rec->ptr, (unsigned long)rec->size, - (int)merr, errstr ? errstr : ""); - return -1; - } - - /* 4. Re-register in tracker at the original VA. */ - tracker.track_alloc((uintptr_t)rec->ptr, rec->size, rec->device, - static_cast(rec->alloc_type)); - - NVSNAP_GPU_LOG_INFO("restore: seq=%u ptr=0x%lx size=%lu H2D OK", - rec->seq_num, (unsigned long)rec->ptr, - (unsigned long)rec->size); - return 0; -} - -/* ═══════════════════════════════════════════════════════════════════════ - * Checkpoint Restore - * ═══════════════════════════════════════════════════════════════════════ */ - -/* Activate a specific GPU device for VMM operations. */ -static int activate_gpu_device(int device) -{ - /* Always re-resolve — after CRIU restore, static locals are stale. */ - CUresult (*real_cuCtxSetCurrent)(CUcontext) = nullptr; - real_cuCtxSetCurrent = (CUresult (*)(CUcontext)) - dlsym(RTLD_NEXT, "cuCtxSetCurrent"); - if (!real_cuCtxSetCurrent) - real_cuCtxSetCurrent = (CUresult (*)(CUcontext)) - nvsnap_resolve_real("cuCtxSetCurrent", "libcuda.so.1"); - - if (!real_ckpt_cuDevicePrimaryCtxRetain) { - NVSNAP_GPU_LOG_ERROR("checkpoint: cuDevicePrimaryCtxRetain not resolved"); - return -1; - } - if (!real_cuCtxSetCurrent) { - NVSNAP_GPU_LOG_ERROR("checkpoint: cuCtxSetCurrent not resolved"); - return -1; - } - - CUcontext ctx = nullptr; - CUresult cres = real_ckpt_cuDevicePrimaryCtxRetain(&ctx, (CUdevice)device); - if (cres != CUDA_SUCCESS) { - NVSNAP_GPU_LOG_ERROR("checkpoint: cuDevicePrimaryCtxRetain(device=%d) failed: %d", - device, (int)cres); - return -1; - } - - cres = real_cuCtxSetCurrent(ctx); - if (cres != CUDA_SUCCESS) { - NVSNAP_GPU_LOG_ERROR("checkpoint: cuCtxSetCurrent failed: %d", (int)cres); - return -1; - } - - NVSNAP_GPU_LOG_INFO("checkpoint: activated GPU context for device %d", device); - return 0; -} - -/* - * Restore one gpu-/ subdirectory. - * Returns 0 on success, -2 on VA mismatches, -1 on hard error. - */ -static int restore_one_gpu_dir(const char *restore_dir) -{ - NVSNAP_GPU_LOG_INFO("checkpoint: restoring from %s", restore_dir); - - char meta_path[PATH_MAX], data_path[PATH_MAX]; - snprintf(meta_path, sizeof(meta_path), "%s/meta.bin", restore_dir); - snprintf(data_path, sizeof(data_path), "%s/gpu_data.bin", restore_dir); - - FILE *meta_fp = fopen(meta_path, "rb"); - if (!meta_fp) { - NVSNAP_GPU_LOG_ERROR("checkpoint: cannot open %s: %s", meta_path, - strerror(errno)); - return -1; - } - - /* 2. Read and validate header. */ - NvSnapCheckpointHeader header = {}; - if (fread(&header, sizeof(header), 1, meta_fp) != 1) { - NVSNAP_GPU_LOG_ERROR("checkpoint: failed to read header from %s", - meta_path); - fclose(meta_fp); - return -1; - } - - /* Skip empty checkpoints (non-GPU processes). */ - if (header.num_allocations == 0) { - NVSNAP_GPU_LOG_INFO("checkpoint: %s has 0 allocations, skipping", restore_dir); - fclose(meta_fp); - return 0; - } - - /* Determine the correct GPU device for this checkpoint. - * - * Primary source: header.source_device — set during checkpoint and - * records which GPU each worker was actually using. - * - * cudaGetDevice() is NOT reliable after CRIU restore: it returns 0 - * for all processes regardless of which device they were using. */ - int device_to_use = (int)header.source_device; - { - /* Override with LOCAL_RANK if set (torchrun/vLLM). */ - const char *lr = getenv("LOCAL_RANK"); - if (lr) device_to_use = atoi(lr); - - NVSNAP_GPU_LOG_INFO("checkpoint: %s: using device %d " - "(header=%u, LOCAL_RANK=%s)", - restore_dir, device_to_use, - header.source_device, - getenv("LOCAL_RANK") ? getenv("LOCAL_RANK") : "unset"); - - if (activate_gpu_device(device_to_use) < 0) { - NVSNAP_GPU_LOG_ERROR("checkpoint: failed to activate device %d for %s", - device_to_use, restore_dir); - fclose(meta_fp); - return -1; - } - } - - if (header.magic != NVSNAP_CHECKPOINT_MAGIC) { - NVSNAP_GPU_LOG_ERROR("checkpoint: bad magic 0x%08x (expected 0x%08x)", - header.magic, NVSNAP_CHECKPOINT_MAGIC); - fclose(meta_fp); - return -1; - } - - if (header.version != NVSNAP_CHECKPOINT_VERSION) { - NVSNAP_GPU_LOG_ERROR("checkpoint: unsupported version %u (expected %u)", - header.version, NVSNAP_CHECKPOINT_VERSION); - fclose(meta_fp); - return -1; - } - - /* 3. Read allocation records. */ - auto *alloc_recs = (NvSnapCheckpointAlloc *)calloc( - header.num_allocations, sizeof(NvSnapCheckpointAlloc)); - if (header.num_allocations > 0 && !alloc_recs) { - NVSNAP_GPU_LOG_ERROR("checkpoint: failed to allocate alloc records"); - fclose(meta_fp); - return -1; - } - if (header.num_allocations > 0) { - size_t n = fread(alloc_recs, sizeof(NvSnapCheckpointAlloc), - header.num_allocations, meta_fp); - if (n != header.num_allocations) { - NVSNAP_GPU_LOG_ERROR("checkpoint: short read on alloc records"); - free(alloc_recs); - fclose(meta_fp); - return -1; - } - } - - /* 4. Read stream records. */ - auto *stream_recs = (NvSnapCheckpointStream *)calloc( - header.num_streams, sizeof(NvSnapCheckpointStream)); - if (header.num_streams > 0) { - if (!stream_recs || - fread(stream_recs, sizeof(NvSnapCheckpointStream), - header.num_streams, meta_fp) != header.num_streams) { - NVSNAP_GPU_LOG_ERROR("checkpoint: failed to read stream records"); - free(alloc_recs); - free(stream_recs); - fclose(meta_fp); - return -1; - } - } - - /* 5. Read NCCL comm records. */ - auto *nccl_recs = (NvSnapCheckpointNcclComm *)calloc( - header.num_nccl_comms, sizeof(NvSnapCheckpointNcclComm)); - if (header.num_nccl_comms > 0) { - if (!nccl_recs || - fread(nccl_recs, sizeof(NvSnapCheckpointNcclComm), - header.num_nccl_comms, meta_fp) != header.num_nccl_comms) { - NVSNAP_GPU_LOG_ERROR("checkpoint: failed to read NCCL records"); - free(alloc_recs); - free(stream_recs); - free(nccl_recs); - fclose(meta_fp); - return -1; - } - } - - fclose(meta_fp); - - /* 6. Open GPU data file. */ - FILE *data_fp = nullptr; - if (header.total_gpu_bytes > 0) { - data_fp = fopen(data_path, "rb"); - if (!data_fp) { - NVSNAP_GPU_LOG_ERROR("checkpoint: cannot open %s: %s", data_path, - strerror(errno)); - free(alloc_recs); - free(stream_recs); - free(nccl_recs); - return -1; - } - } - - /* 7. Restore allocations via cudaMalloc replay for deterministic VAs. - * Records are already sorted by seq_num in meta.bin (#68). - * cudaSetDevice selects the target GPU for all subsequent cudaMalloc calls. */ - auto &tracker = nvsnap::GpuTracker::instance(); - - int va_mismatches = 0; - int restore_errors = 0; - - /* Set the CUDA runtime device for H2D copies. */ - if (real_ckpt_cudaSetDevice) { - NVSNAP_GPU_LOG_INFO("checkpoint: cudaSetDevice(%d)...", device_to_use); - cudaError_t sderr = real_ckpt_cudaSetDevice(device_to_use); - NVSNAP_GPU_LOG_INFO("checkpoint: cudaSetDevice(%d) = %d", device_to_use, (int)sderr); - if (sderr != 0) { - NVSNAP_GPU_LOG_ERROR("checkpoint: cudaSetDevice(%d) failed: %d", - device_to_use, (int)sderr); - if (data_fp) fclose(data_fp); - free(alloc_recs); - free(stream_recs); - free(nccl_recs); - return -1; - } - } - - for (uint32_t i = 0; i < header.num_allocations; i++) { - NvSnapCheckpointAlloc *rec = &alloc_recs[i]; - - /* Override the per-allocation device with the correct one. - * The saved device field may be wrong (same bug as header). */ - if (i == 0) { - NVSNAP_GPU_LOG_INFO("checkpoint: overriding alloc device %d -> %d " - "for %u allocations", - rec->device, device_to_use, - header.num_allocations); - } - rec->device = (int32_t)device_to_use; - - bool has_data = (rec->data_offset != UINT64_MAX); - if (!has_data) - continue; - - if (!real_ckpt_cudaMemcpy) { - NVSNAP_GPU_LOG_ERROR("checkpoint: cudaMemcpy not resolved — cannot H2D"); - restore_errors++; - continue; - } - - int rr = restore_one_alloc(rec, data_fp, tracker); - if (rr == 1) { - va_mismatches++; /* data restored but VA differs */ - } else if (rr < 0) { - restore_errors++; /* hard error — allocation lost */ - } - } - - /* 8. Recreate streams. */ - for (uint32_t i = 0; i < header.num_streams; i++) { - if (real_ckpt_cudaStreamCreate) { - cudaStream_t new_stream = nullptr; - cudaError_t err = real_ckpt_cudaStreamCreate(&new_stream); - if (err == cudaSuccess && new_stream) { - tracker.track_stream(new_stream, stream_recs[i].device, - stream_recs[i].flags); - } - } - } - - /* 9. Recreate events. */ - for (uint32_t i = 0; i < header.num_events; i++) { - if (real_ckpt_cudaEventCreateWithFlags) { - cudaEvent_t new_event = nullptr; - cudaError_t err = - real_ckpt_cudaEventCreateWithFlags(&new_event, 0); - if (err == cudaSuccess && new_event) { - tracker.track_event(new_event, header.source_device, 0); - } - } - } - - if (data_fp) - fclose(data_fp); - - free(alloc_recs); - free(stream_recs); - free(nccl_recs); - - NVSNAP_GPU_LOG_INFO("checkpoint: restored %u allocs (%lu bytes), " - "%u streams from %s (va_mismatches=%d, errors=%d)", - header.num_allocations, - (unsigned long)header.total_gpu_bytes, - header.num_streams, restore_dir, - va_mismatches, restore_errors); - - if (restore_errors > 0) - return -1; - if (va_mismatches > 0) - return 1; /* VA mismatches — data restored but pointers differ */ - return 0; -} - -/* ═══════════════════════════════════════════════════════════════════════ - * Checkpoint Restore — iterate all gpu-* subdirs - * ═══════════════════════════════════════════════════════════════════════ */ - -extern "C" int nvsnap_checkpoint_restore(const char *checkpoint_dir) -{ - if (!checkpoint_dir) { - NVSNAP_GPU_LOG_ERROR("checkpoint_restore: NULL directory"); - return -1; - } - - force_resolve_real_functions(); - - /* Initialize CUDA driver. */ - if (real_ckpt_cuInit) { - CUresult ires = real_ckpt_cuInit(0); - if (ires != CUDA_SUCCESS) { - NVSNAP_GPU_LOG_ERROR("checkpoint: cuInit(0) failed: %d", (int)ires); - return -1; - } - } - - /* - * Scan checkpoint_dir for gpu-* subdirectories containing meta.bin. - * Each subdir holds one GPU's allocations from the original process. - * The restore helper runs as a single process restoring ALL GPUs. - */ - DIR *dir = opendir(checkpoint_dir); - if (!dir) { - NVSNAP_GPU_LOG_ERROR("checkpoint: cannot open directory %s: %s", - checkpoint_dir, strerror(errno)); - return -1; - } - - int total_restored = 0; - int total_errors = 0; - int total_va_failures = 0; - - struct dirent *ent; - while ((ent = readdir(dir)) != NULL) { - /* Match gpu-* directories. */ - if (strncmp(ent->d_name, "gpu-", 4) != 0) - continue; - - /* Check if meta.bin exists in this subdir. */ - char subdir[PATH_MAX]; - snprintf(subdir, sizeof(subdir), "%s/%s", checkpoint_dir, ent->d_name); - - char meta_check[PATH_MAX]; - snprintf(meta_check, sizeof(meta_check), "%s/meta.bin", subdir); - if (access(meta_check, R_OK) != 0) - continue; - - int rr = restore_one_gpu_dir(subdir); - if (rr == 1) { - total_va_failures++; - NVSNAP_GPU_LOG_ERROR("checkpoint: VA mismatches restoring %s", subdir); - } else if (rr < 0) { - total_errors++; - NVSNAP_GPU_LOG_ERROR("checkpoint: failed to restore %s", subdir); - } - total_restored++; - } - closedir(dir); - - NVSNAP_GPU_LOG_INFO("checkpoint: restored %d GPU subdirs (%d errors, %d VA failures)", - total_restored, total_errors, total_va_failures); - - if (total_restored == 0) { - NVSNAP_GPU_LOG_ERROR("checkpoint: no gpu-*/meta.bin found in %s", checkpoint_dir); - return -1; - } - if (total_errors > 0) - return -1; - if (total_va_failures > 0) - return 1; - return 0; -} - -/* ═══════════════════════════════════════════════════════════════════════ - * Checkpoint Restore — per-process (called from inside restored process) - * ═══════════════════════════════════════════════════════════════════════ */ - -extern "C" int nvsnap_checkpoint_restore_self(const char *checkpoint_dir) -{ - /* Ensure NvSnap logging is enabled during restore — log_level may be 0 - * (default) if the original process never set NVSNAP_GPU_LOG_LEVEL. */ - if (nvsnap_gpu_log_level < NVSNAP_GPU_LOG_INFO) - nvsnap_gpu_log_level = NVSNAP_GPU_LOG_INFO; - - NVSNAP_GPU_LOG_INFO("checkpoint_restore_self: ENTERED dir=%s pid=%d", - checkpoint_dir ? checkpoint_dir : "NULL", (int)getpid()); - - if (!checkpoint_dir) { - NVSNAP_GPU_LOG_ERROR("checkpoint_restore_self: NULL directory"); - return -1; - } - - /* Force re-resolve ALL function pointers. After CRIU restore, cached - * pointers from the original run are stale (ASLR, library relocation). */ - force_resolve_real_functions(); - - /* DO NOT call cuInit(0) here. After CRIU restore, the CUDA driver state - * is preserved (process has open /dev/nvidia* fds, RM objects exist). - * Calling cuInit would try to RE-initialize, which can fail or destroy - * the existing state. The restored process already has a working CUDA - * context from before checkpoint. Just re-resolve functions and proceed - * with H2D copy. */ - - /* Build per-PID path. Only restore THIS process's GPU data. - * DO NOT fall back to iterating all subdirs — that would attempt - * H2D to VAs owned by OTHER processes (different CUDA contexts). */ - char restore_dir[PATH_MAX]; - snprintf(restore_dir, sizeof(restore_dir), "%s/gpu-%d", - checkpoint_dir, (int)getpid()); - - char meta_check[PATH_MAX]; - snprintf(meta_check, sizeof(meta_check), "%s/meta.bin", restore_dir); - - if (access(meta_check, R_OK) != 0) { - NVSNAP_GPU_LOG_INFO("checkpoint: gpu-%d/ not found — this process has no GPU data to restore", - (int)getpid()); - return 0; /* Not an error — non-GPU processes have no data */ - } - - NVSNAP_GPU_LOG_INFO("checkpoint: restore_self restoring gpu-%d/", (int)getpid()); - int ret = restore_one_gpu_dir(restore_dir); - NVSNAP_GPU_LOG_INFO("checkpoint: restore_self result=%d", ret); - return ret; -} - -/* ═══════════════════════════════════════════════════════════════════════ - * Pre-checkpoint quiesce: destroy NCCL + disable P2P - * ═══════════════════════════════════════════════════════════════════════ */ - -extern "C" int nvsnap_pre_checkpoint_quiesce(void) -{ - /* Early exit for non-GPU processes. */ - /* This function exists for multi-GPU: abort NCCL comms + disable P2P. - * Skip entirely if no NCCL comms — avoids cudaDeviceSynchronize hanging - * on pending CUDA graphs/async ops (single-GPU TinyLlama, non-TP workers). */ - auto &tracker_check = nvsnap::GpuTracker::instance(); - if (tracker_check.snapshot_nccl_comms().empty()) { - NVSNAP_GPU_LOG_INFO("quiesce: no NCCL comms, skipping pre-checkpoint (pid=%d)", (int)getpid()); - return 0; - } - - NVSNAP_GPU_LOG_INFO("quiesce: starting pre-checkpoint cleanup (pid=%d)", (int)getpid()); - - /* Resolve real CUDA functions (not our wrappers). */ - static cudaError_t (*real_cudaDeviceGetCount)(int *) = nullptr; - static cudaError_t (*real_cudaSetDevice)(int) = nullptr; - static cudaError_t (*real_cudaDeviceCanAccessPeer)(int *, int, int) = nullptr; - static cudaError_t (*real_cudaDeviceDisablePeerAccess)(int) = nullptr; - static cudaError_t (*real_cudaDeviceSynchronize)(void) = nullptr; - static ncclResult_t (*real_ncclCommAbort)(void *) = nullptr; - - if (!real_cudaDeviceGetCount) { - real_cudaDeviceGetCount = (cudaError_t (*)(int *)) - nvsnap_resolve_real("cudaDeviceGetCount", "libcudart.so"); - real_cudaSetDevice = (cudaError_t (*)(int)) - nvsnap_resolve_real("cudaSetDevice", "libcudart.so"); - real_cudaDeviceCanAccessPeer = (cudaError_t (*)(int *, int, int)) - nvsnap_resolve_real("cudaDeviceCanAccessPeer", "libcudart.so"); - real_cudaDeviceDisablePeerAccess = (cudaError_t (*)(int)) - nvsnap_resolve_real("cudaDeviceDisablePeerAccess", "libcudart.so"); - real_cudaDeviceSynchronize = (cudaError_t (*)(void)) - nvsnap_resolve_real("cudaDeviceSynchronize", "libcudart.so"); - real_ncclCommAbort = (ncclResult_t (*)(void *)) - nvsnap_resolve_real("ncclCommAbort", "libnccl.so.2"); - } - - /* 1. Abort all tracked NCCL communicators. - * Use ncclCommAbort (immediate, non-blocking), NOT ncclCommDestroy - * (blocks on pending ops → deadlock when all ranks call simultaneously). */ - if (real_ncclCommAbort) { - auto &tracker = nvsnap::GpuTracker::instance(); - auto comms = tracker.snapshot_nccl_comms(); - for (auto &c : comms) { - NVSNAP_GPU_LOG_INFO("quiesce: aborting NCCL comm %p (rank=%d, nranks=%d)", - (void *)c.comm, c.rank, c.nranks); - real_ncclCommAbort((void *)c.comm); - tracker.untrack_nccl_comm((void *)c.comm); - } - NVSNAP_GPU_LOG_INFO("quiesce: aborted %zu NCCL communicators", comms.size()); - } - - /* 2. Disable P2P access between all GPU pairs. */ - if (real_cudaDeviceGetCount && real_cudaSetDevice && - real_cudaDeviceCanAccessPeer && real_cudaDeviceDisablePeerAccess) { - - int num_devices = 0; - real_cudaDeviceGetCount(&num_devices); - - int current_device = nvsnap_get_current_device(); - int p2p_disabled = 0; - - /* Only disable P2P FROM this process's device. Do NOT call - * cudaSetDevice() for other GPUs — that creates new primary - * contexts on GPUs this process doesn't own, adding cross-GPU - * driver state that makes cuCheckpointProcessLock hang. */ - if (current_device >= 0) { - real_cudaSetDevice(current_device); - for (int j = 0; j < num_devices; j++) { - if (j == current_device) continue; - int can_access = 0; - real_cudaDeviceCanAccessPeer(&can_access, current_device, j); - if (can_access) { - cudaError_t err = real_cudaDeviceDisablePeerAccess(j); - if (err == 0 /* cudaSuccess */) { - p2p_disabled++; - } - /* err=704 (cudaErrorPeerAccessNotEnabled) is fine */ - } - } - } - - NVSNAP_GPU_LOG_INFO("quiesce: disabled %d P2P pairs from device %d (%d total devices)", - p2p_disabled, current_device, num_devices); - } - - /* 3. Synchronize all devices. */ - if (real_cudaDeviceSynchronize) { - real_cudaDeviceSynchronize(); - } - - NVSNAP_GPU_LOG_INFO("quiesce: pre-checkpoint cleanup complete"); - return 0; -} - -/* ═══════════════════════════════════════════════════════════════════════ - * Post-restore resume: re-enable P2P - * ═══════════════════════════════════════════════════════════════════════ */ - -extern "C" int nvsnap_post_restore_resume(void) -{ - NVSNAP_GPU_LOG_INFO("resume: re-enabling P2P access (pid=%d)", (int)getpid()); - - static cudaError_t (*real_cudaDeviceGetCount)(int *) = nullptr; - static cudaError_t (*real_cudaSetDevice)(int) = nullptr; - static cudaError_t (*real_cudaDeviceCanAccessPeer)(int *, int, int) = nullptr; - static cudaError_t (*real_cudaDeviceEnablePeerAccess)(int, unsigned) = nullptr; - - if (!real_cudaDeviceGetCount) { - real_cudaDeviceGetCount = (cudaError_t (*)(int *)) - nvsnap_resolve_real("cudaDeviceGetCount", "libcudart.so"); - real_cudaSetDevice = (cudaError_t (*)(int)) - nvsnap_resolve_real("cudaSetDevice", "libcudart.so"); - real_cudaDeviceCanAccessPeer = (cudaError_t (*)(int *, int, int)) - nvsnap_resolve_real("cudaDeviceCanAccessPeer", "libcudart.so"); - real_cudaDeviceEnablePeerAccess = (cudaError_t (*)(int, unsigned)) - nvsnap_resolve_real("cudaDeviceEnablePeerAccess", "libcudart.so"); - } - - if (!real_cudaDeviceGetCount || !real_cudaDeviceEnablePeerAccess) - return -1; - - int num_devices = 0; - real_cudaDeviceGetCount(&num_devices); - - int current_device = nvsnap_get_current_device(); - int p2p_enabled = 0; - - /* Only re-enable P2P FROM this process's device. */ - if (current_device >= 0) { - real_cudaSetDevice(current_device); - for (int j = 0; j < num_devices; j++) { - if (j == current_device) continue; - int can_access = 0; - real_cudaDeviceCanAccessPeer(&can_access, current_device, j); - if (can_access) { - cudaError_t err = real_cudaDeviceEnablePeerAccess(j, 0); - if (err == 0) p2p_enabled++; - /* err=704 (already enabled) is fine */ - } - } - } - - NVSNAP_GPU_LOG_INFO("resume: enabled %d P2P pairs from device %d", p2p_enabled, current_device); - return 0; -} - -/* ═══════════════════════════════════════════════════════════════════════ - * Query functions (R3) - * ═══════════════════════════════════════════════════════════════════════ */ - -extern "C" int nvsnap_get_alloc_count(void) -{ - auto stats = nvsnap::GpuTracker::instance().get_stats(); - return (int)stats.live_alloc_count; -} - -extern "C" uint64_t nvsnap_get_total_bytes(void) -{ - auto stats = nvsnap::GpuTracker::instance().get_stats(); - return stats.total_allocated_bytes; -} diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/config.c b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/config.c deleted file mode 100644 index 42077dd672..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/config.c +++ /dev/null @@ -1,102 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * NvSnap — Configuration loading from environment variables. - */ -#include "nvsnap/gpu/config.h" - -#include -#include -#include - -static NvSnapConfig g_config; -static int g_initialized = 0; - -/* Parse a size string like "4G", "512M", "1024K", or plain bytes. */ -size_t nvsnap_parse_size(const char *str) -{ - if (!str || !*str) - return 0; - - char *end = NULL; - unsigned long long val = strtoull(str, &end, 10); - - if (end && *end) { - switch (toupper((unsigned char)*end)) { - case 'K': val *= 1024ULL; break; - case 'M': val *= 1024ULL * 1024ULL; break; - case 'G': val *= 1024ULL * 1024ULL * 1024ULL; break; - case 'T': val *= 1024ULL * 1024ULL * 1024ULL * 1024ULL; break; - default: break; - } - } - return (size_t)val; -} - -static int env_int(const char *name, int def) -{ - const char *v = getenv(name); - if (!v || !*v) - return def; - return atoi(v); -} - -static double env_double(const char *name, double def) -{ - const char *v = getenv(name); - if (!v || !*v) - return def; - return atof(v); -} - -static size_t env_size(const char *name, size_t def) -{ - const char *v = getenv(name); - if (!v || !*v) - return def; - return nvsnap_parse_size(v); -} - -static void env_str(const char *name, char *dst, size_t dst_size, const char *def) -{ - const char *v = getenv(name); - if (!v || !*v) - v = def; - if (v) { - strncpy(dst, v, dst_size - 1); - dst[dst_size - 1] = '\0'; - } else { - dst[0] = '\0'; - } -} - -void nvsnap_config_init(void) -{ - if (g_initialized) - return; - - memset(&g_config, 0, sizeof(g_config)); - - g_config.log_level = env_int("NVSNAP_GPU_LOG_LEVEL", 0); - g_config.metrics_enabled = env_int("NVSNAP_METRICS", 1); - g_config.fault_injection_enabled = env_int("NVSNAP_FAULT_INJECTION", 0); - g_config.host_pool_size = env_size("NVSNAP_HOST_POOL_SIZE", 0); - g_config.oversubscription_ratio = env_double("NVSNAP_OVERSUBSCRIPTION_RATIO", 1.0); - g_config.detailed_tracing = env_int("NVSNAP_DETAILED_TRACING", 0); - - env_str("NVSNAP_AGENT_SOCKET", g_config.agent_socket_path, - sizeof(g_config.agent_socket_path), - "/tmp/nvsnap_agent.sock"); - - g_initialized = 1; -} - -const NvSnapConfig *nvsnap_config_get(void) -{ - if (!g_initialized) - nvsnap_config_init(); - return &g_config; -} diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/init.c b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/init.c deleted file mode 100644 index c941c717a5..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/init.c +++ /dev/null @@ -1,229 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * NvSnap — Library initialization and cleanup. - */ -#define _GNU_SOURCE -#include -#include -#include -#include -#include - -#include "nvsnap/gpu/config.h" -#include "nvsnap/gpu/metrics.h" -#include "nvsnap/gpu/interpose.h" -#include "nvsnap/gpu/tracker.h" - -/* Global log level — referenced by NVSNAP_GPU_LOG macro. */ -int nvsnap_gpu_log_level = 0; - -/* Thread-safe logging using write() — no FILE* buffering, no locks. - * Atomic for messages < 4096 bytes (PIPE_BUF on Linux). */ -#include -void nvsnap_log_write(int level, const char *fmt, ...) -{ - char buf[4096]; - int off = snprintf(buf, sizeof(buf), "[nvsnap:%d] ", level); - if (off < 0 || off >= (int)sizeof(buf)) return; - - va_list ap; - va_start(ap, fmt); - int n = vsnprintf(buf + off, sizeof(buf) - (size_t)off, fmt, ap); - va_end(ap); - if (n < 0) return; - off += n; - if (off >= (int)sizeof(buf) - 1) off = (int)sizeof(buf) - 2; - buf[off++] = '\n'; - - /* write() is async-signal-safe and atomic for < PIPE_BUF bytes. */ - ssize_t __attribute__((unused)) wr = write(STDERR_FILENO, buf, (size_t)off); -} - -/* - * Resolve a real function pointer when RTLD_NEXT fails. - * - * This happens when the library containing the symbol (e.g., libnccl.so) - * was loaded via dlopen AFTER our LD_PRELOAD library. RTLD_NEXT only - * searches libraries loaded after us in the link chain, but dlopen'd - * libraries are not in that chain. - * - * Strategy: try to dlopen the likely library with RTLD_NOLOAD (returns - * handle if already loaded, NULL if not) then dlsym from that handle. - */ -static const char *g_lib_search[] = { - "libnccl.so.2", - "libnccl.so", - "libcudart.so", - "libcudart.so.12", - "libcuda.so.1", - "libcuda.so", - NULL -}; - -/* - * Cached handles to real GPU libraries (populated on first resolve). - * We dlopen with RTLD_NOLOAD first (already loaded by PyTorch) to get - * a handle, then use dlvsym or iterate to find the real symbol. - */ -static void *g_real_lib_handles[16] = {0}; -static int g_handles_count = 0; - -static void nvsnap_populate_lib_handles(void) -{ - if (g_handles_count > 0) return; - - for (const char **lib = g_lib_search; *lib; lib++) { - /* RTLD_NOLOAD: only returns handle if already loaded. */ - void *h = dlopen(*lib, RTLD_LAZY | RTLD_NOLOAD); - if (h && g_handles_count < 15) { - g_real_lib_handles[g_handles_count++] = h; - NVSNAP_GPU_LOG_DEBUG("resolve: found loaded library %s", *lib); - } - } - - /* Also try loading them explicitly if not already loaded. */ - for (const char **lib = g_lib_search; *lib; lib++) { - void *h = dlopen(*lib, RTLD_LAZY | RTLD_GLOBAL); - if (h && g_handles_count < 15) { - /* Check if we already have this handle. */ - int dup = 0; - for (int i = 0; i < g_handles_count; i++) { - if (g_real_lib_handles[i] == h) { dup = 1; break; } - } - if (!dup) { - g_real_lib_handles[g_handles_count++] = h; - } - } - } -} - -void *nvsnap_resolve_real(const char *func_name, const char *lib_hint) -{ - void *sym = NULL; - - /* If caller provides a specific library, try that first. */ - if (lib_hint) { - void *h = dlopen(lib_hint, RTLD_LAZY | RTLD_NOLOAD); - if (!h) h = dlopen(lib_hint, RTLD_LAZY); - if (h) { - sym = dlsym(h, func_name); - /* Check it's not our own wrapper (LD_PRELOAD can cause this). */ - Dl_info info; - if (sym && dladdr(sym, &info) && info.dli_fname && - !strstr(info.dli_fname, "libnvsnap_intercept")) { - return sym; - } - } - } - - /* Search cached library handles. */ - nvsnap_populate_lib_handles(); - for (int i = 0; i < g_handles_count; i++) { - sym = dlsym(g_real_lib_handles[i], func_name); - if (sym) { - /* Verify it's not our own wrapper. */ - Dl_info info; - if (dladdr(sym, &info) && info.dli_fname && - !strstr(info.dli_fname, "libnvsnap_intercept")) { - return sym; - } - } - } - - /* Try RTLD_NEXT one more time (library may have been loaded since). */ - sym = dlsym(RTLD_NEXT, func_name); - if (sym) { - Dl_info info; - if (dladdr(sym, &info) && info.dli_fname && - !strstr(info.dli_fname, "libnvsnap_intercept")) { - return sym; - } - } - - return NULL; -} - -/* Thread-local current device. */ -_Thread_local int nvsnap_current_device = 0; -/* Thread-local flag: kept for API compatibility but no longer - * set by stream capture wrappers (which were removed). */ -_Thread_local int nvsnap_in_graph_capture = 0; - -/* - * Fork safety: reset NvSnap state in child process after fork(). - * - * vLLM and other multi-process GPU frameworks use fork() extensively. - * After fork(), the child inherits parent's mutexes (potentially locked), - * shared memory (write positions corrupted), and semaphores (undefined). - * We must reinitialize everything in the child. - */ -static void nvsnap_atfork_child(void) -{ - /* 1. Reset tracker — reconstructs mutexes, clears stale parent state. */ - nvsnap_tracker_reset_after_fork(); - - /* 2. Reinit metrics — child gets its own shared memory segment. */ - nvsnap_metrics_destroy(); - nvsnap_metrics_init(); - - NVSNAP_GPU_LOG_INFO("NvSnap reinitialized after fork (child pid=%d)", - (int)getpid()); -} - -/* Defined in src/self_disable.c. Declared here rather than pulling in - * nvsnap_intercept.h, which this GPU sub-module otherwise does not use. */ -int nvsnap_self_disabled(void); - -__attribute__((constructor(102))) /* After NvSnap init (101), before NvSnap atfork (103) */ -static void nvsnap_gpu_init(void) -{ - if (nvsnap_self_disabled()) - return; - - nvsnap_config_init(); - - const NvSnapConfig *cfg = nvsnap_config_get(); - nvsnap_gpu_log_level = cfg->log_level; - - if (cfg->metrics_enabled) { - nvsnap_metrics_init(); - } - - /* Register fork handler — critical for vLLM and other multi-process frameworks. */ - pthread_atfork(NULL, NULL, nvsnap_atfork_child); - - /* - * Post-CRIU restore detection. - * - * If NVSNAP_RESTORE_DIR is set, this process was checkpointed and is - * now being restored by CRIU. We need to restore GPU memory from the - * checkpoint before the application resumes. - * - * The restore-entrypoint sets this env var before CRIU restore. After - * CRIU unfreezes the process, libnvsnap_intercept.so's constructor re-runs - * (CRIU re-executes constructors on restore for LD_PRELOAD libraries). - * - * Actually — CRIU does NOT re-run constructors. The process resumes - * from where it was frozen. So we need a different detection method. - * - * Instead, NvSnap's libnvsnap_intercept.so will call nvsnap_checkpoint_restore() - * directly when it detects the /run/criu-restored marker. NvSnap just - * needs to expose the function — which it already does. - * - * For standalone (non-NvSnap) restore, use the nvsnap-gpu-restore helper - * or call the C API from the restore orchestrator. - */ - - NVSNAP_GPU_LOG_INFO("NvSnap loaded (pid=%d)", (int)getpid()); -} - -__attribute__((destructor)) -static void nvsnap_fini(void) -{ - NVSNAP_GPU_LOG_INFO("NvSnap unloading (pid=%d)", (int)getpid()); - nvsnap_metrics_destroy(); -} diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/interpose_cudart.c b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/interpose_cudart.c deleted file mode 100644 index c453c20618..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/interpose_cudart.c +++ /dev/null @@ -1,268 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * NvSnap — Minimal CUDA Runtime API interception for checkpoint/restore. - * - * Only the functions needed for allocation tracking, device management, - * and checkpoint quiesce are intercepted here. All pass-through wrappers - * have been removed to avoid stack corruption from vendored type mismatches. - */ -#define _GNU_SOURCE -#include -#include -#include -#include -#include -#include - -#include "nvsnap/gpu/cuda_types.h" -#include "nvsnap/gpu/interpose.h" -#include "nvsnap/gpu/tracker.h" -#include "nvsnap/gpu/metrics.h" -#include "nvsnap/gpu/checkpoint.h" - -/* ═══════════════════════════════════════════════════════════════════════ - * Reliable device query - * - * nvsnap_current_device (thread-local) can be wrong if cudaSetDevice - * was called via a path we don't intercept (driver API, before our - * library loaded, from a different thread). Query the CUDA runtime - * for the actual current device. - * ═══════════════════════════════════════════════════════════════════════ */ -int nvsnap_get_current_device(void) -{ - static cudaError_t (*real_get)(int *) = NULL; - if (!real_get) { - real_get = (cudaError_t (*)(int *))dlsym(RTLD_NEXT, "cudaGetDevice"); - if (!real_get) - real_get = (cudaError_t (*)(int *)) - nvsnap_resolve_real("cudaGetDevice", "libcudart.so"); - } - if (real_get) { - int dev = -1; - if (real_get(&dev) == 0 && dev >= 0) { - nvsnap_current_device = dev; /* keep thread-local in sync */ - return dev; - } - } - return nvsnap_current_device; /* fallback */ -} - -/* - * No auto-detection of post-CRIU restore. NvSnap's libnvsnap_intercept.so - * calls nvsnap_checkpoint_restore_self() directly from its reinit - * handler when it detects the CRIU restore marker. - * - * This avoids: - * - Signal handler conflicts (SIGUSR2) - * - Marker file access() syscalls during checkpoint save - * - In-memory state that CRIU preserves incorrectly - * - Race conditions between save and deferred restore - * - * NvSnap's job: expose nvsnap_checkpoint_restore_self(). - * NvSnap's job: call it at the right time. - */ - -/* Kept for ABI compatibility — no-op. */ -void nvsnap_reset_restore_state(void) { } -void nvsnap_deferred_restore(void) { } - -/* ═══════════════════════════════════════════════════════════════════════ - * Memory Allocation — tracked for checkpoint/restore - * ═══════════════════════════════════════════════════════════════════════ */ - -/* ─── cudaMalloc ─────────────────────────────────────────────────────── */ - -NVSNAP_DECLARE_REAL(cudaError_t, cudaMalloc, void **, size_t); - -cudaError_t cudaMalloc(void **devPtr, size_t size) -{ - NVSNAP_LOAD_REAL(cudaMalloc); - nvsnap_deferred_restore(); - - cudaError_t err = real_cudaMalloc(devPtr, size); - - if (err == cudaSuccess && devPtr && *devPtr) { - nvsnap_tracker_track_alloc((uintptr_t)*devPtr, size, - nvsnap_get_current_device(), 0 /* DEVICE */); - nvsnap_tracker_add_allocated_bytes((int64_t)size); - nvsnap_metrics_write(NVSNAP_METRIC_ALLOC, - nvsnap_get_current_device(), size, - (uint64_t)(uintptr_t)*devPtr); - NVSNAP_GPU_LOG_DEBUG("cudaMalloc(%zu) -> %p", size, *devPtr); - } - return err; -} - -/* ─── cudaFree ───────────────────────────────────────────────────────── */ - -NVSNAP_DECLARE_REAL(cudaError_t, cudaFree, void *); - -cudaError_t cudaFree(void *devPtr) -{ - NVSNAP_LOAD_REAL(cudaFree); - - if (devPtr) { - nvsnap_tracker_untrack_alloc((uintptr_t)devPtr); - nvsnap_metrics_write(NVSNAP_METRIC_FREE, - nvsnap_get_current_device(), 0, - (uint64_t)(uintptr_t)devPtr); - NVSNAP_GPU_LOG_DEBUG("cudaFree(%p)", devPtr); - } - - return real_cudaFree(devPtr); -} - -/* ─── cudaMallocManaged ──────────────────────────────────────────────── */ - -NVSNAP_DECLARE_REAL(cudaError_t, cudaMallocManaged, void **, size_t, unsigned); - -cudaError_t cudaMallocManaged(void **devPtr, size_t size, unsigned flags) -{ - NVSNAP_LOAD_REAL(cudaMallocManaged); - nvsnap_deferred_restore(); - - cudaError_t err = real_cudaMallocManaged(devPtr, size, flags); - - if (err == cudaSuccess && devPtr && *devPtr) { - nvsnap_tracker_track_alloc((uintptr_t)*devPtr, size, - nvsnap_get_current_device(), 1 /* MANAGED */); - nvsnap_tracker_add_allocated_bytes((int64_t)size); - nvsnap_metrics_write(NVSNAP_METRIC_ALLOC, - nvsnap_get_current_device(), size, - (uint64_t)(uintptr_t)*devPtr); - } - return err; -} - -/* ─── cudaHostAlloc (pinned memory) ──────────────────────────────────── */ - -NVSNAP_DECLARE_REAL(cudaError_t, cudaHostAlloc, void **, size_t, unsigned); - -cudaError_t cudaHostAlloc(void **pHost, size_t size, unsigned flags) -{ - NVSNAP_LOAD_REAL(cudaHostAlloc); - nvsnap_deferred_restore(); - - cudaError_t err = real_cudaHostAlloc(pHost, size, flags); - - if (err == cudaSuccess && pHost && *pHost) { - nvsnap_tracker_track_alloc((uintptr_t)*pHost, size, - nvsnap_get_current_device(), 2 /* HOST_PINNED */); - nvsnap_tracker_add_allocated_bytes((int64_t)size); - nvsnap_metrics_write(NVSNAP_METRIC_ALLOC, - nvsnap_get_current_device(), size, - (uint64_t)(uintptr_t)*pHost); - } - return err; -} - -/* ─── cudaFreeHost ───────────────────────────────────────────────────── */ - -NVSNAP_DECLARE_REAL(cudaError_t, cudaFreeHost, void *); - -cudaError_t cudaFreeHost(void *ptr) -{ - NVSNAP_LOAD_REAL(cudaFreeHost); - - if (ptr) { - nvsnap_tracker_untrack_alloc((uintptr_t)ptr); - nvsnap_metrics_write(NVSNAP_METRIC_FREE, - nvsnap_get_current_device(), 0, - (uint64_t)(uintptr_t)ptr); - } - - return real_cudaFreeHost(ptr); -} - -/* ─── cudaMallocAsync ────────────────────────────────────────────────── */ - -NVSNAP_DECLARE_REAL(cudaError_t, cudaMallocAsync, void **, size_t, cudaStream_t); - -cudaError_t cudaMallocAsync(void **devPtr, size_t size, cudaStream_t stream) -{ - NVSNAP_LOAD_REAL(cudaMallocAsync); - nvsnap_deferred_restore(); - - cudaError_t err = real_cudaMallocAsync(devPtr, size, stream); - - if (err == cudaSuccess && devPtr && *devPtr) { - nvsnap_tracker_track_alloc((uintptr_t)*devPtr, size, - nvsnap_get_current_device(), 0); - nvsnap_tracker_add_allocated_bytes((int64_t)size); - nvsnap_metrics_write(NVSNAP_METRIC_ALLOC, - nvsnap_get_current_device(), size, - (uint64_t)(uintptr_t)*devPtr); - } - return err; -} - -/* ─── cudaFreeAsync ──────────────────────────────────────────────────── */ - -NVSNAP_DECLARE_REAL(cudaError_t, cudaFreeAsync, void *, cudaStream_t); - -cudaError_t cudaFreeAsync(void *devPtr, cudaStream_t stream) -{ - NVSNAP_LOAD_REAL(cudaFreeAsync); - - if (devPtr) { - nvsnap_tracker_untrack_alloc((uintptr_t)devPtr); - nvsnap_metrics_write(NVSNAP_METRIC_FREE, - nvsnap_get_current_device(), 0, - (uint64_t)(uintptr_t)devPtr); - } - - return real_cudaFreeAsync(devPtr, stream); -} - -/* ═══════════════════════════════════════════════════════════════════════ - * Device Management — needed for checkpoint coordination - * ═══════════════════════════════════════════════════════════════════════ */ - -NVSNAP_DECLARE_REAL(cudaError_t, cudaSetDevice, int); - -cudaError_t cudaSetDevice(int device) -{ - NVSNAP_LOAD_REAL(cudaSetDevice); - nvsnap_deferred_restore(); - - cudaError_t err = real_cudaSetDevice(device); - - if (err == cudaSuccess) { - nvsnap_current_device = device; - NVSNAP_GPU_LOG_DEBUG("cudaSetDevice(%d)", device); - } - return err; -} - -NVSNAP_DECLARE_REAL(cudaError_t, cudaGetDevice, int *); - -cudaError_t cudaGetDevice(int *device) -{ - NVSNAP_LOAD_REAL(cudaGetDevice); - return real_cudaGetDevice(device); -} - -/* ─── cudaDeviceSynchronize — needed for checkpoint quiesce ──────────── */ - -NVSNAP_DECLARE_REAL(cudaError_t, cudaDeviceSynchronize, void); - -cudaError_t cudaDeviceSynchronize(void) -{ - NVSNAP_LOAD_REAL(cudaDeviceSynchronize); - nvsnap_deferred_restore(); - return real_cudaDeviceSynchronize(); -} - -/* ─── cudaMemGetInfo — pass through, vLLM queries this ───────────────── */ - -NVSNAP_DECLARE_REAL(cudaError_t, cudaMemGetInfo, size_t *, size_t *); - -cudaError_t cudaMemGetInfo(size_t *free_mem, size_t *total_mem) -{ - NVSNAP_LOAD_REAL(cudaMemGetInfo); - return real_cudaMemGetInfo(free_mem, total_mem); -} diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/interpose_cudrv.c b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/interpose_cudrv.c deleted file mode 100644 index aa7a34326f..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/interpose_cudrv.c +++ /dev/null @@ -1,212 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * NvSnap — Minimal CUDA Driver API interception for checkpoint/restore. - * - * Only memory allocation tracking and VMM functions needed for restore - * are intercepted. All pass-through wrappers removed. - */ -#define _GNU_SOURCE -#include -#include - -#include "nvsnap/gpu/cuda_types.h" -#include "nvsnap/gpu/interpose.h" -#include "nvsnap/gpu/tracker.h" -#include "nvsnap/gpu/metrics.h" - -/* Defined in interpose_cudart.c */ -extern void nvsnap_deferred_restore(void); -extern int nvsnap_get_current_device(void); - -/* ═══════════════════════════════════════════════════════════════════════ - * Memory Allocation — tracked for checkpoint/restore - * ═══════════════════════════════════════════════════════════════════════ */ - -/* ─── cuMemAlloc_v2 ──────────────────────────────────────────────────── */ - -NVSNAP_DECLARE_REAL(CUresult, cuMemAlloc_v2, CUdeviceptr *, size_t); - -CUresult cuMemAlloc_v2(CUdeviceptr *dptr, size_t bytesize) -{ - NVSNAP_LOAD_REAL(cuMemAlloc_v2); - nvsnap_deferred_restore(); - - CUresult err = real_cuMemAlloc_v2(dptr, bytesize); - - if (err == CUDA_SUCCESS && dptr) { - nvsnap_tracker_track_alloc((uintptr_t)*dptr, bytesize, - nvsnap_get_current_device(), 0); - nvsnap_tracker_add_allocated_bytes((int64_t)bytesize); - nvsnap_metrics_write(NVSNAP_METRIC_ALLOC, - nvsnap_get_current_device(), bytesize, - (uint64_t)*dptr); - NVSNAP_GPU_LOG_DEBUG("cuMemAlloc_v2(%zu) -> 0x%llx", bytesize, - (unsigned long long)*dptr); - } - return err; -} - -CUresult cuMemAlloc(CUdeviceptr *dptr, size_t bytesize) -{ - return cuMemAlloc_v2(dptr, bytesize); -} - -/* ─── cuMemFree_v2 ───────────────────────────────────────────────────── */ - -NVSNAP_DECLARE_REAL(CUresult, cuMemFree_v2, CUdeviceptr); - -CUresult cuMemFree_v2(CUdeviceptr dptr) -{ - NVSNAP_LOAD_REAL(cuMemFree_v2); - - if (dptr) { - nvsnap_tracker_untrack_alloc((uintptr_t)dptr); - nvsnap_metrics_write(NVSNAP_METRIC_FREE, - nvsnap_get_current_device(), 0, (uint64_t)dptr); - NVSNAP_GPU_LOG_DEBUG("cuMemFree_v2(0x%llx)", (unsigned long long)dptr); - } - - return real_cuMemFree_v2(dptr); -} - -CUresult cuMemFree(CUdeviceptr dptr) -{ - return cuMemFree_v2(dptr); -} - -/* ─── cuMemAllocManaged ──────────────────────────────────────────────── */ - -NVSNAP_DECLARE_REAL(CUresult, cuMemAllocManaged, CUdeviceptr *, size_t, unsigned); - -CUresult cuMemAllocManaged(CUdeviceptr *dptr, size_t bytesize, unsigned flags) -{ - NVSNAP_LOAD_REAL(cuMemAllocManaged); - nvsnap_deferred_restore(); - - CUresult err = real_cuMemAllocManaged(dptr, bytesize, flags); - - if (err == CUDA_SUCCESS && dptr) { - nvsnap_tracker_track_alloc((uintptr_t)*dptr, bytesize, - nvsnap_get_current_device(), 1); - nvsnap_tracker_add_allocated_bytes((int64_t)bytesize); - nvsnap_metrics_write(NVSNAP_METRIC_ALLOC, - nvsnap_get_current_device(), bytesize, - (uint64_t)*dptr); - } - return err; -} - -/* ═══════════════════════════════════════════════════════════════════════ - * VMM (Virtual Memory Management) — needed for restore - * ═══════════════════════════════════════════════════════════════════════ */ - -NVSNAP_DECLARE_REAL(CUresult, cuMemAddressReserve, CUdeviceptr *, size_t, - size_t, CUdeviceptr, unsigned long long); - -CUresult cuMemAddressReserve(CUdeviceptr *ptr, size_t size, size_t alignment, - CUdeviceptr addr, unsigned long long flags) -{ - NVSNAP_LOAD_REAL(cuMemAddressReserve); - CUresult err = real_cuMemAddressReserve(ptr, size, alignment, addr, flags); - if (err == CUDA_SUCCESS && ptr) { - NVSNAP_GPU_LOG_DEBUG("cuMemAddressReserve(%zu) -> 0x%llx", size, - (unsigned long long)*ptr); - } - return err; -} - -NVSNAP_DECLARE_REAL(CUresult, cuMemCreate, CUmemGenericAllocationHandle *, - size_t, const CUmemAllocationProp *, unsigned long long); - -CUresult cuMemCreate(CUmemGenericAllocationHandle *handle, size_t size, - const CUmemAllocationProp *prop, unsigned long long flags) -{ - NVSNAP_LOAD_REAL(cuMemCreate); - - CUresult err = real_cuMemCreate(handle, size, prop, flags); - if (err == CUDA_SUCCESS) { - NVSNAP_GPU_LOG_DEBUG("cuMemCreate(%zu) -> handle=%llu", size, - handle ? (unsigned long long)*handle : 0ULL); - } - return err; -} - -NVSNAP_DECLARE_REAL(CUresult, cuMemMap, CUdeviceptr, size_t, size_t, - CUmemGenericAllocationHandle, unsigned long long); - -CUresult cuMemMap(CUdeviceptr ptr, size_t size, size_t offset, - CUmemGenericAllocationHandle handle, unsigned long long flags) -{ - NVSNAP_LOAD_REAL(cuMemMap); - CUresult err = real_cuMemMap(ptr, size, offset, handle, flags); - - if (err == CUDA_SUCCESS) { - nvsnap_tracker_track_vmm_mapping((uintptr_t)ptr, size, - (uint64_t)handle, - nvsnap_get_current_device()); - nvsnap_metrics_write(NVSNAP_METRIC_VMM_MAP, - nvsnap_get_current_device(), size, (uint64_t)ptr); - NVSNAP_GPU_LOG_DEBUG("cuMemMap(0x%llx, %zu)", (unsigned long long)ptr, size); - } - return err; -} - -NVSNAP_DECLARE_REAL(CUresult, cuMemSetAccess, CUdeviceptr, size_t, - const CUmemAccessDesc *, size_t); - -CUresult cuMemSetAccess(CUdeviceptr ptr, size_t size, - const CUmemAccessDesc *desc, size_t count) -{ - NVSNAP_LOAD_REAL(cuMemSetAccess); - return real_cuMemSetAccess(ptr, size, desc, count); -} - -NVSNAP_DECLARE_REAL(CUresult, cuMemUnmap, CUdeviceptr, size_t); - -CUresult cuMemUnmap(CUdeviceptr ptr, size_t size) -{ - NVSNAP_LOAD_REAL(cuMemUnmap); - - nvsnap_tracker_untrack_vmm_mapping((uintptr_t)ptr); - nvsnap_metrics_write(NVSNAP_METRIC_VMM_UNMAP, - nvsnap_get_current_device(), size, (uint64_t)ptr); - - return real_cuMemUnmap(ptr, size); -} - -NVSNAP_DECLARE_REAL(CUresult, cuMemRelease, CUmemGenericAllocationHandle); - -CUresult cuMemRelease(CUmemGenericAllocationHandle handle) -{ - NVSNAP_LOAD_REAL(cuMemRelease); - NVSNAP_GPU_LOG_DEBUG("cuMemRelease(handle=%llu)", (unsigned long long)handle); - return real_cuMemRelease(handle); -} - -NVSNAP_DECLARE_REAL(CUresult, cuMemAddressFree, CUdeviceptr, size_t); - -CUresult cuMemAddressFree(CUdeviceptr ptr, size_t size) -{ - NVSNAP_LOAD_REAL(cuMemAddressFree); - NVSNAP_GPU_LOG_DEBUG("cuMemAddressFree(0x%llx, %zu)", (unsigned long long)ptr, size); - return real_cuMemAddressFree(ptr, size); -} - -/* ─── cuMemGetInfo_v2 — pass through, vLLM queries this ─────────────── */ - -NVSNAP_DECLARE_REAL(CUresult, cuMemGetInfo_v2, size_t *, size_t *); - -CUresult cuMemGetInfo_v2(size_t *free_mem, size_t *total_mem) -{ - NVSNAP_LOAD_REAL(cuMemGetInfo_v2); - return real_cuMemGetInfo_v2(free_mem, total_mem); -} - -CUresult cuMemGetInfo(size_t *free_mem, size_t *total_mem) -{ - return cuMemGetInfo_v2(free_mem, total_mem); -} diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/interpose_nccl.c b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/interpose_nccl.c deleted file mode 100644 index c759aec25f..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/interpose_nccl.c +++ /dev/null @@ -1,104 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * NvSnap — Minimal NCCL interception for checkpoint/restore. - * - * Only communicator lifecycle tracking is needed. Collective operations - * are not intercepted — NvSnap handles NCCL quiesce separately. - */ -#define _GNU_SOURCE -#include -#include - -#include "nvsnap/gpu/cuda_types.h" -#include "nvsnap/gpu/interpose.h" -#include "nvsnap/gpu/tracker.h" - -/* ═══════════════════════════════════════════════════════════════════════ - * Communicator lifecycle — tracked for checkpoint/restore - * ═══════════════════════════════════════════════════════════════════════ */ - -NVSNAP_DECLARE_REAL(ncclResult_t, ncclCommInitRank, ncclComm_t *, int, ncclUniqueId, int); - -ncclResult_t ncclCommInitRank(ncclComm_t *comm, int nranks, - ncclUniqueId commId, int rank) -{ - NVSNAP_LOAD_REAL(ncclCommInitRank); - - ncclResult_t err = real_ncclCommInitRank(comm, nranks, commId, rank); - - if (err == ncclSuccess && comm && *comm) { - nvsnap_tracker_track_nccl_comm( - *comm, nranks, rank, - (const unsigned char *)commId.internal, - nvsnap_current_device); - NVSNAP_GPU_LOG_INFO("ncclCommInitRank(nranks=%d, rank=%d) -> %p", - nranks, rank, *comm); - } - return err; -} - -NVSNAP_DECLARE_REAL(ncclResult_t, ncclCommInitRankConfig, ncclComm_t *, int, - ncclUniqueId, int, void *); - -ncclResult_t ncclCommInitRankConfig(ncclComm_t *comm, int nranks, - ncclUniqueId commId, int rank, - void *config) -{ - NVSNAP_LOAD_REAL(ncclCommInitRankConfig); - - ncclResult_t err = real_ncclCommInitRankConfig(comm, nranks, commId, - rank, config); - - if (err == ncclSuccess && comm && *comm) { - nvsnap_tracker_track_nccl_comm( - *comm, nranks, rank, - (const unsigned char *)commId.internal, - nvsnap_current_device); - NVSNAP_GPU_LOG_INFO("ncclCommInitRankConfig(nranks=%d, rank=%d) -> %p", - nranks, rank, *comm); - } - return err; -} - -NVSNAP_DECLARE_REAL(ncclResult_t, ncclCommInitAll, ncclComm_t *, int, const int *); - -ncclResult_t ncclCommInitAll(ncclComm_t *comms, int ndev, const int *devlist) -{ - NVSNAP_LOAD_REAL(ncclCommInitAll); - - ncclResult_t err = real_ncclCommInitAll(comms, ndev, devlist); - - if (err == ncclSuccess && comms) { - unsigned char zero_id[128] = {0}; - for (int i = 0; i < ndev; i++) { - int dev = devlist ? devlist[i] : i; - nvsnap_tracker_track_nccl_comm(comms[i], ndev, i, zero_id, dev); - } - NVSNAP_GPU_LOG_INFO("ncclCommInitAll(ndev=%d)", ndev); - } - return err; -} - -NVSNAP_DECLARE_REAL(ncclResult_t, ncclCommDestroy, ncclComm_t); - -ncclResult_t ncclCommDestroy(ncclComm_t comm) -{ - NVSNAP_LOAD_REAL(ncclCommDestroy); - NVSNAP_GPU_LOG_INFO("ncclCommDestroy(%p)", comm); - nvsnap_tracker_untrack_nccl_comm(comm); - return real_ncclCommDestroy(comm); -} - -NVSNAP_DECLARE_REAL(ncclResult_t, ncclCommAbort, ncclComm_t); - -ncclResult_t ncclCommAbort(ncclComm_t comm) -{ - NVSNAP_LOAD_REAL(ncclCommAbort); - NVSNAP_GPU_LOG_INFO("ncclCommAbort(%p)", comm); - nvsnap_tracker_untrack_nccl_comm(comm); - return real_ncclCommAbort(comm); -} diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/metrics.c b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/metrics.c deleted file mode 100644 index ebef07a101..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/metrics.c +++ /dev/null @@ -1,110 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * NvSnap — Shared-memory ring buffer for zero-syscall metrics. - * - * Uses MAP_ANONYMOUS instead of shm_open so CRIU can checkpoint/restore - * the mapping natively without needing /dev/shm/ files to exist. - */ -#include "nvsnap/gpu/metrics.h" -#include "nvsnap/gpu/config.h" - -#include -#include -#include -#include -#include - -static NvSnapMetricsRingBuffer *g_ring = NULL; - -static uint64_t metrics_now_ns(void) -{ - struct timespec ts; - clock_gettime(CLOCK_MONOTONIC, &ts); - return (uint64_t)ts.tv_sec * 1000000000ULL + (uint64_t)ts.tv_nsec; -} - -int nvsnap_metrics_init(void) -{ - if (g_ring) - return 0; /* already initialized */ - - /* Anonymous mmap — no /dev/shm file, no path for CRIU to track. - * CRIU handles MAP_ANONYMOUS | MAP_SHARED natively. */ - g_ring = (NvSnapMetricsRingBuffer *)mmap( - NULL, sizeof(NvSnapMetricsRingBuffer), - PROT_READ | PROT_WRITE, - MAP_ANONYMOUS | MAP_SHARED, -1, 0); - - if (g_ring == MAP_FAILED) { - g_ring = NULL; - return -1; - } - - memset(g_ring, 0, sizeof(NvSnapMetricsRingBuffer)); - return 0; -} - -void nvsnap_metrics_write(NvSnapMetricType type, int device, - uint64_t value, uint64_t ptr) -{ - NvSnapMetricsRingBuffer *ring = g_ring; - if (__builtin_expect(ring == NULL, 0)) - return; - - /* Atomically claim a slot. */ - uint64_t pos = atomic_fetch_add_explicit(&ring->write_pos, 1, - memory_order_relaxed); - uint64_t idx = pos & (NVSNAP_METRICS_RING_SIZE - 1); - - /* Check for overflow (writer lapping reader). */ - uint64_t rpos = atomic_load_explicit(&ring->read_pos, memory_order_relaxed); - if (pos - rpos >= NVSNAP_METRICS_RING_SIZE) { - atomic_fetch_add_explicit(&ring->overflow_count, 1, - memory_order_relaxed); - } - - NvSnapMetricsEntry *e = &ring->entries[idx]; - e->type = (uint32_t)type; - e->pid = (uint32_t)getpid(); - e->device = (int32_t)device; - e->timestamp = metrics_now_ns(); - e->value = value; - e->ptr = ptr; -} - -int nvsnap_metrics_read(NvSnapMetricsEntry *out) -{ - NvSnapMetricsRingBuffer *ring = g_ring; - if (__builtin_expect(ring == NULL, 0)) - return -1; - - uint64_t rpos = atomic_load_explicit(&ring->read_pos, memory_order_relaxed); - uint64_t wpos = atomic_load_explicit(&ring->write_pos, memory_order_acquire); - - if (rpos >= wpos) - return -1; /* empty */ - - uint64_t idx = rpos & (NVSNAP_METRICS_RING_SIZE - 1); - if (out) - *out = ring->entries[idx]; - - atomic_fetch_add_explicit(&ring->read_pos, 1, memory_order_release); - return 0; -} - -NvSnapMetricsRingBuffer *nvsnap_metrics_get_buffer(void) -{ - return g_ring; -} - -void nvsnap_metrics_destroy(void) -{ - if (g_ring) { - munmap(g_ring, sizeof(NvSnapMetricsRingBuffer)); - g_ring = NULL; - } -} diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/symbol_table.c b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/symbol_table.c deleted file mode 100644 index 7eb3c716c1..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/symbol_table.c +++ /dev/null @@ -1,122 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * NvSnap — dlsym routing table. - * - * Maps symbol names to our wrapper functions. Only includes the minimal - * set of functions needed for checkpoint/restore. - * - * LD_PRELOAD gives us symbol precedence for all exported functions. - */ -#define _GNU_SOURCE -#include -#include -#include -#include - -#include "nvsnap/gpu/cuda_types.h" -#include "nvsnap/gpu/interpose.h" - -/* ── Forward declarations of all wrapper functions ───────────────────── */ - -/* CUDA Runtime API — interpose_cudart.c */ -extern cudaError_t cudaMalloc(void **, size_t); -extern cudaError_t cudaFree(void *); -extern cudaError_t cudaMallocManaged(void **, size_t, unsigned); -extern cudaError_t cudaHostAlloc(void **, size_t, unsigned); -extern cudaError_t cudaFreeHost(void *); -extern cudaError_t cudaMallocAsync(void **, size_t, cudaStream_t); -extern cudaError_t cudaFreeAsync(void *, cudaStream_t); -extern cudaError_t cudaSetDevice(int); -extern cudaError_t cudaGetDevice(int *); -extern cudaError_t cudaDeviceSynchronize(void); -extern cudaError_t cudaMemGetInfo(size_t *, size_t *); - -/* CUDA Driver API — interpose_cudrv.c */ -extern CUresult cuMemAlloc_v2(CUdeviceptr *, size_t); -extern CUresult cuMemAlloc(CUdeviceptr *, size_t); -extern CUresult cuMemFree_v2(CUdeviceptr); -extern CUresult cuMemFree(CUdeviceptr); -extern CUresult cuMemAllocManaged(CUdeviceptr *, size_t, unsigned); -extern CUresult cuMemAddressReserve(CUdeviceptr *, size_t, size_t, CUdeviceptr, unsigned long long); -extern CUresult cuMemCreate(CUmemGenericAllocationHandle *, size_t, const CUmemAllocationProp *, unsigned long long); -extern CUresult cuMemMap(CUdeviceptr, size_t, size_t, CUmemGenericAllocationHandle, unsigned long long); -extern CUresult cuMemSetAccess(CUdeviceptr, size_t, const CUmemAccessDesc *, size_t); -extern CUresult cuMemUnmap(CUdeviceptr, size_t); -extern CUresult cuMemRelease(CUmemGenericAllocationHandle); -extern CUresult cuMemAddressFree(CUdeviceptr, size_t); -extern CUresult cuMemGetInfo_v2(size_t *, size_t *); -extern CUresult cuMemGetInfo(size_t *, size_t *); - -/* NCCL — interpose_nccl.c */ -extern ncclResult_t ncclCommInitRank(ncclComm_t *, int, ncclUniqueId, int); -extern ncclResult_t ncclCommInitRankConfig(ncclComm_t *, int, ncclUniqueId, int, void *); -extern ncclResult_t ncclCommInitAll(ncclComm_t *, int, const int *); -extern ncclResult_t ncclCommDestroy(ncclComm_t); -extern ncclResult_t ncclCommAbort(ncclComm_t); - -/* ── Symbol routing table ────────────────────────────────────────────── */ - -typedef struct { - const char *name; - void *wrapper; -} SymbolEntry; - -#define SYM(func) { #func, (void *)(func) } - -static const SymbolEntry g_symbol_table[] = { - /* CUDA Runtime API */ - SYM(cudaMalloc), - SYM(cudaFree), - SYM(cudaMallocManaged), - SYM(cudaHostAlloc), - SYM(cudaFreeHost), - SYM(cudaMallocAsync), - SYM(cudaFreeAsync), - SYM(cudaSetDevice), - SYM(cudaGetDevice), - SYM(cudaDeviceSynchronize), - SYM(cudaMemGetInfo), - - /* CUDA Driver API */ - SYM(cuMemAlloc_v2), - SYM(cuMemAlloc), - SYM(cuMemFree_v2), - SYM(cuMemFree), - SYM(cuMemAllocManaged), - SYM(cuMemAddressReserve), - SYM(cuMemCreate), - SYM(cuMemMap), - SYM(cuMemSetAccess), - SYM(cuMemUnmap), - SYM(cuMemRelease), - SYM(cuMemAddressFree), - SYM(cuMemGetInfo_v2), - SYM(cuMemGetInfo), - - /* NCCL */ - SYM(ncclCommInitRank), - SYM(ncclCommInitRankConfig), - SYM(ncclCommInitAll), - SYM(ncclCommDestroy), - SYM(ncclCommAbort), - - { NULL, NULL } -}; - -#undef SYM - -/* ── Lookup (used by internal code, e.g., tests) ───────────────────── */ - -void *nvsnap_lookup_symbol(const char *name) -{ - for (const SymbolEntry *e = g_symbol_table; e->name != NULL; e++) { - if (strcmp(e->name, name) == 0) { - return e->wrapper; - } - } - return NULL; -} diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/tracker.cpp b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/tracker.cpp deleted file mode 100644 index 1503d4f899..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu/tracker.cpp +++ /dev/null @@ -1,363 +0,0 @@ -/* - * NvSnap — GpuTracker implementation. - */ -#include "nvsnap/gpu/tracker.h" - -#include -#include -#include - -namespace nvsnap { - -static uint64_t now_ns() -{ - using namespace std::chrono; - return (uint64_t)duration_cast( - steady_clock::now().time_since_epoch()) - .count(); -} - -GpuTracker &GpuTracker::instance() -{ - static GpuTracker s_instance; - return s_instance; -} - -/* ── Allocations ─────────────────────────────────────────────────────── */ - -void GpuTracker::track_alloc(uintptr_t ptr, size_t size, int device, AllocType type) -{ - AllocInfo info{ptr, size, device, type, now_ns(), - next_seq_num_.fetch_add(1, std::memory_order_relaxed)}; - std::unique_lock lock(alloc_mu_); - allocs_[ptr] = info; -} - -bool GpuTracker::untrack_alloc(uintptr_t ptr) -{ - std::unique_lock lock(alloc_mu_); - auto it = allocs_.find(ptr); - if (it == allocs_.end()) - return false; - size_t size = it->second.size; - allocs_.erase(it); - /* Decrement live bytes (saturating). */ - uint64_t old = total_allocated_bytes_.load(std::memory_order_relaxed); - while (old >= size && - !total_allocated_bytes_.compare_exchange_weak( - old, old - size, std::memory_order_relaxed)) - ; - return true; -} - -bool GpuTracker::lookup_alloc(uintptr_t ptr, AllocInfo *out) const -{ - std::shared_lock lock(alloc_mu_); - auto it = allocs_.find(ptr); - if (it == allocs_.end()) - return false; - if (out) - *out = it->second; - return true; -} - -std::vector GpuTracker::snapshot_allocs() const -{ - std::shared_lock lock(alloc_mu_); - std::vector v; - v.reserve(allocs_.size()); - for (auto &kv : allocs_) - v.push_back(kv.second); - return v; -} - -/* ── Streams ─────────────────────────────────────────────────────────── */ - -void GpuTracker::track_stream(void *handle, int device, unsigned flags) -{ - StreamInfo info{handle, device, flags, 0}; - std::unique_lock lock(stream_mu_); - streams_[(uintptr_t)handle] = info; -} - -bool GpuTracker::untrack_stream(void *handle) -{ - std::unique_lock lock(stream_mu_); - return streams_.erase((uintptr_t)handle) > 0; -} - -std::vector GpuTracker::snapshot_streams() const -{ - std::shared_lock lock(stream_mu_); - std::vector v; - v.reserve(streams_.size()); - for (auto &kv : streams_) - v.push_back(kv.second); - return v; -} - -void GpuTracker::inc_stream_kernel_count(void *stream) -{ - if (!stream) - return; /* default stream — skip map lookup on hot path */ - std::shared_lock lock(stream_mu_); - auto it = streams_.find((uintptr_t)stream); - if (it != streams_.end()) { - /* Not perfectly atomic, but good enough for stats. */ - it->second.kernel_count++; - } -} - -/* ── Events ──────────────────────────────────────────────────────────── */ - -void GpuTracker::track_event(void *handle, int device, unsigned flags) -{ - EventInfo info{handle, device, flags}; - std::unique_lock lock(event_mu_); - events_[(uintptr_t)handle] = info; -} - -bool GpuTracker::untrack_event(void *handle) -{ - std::unique_lock lock(event_mu_); - return events_.erase((uintptr_t)handle) > 0; -} - -/* ── NCCL communicators ──────────────────────────────────────────────── */ - -void GpuTracker::track_nccl_comm(void *comm, int nranks, int rank, - const uint8_t unique_id[128], int device) -{ - NcclCommInfo info{}; - info.comm = comm; - info.nranks = nranks; - info.rank = rank; - info.device = device; - info.collective_count = 0; - info.bytes_transferred = 0; - if (unique_id) - std::memcpy(info.unique_id, unique_id, 128); - std::unique_lock lock(nccl_mu_); - nccl_comms_[(uintptr_t)comm] = info; -} - -bool GpuTracker::untrack_nccl_comm(void *comm) -{ - std::unique_lock lock(nccl_mu_); - return nccl_comms_.erase((uintptr_t)comm) > 0; -} - -void GpuTracker::record_nccl_collective(void *comm, uint64_t bytes) -{ - std::shared_lock lock(nccl_mu_); - auto it = nccl_comms_.find((uintptr_t)comm); - if (it != nccl_comms_.end()) { - it->second.collective_count++; - it->second.bytes_transferred += bytes; - } -} - -std::vector GpuTracker::snapshot_nccl_comms() const -{ - std::shared_lock lock(nccl_mu_); - std::vector v; - v.reserve(nccl_comms_.size()); - for (auto &kv : nccl_comms_) - v.push_back(kv.second); - return v; -} - -/* ── VMM mappings ────────────────────────────────────────────────────── */ - -void GpuTracker::track_vmm_mapping(uintptr_t va, size_t size, - uint64_t handle, int device) -{ - VmmMappingInfo info{va, size, handle, device}; - std::unique_lock lock(vmm_mu_); - vmm_mappings_[va] = info; -} - -bool GpuTracker::untrack_vmm_mapping(uintptr_t va) -{ - std::unique_lock lock(vmm_mu_); - return vmm_mappings_.erase(va) > 0; -} - -std::vector GpuTracker::snapshot_vmm_mappings() const -{ - std::shared_lock lock(vmm_mu_); - std::vector v; - v.reserve(vmm_mappings_.size()); - for (auto &kv : vmm_mappings_) - v.push_back(kv.second); - return v; -} - -/* ── Fork safety ─────────────────────────────────────────────────────── */ - -void GpuTracker::reset_after_fork() -{ - /* After fork(), std::shared_mutex is in undefined state if the parent - * held any lock. We reconstruct them in-place and clear all tracked data - * (child process starts with a clean GPU state). */ - new (&alloc_mu_) std::shared_mutex(); - new (&stream_mu_) std::shared_mutex(); - new (&event_mu_) std::shared_mutex(); - new (&nccl_mu_) std::shared_mutex(); - new (&vmm_mu_) std::shared_mutex(); - - allocs_.clear(); - streams_.clear(); - events_.clear(); - nccl_comms_.clear(); - vmm_mappings_.clear(); - - next_seq_num_.store(0, std::memory_order_relaxed); - total_allocated_bytes_.store(0, std::memory_order_relaxed); - total_kernel_launches_.store(0, std::memory_order_relaxed); - total_memcpy_bytes_.store(0, std::memory_order_relaxed); -} - -/* ── Atomic counters ─────────────────────────────────────────────────── */ - -void GpuTracker::add_allocated_bytes(int64_t delta) -{ - if (delta > 0) { - total_allocated_bytes_.fetch_add((uint64_t)delta, - std::memory_order_relaxed); - } - /* Decrement is handled in untrack_alloc for correctness. */ -} - -void GpuTracker::inc_kernel_launches() -{ - total_kernel_launches_.fetch_add(1, std::memory_order_relaxed); -} - -void GpuTracker::add_memcpy_bytes(uint64_t bytes) -{ - total_memcpy_bytes_.fetch_add(bytes, std::memory_order_relaxed); -} - -GpuStats GpuTracker::get_stats() const -{ - GpuStats s{}; - s.total_allocated_bytes = total_allocated_bytes_.load(std::memory_order_relaxed); - s.total_kernel_launches = total_kernel_launches_.load(std::memory_order_relaxed); - s.total_memcpy_bytes = total_memcpy_bytes_.load(std::memory_order_relaxed); - - { - std::shared_lock lock(alloc_mu_); - s.live_alloc_count = allocs_.size(); - } - { - std::shared_lock lock(stream_mu_); - s.live_stream_count = streams_.size(); - } - { - std::shared_lock lock(event_mu_); - s.live_event_count = events_.size(); - } - { - std::shared_lock lock(nccl_mu_); - s.live_nccl_comm_count = nccl_comms_.size(); - } - { - std::shared_lock lock(vmm_mu_); - s.live_vmm_mapping_count = vmm_mappings_.size(); - } - return s; -} - -} /* namespace nvsnap */ - -/* ═══════════════════════════════════════════════════════════════════════ - * C API — thin wrappers around the singleton - * ═══════════════════════════════════════════════════════════════════════ */ - -extern "C" { - -void nvsnap_tracker_track_alloc(uintptr_t ptr, size_t size, int device, int type) -{ - nvsnap::GpuTracker::instance().track_alloc( - ptr, size, device, static_cast(type)); -} - -void nvsnap_tracker_untrack_alloc(uintptr_t ptr) -{ - nvsnap::GpuTracker::instance().untrack_alloc(ptr); -} - -void nvsnap_tracker_add_allocated_bytes(int64_t delta) -{ - nvsnap::GpuTracker::instance().add_allocated_bytes(delta); -} - -void nvsnap_tracker_inc_kernel_launches(void) -{ - nvsnap::GpuTracker::instance().inc_kernel_launches(); -} - -void nvsnap_tracker_add_memcpy_bytes(uint64_t bytes) -{ - nvsnap::GpuTracker::instance().add_memcpy_bytes(bytes); -} - -void nvsnap_tracker_track_stream(void *handle, int device, unsigned flags) -{ - nvsnap::GpuTracker::instance().track_stream(handle, device, flags); -} - -void nvsnap_tracker_untrack_stream(void *handle) -{ - nvsnap::GpuTracker::instance().untrack_stream(handle); -} - -void nvsnap_tracker_inc_stream_kernel_count(void *stream) -{ - nvsnap::GpuTracker::instance().inc_stream_kernel_count(stream); -} - -void nvsnap_tracker_track_event(void *handle, int device, unsigned flags) -{ - nvsnap::GpuTracker::instance().track_event(handle, device, flags); -} - -void nvsnap_tracker_untrack_event(void *handle) -{ - nvsnap::GpuTracker::instance().untrack_event(handle); -} - -void nvsnap_tracker_track_nccl_comm(void *comm, int nranks, int rank, - const unsigned char unique_id[128], int device) -{ - nvsnap::GpuTracker::instance().track_nccl_comm(comm, nranks, rank, unique_id, device); -} - -void nvsnap_tracker_untrack_nccl_comm(void *comm) -{ - nvsnap::GpuTracker::instance().untrack_nccl_comm(comm); -} - -void nvsnap_tracker_record_nccl_collective(void *comm, uint64_t bytes) -{ - nvsnap::GpuTracker::instance().record_nccl_collective(comm, bytes); -} - -void nvsnap_tracker_track_vmm_mapping(uintptr_t va, size_t size, - uint64_t handle, int device) -{ - nvsnap::GpuTracker::instance().track_vmm_mapping(va, size, handle, device); -} - -void nvsnap_tracker_untrack_vmm_mapping(uintptr_t va) -{ - nvsnap::GpuTracker::instance().untrack_vmm_mapping(va); -} - -void nvsnap_tracker_reset_after_fork(void) -{ - nvsnap::GpuTracker::instance().reset_after_fork(); -} - -} /* extern "C" */ diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu_checkpoint.c b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu_checkpoint.c deleted file mode 100644 index 7aac7fbdfd..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/gpu_checkpoint.c +++ /dev/null @@ -1,145 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * GPU checkpoint/restore helpers for libnvsnap_intercept.so - * - * Provides P2P disable/enable and device sync for multi-GPU checkpoint. - * Called from the quiesce path (AFTER NCCL destroy, BEFORE CUDA Checkpoint API). - * - * All CUDA functions resolved via dlsym at runtime — no link dependency on CUDA. - */ -#define _GNU_SOURCE -#include -#include -#include -#include - -/* Logging via nvsnap intercept */ -extern void nvsnap_log(int level, const char *func, const char *fmt, ...); -#define NVSNAP_INFO(fmt, ...) nvsnap_log(3, __func__, fmt, ##__VA_ARGS__) -#define NVSNAP_WARN(fmt, ...) nvsnap_log(2, __func__, fmt, ##__VA_ARGS__) - -typedef int cudaError_t; - -/* Resolve a symbol from a specific library, bypassing our own wrappers. */ -static void *resolve_real(const char *sym, const char *lib) { - void *h = dlopen(lib, RTLD_LAZY | RTLD_NOLOAD); - if (!h) h = dlopen(lib, RTLD_LAZY); - if (!h) return NULL; - return dlsym(h, sym); -} - -/* Cached function pointers */ -static cudaError_t (*fn_getDeviceCount)(int *) = NULL; -static cudaError_t (*fn_setDevice)(int) = NULL; -static cudaError_t (*fn_getDevice)(int *) = NULL; -static cudaError_t (*fn_canAccessPeer)(int *, int, int) = NULL; -static cudaError_t (*fn_disablePeerAccess)(int) = NULL; -static cudaError_t (*fn_enablePeerAccess)(int, unsigned) = NULL; -static cudaError_t (*fn_deviceSync)(void) = NULL; -static int fns_resolved = 0; - -static void resolve_all(void) { - if (fns_resolved) return; - fn_getDeviceCount = resolve_real("cudaGetDeviceCount", "libcudart.so"); - fn_setDevice = resolve_real("cudaSetDevice", "libcudart.so"); - fn_getDevice = resolve_real("cudaGetDevice", "libcudart.so"); - fn_canAccessPeer = resolve_real("cudaDeviceCanAccessPeer", "libcudart.so"); - fn_disablePeerAccess = resolve_real("cudaDeviceDisablePeerAccess", "libcudart.so"); - fn_enablePeerAccess = resolve_real("cudaDeviceEnablePeerAccess", "libcudart.so"); - fn_deviceSync = resolve_real("cudaDeviceSynchronize", "libcudart.so"); - fns_resolved = 1; -} - -/* - * Disable P2P access between all GPU pairs and sync all devices. - * Called during quiesce, AFTER NCCL comms are destroyed. - * Returns: number of P2P pairs disabled, or -1 on error. - */ -int nvsnap_gpu_pre_checkpoint(void) { - resolve_all(); - - if (!fn_getDeviceCount || !fn_setDevice || !fn_getDevice) { - NVSNAP_WARN("CUDA runtime not available, skipping P2P disable"); - return 0; - } - - int num_devices = 0; - fn_getDeviceCount(&num_devices); - if (num_devices <= 1) return 0; /* Single GPU, nothing to do */ - - int current_device = -1; - fn_getDevice(¤t_device); - - /* Disable P2P access FROM this process's device to all other devices. - * - * CRITICAL: Do NOT call cudaSetDevice() for other GPUs. That creates new - * CUDA primary contexts on GPUs this process doesn't own, adding cross-GPU - * driver state that makes cuCheckpointProcessLock/Checkpoint hang. - * Each TP worker only owns one GPU — only touch that one. */ - int p2p_disabled = 0; - if (current_device >= 0 && fn_canAccessPeer && fn_disablePeerAccess) { - fn_setDevice(current_device); - for (int j = 0; j < num_devices; j++) { - if (j == current_device) continue; - int can_access = 0; - fn_canAccessPeer(&can_access, current_device, j); - if (can_access) { - cudaError_t err = fn_disablePeerAccess(j); - if (err == 0) { - p2p_disabled++; - } - /* err=704 (not enabled) is fine — skip silently */ - } - } - } - - /* Sync this device only */ - if (fn_deviceSync && current_device >= 0) { - fn_setDevice(current_device); - fn_deviceSync(); - } - - NVSNAP_INFO("P2P disabled: %d pairs across %d devices", p2p_disabled, num_devices); - return p2p_disabled; -} - -/* - * Re-enable P2P access between all GPU pairs. - * Called during post-restore reinit. - * Returns: number of P2P pairs enabled, or -1 on error. - */ -int nvsnap_gpu_post_restore(void) { - resolve_all(); - - if (!fn_getDeviceCount || !fn_setDevice || !fn_getDevice) return 0; - - int num_devices = 0; - fn_getDeviceCount(&num_devices); - if (num_devices <= 1) return 0; - - int current_device = -1; - fn_getDevice(¤t_device); - - /* Re-enable P2P FROM this device only — same logic as pre_checkpoint. */ - int p2p_enabled = 0; - if (current_device >= 0 && fn_canAccessPeer && fn_enablePeerAccess) { - fn_setDevice(current_device); - for (int j = 0; j < num_devices; j++) { - if (j == current_device) continue; - int can_access = 0; - fn_canAccessPeer(&can_access, current_device, j); - if (can_access) { - cudaError_t err = fn_enablePeerAccess(j, 0); - if (err == 0) p2p_enabled++; - /* err=704 (already enabled) is fine */ - } - } - } - - NVSNAP_INFO("P2P re-enabled: %d pairs across %d devices", p2p_enabled, num_devices); - return p2p_enabled; -} diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/init.c b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/init.c deleted file mode 100644 index 16775d89b1..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/init.c +++ /dev/null @@ -1,347 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * NVSNAP Interception Library - Initialization - * - * This library intercepts io_uring and libuv for checkpoint/restore support. - * CUDA/GPU state is handled externally by cuda-checkpoint (NVIDIA's tool). - * - * Initialization happens automatically via __attribute__((constructor)). - */ - -#define _GNU_SOURCE -#include -#include -#include -#include -#include -#include -#include -#include - -#include - -#include "nvsnap_intercept.h" -#include - -/* From io_uring_intercept.c */ -void nvsnap_install_uvloop_hook_async(void); -/* From zmq_intercept.c */ -void nvsnap_zmq_reinit_all_if_restored(void); -/* From quiesce.c */ -void nvsnap_quiesce_init(void); -void nvsnap_start_quiesce_worker(void); -void nvsnap_perform_restore_reinit(void); - -static pthread_once_t g_restore_watch_once = PTHREAD_ONCE_INIT; - -static void *nvsnap_restore_watch_thread(void *arg) { - (void)arg; - for (;;) { - /* - * Marker locations (nvsnap#186): - * /var/run/nvsnap/.restored — canonical, always writable. - * Matches NVSNAP_RESTORE_MARKER in - * io_uring_intercept.c + - * libuv_intercept.c. This is the - * one that works when /nvsnap and - * /nvsnap-lib are hostPath read-only - * (the v0.0.20+ production webhook - * injection path — see - * internal/webhook/restore_entrypoint.go). - * /nvsnap-lib/.restored, - * /nvsnap/.restored — legacy, written by restore-entrypoint - * when those mounts are writable - * (emptyDir test-workload pattern). - * Keep checking them so existing - * test workloads aren't regressed. - * /run/criu-restored — historical alternate. - */ - if (access("/var/run/nvsnap/.restored", F_OK) == 0 || - access("/nvsnap-lib/.restored", F_OK) == 0 || access("/nvsnap/.restored", F_OK) == 0 || - access("/run/criu-restored", F_OK) == 0) { - NVSNAP_WARN("Restore watch detected marker, triggering full reinit"); - nvsnap_zmq_reinit_all_if_restored(); - nvsnap_perform_restore_reinit(); - break; - } - usleep(200 * 1000); - } - return NULL; -} - -static void nvsnap_start_restore_watch(void) { - pthread_t tid; - if (pthread_create(&tid, NULL, nvsnap_restore_watch_thread, NULL) == 0) { - pthread_detach(tid); - } else { - NVSNAP_WARN("Restore watch thread failed to start"); - } -} - -/* - * ============================================================================= - * GLOBAL STATE - * ============================================================================= - */ - -static nvsnap_state_t g_state = { - .initialized = false, - .enabled = true, - .init_mutex = PTHREAD_MUTEX_INITIALIZER, - .log_mutex = PTHREAD_MUTEX_INITIALIZER, - .log_level = NVSNAP_LOG_INFO, - .log_file = NULL, -}; - -nvsnap_state_t* nvsnap_get_state(void) { - return &g_state; -} - -/* - * ============================================================================= - * LOGGING - * ============================================================================= - */ - -static const char* log_level_str[] = { - "OFF", "ERROR", "WARN", "INFO", "DEBUG", "TRACE" -}; - -void nvsnap_log(nvsnap_log_level_t level, const char* func, const char* fmt, ...) { - nvsnap_state_t* state = &g_state; - - if (level > state->log_level) { - return; - } - - pthread_mutex_lock(&state->log_mutex); - - FILE* out = state->log_file ? state->log_file : stderr; - - /* Timestamp */ - struct timespec ts; - clock_gettime(CLOCK_REALTIME, &ts); - struct tm tm; - localtime_r(&ts.tv_sec, &tm); - - fprintf(out, "[%02d:%02d:%02d.%03ld] [%s] [%s] ", - tm.tm_hour, tm.tm_min, tm.tm_sec, ts.tv_nsec / 1000000, - log_level_str[level], func); - - va_list args; - va_start(args, fmt); - vfprintf(out, fmt, args); - va_end(args); - - fprintf(out, "\n"); - fflush(out); - - pthread_mutex_unlock(&state->log_mutex); -} - -/* - * ============================================================================= - * INITIALIZATION - * ============================================================================= - */ - -static void init_logging(nvsnap_state_t* state) { - /* Log level */ - const char* level_str = getenv(NVSNAP_ENV_LOG_LEVEL); - if (level_str) { - state->log_level = atoi(level_str); - if (state->log_level > NVSNAP_LOG_TRACE) { - state->log_level = NVSNAP_LOG_TRACE; - } - } else { - state->log_level = NVSNAP_LOG_INFO; /* Default */ - } - - /* Log file */ - const char* log_file = getenv(NVSNAP_ENV_LOG_FILE); - if (log_file && strcmp(log_file, "stderr") != 0) { - state->log_file = fopen(log_file, "a"); - if (!state->log_file) { - fprintf(stderr, "NVSNAP: Failed to open log file %s, using stderr\n", log_file); - state->log_file = NULL; - } - } -} - -/* Check if interception is enabled */ -static bool check_enabled(void) { - const char* enabled = getenv(NVSNAP_ENV_ENABLED); - if (enabled && strcmp(enabled, "0") == 0) { - return false; - } - /* Skip init for processes that must not be intercepted. - * cuda-checkpoint is called by the CRIU CUDA plugin during restore. - * The plugin parses cuda-checkpoint's stdout/stderr to extract the - * restore thread ID. If the intercept library initializes, its log - * messages corrupt the output, causing GPU resume to fail (tid=0). */ - const char* comm = NULL; - char buf[256]; - FILE* f = fopen("/proc/self/comm", "r"); - if (f) { - if (fgets(buf, sizeof(buf), f)) { - /* Strip trailing newline */ - char* nl = strchr(buf, '\n'); - if (nl) *nl = '\0'; - comm = buf; - } - fclose(f); - } - if (comm && (strstr(comm, "cuda-checkpoint") != NULL || - strstr(comm, "cuda_checkpoint") != NULL)) { - return false; - } - return true; -} - -/* Check if seccomp interception is enabled */ -static bool check_seccomp_enabled(void) { - const char* enabled = getenv(NVSNAP_ENV_SECCOMP_ENABLED); - if (enabled && strcmp(enabled, "1") == 0) { - return true; - } - return false; -} - -/* Check if we're in post-restore mode */ -static bool check_post_restore(void) { - const char* post_restore = getenv(NVSNAP_ENV_POST_RESTORE); - if (post_restore && strcmp(post_restore, "1") == 0) { - return true; - } - return false; -} - -/* No-op signal handler for SIGUSR2 — causes EINTR in blocked syscalls */ -static void nvsnap_sigusr2_noop(int sig) { (void)sig; } - -int nvsnap_init_explicit(void) { - nvsnap_state_t* state = &g_state; - - pthread_mutex_lock(&state->init_mutex); - - if (state->initialized) { - pthread_mutex_unlock(&state->init_mutex); - return 0; - } - - /* Check if enabled - allow complete disable for debugging */ - state->enabled = check_enabled(); - if (!state->enabled) { - /* Silent disable - don't even print anything */ - state->initialized = true; - pthread_mutex_unlock(&state->init_mutex); - return 0; - } - - /* Initialize logging first */ - init_logging(state); - - NVSNAP_INFO("=== NVSNAP Interception Library Initializing ==="); - NVSNAP_INFO("PID: %d, PPID: %d", getpid(), getppid()); - NVSNAP_INFO("Purpose: io_uring draining + libuv reinit for CRIU checkpoint/restore"); - NVSNAP_INFO("Note: GPU state is handled by cuda-checkpoint (NVIDIA)"); - - /* Install no-op SIGUSR2 handler for CRIU restore wakeup. - * After CRIU restore, library I/O threads are stuck in epoll_wait(). - * The restore-entrypoint sends SIGUSR2 to all threads, causing EINTR, - * which lets the library event loops continue and detect the restore. - * SA_RESTART is NOT set — we WANT EINTR to interrupt blocking syscalls. */ - { - struct sigaction sa; - memset(&sa, 0, sizeof(sa)); - sa.sa_handler = nvsnap_sigusr2_noop; - sa.sa_flags = 0; - sigemptyset(&sa.sa_mask); - if (sigaction(SIGUSR2, &sa, NULL) == 0) { - NVSNAP_INFO("Installed SIGUSR2 handler for CRIU restore wakeup"); - } - } - - /* Install uvloop hook when Python initializes */ - nvsnap_install_uvloop_hook_async(); - - /* Ensure we trigger restore reinit even if no intercepts fire */ - pthread_once(&g_restore_watch_once, nvsnap_start_restore_watch); - - /* If we were restored, kick off reinit (uvloop handles its own fork via native patch). - * /var/run/nvsnap/.restored is the canonical marker that always works — the legacy - * /nvsnap and /nvsnap-lib paths are checked too so existing test workloads (which use - * writable emptyDirs there) stay compatible. See nvsnap_restore_watch_thread above - * for the full rationale (nvsnap#186). */ - if (access("/var/run/nvsnap/.restored", F_OK) == 0 || - access("/nvsnap-lib/.restored", F_OK) == 0 || access("/nvsnap/.restored", F_OK) == 0) { - nvsnap_zmq_reinit_all_if_restored(); - } - - /* Check if seccomp interception is enabled */ - if (check_seccomp_enabled()) { - NVSNAP_INFO("seccomp interception enabled via %s", NVSNAP_ENV_SECCOMP_ENABLED); - - if (nvsnap_seccomp_install_filter() == 0) { - NVSNAP_INFO("seccomp filter installed successfully"); - - /* Check if post-restore mode */ - if (check_post_restore()) { - NVSNAP_INFO("Post-restore mode detected - will monitor io_uring for healing"); - nvsnap_seccomp_set_post_restore(true); - } - } else { - NVSNAP_ERROR("Failed to install seccomp filter"); - } - } else { - NVSNAP_DEBUG("seccomp interception disabled (set %s=1 to enable)", - NVSNAP_ENV_SECCOMP_ENABLED); - } - - /* Initialize quiesce module: installs SIGUSR1/SIGUSR2 handlers (if - * NVSNAP_QUIESCE_SIGNALS=1), sets up ACK pipe, starts worker thread. - * Without this, SIGUSR2 resume never fires and quiesce spin-loops forever. */ - nvsnap_quiesce_init(); - - /* Ensure quiesce worker thread runs even if NVSNAP_QUIESCE_SIGNALS is not set. - * File-based triggers work without signals. */ - nvsnap_start_quiesce_worker(); - - state->initialized = true; - - NVSNAP_INFO("=== NVSNAP Initialization Complete ==="); - - pthread_mutex_unlock(&state->init_mutex); - return 0; -} - -/* Automatic initialization via constructor */ -__attribute__((constructor(101))) /* Priority 101 = run early but after libc */ -void nvsnap_init(void) { - if (nvsnap_self_disabled()) - return; - nvsnap_init_explicit(); -} - -/* Cleanup on unload */ -__attribute__((destructor)) -void nvsnap_fini(void) { - nvsnap_state_t* state = &g_state; - - if (!state->initialized) { - return; - } - - NVSNAP_INFO("=== NVSNAP Shutting Down ==="); - - /* Print quiesce state if any io_uring or libuv was tracked */ - nvsnap_dump_quiesce_state(stderr); - - if (state->log_file && state->log_file != stderr) { - fclose(state->log_file); - } -} diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/io_uring_intercept.c b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/io_uring_intercept.c deleted file mode 100644 index 19eeb2546c..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/io_uring_intercept.c +++ /dev/null @@ -1,1284 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * io_uring Interception for NVSNAP - * - * This module intercepts io_uring syscalls to track all io_uring instances. - * On quiesce signal, we can drain all rings before CRIU checkpoint. - * - * Why intercept at syscall level: - * - io_uring has no library - apps use raw syscalls or liburing - * - We can't dlsym a syscall - * - Solution: intercept via libc's syscall() function wrapper - * - * TRANSPARENT REINIT ARCHITECTURE: - * After CRIU restore, io_uring fds may be invalidated (e.g., overwritten by BPF). - * Instead of debugging WHO breaks the fd, we DETECT and FIX it transparently: - * - * 1. On first io_uring_enter after restore, verify fd is still io_uring - * 2. If not (EBADF or wrong fd type), transparently recreate: - * - Create new io_uring with saved params - * - Close whatever is at the expected fd - * - dup2 new io_uring to expected fd - * - mmap ring memory at the same addresses app expects - * - Sync ring indices - * 3. Continue - app never knows anything happened - * - * This is GENERIC - works for any application using io_uring without modification. - * - * NOTE: uvloop/libuv kernel state reinit (uv_loop_fork) is handled natively - * by the patched uvloop, not by this module. See uvloop checkpoint-restore-v1. - * - * Build: - * Part of libnvsnap_intercept.so - */ - -#define _GNU_SOURCE -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include /* for backtrace() - debugging */ -#include -#include - -#include "nvsnap_intercept.h" - -/* io_uring syscall numbers */ -#ifndef __NR_io_uring_setup -#define __NR_io_uring_setup 425 -#endif -#ifndef __NR_io_uring_enter -#define __NR_io_uring_enter 426 -#endif -#ifndef __NR_io_uring_register -#define __NR_io_uring_register 427 -#endif - -/* io_uring setup flags */ -#define IORING_SETUP_SQPOLL (1U << 1) -#define IORING_SETUP_SQE128 (1U << 10) -#define IORING_SETUP_CQE32 (1U << 11) - -/* io_uring features */ -#define IORING_FEAT_SINGLE_MMAP (1U << 0) - -/* io_uring enter flags */ -#define IORING_ENTER_GETEVENTS (1U << 0) -#define IORING_ENTER_SQ_WAKEUP (1U << 1) - -/* io_uring sq flags */ -#define IORING_SQ_NEED_WAKEUP (1U << 0) - -/* mmap offsets for io_uring */ -#define IORING_OFF_SQ_RING 0ULL -#define IORING_OFF_CQ_RING 0x8000000ULL -#define IORING_OFF_SQES 0x10000000ULL - -/* External functions from quiesce.c */ -extern int nvsnap_is_restored(void); -extern int nvsnap_perform_quiescence(void); - -/* Real syscall function */ -static long (*real_syscall)(long number, ...) = NULL; - -/* Real mmap and close functions - needed for reinit to avoid recursive interception */ -typedef void* (*mmap_fn)(void *addr, size_t length, int prot, int flags, int fd, off_t offset); -typedef int (*close_fn)(int fd); -static mmap_fn real_mmap = NULL; -static close_fn real_close = NULL; - -/* Restore state tracking - * Cross-thread flags use _Atomic int to avoid undefined behavior. - */ -static _Atomic int g_restore_checked = 0; -static _Atomic int g_is_restored = 0; -/* g_io_uring_calls_since_restore removed - was used by uvloop fork tracking */ - -/* Config values: only written once during init or after restore reset, read-mostly */ -static int g_disable_io_uring_reinit = -1; -static int g_debug_io_uring = -1; - -/* Generic CRIU restore marker - used by all patched components */ -#define CRIU_RESTORE_MARKER "/run/criu-restored" - -/* Legacy NVSNAP-specific markers (for backwards compatibility) */ -#define NVSNAP_RESTORE_MARKER "/var/run/nvsnap/.restored" - -/* - * ============================================================================= - * IO_URING TRACKING WITH FULL PARAMS FOR TRANSPARENT REINIT - * ============================================================================= - */ - -/* - * io_uring_params structure (from linux/io_uring.h) - * Full definition for capturing setup params - */ -struct io_uring_params { - uint32_t sq_entries; - uint32_t cq_entries; - uint32_t flags; - uint32_t sq_thread_cpu; - uint32_t sq_thread_idle; - uint32_t features; - uint32_t wq_fd; - uint32_t resv[3]; - struct io_sqring_offsets { - uint32_t head; - uint32_t tail; - uint32_t ring_mask; - uint32_t ring_entries; - uint32_t flags; - uint32_t dropped; - uint32_t array; - uint32_t resv1; - uint64_t resv2; - } sq_off; - struct io_cqring_offsets { - uint32_t head; - uint32_t tail; - uint32_t ring_mask; - uint32_t ring_entries; - uint32_t overflow; - uint32_t cqes; - uint32_t flags; - uint32_t resv1; - uint64_t resv2; - } cq_off; -}; - -/* Maximum tracked instances */ -#define MAX_IO_URING_INSTANCES 256 - -/* - * Extended tracking structure with all info needed for reinit - */ -typedef struct { - int fd; /* io_uring file descriptor */ - int active; /* Is this slot in use? */ - int validated_after_restore; /* Has been verified working post-restore */ - int reinit_succeeded; /* Did reinit succeed? */ - int reinit_attempts; /* How many reinit attempts (max 3) */ - int invalidated_after_restore; /* Have we invalidated this fd post-restore */ - - /* Setup params - needed to recreate */ - uint32_t setup_entries; /* entries arg to io_uring_setup */ - struct io_uring_params params; /* Full params returned by kernel */ - - /* Memory mapping addresses - needed to remap at same locations */ - unsigned long sq_ring_addr; - unsigned long cq_ring_addr; /* May be same as sq_ring_addr if SINGLE_MMAP */ - unsigned long sqe_addr; - size_t sq_ring_size; - size_t cq_ring_size; - size_t sqe_size; - - /* Ring indices at last known good state */ - uint32_t sq_head; - uint32_t sq_tail; - uint32_t cq_head; - uint32_t cq_tail; - - pthread_t owner_thread; -} nvsnap_io_uring_instance_t; - -static nvsnap_io_uring_instance_t g_io_urings[MAX_IO_URING_INSTANCES]; -static int g_io_uring_count = 0; -static pthread_mutex_t g_io_uring_mutex = PTHREAD_MUTEX_INITIALIZER; - -/* Forward declarations for state flags defined later */ -static _Atomic int g_proactive_reinit_done; -static _Atomic int g_pre_restore_check_done; - -/* Check if we're in a restored process. - * IMPORTANT: We ALWAYS check the marker file because CRIU preserves - * memory state from before checkpoint, including all state flags. - * So we cannot rely on a one-time check; we must check on every call - * until we detect restore. - * - * When restore is detected, we also RESET all state flags to ensure - * the reinit logic runs. - */ -static int check_restored(void) { - /* If already detected as restored, return immediately */ - if (g_is_restored) return 1; - - int detected = 0; - - /* Check generic CRIU restore marker (preferred) */ - if (access(CRIU_RESTORE_MARKER, F_OK) == 0) { - NVSNAP_WARN("=== RESTORE DETECTED via %s ===", CRIU_RESTORE_MARKER); - detected = 1; - } - /* Check legacy NVSNAP-specific markers */ - else if (access("/nvsnap-lib/.restored", F_OK) == 0) { - NVSNAP_WARN("=== RESTORE DETECTED via /nvsnap-lib/.restored ==="); - detected = 1; - } else if (access("/nvsnap/.restored", F_OK) == 0) { - NVSNAP_WARN("=== RESTORE DETECTED via /nvsnap/.restored ==="); - detected = 1; - } else if (access(NVSNAP_RESTORE_MARKER, F_OK) == 0) { - NVSNAP_WARN("=== RESTORE DETECTED via %s ===", NVSNAP_RESTORE_MARKER); - detected = 1; - } - - /* Check environment variable */ - if (!detected && getenv("NVSNAP_RESTORED")) { - NVSNAP_WARN("=== RESTORE DETECTED via NVSNAP_RESTORED env ==="); - detected = 1; - } - - /* Check via quiesce.c export */ - if (!detected && nvsnap_is_restored()) { - NVSNAP_WARN("=== RESTORE DETECTED via quiesce ==="); - detected = 1; - } - - if (detected) { - g_is_restored = 1; - /* CRITICAL: Reset all state flags that were preserved from checkpoint. - * These flags would prevent reinit logic from running. */ - g_proactive_reinit_done = 0; - g_pre_restore_check_done = 0; - g_restore_checked = 0; - g_debug_io_uring = -1; - g_disable_io_uring_reinit = -1; - NVSNAP_WARN("Reset state flags: proactive_reinit_done=0, pre_restore_check_done=0"); - nvsnap_zmq_reinit_all_if_restored(); - return 1; - } - - return 0; -} - -static int is_debug_io_uring_enabled(void) { - if (g_debug_io_uring >= 0) { - return g_debug_io_uring; - } - const char *env = getenv("NVSNAP_DEBUG_IO_URING"); - if (env && (!strcmp(env, "1") || !strcasecmp(env, "true") || !strcasecmp(env, "yes"))) { - g_debug_io_uring = 1; - return 1; - } - if (access("/nvsnap-lib/.debug_io_uring", F_OK) == 0) { - g_debug_io_uring = 1; - return 1; - } - g_debug_io_uring = 0; - return 0; -} - -static int is_io_uring_reinit_disabled(void) { - if (g_disable_io_uring_reinit >= 0) { - return g_disable_io_uring_reinit; - } - const char *env = getenv("NVSNAP_DISABLE_IO_URING_REINIT"); - if (env && (!strcmp(env, "1") || !strcasecmp(env, "true") || !strcasecmp(env, "yes"))) { - g_disable_io_uring_reinit = 1; - return 1; - } - if (access("/nvsnap-lib/.disable_io_uring_reinit", F_OK) == 0) { - g_disable_io_uring_reinit = 1; - return 1; - } - g_disable_io_uring_reinit = 0; - return 0; -} - -/* - * Initialize real syscall pointer - */ -static void init_real_syscall(void) { - if (!real_syscall) { - real_syscall = dlsym(RTLD_NEXT, "syscall"); - if (!real_syscall) { - /* Fallback to libc */ - real_syscall = dlsym(RTLD_DEFAULT, "syscall"); - } - } -} - -/* - * ============================================================================= - * FD VERIFICATION AND INFO PARSING - * ============================================================================= - */ - -/* Verify an fd is actually an io_uring by checking /proc/self/fd link */ -static int is_valid_io_uring_fd(int fd) { - char link_path[64]; - char target[256]; - - snprintf(link_path, sizeof(link_path), "/proc/self/fd/%d", fd); - ssize_t len = readlink(link_path, target, sizeof(target) - 1); - - if (len <= 0) { - return 0; /* fd doesn't exist */ - } - target[len] = '\0'; - - /* Check if it's an io_uring */ - return (strstr(target, "io_uring") != NULL); -} - -/* Get what an fd points to */ -static int get_fd_type(int fd, char *buf, size_t buflen) { - char link_path[64]; - snprintf(link_path, sizeof(link_path), "/proc/self/fd/%d", fd); - ssize_t len = readlink(link_path, buf, buflen - 1); - if (len > 0) { - buf[len] = '\0'; - return 0; - } - return -1; -} - -/* Parse io_uring params from fdinfo */ -static int parse_io_uring_fdinfo(int fd, uint32_t *sq_entries, uint32_t *cq_entries, - uint32_t *sq_head, uint32_t *sq_tail, - uint32_t *cq_head, uint32_t *cq_tail) { - char fdinfo_path[64]; - char line[256]; - snprintf(fdinfo_path, sizeof(fdinfo_path), "/proc/self/fdinfo/%d", fd); - - FILE *f = fopen(fdinfo_path, "r"); - if (!f) return 0; - - int is_io_uring = 0; - *sq_entries = 0; *cq_entries = 0; - *sq_head = 0; *sq_tail = 0; *cq_head = 0; *cq_tail = 0; - - while (fgets(line, sizeof(line), f)) { - unsigned int val; - if (strncmp(line, "SqMask:", 7) == 0) { - is_io_uring = 1; - if (sscanf(line, "SqMask: %u", &val) == 1) { - *sq_entries = val + 1; - } - } else if (strncmp(line, "CqMask:", 7) == 0) { - if (sscanf(line, "CqMask: %u", &val) == 1) { - *cq_entries = val + 1; - } - } else if (strncmp(line, "SqHead:", 7) == 0) { - sscanf(line, "SqHead: %u", sq_head); - } else if (strncmp(line, "SqTail:", 7) == 0) { - sscanf(line, "SqTail: %u", sq_tail); - } else if (strncmp(line, "CqHead:", 7) == 0) { - sscanf(line, "CqHead: %u", cq_head); - } else if (strncmp(line, "CqTail:", 7) == 0) { - sscanf(line, "CqTail: %u", cq_tail); - } - } - - fclose(f); - return is_io_uring; -} - -static void dump_io_uring_fdinfo(int fd, const char *tag) { - uint32_t sq_entries = 0; - uint32_t cq_entries = 0; - uint32_t sq_head = 0; - uint32_t sq_tail = 0; - uint32_t cq_head = 0; - uint32_t cq_tail = 0; - - int is_ring = parse_io_uring_fdinfo(fd, &sq_entries, &cq_entries, - &sq_head, &sq_tail, &cq_head, &cq_tail); - if (!is_ring) { - NVSNAP_WARN("[%s] fd=%d is not io_uring or fdinfo missing", tag, fd); - return; - } - - NVSNAP_WARN("[%s] fd=%d sq_entries=%u cq_entries=%u sq=%u/%u cq=%u/%u", - tag, fd, sq_entries, cq_entries, sq_head, sq_tail, cq_head, cq_tail); -} - -/* Parse io_uring mmap addresses from /proc/self/maps */ -static void capture_io_uring_mmap_addrs(int fd, nvsnap_io_uring_instance_t *inst) { - char line[512]; - char fd_path[64]; - char fd_link[256]; - - snprintf(fd_path, sizeof(fd_path), "/proc/self/fd/%d", fd); - ssize_t link_len = readlink(fd_path, fd_link, sizeof(fd_link) - 1); - if (link_len <= 0) return; - fd_link[link_len] = '\0'; - - FILE *f = fopen("/proc/self/maps", "r"); - if (!f) return; - - while (fgets(line, sizeof(line), f)) { - if (strstr(line, fd_link) || strstr(line, "io_uring")) { - unsigned long start, end, offset; - char perms[8]; - - if (sscanf(line, "%lx-%lx %7s %lx", &start, &end, perms, &offset) >= 4) { - size_t size = end - start; - - /* Identify which mapping this is based on offset */ - if (offset == IORING_OFF_SQ_RING || offset == 0) { - if (inst->sq_ring_addr == 0) { - inst->sq_ring_addr = start; - inst->sq_ring_size = size; - NVSNAP_DEBUG(" Captured SQ ring: addr=0x%lx size=%zu", start, size); - } - } else if (offset == IORING_OFF_CQ_RING || offset == 0x8000000) { - if (inst->cq_ring_addr == 0 || inst->cq_ring_addr == inst->sq_ring_addr) { - inst->cq_ring_addr = start; - inst->cq_ring_size = size; - NVSNAP_DEBUG(" Captured CQ ring: addr=0x%lx size=%zu", start, size); - } - } else if (offset == IORING_OFF_SQES || offset == 0x10000000) { - if (inst->sqe_addr == 0) { - inst->sqe_addr = start; - inst->sqe_size = size; - NVSNAP_DEBUG(" Captured SQEs: addr=0x%lx size=%zu", start, size); - } - } - } - } - } - - fclose(f); -} - -/* - * ============================================================================= - * TRANSPARENT IO_URING REINIT - * ============================================================================= - */ - -/* Page-align a size */ -static inline unsigned long align_up(unsigned long size, unsigned long align) { - return (size + align - 1) & ~(align - 1); -} - -/* - * Transparently recreate an io_uring at the expected fd and mmap addresses. - * This is the core of the self-healing architecture. - */ -static int reinit_io_uring_at_fd(nvsnap_io_uring_instance_t *inst) { - init_real_syscall(); - - NVSNAP_WARN("=== TRANSPARENT IO_URING REINIT for fd=%d ===", inst->fd); - NVSNAP_INFO(" Original params: entries=%u flags=0x%x", - inst->setup_entries, inst->params.flags); - NVSNAP_INFO(" Saved addrs: sq=0x%lx cq=0x%lx sqe=0x%lx", - inst->sq_ring_addr, inst->cq_ring_addr, inst->sqe_addr); - - /* Prepare params for new io_uring - strip SQPOLL for restore compatibility */ - struct io_uring_params new_params = {0}; - new_params.flags = inst->params.flags & ~IORING_SETUP_SQPOLL; /* Strip SQPOLL */ - - /* Ensure we have real_close for cleanup */ - if (!real_close) { - real_close = dlsym(RTLD_NEXT, "close"); - } - - /* Create new io_uring */ - int new_fd = real_syscall(__NR_io_uring_setup, inst->setup_entries, &new_params); - if (new_fd < 0) { - NVSNAP_WARN(" io_uring_setup failed: %s", strerror(errno)); - return -1; - } - - NVSNAP_INFO(" Created new io_uring: fd=%d sq=%u cq=%u features=0x%x", - new_fd, new_params.sq_entries, new_params.cq_entries, new_params.features); - - int single_mmap = (new_params.features & IORING_FEAT_SINGLE_MMAP) != 0; - - /* Calculate mmap sizes from kernel params */ - size_t sq_ring_size = align_up(new_params.sq_off.array + - new_params.sq_entries * sizeof(uint32_t), 4096); - size_t sqe_size = align_up(new_params.sq_entries * 64, 4096); /* sizeof(io_uring_sqe) = 64 */ - size_t cq_ring_size = align_up(new_params.cq_off.cqes + - new_params.cq_entries * 16, 4096); /* sizeof(io_uring_cqe) = 16 */ - - if (single_mmap && cq_ring_size > sq_ring_size) { - sq_ring_size = cq_ring_size; - } - - /* - * Map rings using real_mmap to avoid recursive interception! - */ - if (!real_mmap) { - real_mmap = dlsym(RTLD_NEXT, "mmap"); - } - - /* Map SQ ring at the original address */ - if (inst->sq_ring_addr) { - /* First unmap anything at target address */ - munmap((void *)inst->sq_ring_addr, inst->sq_ring_size > 0 ? inst->sq_ring_size : sq_ring_size); - - void *sq_ptr = real_mmap((void *)inst->sq_ring_addr, sq_ring_size, - PROT_READ | PROT_WRITE, MAP_SHARED | MAP_FIXED, - new_fd, IORING_OFF_SQ_RING); - if (sq_ptr == MAP_FAILED) { - NVSNAP_WARN(" Failed to mmap SQ ring at 0x%lx: %s", - inst->sq_ring_addr, strerror(errno)); - real_close(new_fd); - return -1; - } - NVSNAP_INFO(" Mapped SQ ring at 0x%lx (size=%zu)", inst->sq_ring_addr, sq_ring_size); - - /* Map CQ ring if not single_mmap */ - if (!single_mmap && inst->cq_ring_addr && inst->cq_ring_addr != inst->sq_ring_addr) { - munmap((void *)inst->cq_ring_addr, inst->cq_ring_size > 0 ? inst->cq_ring_size : cq_ring_size); - - void *cq_ptr = real_mmap((void *)inst->cq_ring_addr, cq_ring_size, - PROT_READ | PROT_WRITE, MAP_SHARED | MAP_FIXED, - new_fd, IORING_OFF_CQ_RING); - if (cq_ptr == MAP_FAILED) { - NVSNAP_WARN(" Failed to mmap CQ ring at 0x%lx: %s", - inst->cq_ring_addr, strerror(errno)); - /* Non-fatal - continue */ - } else { - NVSNAP_INFO(" Mapped CQ ring at 0x%lx (size=%zu)", inst->cq_ring_addr, cq_ring_size); - } - } - } - - /* Map SQEs at the original address */ - if (inst->sqe_addr) { - munmap((void *)inst->sqe_addr, inst->sqe_size > 0 ? inst->sqe_size : sqe_size); - - void *sqe_ptr = real_mmap((void *)inst->sqe_addr, sqe_size, - PROT_READ | PROT_WRITE, MAP_SHARED | MAP_FIXED, - new_fd, IORING_OFF_SQES); - if (sqe_ptr == MAP_FAILED) { - NVSNAP_WARN(" Failed to mmap SQEs at 0x%lx: %s", - inst->sqe_addr, strerror(errno)); - real_close(new_fd); - return -1; - } - NVSNAP_INFO(" Mapped SQEs at 0x%lx (size=%zu)", inst->sqe_addr, sqe_size); - } - - /* Sync ring indices */ - if (inst->sq_ring_addr) { - unsigned long cq_base = single_mmap ? inst->sq_ring_addr : - (inst->cq_ring_addr ? inst->cq_ring_addr : inst->sq_ring_addr); - - volatile uint32_t *sq_head = (uint32_t *)((char *)inst->sq_ring_addr + new_params.sq_off.head); - volatile uint32_t *sq_tail = (uint32_t *)((char *)inst->sq_ring_addr + new_params.sq_off.tail); - volatile uint32_t *cq_head = (uint32_t *)((char *)cq_base + new_params.cq_off.head); - volatile uint32_t *cq_tail = (uint32_t *)((char *)cq_base + new_params.cq_off.tail); - volatile uint32_t *sqflags = (uint32_t *)((char *)inst->sq_ring_addr + new_params.sq_off.flags); - - /* Sync to saved indices (make rings appear empty at app's expected position) */ - *sq_head = inst->sq_tail; - *sq_tail = inst->sq_tail; - *cq_head = inst->cq_head; - *cq_tail = inst->cq_head; - - NVSNAP_INFO(" Synced indices: sq=%u/%u cq=%u/%u", *sq_head, *sq_tail, *cq_head, *cq_tail); - - /* Set NEED_WAKEUP since we stripped SQPOLL */ - *sqflags |= IORING_SQ_NEED_WAKEUP; - NVSNAP_INFO(" Set IORING_SQ_NEED_WAKEUP in sqflags"); - } - - /* Move new fd to expected fd number */ - if (new_fd != inst->fd) { - if (!real_close) { - real_close = dlsym(RTLD_NEXT, "close"); - } - - /* Close whatever is at the target fd */ - real_close(inst->fd); - - if (dup2(new_fd, inst->fd) < 0) { - NVSNAP_WARN(" dup2(%d, %d) failed: %s", new_fd, inst->fd, strerror(errno)); - real_close(new_fd); - return -1; - } - real_close(new_fd); - NVSNAP_INFO(" Moved io_uring from fd %d to fd %d", new_fd, inst->fd); - } - - /* Update tracking with new params */ - inst->params = new_params; - inst->reinit_succeeded = 1; - inst->reinit_attempts++; - inst->validated_after_restore = 1; - - NVSNAP_WARN("=== TRANSPARENT REINIT COMPLETE for fd=%d ===", inst->fd); - return 0; -} - -/* - * ============================================================================= - * PROACTIVE REINIT - * ============================================================================= - */ - -static _Atomic int g_proactive_reinit_done = 0; -static _Atomic int g_pre_restore_check_done = 0; - -/* - * Force reinit mode - set to 1 when we detect EBADF (indicating restore) - */ -static _Atomic int g_force_reinit_mode = 0; - -static void proactive_reinit_all_io_urings(void) { - if (g_proactive_reinit_done) return; - g_proactive_reinit_done = 1; - - if (g_force_reinit_mode) { - NVSNAP_WARN("=== PROACTIVE IO_URING REINIT (FORCED - post-restore) ==="); - } else { - NVSNAP_INFO("=== PROACTIVE IO_URING REINIT (checking validity) ==="); - } - - pthread_mutex_lock(&g_io_uring_mutex); - - int reinit_count = 0; - int success_count = 0; - int skipped_valid = 0; - - for (int i = 0; i < MAX_IO_URING_INSTANCES; i++) { - if (!g_io_urings[i].active) continue; - - nvsnap_io_uring_instance_t *inst = &g_io_urings[i]; - - if (inst->validated_after_restore || inst->reinit_succeeded || inst->reinit_attempts >= 3) continue; - - NVSNAP_INFO("Checking io_uring fd=%d (entries=%u, flags=0x%x, SQPOLL=%d)", - inst->fd, inst->setup_entries, inst->params.flags, - (inst->params.flags & IORING_SETUP_SQPOLL) != 0); - - if (!g_force_reinit_mode && is_valid_io_uring_fd(inst->fd)) { - inst->validated_after_restore = 1; - NVSNAP_INFO(" fd=%d is valid io_uring - skipping reinit", inst->fd); - skipped_valid++; - continue; - } - - char fd_type[256] = "unknown"; - get_fd_type(inst->fd, fd_type, sizeof(fd_type)); - NVSNAP_WARN(" fd=%d (now: %s) - reinitializing", inst->fd, fd_type); - - reinit_count++; - - if (inst->sq_ring_addr == 0 || inst->sqe_addr == 0) { - NVSNAP_WARN(" Cannot reinit fd=%d - missing mmap addresses (sq=0x%lx sqe=0x%lx)", - inst->fd, inst->sq_ring_addr, inst->sqe_addr); - inst->reinit_attempts = 3; - continue; - } - - if (reinit_io_uring_at_fd(inst) == 0) { - success_count++; - NVSNAP_INFO(" fd=%d reinit SUCCESS", inst->fd); - } else { - NVSNAP_WARN(" fd=%d reinit FAILED", inst->fd); - } - } - - pthread_mutex_unlock(&g_io_uring_mutex); - - NVSNAP_WARN("=== PROACTIVE REINIT COMPLETE: %d/%d reinited, %d skipped (valid) ===", - success_count, reinit_count, skipped_valid); - - if (success_count > 0 || g_force_reinit_mode) { - /* io_uring reinit complete. uvloop handles uv_loop_fork() natively now. */ - NVSNAP_INFO("Manual io_uring reinit complete (uvloop handles uv_loop_fork natively)"); - } -} - -/* - * ============================================================================= - * TRACKING FUNCTIONS (exported to quiesce.c) - * ============================================================================= - */ - -static void nvsnap_write_io_uring_map_locked(void) { - FILE *f = fopen("/nvsnap-lib/.io_uring_map", "w"); - if (!f) { - NVSNAP_WARN("Failed to write /nvsnap-lib/.io_uring_map: %s", strerror(errno)); - return; - } - - for (int i = 0; i < MAX_IO_URING_INSTANCES; i++) { - if (!g_io_urings[i].active) - continue; - fprintf(f, "%d %u %u 0x%x\n", - g_io_urings[i].fd, - g_io_urings[i].params.sq_entries, - g_io_urings[i].params.cq_entries, - g_io_urings[i].params.flags); - } - fflush(f); - fsync(fileno(f)); - fclose(f); -} - -/* Track a new io_uring instance */ -int nvsnap_track_io_uring(int fd, uint32_t sq_entries, uint32_t cq_entries, - uint32_t flags) { - pthread_mutex_lock(&g_io_uring_mutex); - - if (g_io_uring_count >= MAX_IO_URING_INSTANCES) { - NVSNAP_WARN("Too many io_uring instances (%d), can't track fd=%d", - g_io_uring_count, fd); - pthread_mutex_unlock(&g_io_uring_mutex); - return -1; - } - - /* Find empty slot */ - int slot = -1; - for (int i = 0; i < MAX_IO_URING_INSTANCES; i++) { - if (!g_io_urings[i].active) { - slot = i; - break; - } - } - - if (slot < 0) { - NVSNAP_WARN("No empty io_uring slots"); - pthread_mutex_unlock(&g_io_uring_mutex); - return -1; - } - - memset(&g_io_urings[slot], 0, sizeof(nvsnap_io_uring_instance_t)); - g_io_urings[slot].fd = fd; - g_io_urings[slot].active = 1; - g_io_urings[slot].setup_entries = sq_entries; - g_io_urings[slot].params.sq_entries = sq_entries; - g_io_urings[slot].params.cq_entries = cq_entries; - g_io_urings[slot].params.flags = flags; - g_io_urings[slot].owner_thread = pthread_self(); - g_io_uring_count++; - - NVSNAP_INFO("Tracked io_uring fd=%d entries=%u/%u flags=0x%x (SQPOLL=%d)", - fd, sq_entries, cq_entries, flags, (flags & IORING_SETUP_SQPOLL) != 0); - - nvsnap_write_io_uring_map_locked(); - - pthread_mutex_unlock(&g_io_uring_mutex); - return 0; -} - -/* Update mmap addresses for a tracked io_uring (call after mmap is done) */ -int nvsnap_update_io_uring_addrs(int fd) { - pthread_mutex_lock(&g_io_uring_mutex); - - for (int i = 0; i < MAX_IO_URING_INSTANCES; i++) { - if (g_io_urings[i].active && g_io_urings[i].fd == fd) { - capture_io_uring_mmap_addrs(fd, &g_io_urings[i]); - - /* Also capture current ring indices */ - parse_io_uring_fdinfo(fd, - &g_io_urings[i].params.sq_entries, - &g_io_urings[i].params.cq_entries, - &g_io_urings[i].sq_head, - &g_io_urings[i].sq_tail, - &g_io_urings[i].cq_head, - &g_io_urings[i].cq_tail); - - NVSNAP_DEBUG("Updated io_uring fd=%d addrs: sq=0x%lx sqe=0x%lx indices: sq=%u/%u cq=%u/%u", - fd, g_io_urings[i].sq_ring_addr, g_io_urings[i].sqe_addr, - g_io_urings[i].sq_head, g_io_urings[i].sq_tail, - g_io_urings[i].cq_head, g_io_urings[i].cq_tail); - - pthread_mutex_unlock(&g_io_uring_mutex); - return 0; - } - } - - pthread_mutex_unlock(&g_io_uring_mutex); - return -1; -} - -/* Untrack an io_uring instance (on close) */ -int nvsnap_untrack_io_uring(int fd) { - pthread_mutex_lock(&g_io_uring_mutex); - - for (int i = 0; i < MAX_IO_URING_INSTANCES; i++) { - if (g_io_urings[i].active && g_io_urings[i].fd == fd) { - NVSNAP_DEBUG("Untracked io_uring fd=%d", fd); - g_io_urings[i].active = 0; - g_io_uring_count--; - pthread_mutex_unlock(&g_io_uring_mutex); - return 0; - } - } - - pthread_mutex_unlock(&g_io_uring_mutex); - return -1; -} - -/* Check if an io_uring fd needs reinit after restore */ -int nvsnap_io_uring_needs_reinit(int fd) { - pthread_mutex_lock(&g_io_uring_mutex); - - for (int i = 0; i < MAX_IO_URING_INSTANCES; i++) { - if (g_io_urings[i].active && g_io_urings[i].fd == fd) { - int needs = !g_io_urings[i].validated_after_restore && - !g_io_urings[i].reinit_succeeded && - g_io_urings[i].reinit_attempts < 3; - pthread_mutex_unlock(&g_io_uring_mutex); - return needs; - } - } - - pthread_mutex_unlock(&g_io_uring_mutex); - return 0; -} - -/* Mark an io_uring fd as validated after restore */ -void nvsnap_mark_io_uring_validated(int fd) { - pthread_mutex_lock(&g_io_uring_mutex); - - for (int i = 0; i < MAX_IO_URING_INSTANCES; i++) { - if (g_io_urings[i].active && g_io_urings[i].fd == fd) { - g_io_urings[i].validated_after_restore = 1; - NVSNAP_DEBUG("io_uring fd=%d marked as validated after restore", fd); - break; - } - } - - pthread_mutex_unlock(&g_io_uring_mutex); -} - -/* - * ============================================================================= - * SYSCALL INTERCEPTION - * ============================================================================= - */ - -long syscall(long number, ...) { - init_real_syscall(); - - /* Check for quiesce request on any syscall (opportunistic) */ - nvsnap_perform_quiescence(); - - va_list ap; - va_start(ap, number); - - long ret; - - switch (number) { - case __NR_io_uring_setup: { - /* io_uring_setup(entries, params) */ - unsigned int entries = va_arg(ap, unsigned int); - struct io_uring_params* params = va_arg(ap, struct io_uring_params*); - va_end(ap); - - NVSNAP_DEBUG("io_uring_setup(entries=%u, params=%p)", entries, params); - - ret = real_syscall(__NR_io_uring_setup, entries, params); - - if (ret >= 0) { - uint32_t sq_entries = params ? params->sq_entries : entries; - uint32_t cq_entries = params ? params->cq_entries : entries * 2; - uint32_t flags = params ? params->flags : 0; - - nvsnap_track_io_uring((int)ret, sq_entries, cq_entries, flags); - - /* Store full params */ - pthread_mutex_lock(&g_io_uring_mutex); - for (int i = 0; i < MAX_IO_URING_INSTANCES; i++) { - if (g_io_urings[i].active && g_io_urings[i].fd == (int)ret) { - g_io_urings[i].setup_entries = entries; - if (params) { - g_io_urings[i].params = *params; - } - break; - } - } - pthread_mutex_unlock(&g_io_uring_mutex); - - if (flags & IORING_SETUP_SQPOLL) { - NVSNAP_WARN("io_uring fd=%ld has SQPOLL - kernel thread created", ret); - } - } - - return ret; - } - - case __NR_io_uring_enter: { - /* io_uring_enter(fd, to_submit, min_complete, flags, sig) */ - unsigned int fd = va_arg(ap, unsigned int); - unsigned int to_submit = va_arg(ap, unsigned int); - unsigned int min_complete = va_arg(ap, unsigned int); - unsigned int flags = va_arg(ap, unsigned int); - void* sig = va_arg(ap, void*); - va_end(ap); - - NVSNAP_TRACE("io_uring_enter(fd=%u, submit=%u, complete=%u, flags=0x%x)", - fd, to_submit, min_complete, flags); - if (is_debug_io_uring_enabled()) { - NVSNAP_INFO("io_uring_enter intercepted: fd=%u submit=%u complete=%u flags=0x%x", - fd, to_submit, min_complete, flags); - } - - if (is_io_uring_reinit_disabled()) { - ret = real_syscall(__NR_io_uring_enter, fd, to_submit, min_complete, - flags, sig); - if (is_debug_io_uring_enabled()) { - NVSNAP_INFO("io_uring_enter(fd=%u) -> ret=%ld errno=%d (%s)", - fd, ret, errno, strerror(errno)); - } - return ret; - } - - /* Pre-restore check: detect restore via marker files */ - if (!g_pre_restore_check_done && g_io_uring_count > 0) { - int is_restored = 0; - if (access(CRIU_RESTORE_MARKER, F_OK) == 0) { - NVSNAP_WARN("Restore marker found: %s", CRIU_RESTORE_MARKER); - is_restored = 1; - } else if (access(NVSNAP_RESTORE_MARKER, F_OK) == 0) { - NVSNAP_WARN("Restore marker found: %s", NVSNAP_RESTORE_MARKER); - is_restored = 1; - } else if (access("/nvsnap-lib/.restored", F_OK) == 0) { - NVSNAP_WARN("Restore marker found: /nvsnap-lib/.restored"); - is_restored = 1; - } else { - /* Also check by detecting broken io_uring */ - char fd_type[256]; - if (get_fd_type(fd, fd_type, sizeof(fd_type)) == 0) { - if (strstr(fd_type, "io_uring") == NULL) { - NVSNAP_WARN("io_uring fd=%u is now: %s - RESTORE DETECTED!", fd, fd_type); - is_restored = 1; - } - } - } - - if (is_restored) { - NVSNAP_WARN("=== RESTORE DETECTED (pre-syscall) ==="); - /* uvloop handles uv_loop_fork natively now - just do io_uring reinit */ - g_pre_restore_check_done = 1; - } - g_pre_restore_check_done = 1; - } - - /* Proactive reinit on first io_uring_enter after restore */ - if (!g_proactive_reinit_done && g_io_uring_count > 0) { - char fd_type[256]; - int fd_is_invalid = 0; - if (get_fd_type(fd, fd_type, sizeof(fd_type)) == 0) { - if (strstr(fd_type, "io_uring") == NULL) { - NVSNAP_WARN("io_uring fd=%u is now: %s - RESTORE DETECTED!", fd, fd_type); - fd_is_invalid = 1; - } - } - - if (fd_is_invalid) { - g_proactive_reinit_done = 1; - /* uvloop handles uv_loop_fork natively - just reinit io_urings */ - } - - proactive_reinit_all_io_urings(); - } - - /* Execute the actual syscall */ - ret = real_syscall(__NR_io_uring_enter, fd, to_submit, min_complete, - flags, sig); - if (is_debug_io_uring_enabled()) { - NVSNAP_INFO("io_uring_enter(fd=%u) -> ret=%ld errno=%d (%s)", - fd, ret, errno, strerror(errno)); - } - - if (is_debug_io_uring_enabled() && (flags & IORING_ENTER_GETEVENTS)) { - if (ret == -1 && errno == EINTR) { - NVSNAP_WARN("io_uring_enter(fd=%u) EINTR: submit=%u min=%u flags=0x%x", - fd, to_submit, min_complete, flags); - } - - if (ret >= 0 && min_complete > 0 && ret != (long)min_complete) { - NVSNAP_WARN("io_uring_enter(fd=%u) short getevents: ret=%ld expected=%u submit=%u flags=0x%x", - fd, ret, min_complete, to_submit, flags); - dump_io_uring_fdinfo(fd, "short-getevents"); - void *bt[32]; - int bt_size = backtrace(bt, (int)(sizeof(bt) / sizeof(bt[0]))); - backtrace_symbols_fd(bt, bt_size, STDERR_FILENO); - } - } - - /* If EBADF or ENOTSUP, try reinit and retry */ - if (ret < 0 && (errno == EBADF || errno == ENOTSUP || errno == EOPNOTSUPP)) { - int saved_errno = errno; - - NVSNAP_WARN("io_uring_enter(fd=%u) got %s - this indicates CRIU restore!", - fd, strerror(errno)); - - char fd_type2[256]; - if (get_fd_type(fd, fd_type2, sizeof(fd_type2)) == 0) { - NVSNAP_WARN(" fd=%u is currently: %s", fd, fd_type2); - } - - /* Reset and force reinit all io_urings */ - NVSNAP_WARN("Resetting proactive reinit flag and enabling FORCE mode"); - g_proactive_reinit_done = 0; - g_force_reinit_mode = 1; - - /* Reset validation state for ALL tracked io_urings */ - pthread_mutex_lock(&g_io_uring_mutex); - for (int i = 0; i < MAX_IO_URING_INSTANCES; i++) { - if (g_io_urings[i].active) { - g_io_urings[i].validated_after_restore = 0; - g_io_urings[i].reinit_succeeded = 0; - g_io_urings[i].reinit_attempts = 0; - } - } - pthread_mutex_unlock(&g_io_uring_mutex); - - NVSNAP_WARN("Running proactive reinit on ALL tracked io_urings"); - proactive_reinit_all_io_urings(); - - nvsnap_perform_restore_reinit(); - - /* Retry the syscall */ - ret = real_syscall(__NR_io_uring_enter, fd, to_submit, min_complete, - flags, sig); - - if (ret >= 0) { - NVSNAP_INFO("io_uring_enter(fd=%u) SUCCEEDED after proactive reinit!", fd); - } else { - NVSNAP_WARN("io_uring_enter(fd=%u) still failed after reinit: %s", - fd, strerror(errno)); - errno = saved_errno; - } - } - - return ret; - } - - case __NR_io_uring_register: { - unsigned int fd = va_arg(ap, unsigned int); - unsigned int opcode = va_arg(ap, unsigned int); - void* arg = va_arg(ap, void*); - unsigned int nr_args = va_arg(ap, unsigned int); - va_end(ap); - - NVSNAP_DEBUG("io_uring_register(fd=%u, opcode=%u, arg=%p, nr=%u)", - fd, opcode, arg, nr_args); - - ret = real_syscall(__NR_io_uring_register, fd, opcode, arg, nr_args); - return ret; - } - - case __NR_close: { - int fd = va_arg(ap, int); - va_end(ap); - - if (check_restored() && nvsnap_io_uring_needs_reinit(fd)) { - NVSNAP_WARN("CLOSE (via syscall) called on io_uring fd=%d AFTER RESTORE!", fd); - } - - nvsnap_untrack_io_uring(fd); - ret = real_syscall(__NR_close, fd); - return ret; - } - - default: { - long a1 = va_arg(ap, long); - long a2 = va_arg(ap, long); - long a3 = va_arg(ap, long); - long a4 = va_arg(ap, long); - long a5 = va_arg(ap, long); - long a6 = va_arg(ap, long); - va_end(ap); - - return real_syscall(number, a1, a2, a3, a4, a5, a6); - } - } -} - -/* - * ============================================================================= - * LIBURING INTERCEPTION (higher-level library) - * ============================================================================= - */ - -typedef int (*io_uring_queue_init_fn)(unsigned entries, void* ring, unsigned flags); -static io_uring_queue_init_fn real_io_uring_queue_init = NULL; - -int io_uring_queue_init(unsigned entries, void* ring, unsigned flags) { - if (!real_io_uring_queue_init) { - real_io_uring_queue_init = dlsym(RTLD_NEXT, "io_uring_queue_init"); - if (!real_io_uring_queue_init) { - NVSNAP_WARN("io_uring_queue_init not found - liburing not loaded?"); - errno = ENOSYS; - return -1; - } - } - - NVSNAP_DEBUG("io_uring_queue_init(entries=%u, ring=%p, flags=0x%x)", - entries, ring, flags); - - int ret = real_io_uring_queue_init(entries, ring, flags); - - if (ret == 0) { - int ring_fd = *(int*)ring; - nvsnap_track_io_uring(ring_fd, entries, entries * 2, flags); - nvsnap_update_io_uring_addrs(ring_fd); - - NVSNAP_INFO("io_uring_queue_init succeeded: fd=%d entries=%u flags=0x%x", - ring_fd, entries, flags); - } - - return ret; -} - -typedef int (*io_uring_queue_init_params_fn)(unsigned entries, void* ring, - struct io_uring_params* p); -static io_uring_queue_init_params_fn real_io_uring_queue_init_params = NULL; - -int io_uring_queue_init_params(unsigned entries, void* ring, struct io_uring_params* p) { - if (!real_io_uring_queue_init_params) { - real_io_uring_queue_init_params = dlsym(RTLD_NEXT, "io_uring_queue_init_params"); - if (!real_io_uring_queue_init_params) { - NVSNAP_WARN("io_uring_queue_init_params not found"); - errno = ENOSYS; - return -1; - } - } - - NVSNAP_DEBUG("io_uring_queue_init_params(entries=%u, ring=%p, params=%p)", - entries, ring, p); - - int ret = real_io_uring_queue_init_params(entries, ring, p); - - if (ret == 0) { - int ring_fd = *(int*)ring; - uint32_t sq_entries = p ? p->sq_entries : entries; - uint32_t cq_entries = p ? p->cq_entries : entries * 2; - uint32_t flags = p ? p->flags : 0; - - nvsnap_track_io_uring(ring_fd, sq_entries, cq_entries, flags); - - pthread_mutex_lock(&g_io_uring_mutex); - for (int i = 0; i < MAX_IO_URING_INSTANCES; i++) { - if (g_io_urings[i].active && g_io_urings[i].fd == ring_fd && p) { - g_io_urings[i].params = *p; - break; - } - } - pthread_mutex_unlock(&g_io_uring_mutex); - - nvsnap_update_io_uring_addrs(ring_fd); - - NVSNAP_INFO("io_uring_queue_init_params succeeded: fd=%d sq=%u cq=%u flags=0x%x", - ring_fd, sq_entries, cq_entries, flags); - } - - return ret; -} - -typedef void (*io_uring_queue_exit_fn)(void* ring); -static io_uring_queue_exit_fn real_io_uring_queue_exit = NULL; - -void io_uring_queue_exit(void* ring) { - if (!real_io_uring_queue_exit) { - real_io_uring_queue_exit = dlsym(RTLD_NEXT, "io_uring_queue_exit"); - if (!real_io_uring_queue_exit) { - return; - } - } - - if (ring) { - int ring_fd = *(int*)ring; - NVSNAP_DEBUG("io_uring_queue_exit(ring=%p fd=%d)", ring, ring_fd); - nvsnap_untrack_io_uring(ring_fd); - } - - real_io_uring_queue_exit(ring); -} - -/* - * ============================================================================= - * DIRECT close() INTERCEPTION - * ============================================================================= - */ - -int close(int fd) { - if (!real_close) { - real_close = dlsym(RTLD_NEXT, "close"); - if (!real_close) { - errno = EBADF; - return -1; - } - } - - if (check_restored() && nvsnap_io_uring_needs_reinit(fd)) { - NVSNAP_WARN("close(%d) called on io_uring fd AFTER RESTORE", fd); - } - - nvsnap_untrack_io_uring(fd); - return real_close(fd); -} - -/* - * ============================================================================= - * MMAP INTERCEPTION - Capture io_uring ring addresses - * ============================================================================= - */ - -void *mmap(void *addr, size_t length, int prot, int flags, int fd, off_t offset) { - if (!real_mmap) { - real_mmap = dlsym(RTLD_NEXT, "mmap"); - if (!real_mmap) { - errno = ENOMEM; - return MAP_FAILED; - } - } - - void *result = real_mmap(addr, length, prot, flags, fd, offset); - - /* If this is an io_uring mmap, capture the address */ - if (result != MAP_FAILED && fd >= 0) { - pthread_mutex_lock(&g_io_uring_mutex); - for (int i = 0; i < MAX_IO_URING_INSTANCES; i++) { - if (g_io_urings[i].active && g_io_urings[i].fd == fd) { - if (offset == IORING_OFF_SQ_RING || offset == 0) { - g_io_urings[i].sq_ring_addr = (unsigned long)result; - g_io_urings[i].sq_ring_size = length; - NVSNAP_DEBUG("Captured io_uring fd=%d SQ ring mmap: addr=%p size=%zu", - fd, result, length); - } else if (offset == IORING_OFF_CQ_RING) { - g_io_urings[i].cq_ring_addr = (unsigned long)result; - g_io_urings[i].cq_ring_size = length; - NVSNAP_DEBUG("Captured io_uring fd=%d CQ ring mmap: addr=%p size=%zu", - fd, result, length); - } else if (offset == IORING_OFF_SQES) { - g_io_urings[i].sqe_addr = (unsigned long)result; - g_io_urings[i].sqe_size = length; - NVSNAP_DEBUG("Captured io_uring fd=%d SQEs mmap: addr=%p size=%zu", - fd, result, length); - } - break; - } - } - pthread_mutex_unlock(&g_io_uring_mutex); - } - - return result; -} - -/* Also intercept mmap64 which some libraries use */ -void *mmap64(void *addr, size_t length, int prot, int flags, int fd, off_t offset) { - return mmap(addr, length, prot, flags, fd, offset); -} - -/* Stubs for removed uvloop functions (still referenced from header) */ -void nvsnap_dump_uvloop_metadata(void) { - /* No-op: uvloop handles uv_loop_fork natively now */ -} - -void nvsnap_install_uvloop_hook_async(void) { - /* No-op: uvloop handles uv_loop_fork natively now */ -} diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/libuv_intercept.c b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/libuv_intercept.c deleted file mode 100644 index dc43e060ca..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/libuv_intercept.c +++ /dev/null @@ -1,1043 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * libuv Interception for NVSNAP - * - * This module intercepts libuv calls to: - * 1. Track all uv_loop_t instances AND all uv_handle_t instances - * 2. Call uv_loop_fork() after CRIU restore - * 3. Reinitialize individual handles that become stale after restore - * - * The problem with uvloop after CRIU restore: - * - libuv's internal state (epoll fd, signal handlers) is stale - * - uvloop's Python objects contain cached C pointers to uv_handle_t - * - uv_loop_fork() fixes loop backend BUT NOT individual handles - * - * Our solution: - * - Track ALL handles (not just loops) as they're created - * - On restore, mark all handles as "stale" - * - Before any handle operation, check if handle needs reinit - * - Reinit handles lazily on first post-restore use - * - * Key insight: libuv handles have a consistent structure where the first - * field is always a pointer to the loop. We can use this to validate handles. - * - * Build: - * Part of libnvsnap_intercept.so - */ - -#define _GNU_SOURCE -#include -#include -#include -#include -#include -#include -#include -#include - -#include "nvsnap_intercept.h" - -/* External functions from quiesce.c */ -extern int nvsnap_track_libuv_loop(void* loop); -extern int nvsnap_ensure_libuv_loop_ready(void* loop); -extern int nvsnap_perform_quiescence(void); - -/* - * ============================================================================= - * HANDLE TRACKING - * ============================================================================= - * - * We track ALL libuv handles to enable post-restore reinitialization. - * Each handle type needs different reinit logic. - */ - -typedef enum { - NVSNAP_UV_UNKNOWN = 0, - NVSNAP_UV_ASYNC, - NVSNAP_UV_CHECK, - NVSNAP_UV_FS_EVENT, - NVSNAP_UV_FS_POLL, - NVSNAP_UV_IDLE, - NVSNAP_UV_PIPE, - NVSNAP_UV_POLL, - NVSNAP_UV_PREPARE, - NVSNAP_UV_PROCESS, - NVSNAP_UV_SIGNAL, - NVSNAP_UV_TCP, - NVSNAP_UV_TIMER, - NVSNAP_UV_TTY, - NVSNAP_UV_UDP, -} nvsnap_uv_handle_type_t; - -typedef struct nvsnap_uv_handle { - void* handle; /* uv_handle_t* pointer */ - void* loop; /* Owning loop */ - nvsnap_uv_handle_type_t type; /* Handle type for reinit dispatch */ - uint64_t generation; /* Generation when created */ - int needs_reinit; /* Flag: needs reinit after restore */ - int has_signum; /* Signal signum known */ - int signum; /* Signal number (if has_signum) */ - struct nvsnap_uv_handle* next; /* Linked list */ -} nvsnap_uv_handle_t; - -#define MAX_TRACKED_HANDLES 4096 -static nvsnap_uv_handle_t* g_handles = NULL; -static int g_handle_count = 0; -static pthread_mutex_t g_handle_mutex = PTHREAD_MUTEX_INITIALIZER; -static uint64_t g_generation = 1; /* Increments on each restore */ - -/* Forward declaration - full definition with restore detection below */ -static int g_libuv_restored = 0; -static void nvsnap_dump_handles(FILE* out); - -/* Minimal public prefix of uv_handle_t (data, loop) */ -typedef struct { - void* data; - void* loop; -} nvsnap_uv_handle_public_t; - -static void* get_handle_loop(void* handle) { - if (!handle) return NULL; - nvsnap_uv_handle_public_t* pub = (nvsnap_uv_handle_public_t*)handle; - return pub->loop; -} - -/* Track a handle */ -static void track_handle(void* handle, void* loop, nvsnap_uv_handle_type_t type) { - if (!handle) return; - - pthread_mutex_lock(&g_handle_mutex); - - /* Check if already tracked */ - for (nvsnap_uv_handle_t* h = g_handles; h; h = h->next) { - if (h->handle == handle) { - pthread_mutex_unlock(&g_handle_mutex); - return; - } - } - - /* Add new entry */ - nvsnap_uv_handle_t* entry = malloc(sizeof(nvsnap_uv_handle_t)); - if (!entry) { - pthread_mutex_unlock(&g_handle_mutex); - return; - } - - entry->handle = handle; - entry->loop = loop; - entry->type = type; - entry->generation = g_generation; - entry->needs_reinit = 0; - entry->has_signum = 0; - entry->signum = 0; - entry->next = g_handles; - g_handles = entry; - g_handle_count++; - - NVSNAP_DEBUG("Tracked handle %p type=%d loop=%p (total=%d)", - handle, type, loop, g_handle_count); - - pthread_mutex_unlock(&g_handle_mutex); -} - -/* Untrack a handle (on close) */ -static void untrack_handle(void* handle) { - if (!handle) return; - - pthread_mutex_lock(&g_handle_mutex); - - nvsnap_uv_handle_t** pp = &g_handles; - while (*pp) { - if ((*pp)->handle == handle) { - nvsnap_uv_handle_t* to_free = *pp; - *pp = (*pp)->next; - free(to_free); - g_handle_count--; - NVSNAP_DEBUG("Untracked handle %p (total=%d)", handle, g_handle_count); - pthread_mutex_unlock(&g_handle_mutex); - return; - } - pp = &(*pp)->next; - } - - pthread_mutex_unlock(&g_handle_mutex); -} - -/* Find a tracked handle */ -static nvsnap_uv_handle_t* find_handle(void* handle) { - for (nvsnap_uv_handle_t* h = g_handles; h; h = h->next) { - if (h->handle == handle) { - return h; - } - } - return NULL; -} - -static void set_handle_signum(void* handle, int signum) { - if (!handle) return; - pthread_mutex_lock(&g_handle_mutex); - nvsnap_uv_handle_t* h = find_handle(handle); - if (h) { - h->has_signum = 1; - h->signum = signum; - NVSNAP_DEBUG("Signal handle %p updated signum=%d", handle, signum); - } else { - NVSNAP_DEBUG("Signal handle %p not tracked; cannot record signum=%d", - handle, signum); - } - pthread_mutex_unlock(&g_handle_mutex); -} - -/* Mark all handles as needing reinit (called on restore detection) */ -void nvsnap_mark_handles_for_reinit(void) { - pthread_mutex_lock(&g_handle_mutex); - - g_generation++; - int count = 0; - - for (nvsnap_uv_handle_t* h = g_handles; h; h = h->next) { - h->needs_reinit = 1; - count++; - } - - NVSNAP_INFO("Marked %d handles for reinit (generation=%lu)", count, g_generation); - - pthread_mutex_unlock(&g_handle_mutex); -} - -/* - * ============================================================================= - * HANDLE REINITIALIZATION - * ============================================================================= - * - * Different handle types need different reinit strategies. - * The key insight is that most handles just need their internal - * state reset - the Python/application-level state is still valid. - */ - -/* (Function pointers are now obtained dynamically via get_real_libuv_func) */ - -/* - * Reinitialize a handle after restore. - * - * Strategy varies by type: - * - Timer: Just needs loop reinit, timer state preserved - * - Async: Needs loop reinit, internal eventfd/pipe recreated by uv_loop_fork - * - Process: Most complex - child process relationship may need verification - * - Signal: uv_loop_fork should handle this - * - TCP/UDP/Pipe: Socket state needs validation - */ -static int reinit_handle(nvsnap_uv_handle_t* h) { - if (!h || !h->needs_reinit) return 0; - - NVSNAP_DEBUG("Reinitializing handle %p type=%d", h->handle, h->type); - - /* First, ensure the loop is reinitialized */ - nvsnap_ensure_libuv_loop_ready(h->loop); - - /* - * For most handle types, uv_loop_fork() has already done the heavy lifting. - * The handles should work as long as: - * 1. The loop is valid - * 2. The handle's loop pointer matches - * - * We validate this by checking if the handle's loop pointer (first field - * in all uv_handle_t structures) matches what we expect. - */ - - /* uv_handle_t public prefix: data, loop */ - void* handle_loop = get_handle_loop(h->handle); - - if (handle_loop != h->loop) { - NVSNAP_WARN("Handle %p loop pointer mismatch: expected %p, got %p", - h->handle, h->loop, handle_loop); - /* This is bad - handle is corrupted or loop was replaced */ - /* For now, we can't fix this without more invasive changes */ - return -1; - } - - switch (h->type) { - case NVSNAP_UV_PROCESS: - /* - * Process handles are tricky. After restore: - * - The child PID field should still be valid - * - The child process was also restored (same process tree) - * - The exit callback and status fields should be intact - * - * Main concern: signal handling for SIGCHLD - * uv_loop_fork() should reinit the signal infrastructure. - * - * We mark it as reinited and hope for the best. - * If it fails, we'll catch it in the error handling. - */ - NVSNAP_INFO("Process handle %p: marked as reinited (child should be valid)", - h->handle); - break; - - case NVSNAP_UV_SIGNAL: - /* - * Signal handles need the loop's signal infrastructure. - * uv_loop_fork() should have reinited this. - */ - NVSNAP_DEBUG("Signal handle %p: relying on uv_loop_fork reinit", h->handle); - break; - - case NVSNAP_UV_TIMER: - case NVSNAP_UV_IDLE: - case NVSNAP_UV_PREPARE: - case NVSNAP_UV_CHECK: - /* - * These are simple handles that just need the loop to work. - * uv_loop_fork() is sufficient. - */ - NVSNAP_DEBUG("Simple handle %p type=%d: loop reinit sufficient", - h->handle, h->type); - break; - - case NVSNAP_UV_ASYNC: - /* - * Async handles use an eventfd or pipe pair. - * uv_loop_fork() recreates these. - * The handle should work after loop reinit. - */ - NVSNAP_DEBUG("Async handle %p: relying on uv_loop_fork", h->handle); - break; - - case NVSNAP_UV_TCP: - case NVSNAP_UV_UDP: - case NVSNAP_UV_PIPE: - /* - * Socket handles are complex: - * - The FD was restored by CRIU - * - But the socket state might need validation - * - * For now, we trust CRIU's socket restoration. - * If there are issues, we'll add specific handling. - */ - NVSNAP_DEBUG("Socket handle %p type=%d: trusting CRIU socket restore", - h->handle, h->type); - break; - - case NVSNAP_UV_POLL: - /* - * Poll handles wrap an external FD. - * The FD should be valid after CRIU restore. - * The poll registration in the loop is handled by uv_loop_fork. - */ - NVSNAP_DEBUG("Poll handle %p: relying on uv_loop_fork", h->handle); - break; - - default: - NVSNAP_DEBUG("Unknown handle type %d: assuming loop reinit sufficient", h->type); - break; - } - - h->needs_reinit = 0; - h->generation = g_generation; - - return 0; -} - -/* - * Ensure a handle is ready for use after restore. - * Called before any operation on the handle. - */ -static void ensure_handle_ready(void* handle) { - if (!g_libuv_restored || !handle) return; - - pthread_mutex_lock(&g_handle_mutex); - - nvsnap_uv_handle_t* h = find_handle(handle); - if (h && h->needs_reinit) { - reinit_handle(h); - } - - pthread_mutex_unlock(&g_handle_mutex); -} - -/* - * Function pointers are obtained dynamically via get_real_libuv_func() - * to avoid issues with dlvsym and versioned symbols. - * Each intercepted function caches its real function pointer locally. - */ - -/* - * Macro for intercepting loop functions. - * Uses get_real_libuv_func() which properly bypasses our dlsym hook. - */ -#define INTERCEPT_LOOP_FN(name, ...) \ - static int (*real_fn)() = NULL; \ - if (!real_fn) { \ - real_fn = get_real_libuv_func(#name); \ - if (!real_fn) { \ - NVSNAP_WARN(#name " not found"); \ - return -1; \ - } \ - } \ - if (check_if_restored()) { \ - nvsnap_ensure_libuv_loop_ready(loop); \ - } \ - return real_fn(__VA_ARGS__) - -/* Macro for void loop functions */ -#define INTERCEPT_LOOP_VOID_FN(name, ...) \ - static void (*real_fn)() = NULL; \ - if (!real_fn) { \ - real_fn = get_real_libuv_func(#name); \ - if (!real_fn) { \ - NVSNAP_WARN(#name " not found"); \ - return; \ - } \ - } \ - if (check_if_restored()) { \ - nvsnap_ensure_libuv_loop_ready(loop); \ - } \ - real_fn(__VA_ARGS__) - -/* - * ============================================================================= - * LOOP LIFECYCLE - * ============================================================================= - */ - -/* - * LIBUV INTERCEPTION STRATEGY - * - * Problem: We need to track libuv loops and reinitialize them after CRIU restore. - * - * Key insight: Our dlsym hook in dlopen_hook.c ONLY intercepts "cu*" and "nccl*" - * symbols. So dlsym(RTLD_NEXT, "uv_*") will pass through to real dlsym and work! - * - * The previous issue was using dlvsym with empty version string (""), which - * fails for unversioned symbols like libuv. We now use dlsym directly. - * - * Restore detection: CRIU preserves the original environment, so NVSNAP_RESTORED=1 - * won't be set in the restored process. Instead, we use a file marker: - * /var/run/nvsnap/.restored (created by restore-entrypoint before CRIU restore) - */ - -/* Real dlsym - obtained at init time to avoid recursion */ -static void* (*real_libuv_dlsym)(void*, const char*) = NULL; -static void* (*real_dlopen)(const char*, int) = NULL; -static void* uvloop_handle = NULL; - -typedef struct { - void* handle; -} nvsnap_uvloop_find_ctx_t; - -static int nvsnap_find_uvloop_cb(struct dl_phdr_info* info, size_t size, void* data) { - (void)size; - if (!info || !info->dlpi_name || !info->dlpi_name[0]) { - return 0; - } - if (strstr(info->dlpi_name, "uvloop") && strstr(info->dlpi_name, ".so")) { - nvsnap_uvloop_find_ctx_t* out = (nvsnap_uvloop_find_ctx_t*)data; - if (real_dlopen) { - out->handle = real_dlopen(info->dlpi_name, RTLD_LAZY | RTLD_NOLOAD); - if (out->handle) { - NVSNAP_DEBUG("Opened uvloop module: %s", info->dlpi_name); - return 1; /* stop iteration */ - } - } - } - return 0; -} - -/* Marker file for restore detection (g_libuv_restored defined at top of file) */ -#define NVSNAP_RESTORE_MARKER "/var/run/nvsnap/.restored" - -/* Handle for libuv library */ -static void* libuv_handle = NULL; - -/* Get real libuv function - use dlopen to get explicit handle */ -static void* get_real_libuv_func(const char* name) { - if (!real_libuv_dlsym) { - /* Bootstrap: get real dlsym via dlvsym (we don't hook dlvsym) */ - real_libuv_dlsym = dlvsym(RTLD_NEXT, "dlsym", "GLIBC_2.34"); - if (!real_libuv_dlsym) { - real_libuv_dlsym = dlvsym(RTLD_NEXT, "dlsym", "GLIBC_2.2.5"); - } - if (!real_libuv_dlsym) { - real_libuv_dlsym = dlvsym(RTLD_NEXT, "dlsym", "GLIBC_2.17"); - } - } - if (!real_libuv_dlsym) { - NVSNAP_WARN("Could not get real dlsym"); - return NULL; - } - - /* - * Strategy: Try to get an explicit handle to libuv, then use dlsym on that handle. - * This avoids finding our own intercepted symbols via RTLD_DEFAULT. - * - * uvloop bundles its own libuv, so we check for that first. - */ - if (!real_dlopen) { - real_dlopen = dlvsym(RTLD_NEXT, "dlopen", "GLIBC_2.34"); - if (!real_dlopen) { - real_dlopen = dlvsym(RTLD_NEXT, "dlopen", "GLIBC_2.2.5"); - } - if (!real_dlopen) { - real_dlopen = dlvsym(RTLD_NEXT, "dlopen", "GLIBC_2.17"); - } - } - - if (!libuv_handle) { - /* Try uvloop's bundled libuv first (it's typically at this path) */ - - if (real_dlopen) { - /* Try uvloop's bundled libuv - use RTLD_NOLOAD to check if already loaded */ - libuv_handle = real_dlopen("libuv.so.1", RTLD_LAZY | RTLD_NOLOAD); - if (!libuv_handle) { - libuv_handle = real_dlopen("libuv.so", RTLD_LAZY | RTLD_NOLOAD); - } - - /* If not loaded yet, try loading directly (may be in LD_LIBRARY_PATH) */ - if (!libuv_handle) { - libuv_handle = real_dlopen("libuv.so.1", RTLD_LAZY); - } - if (!libuv_handle) { - libuv_handle = real_dlopen("libuv.so", RTLD_LAZY); - } - - if (libuv_handle) { - NVSNAP_DEBUG("Opened libuv library: %p", libuv_handle); - } - } - } - - void* func = NULL; - - /* Try using explicit libuv handle first */ - if (libuv_handle) { - func = real_libuv_dlsym(libuv_handle, name); - if (func) { - NVSNAP_DEBUG("Found %s in libuv handle: %p", name, func); - return func; - } - } - - /* Fallback: use RTLD_NEXT to skip our library and find the next one */ - func = real_libuv_dlsym(RTLD_NEXT, name); - if (func) { - NVSNAP_DEBUG("Found %s via RTLD_NEXT: %p", name, func); - return func; - } - - /* Final fallback: resolve from uvloop module if it bundles libuv */ - if (!uvloop_handle && real_dlopen) { - nvsnap_uvloop_find_ctx_t ctx = {0}; - (void)dl_iterate_phdr(nvsnap_find_uvloop_cb, &ctx); - if (ctx.handle) { - uvloop_handle = ctx.handle; - } - } - - if (uvloop_handle) { - func = real_libuv_dlsym(uvloop_handle, name); - if (func) { - NVSNAP_DEBUG("Found %s in uvloop module: %p", name, func); - return func; - } - } - - NVSNAP_DEBUG("Could not find libuv function: %s", name); - return NULL; -} - -/* Check if we're in a restored process */ -static int check_if_restored(void) { - if (g_libuv_restored) return 1; - - /* Check env var (might be set by wrapper scripts) */ - if (getenv("NVSNAP_RESTORED")) { - g_libuv_restored = 1; - return 1; - } - - /* Check marker file */ - if (access(NVSNAP_RESTORE_MARKER, F_OK) == 0) { - g_libuv_restored = 1; - return 1; - } - - return 0; -} - -/* Called by quiesce.c when restore is detected */ -void nvsnap_libuv_enable_interception(void) { - g_libuv_restored = 1; - NVSNAP_INFO("libuv restore mode activated"); -} - -/* - * Note: We use weak symbols so that if libuv is statically linked (like in uvloop), - * our interception functions won't override the internal symbols. - * - * However, this doesn't always work because: - * 1. LD_PRELOAD typically overrides regardless of weak/strong - * 2. Cython extensions may use different symbol resolution - * - * So we also check if we can find the real function, and if not, we need a fallback. - */ - -/* Flag to disable libuv interception if real functions not available */ -static int g_libuv_interception_enabled = -1; /* -1 = not checked yet */ - -static int check_libuv_available(void) { - if (g_libuv_interception_enabled == -1) { - void* test = get_real_libuv_func("uv_loop_init"); - g_libuv_interception_enabled = (test != NULL) ? 1 : 0; - if (!g_libuv_interception_enabled) { - NVSNAP_INFO("libuv not dynamically available (may be statically linked in uvloop) - " - "disabling libuv interception"); - } else { - NVSNAP_INFO("libuv dynamically available - interception enabled"); - } - } - return g_libuv_interception_enabled; -} - -int uv_loop_init(void* loop) { - /* Check if libuv is available for interception */ - if (!check_libuv_available()) { - /* libuv is statically linked or not available. - * This function shouldn't even be called in that case, - * but if it is, we need to fail gracefully. */ - NVSNAP_WARN("uv_loop_init called but libuv not available for interception"); - return -1; /* EPERM - will cause uvloop to error */ - } - - static int (*real_fn)(void*) = NULL; - if (!real_fn) { - real_fn = get_real_libuv_func("uv_loop_init"); - if (!real_fn) { - NVSNAP_WARN("uv_loop_init: could not find real function"); - return -1; - } - } - - NVSNAP_DEBUG("uv_loop_init(%p)", loop); - - int ret = real_fn(loop); - - if (ret == 0) { - /* Track the loop for post-restore reinit */ - nvsnap_track_libuv_loop(loop); - NVSNAP_DEBUG("Tracked libuv loop %p", loop); - } - - return ret; -} - -int uv_loop_close(void* loop) { - static int (*real_fn)(void*) = NULL; - if (!real_fn) { - real_fn = get_real_libuv_func("uv_loop_close"); - if (!real_fn) return -1; - } - - NVSNAP_DEBUG("uv_loop_close(%p)", loop); - - /* No need to untrack - loop will be reused or freed */ - return real_fn(loop); -} - -/* - * ============================================================================= - * LOOP OPERATIONS - These need reinit after restore - * ============================================================================= - */ - -int uv_run(void* loop, int mode) { - INTERCEPT_LOOP_FN(uv_run, loop, mode); -} - -int uv_loop_alive(void* loop) { - INTERCEPT_LOOP_FN(uv_loop_alive, loop); -} - -int uv_backend_fd(void* loop) { - INTERCEPT_LOOP_FN(uv_backend_fd, loop); -} - -int uv_backend_timeout(void* loop) { - INTERCEPT_LOOP_FN(uv_backend_timeout, loop); -} - -void uv_stop(void* loop) { - INTERCEPT_LOOP_VOID_FN(uv_stop, loop); -} - -void uv_update_time(void* loop) { - INTERCEPT_LOOP_VOID_FN(uv_update_time, loop); -} - -uint64_t uv_now(void* loop) { - static uint64_t (*real_fn)(void*) = NULL; - if (!real_fn) { - real_fn = get_real_libuv_func("uv_now"); - if (!real_fn) return 0; - } - if (check_if_restored()) { - nvsnap_ensure_libuv_loop_ready(loop); - } - return real_fn(loop); -} - -/* - * ============================================================================= - * HANDLE INITIALIZATION - Track handles AND bind to loop - * ============================================================================= - */ - -/* Macro for init functions that also track the handle */ -#define INTERCEPT_INIT_FN(name, type, ...) \ - static int (*real_fn)() = NULL; \ - if (!real_fn) { \ - real_fn = get_real_libuv_func(#name); \ - if (!real_fn) { \ - NVSNAP_WARN(#name " not found"); \ - return -1; \ - } \ - } \ - if (check_if_restored()) { \ - nvsnap_ensure_libuv_loop_ready(loop); \ - } \ - int ret = real_fn(__VA_ARGS__); \ - if (ret == 0) { \ - track_handle(handle, loop, type); \ - } \ - return ret - -int uv_tcp_init(void* loop, void* handle) { - INTERCEPT_INIT_FN(uv_tcp_init, NVSNAP_UV_TCP, loop, handle); -} - -int uv_tcp_init_ex(void* loop, void* handle, unsigned int flags) { - INTERCEPT_INIT_FN(uv_tcp_init_ex, NVSNAP_UV_TCP, loop, handle, flags); -} - -int uv_udp_init(void* loop, void* handle) { - INTERCEPT_INIT_FN(uv_udp_init, NVSNAP_UV_UDP, loop, handle); -} - -int uv_udp_init_ex(void* loop, void* handle, unsigned int flags) { - INTERCEPT_INIT_FN(uv_udp_init_ex, NVSNAP_UV_UDP, loop, handle, flags); -} - -int uv_pipe_init(void* loop, void* handle, int ipc) { - INTERCEPT_INIT_FN(uv_pipe_init, NVSNAP_UV_PIPE, loop, handle, ipc); -} - -int uv_tty_init(void* loop, void* handle, int fd, int readable) { - INTERCEPT_INIT_FN(uv_tty_init, NVSNAP_UV_TTY, loop, handle, fd, readable); -} - -int uv_poll_init(void* loop, void* handle, int fd) { - INTERCEPT_INIT_FN(uv_poll_init, NVSNAP_UV_POLL, loop, handle, fd); -} - -int uv_poll_init_socket(void* loop, void* handle, int socket) { - INTERCEPT_INIT_FN(uv_poll_init_socket, NVSNAP_UV_POLL, loop, handle, socket); -} - -int uv_timer_init(void* loop, void* handle) { - INTERCEPT_INIT_FN(uv_timer_init, NVSNAP_UV_TIMER, loop, handle); -} - -int uv_prepare_init(void* loop, void* handle) { - INTERCEPT_INIT_FN(uv_prepare_init, NVSNAP_UV_PREPARE, loop, handle); -} - -int uv_check_init(void* loop, void* handle) { - INTERCEPT_INIT_FN(uv_check_init, NVSNAP_UV_CHECK, loop, handle); -} - -int uv_idle_init(void* loop, void* handle) { - INTERCEPT_INIT_FN(uv_idle_init, NVSNAP_UV_IDLE, loop, handle); -} - -int uv_async_init(void* loop, void* handle, void* cb) { - INTERCEPT_INIT_FN(uv_async_init, NVSNAP_UV_ASYNC, loop, handle, cb); -} - -int uv_signal_init(void* loop, void* handle) { - INTERCEPT_INIT_FN(uv_signal_init, NVSNAP_UV_SIGNAL, loop, handle); -} - -int uv_fs_event_init(void* loop, void* handle) { - INTERCEPT_INIT_FN(uv_fs_event_init, NVSNAP_UV_FS_EVENT, loop, handle); -} - -int uv_fs_poll_init(void* loop, void* handle) { - INTERCEPT_INIT_FN(uv_fs_poll_init, NVSNAP_UV_FS_POLL, loop, handle); -} - -/* - * ============================================================================= - * PROCESS HANDLING - Critical for uvloop subprocesses - * ============================================================================= - * - * uv_spawn is particularly important because: - * - vLLM uses multiprocessing with spawn context - * - uvloop creates subprocess handles (UVProcess) - * - After restore, these handles need to maintain connection to child - * - * Key insight: CRIU restores the entire process tree, so: - * - Child processes ARE restored with their original PIDs - * - The uv_process_t handle's pid field is still valid - * - We just need to ensure signal handling works - */ - -int uv_spawn(void* loop, void* handle, void* options) { - static int (*real_fn)(void*, void*, void*) = NULL; - if (!real_fn) { - real_fn = get_real_libuv_func("uv_spawn"); - if (!real_fn) { - NVSNAP_WARN("uv_spawn not found"); - return -1; - } - } - - if (check_if_restored()) { - nvsnap_ensure_libuv_loop_ready(loop); - } - - NVSNAP_DEBUG("uv_spawn(loop=%p, handle=%p, options=%p)", loop, handle, options); - - int ret = real_fn(loop, handle, options); - - if (ret == 0) { - /* Track this process handle - critical for post-restore reinit */ - track_handle(handle, loop, NVSNAP_UV_PROCESS); - NVSNAP_INFO("uv_spawn succeeded: handle=%p (process handle tracked)", handle); - } else { - NVSNAP_DEBUG("uv_spawn failed: %d", ret); - } - - return ret; -} - -/* - * ============================================================================= - * HANDLE CLOSE - Untrack handles when they're closed - * ============================================================================= - */ - -void uv_close(void* handle, void* close_cb) { - static void (*real_fn)(void*, void*) = NULL; - if (!real_fn) { - real_fn = get_real_libuv_func("uv_close"); - if (!real_fn) { - NVSNAP_WARN("uv_close not found"); - return; - } - } - - /* Ensure handle is reinited before close (in case close does cleanup) */ - if (check_if_restored()) { - ensure_handle_ready(handle); - } - - NVSNAP_DEBUG("uv_close(handle=%p, cb=%p)", handle, close_cb); - - /* Untrack the handle */ - untrack_handle(handle); - - real_fn(handle, close_cb); -} - -/* - * ============================================================================= - * HANDLE OPERATIONS - Ensure handle is ready before use - * ============================================================================= - */ - -/* Intercept handle operations to ensure reinit before use */ - -/* Macro for handle operations */ -#define INTERCEPT_HANDLE_OP(name, ...) \ - static int (*real_fn)() = NULL; \ - if (!real_fn) { \ - real_fn = get_real_libuv_func(#name); \ - if (!real_fn) return -1; \ - } \ - if (check_if_restored()) { \ - ensure_handle_ready(handle); \ - } \ - return real_fn(__VA_ARGS__) - -/* Timer operations */ -int uv_timer_start(void* handle, void* cb, uint64_t timeout, uint64_t repeat) { - INTERCEPT_HANDLE_OP(uv_timer_start, handle, cb, timeout, repeat); -} - -int uv_timer_stop(void* handle) { - INTERCEPT_HANDLE_OP(uv_timer_stop, handle); -} - -/* Signal operations */ -int uv_signal_start(void* handle, void* cb, int signum) { - NVSNAP_DEBUG("uv_signal_start(handle=%p, signum=%d)", handle, signum); - set_handle_signum(handle, signum); - INTERCEPT_HANDLE_OP(uv_signal_start, handle, cb, signum); -} - -int uv_signal_start_oneshot(void* handle, void* cb, int signum) { - set_handle_signum(handle, signum); - INTERCEPT_HANDLE_OP(uv_signal_start_oneshot, handle, cb, signum); -} - -int uv_signal_stop(void* handle) { - INTERCEPT_HANDLE_OP(uv_signal_stop, handle); -} - -/* Async send - important for cross-thread wakeup */ -int uv_async_send(void* handle) { - INTERCEPT_HANDLE_OP(uv_async_send, handle); -} - -/* Process kill - used to send signals to child processes */ -int uv_process_kill(void* handle, int signum) { - NVSNAP_DEBUG("uv_process_kill(handle=%p, signum=%d)", handle, signum); - INTERCEPT_HANDLE_OP(uv_process_kill, handle, signum); -} - -/* - * ============================================================================= - * SPECIAL: uv_loop_fork PASSTHROUGH - * ============================================================================= - * - * We don't intercept uv_loop_fork itself - we call it internally. - * This is a passthrough for apps that call it directly. - */ - -int uv_loop_fork_passthrough(void* loop) { - static int (*real_fn)(void*) = NULL; - if (!real_fn) { - real_fn = get_real_libuv_func("uv_loop_fork"); - if (!real_fn) { - NVSNAP_WARN("uv_loop_fork not found"); - return -1; - } - } - - NVSNAP_DEBUG("uv_loop_fork(%p) passthrough", loop); - return real_fn(loop); -} - -/* - * ============================================================================= - * UVLOOP-SPECIFIC HANDLING - * ============================================================================= - * - * uvloop uses libuv underneath, so our libuv interception catches it. - * However, uvloop also has Python-level cached state that we can't fix - * from C code. - * - * For uvloop, the remaining issues after uv_loop_fork() are: - * - * 1. UVProcess._handle - points to stale uv_process_t - * - This causes crashes when subprocess.wait() is called - * - Solution: Python-level patch in sitecustomize.py - * - * 2. Signal handlers - uvloop caches signal handler state - * - uv_loop_fork() should fix this - * - * 3. Child watchers - for monitoring subprocesses - * - Partially fixed by uv_loop_fork() - * - May need Python-level reinit - * - * The comprehensive solution requires: - * - This C library for libuv/io_uring (what we're doing) - * - Python sitecustomize.py for uvloop Python objects - * - OR: Patching uvloop source (better long-term) - */ - -/* - * Debug helper: Check if we're in a uvloop context - */ -static int is_uvloop_loaded(void) { - /* Check if uvloop's internal symbols are present */ - void* uvloop_run = dlsym(RTLD_DEFAULT, "uvloop_run"); - return uvloop_run != NULL; -} - -/* - * Called when restore is detected - marks all handles for reinit - */ -void nvsnap_libuv_on_restore(void) { - g_libuv_restored = 1; - nvsnap_mark_handles_for_reinit(); - nvsnap_dump_handles(stderr); -} - -/* - * Diagnostic: dump all tracked handles - */ -void nvsnap_dump_handles(FILE* out) { - pthread_mutex_lock(&g_handle_mutex); - - fprintf(out, "\n=== Tracked libuv handles (%d) ===\n", g_handle_count); - - const char* type_names[] = { - "UNKNOWN", "ASYNC", "CHECK", "FS_EVENT", "FS_POLL", - "IDLE", "PIPE", "POLL", "PREPARE", "PROCESS", - "SIGNAL", "TCP", "TIMER", "TTY", "UDP" - }; - - for (nvsnap_uv_handle_t* h = g_handles; h; h = h->next) { - const char* type_name = (h->type < sizeof(type_names)/sizeof(type_names[0])) - ? type_names[h->type] : "?"; - if (h->has_signum) { - fprintf(out, " handle=%p loop=%p type=%s gen=%lu needs_reinit=%d signum=%d\n", - h->handle, h->loop, type_name, h->generation, h->needs_reinit, - h->signum); - } else { - fprintf(out, " handle=%p loop=%p type=%s gen=%lu needs_reinit=%d\n", - h->handle, h->loop, type_name, h->generation, h->needs_reinit); - } - } - - fprintf(out, "================================\n\n"); - - pthread_mutex_unlock(&g_handle_mutex); -} - -__attribute__((constructor(103))) -static void libuv_intercept_init(void) { - if (nvsnap_self_disabled()) - return; - - /* Check if libuv interception is disabled */ - const char* disable_libuv = getenv("NVSNAP_DISABLE_LIBUV"); - if (disable_libuv && strcmp(disable_libuv, "1") == 0) { - NVSNAP_DEBUG("libuv interception disabled via NVSNAP_DISABLE_LIBUV=1"); - return; - } - - /* Check if lightweight mode - skip libuv/io_uring */ - const char* lightweight = getenv("NVSNAP_LIGHTWEIGHT"); - if (lightweight && strcmp(lightweight, "1") == 0) { - NVSNAP_DEBUG("libuv interception disabled (lightweight mode)"); - return; - } - - int uvloop = is_uvloop_loaded(); - - /* Check if we're in a restored process - use file marker */ - g_libuv_restored = check_if_restored(); - - NVSNAP_INFO("libuv intercept initialized (uvloop=%d, restored=%d, generation=%lu)", - uvloop, g_libuv_restored, g_generation); - - if (g_libuv_restored) { - NVSNAP_INFO("Restore detected - handles will be reinited on first use"); - nvsnap_mark_handles_for_reinit(); - } -} diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/nccl_intercept.c b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/nccl_intercept.c deleted file mode 100644 index 6c434cc1af..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/nccl_intercept.c +++ /dev/null @@ -1,714 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * NCCL Interception for Multi-GPU Checkpoint/Restore - * - * Hooks ncclCommInitRank() to track NCCL communicators. On SIGUSR1 quiesce, - * calls cudaDeviceSynchronize() + ncclCommAbort() on all tracked communicators - * to remove cross-GPU NCCL dependencies before cuda-checkpoint runs. - * - * Without this, cuda-checkpoint --action lock deadlocks on multi-GPU workloads - * because locking one GPU freezes its NCCL ring buffers, causing other GPUs to - * block on NCCL collectives that need the frozen GPU. - * - * Enable: NVSNAP_NCCL_INTERCEPT=1 - */ - -#define _GNU_SOURCE -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "nvsnap_intercept.h" - -/* ========================================================================= - * NCCL type definitions (no nccl.h dependency) - * ========================================================================= */ - -#define NCCL_UNIQUE_ID_BYTES 128 - -typedef void *ncclComm_t; - -typedef struct { - char internal[NCCL_UNIQUE_ID_BYTES]; -} ncclUniqueId; - -typedef enum { - ncclSuccess = 0, - ncclUnhandledCudaError = 1, - ncclSystemError = 2, - ncclInternalError = 3, - ncclInvalidArgument = 4, - ncclInvalidUsage = 5, -} ncclResult_t; - -/* CUDA runtime */ -typedef enum { - cudaSuccess = 0, -} cudaError_t; - -/* NCCL data types and reduction ops (opaque — just pass through) */ -typedef int ncclDataType_t; -typedef int ncclRedOp_t; - -/* CUDA stream (opaque pointer) */ -typedef void *cudaStream_t; - -/* ========================================================================= - * Function pointer types - * ========================================================================= */ - -typedef ncclResult_t (*ncclCommInitRank_fn)(ncclComm_t *comm, int nranks, - ncclUniqueId commId, int rank); -typedef ncclResult_t (*ncclCommAbort_fn)(ncclComm_t comm); -typedef ncclResult_t (*ncclCommDestroy_fn)(ncclComm_t comm); -typedef ncclResult_t (*ncclCommFinalize_fn)(ncclComm_t comm); -typedef ncclResult_t (*ncclGetUniqueId_fn)(ncclUniqueId *uniqueId); -typedef cudaError_t (*cudaDeviceSynchronize_fn)(void); - -/* Collective operation function pointer types */ -typedef ncclResult_t (*ncclAllReduce_fn)(const void *, void *, size_t, - ncclDataType_t, ncclRedOp_t, - ncclComm_t, cudaStream_t); -typedef ncclResult_t (*ncclBroadcast_fn)(const void *, void *, size_t, - ncclDataType_t, int, - ncclComm_t, cudaStream_t); -typedef ncclResult_t (*ncclAllGather_fn)(const void *, void *, size_t, - ncclDataType_t, - ncclComm_t, cudaStream_t); -typedef ncclResult_t (*ncclReduceScatter_fn)(const void *, void *, size_t, - ncclDataType_t, ncclRedOp_t, - ncclComm_t, cudaStream_t); -typedef ncclResult_t (*ncclSend_fn)(const void *, size_t, ncclDataType_t, - int, ncclComm_t, cudaStream_t); -typedef ncclResult_t (*ncclRecv_fn)(void *, size_t, ncclDataType_t, - int, ncclComm_t, cudaStream_t); - -/* ========================================================================= - * Global state - * ========================================================================= */ - -#define NVSNAP_NCCL_MAX_COMMS 64 - -typedef struct nvsnap_nccl_comm { - ncclComm_t real_comm; - int rank; - int nranks; - ncclUniqueId unique_id; - int aborted; -} nvsnap_nccl_comm_t; - -static nvsnap_nccl_comm_t g_nccl_comms[NVSNAP_NCCL_MAX_COMMS]; -static int g_nccl_ncomms = 0; -static pthread_mutex_t g_nccl_mutex = PTHREAD_MUTEX_INITIALIZER; -static int g_is_nccl_parent = 0; - -/* Real function pointers */ -static ncclCommInitRank_fn real_ncclCommInitRank = NULL; -static ncclCommAbort_fn real_ncclCommAbort = NULL; -static ncclCommDestroy_fn real_ncclCommDestroy = NULL; -static ncclCommFinalize_fn real_ncclCommFinalize = NULL; -static ncclGetUniqueId_fn real_ncclGetUniqueId = NULL; -static cudaDeviceSynchronize_fn real_cudaDeviceSynchronize = NULL; - -/* Real collective function pointers */ -static ncclAllReduce_fn real_ncclAllReduce = NULL; -static ncclBroadcast_fn real_ncclBroadcast = NULL; -static ncclAllGather_fn real_ncclAllGather = NULL; -static ncclReduceScatter_fn real_ncclReduceScatter = NULL; -static ncclSend_fn real_ncclSend = NULL; -static ncclRecv_fn real_ncclRecv = NULL; - -/* ========================================================================= - * Comm pointer remap table (old checkpoint comm → new restored comm) - * ========================================================================= */ - -#define NVSNAP_NCCL_MAX_REMAP 64 - -static ncclComm_t g_comm_remap_old[NVSNAP_NCCL_MAX_REMAP]; -static ncclComm_t g_comm_remap_new[NVSNAP_NCCL_MAX_REMAP]; -static int g_comm_remap_count = 0; - -static ncclComm_t nvsnap_nccl_remap_comm(ncclComm_t comm) -{ - for (int i = 0; i < g_comm_remap_count; i++) { - if (g_comm_remap_old[i] == comm) - return g_comm_remap_new[i]; - } - return comm; /* no remap found, pass through */ -} - -/* Library handles */ -static void *g_libnccl = NULL; -static void *g_libcudart = NULL; - -/* Use dlvsym to get the real dlsym, bypassing our override in zmq_intercept.c */ -static void *(*g_real_dlsym)(void *, const char *) = NULL; - -static void *nccl_real_dlsym(void *handle, const char *symbol) -{ - if (!g_real_dlsym) - g_real_dlsym = dlvsym(RTLD_NEXT, "dlsym", "GLIBC_2.2.5"); - if (g_real_dlsym) - return g_real_dlsym(handle, symbol); - return NULL; -} - -/* Init state */ -static pthread_once_t g_nccl_once = PTHREAD_ONCE_INIT; -static int g_nccl_enabled = -1; /* -1 = unchecked */ - -/* ========================================================================= - * Fork handling - * ========================================================================= */ - -/* - * Mark this process as the NCCL parent (called via pthread_atfork prepare). - * EngineCore creates NCCL comms then forks TP workers. After fork, both - * parent and children receive quiesce triggers. If the parent calls - * ncclCommFinalize/Destroy, it corrupts shared kernel-side NCCL state - * (proxy sockets, FDs) that workers are actively using. Skip quiesce - * in the parent — children own the comms after fork. - */ -void nvsnap_nccl_mark_parent(void) -{ - if (g_nccl_ncomms > 0) - g_is_nccl_parent = 1; -} - -/* - * Reset NCCL tracking state after fork. Called from quiesce.c's - * pthread_atfork child handler. The child process will create its own - * NCCL communicators — we must not carry stale entries from the parent. - */ -void nvsnap_nccl_atfork_child(void) -{ - g_nccl_once = (pthread_once_t)PTHREAD_ONCE_INIT; - g_nccl_ncomms = 0; - g_nccl_mutex = (pthread_mutex_t)PTHREAD_MUTEX_INITIALIZER; - real_ncclCommInitRank = NULL; - real_ncclCommAbort = NULL; - real_ncclCommDestroy = NULL; - real_ncclCommFinalize = NULL; - real_ncclGetUniqueId = NULL; - real_cudaDeviceSynchronize = NULL; - g_real_dlsym = NULL; - g_nccl_enabled = -1; - /* Don't close g_libnccl / g_libcudart — shared with parent, child inherits */ - g_libnccl = NULL; - g_libcudart = NULL; -} - -/* ========================================================================= - * Helpers - * ========================================================================= */ - -int nvsnap_nccl_is_enabled(void) -{ - if (g_nccl_enabled < 0) { - const char *env = getenv("NVSNAP_NCCL_INTERCEPT"); - g_nccl_enabled = (env && strcmp(env, "1") == 0) ? 1 : 0; - } - return g_nccl_enabled; -} - -static void nvsnap_nccl_load_real(void) -{ - if (!nvsnap_nccl_is_enabled()) - return; - - /* Find libnccl.so — try env var first, then standard names */ - const char *nccl_path = getenv("VLLM_NCCL_SO_PATH"); - if (nccl_path && nccl_path[0]) { - g_libnccl = dlopen(nccl_path, RTLD_LAZY | RTLD_GLOBAL); - if (g_libnccl) - NVSNAP_INFO("Loaded NCCL from VLLM_NCCL_SO_PATH: %s", nccl_path); - } - if (!g_libnccl) { - g_libnccl = dlopen("libnccl.so.2", RTLD_LAZY | RTLD_GLOBAL); - if (g_libnccl) - NVSNAP_INFO("Loaded NCCL from libnccl.so.2"); - } - if (!g_libnccl) { - /* Try RTLD_NEXT — NCCL may already be loaded by the application */ - real_ncclCommInitRank = (ncclCommInitRank_fn)nccl_real_dlsym(RTLD_NEXT, "ncclCommInitRank"); - if (real_ncclCommInitRank) { - NVSNAP_INFO("Found ncclCommInitRank via RTLD_NEXT"); - real_ncclCommAbort = (ncclCommAbort_fn)nccl_real_dlsym(RTLD_NEXT, "ncclCommAbort"); - real_ncclCommDestroy = (ncclCommDestroy_fn)nccl_real_dlsym(RTLD_NEXT, "ncclCommDestroy"); - real_ncclCommFinalize = (ncclCommFinalize_fn)nccl_real_dlsym(RTLD_NEXT, "ncclCommFinalize"); - real_ncclGetUniqueId = (ncclGetUniqueId_fn)nccl_real_dlsym(RTLD_NEXT, "ncclGetUniqueId"); - } else { - NVSNAP_WARN("Could not find libnccl.so — NCCL interception disabled"); - g_nccl_enabled = 0; - return; - } - } - - if (g_libnccl && !real_ncclCommInitRank) { - real_ncclCommInitRank = (ncclCommInitRank_fn)nccl_real_dlsym(g_libnccl, "ncclCommInitRank"); - real_ncclCommAbort = (ncclCommAbort_fn)nccl_real_dlsym(g_libnccl, "ncclCommAbort"); - real_ncclCommDestroy = (ncclCommDestroy_fn)nccl_real_dlsym(g_libnccl, "ncclCommDestroy"); - real_ncclCommFinalize = (ncclCommFinalize_fn)nccl_real_dlsym(g_libnccl, "ncclCommFinalize"); - real_ncclGetUniqueId = (ncclGetUniqueId_fn)nccl_real_dlsym(g_libnccl, "ncclGetUniqueId"); - } - - /* Load collective functions from wherever we found NCCL */ - { - void *h = g_libnccl ? g_libnccl : RTLD_NEXT; - if (!real_ncclAllReduce) - real_ncclAllReduce = (ncclAllReduce_fn)nccl_real_dlsym(h, "ncclAllReduce"); - if (!real_ncclBroadcast) - real_ncclBroadcast = (ncclBroadcast_fn)nccl_real_dlsym(h, "ncclBroadcast"); - if (!real_ncclAllGather) - real_ncclAllGather = (ncclAllGather_fn)nccl_real_dlsym(h, "ncclAllGather"); - if (!real_ncclReduceScatter) - real_ncclReduceScatter = (ncclReduceScatter_fn)nccl_real_dlsym(h, "ncclReduceScatter"); - if (!real_ncclSend) - real_ncclSend = (ncclSend_fn)nccl_real_dlsym(h, "ncclSend"); - if (!real_ncclRecv) - real_ncclRecv = (ncclRecv_fn)nccl_real_dlsym(h, "ncclRecv"); - } - - if (!real_ncclCommInitRank || !real_ncclCommAbort) { - NVSNAP_WARN("Could not resolve NCCL functions — interception disabled"); - g_nccl_enabled = 0; - return; - } - - /* Find cudaDeviceSynchronize */ - g_libcudart = dlopen("libcudart.so", RTLD_LAZY | RTLD_GLOBAL); - if (!g_libcudart) - g_libcudart = dlopen("libcudart.so.12", RTLD_LAZY | RTLD_GLOBAL); - if (g_libcudart) { - real_cudaDeviceSynchronize = (cudaDeviceSynchronize_fn)nccl_real_dlsym( - g_libcudart, "cudaDeviceSynchronize"); - } - if (!real_cudaDeviceSynchronize) { - real_cudaDeviceSynchronize = (cudaDeviceSynchronize_fn)nccl_real_dlsym( - RTLD_NEXT, "cudaDeviceSynchronize"); - } - if (!real_cudaDeviceSynchronize) - NVSNAP_WARN("Could not find cudaDeviceSynchronize — will skip GPU sync before NCCL abort"); - - NVSNAP_INFO("NCCL interception initialized: ncclCommInitRank=%p ncclCommAbort=%p cudaDeviceSync=%p", - (void *)real_ncclCommInitRank, (void *)real_ncclCommAbort, - (void *)real_cudaDeviceSynchronize); -} - -/* ========================================================================= - * Shared memory persistence for comm tracking across fork - * - * Parent writes comm entries to /dev/shm/nvsnap-nccl-comms-. - * After fork, child's nvsnap_nccl_quiesce() calls nvsnap_nccl_recover_comms() - * to load the parent's entries. Comm pointers are valid in children - * because fork() copies the address space. - * ========================================================================= */ - -#define NVSNAP_NCCL_SHM_PREFIX "nvsnap-nccl-comms-" - -typedef struct { - uint64_t comm_ptr; - int rank; - int nranks; -} nvsnap_nccl_shm_entry_t; - -/* Write current comm table to /dev/shm. Caller holds g_nccl_mutex. */ -static void nvsnap_nccl_persist_comms_locked(void) -{ - char path[PATH_MAX]; - snprintf(path, sizeof(path), "/dev/shm/%s%d", NVSNAP_NCCL_SHM_PREFIX, getpid()); - - FILE *f = fopen(path, "w"); - if (!f) { - NVSNAP_WARN("Failed to persist NCCL comms to %s: %s", path, strerror(errno)); - return; - } - - for (int i = 0; i < g_nccl_ncomms; i++) { - nvsnap_nccl_shm_entry_t entry = { - .comm_ptr = (uint64_t)g_nccl_comms[i].real_comm, - .rank = g_nccl_comms[i].rank, - .nranks = g_nccl_comms[i].nranks, - }; - fwrite(&entry, sizeof(entry), 1, f); - } - fclose(f); - NVSNAP_DEBUG("Persisted %d NCCL comms to %s", g_nccl_ncomms, path); -} - -/* Recover comm table from parent's shared memory file. */ -static int nvsnap_nccl_recover_comms(void) -{ - DIR *dir = opendir("/dev/shm"); - if (!dir) - return 0; - - int recovered = 0; - struct dirent *ent; - while ((ent = readdir(dir)) != NULL) { - if (strncmp(ent->d_name, NVSNAP_NCCL_SHM_PREFIX, - strlen(NVSNAP_NCCL_SHM_PREFIX)) != 0) - continue; - - char path[PATH_MAX]; - snprintf(path, sizeof(path), "/dev/shm/%s", ent->d_name); - - FILE *f = fopen(path, "r"); - if (!f) - continue; - - nvsnap_nccl_shm_entry_t entry; - while (fread(&entry, sizeof(entry), 1, f) == 1) { - if (g_nccl_ncomms >= NVSNAP_NCCL_MAX_COMMS) - break; - - /* Skip duplicates */ - int dup = 0; - for (int i = 0; i < g_nccl_ncomms; i++) { - if ((uint64_t)g_nccl_comms[i].real_comm == entry.comm_ptr) { - dup = 1; - break; - } - } - if (dup) - continue; - - g_nccl_comms[g_nccl_ncomms].real_comm = (ncclComm_t)entry.comm_ptr; - g_nccl_comms[g_nccl_ncomms].rank = entry.rank; - g_nccl_comms[g_nccl_ncomms].nranks = entry.nranks; - g_nccl_comms[g_nccl_ncomms].aborted = 0; - g_nccl_ncomms++; - recovered++; - } - fclose(f); - } - closedir(dir); - - if (recovered > 0) - NVSNAP_INFO("Recovered %d NCCL comms from shared memory", recovered); - - return recovered; -} - -/* ncclCommInitRank hook is now in nvsnap_interpose_nccl.c (NvSnap's version). - * NvSnap handles allocation tracking + NCCL lifecycle. - * NvSnap keeps collective hooks (ncclAllReduce, etc.) for restore remapping. */ - -/* ========================================================================= - * Quiesce: abort all communicators before checkpoint - * ========================================================================= */ - -void nvsnap_nccl_quiesce(void) -{ - if (!nvsnap_nccl_is_enabled()) - return; - - pthread_mutex_lock(&g_nccl_mutex); - int ncomms = g_nccl_ncomms; - pthread_mutex_unlock(&g_nccl_mutex); - - if (ncomms == 0) { - NVSNAP_INFO("NCCL quiesce: no communicators tracked, skipping"); - return; - } - - /* No-op: NCCL destroy is handled by NvSnap's nvsnap_pre_checkpoint_quiesce() - * which is called via dlsym from the quiesce path. NvSnap only tracks comms - * for observability — the actual destroy must be done by the library that - * also tracks allocations (NvSnap), because NCCL destroy order matters - * relative to P2P disable and D2H save. - * - * The working 322 GB checkpoint (v0.9.97) used this exact pattern: - * nvsnap_nccl_quiesce = no-op, NvSnap does the real work. */ - NVSNAP_INFO("NCCL quiesce: %d communicator(s) tracked (handled by NvSnap)", ncomms); -} - -/* ========================================================================= - * Restore: recreate NCCL communicators after CRIU restore - * - * All ranks coordinate via /dev/shm: - * - Rank 0 generates a new ncclUniqueId, writes it to /dev/shm/nvsnap-nccl-uid - * - Other ranks poll for /dev/shm/nvsnap-nccl-uid-ready, then read the ID - * - All ranks call ncclCommInitRank() with the new ID - * - Old comm pointers are mapped to new ones in the remap table - * ========================================================================= */ - -#define NVSNAP_NCCL_UID_SHM "/dev/shm/nvsnap-nccl-uid" -#define NVSNAP_NCCL_UID_READY_SHM "/dev/shm/nvsnap-nccl-uid-ready" -#define NVSNAP_NCCL_RESTORE_POLL_US 10000 /* 10ms */ -#define NVSNAP_NCCL_RESTORE_TIMEOUT_S 60 - -void nvsnap_nccl_restore(void) -{ - if (!nvsnap_nccl_is_enabled()) - return; - - pthread_once(&g_nccl_once, nvsnap_nccl_load_real); - - pthread_mutex_lock(&g_nccl_mutex); - int ncomms = g_nccl_ncomms; - pthread_mutex_unlock(&g_nccl_mutex); - - if (ncomms == 0) - return; - - if (!real_ncclCommInitRank || !real_ncclGetUniqueId) { - NVSNAP_ERROR("NCCL restore: missing ncclCommInitRank or ncclGetUniqueId"); - return; - } - - NVSNAP_INFO("NCCL restore: recreating %d communicator(s)", ncomms); - - /* - * Recreate each communicator group. Multiple comm groups (TP, PP, DP) - * are tracked separately in g_nccl_comms[]. We recreate them in order, - * using per-group shm files to coordinate unique IDs across ranks. - */ - pthread_mutex_lock(&g_nccl_mutex); - for (int i = 0; i < g_nccl_ncomms; i++) { - nvsnap_nccl_comm_t *entry = &g_nccl_comms[i]; - ncclUniqueId new_id; - char uid_path[256]; - char ready_path[256]; - - /* Per-comm-group shm paths to allow multiple groups */ - snprintf(uid_path, sizeof(uid_path), - NVSNAP_NCCL_UID_SHM "-%d", i); - snprintf(ready_path, sizeof(ready_path), - NVSNAP_NCCL_UID_READY_SHM "-%d", i); - - if (entry->rank == 0) { - /* Rank 0: generate new unique ID and share via shm */ - ncclResult_t rc = real_ncclGetUniqueId(&new_id); - if (rc != ncclSuccess) { - NVSNAP_ERROR("NCCL restore: ncclGetUniqueId failed rc=%d", (int)rc); - pthread_mutex_unlock(&g_nccl_mutex); - return; - } - - /* Write ID to shm */ - int fd = open(uid_path, O_CREAT | O_WRONLY | O_TRUNC, 0666); - if (fd < 0) { - NVSNAP_ERROR("NCCL restore: cannot create %s: %s", - uid_path, strerror(errno)); - pthread_mutex_unlock(&g_nccl_mutex); - return; - } - ssize_t n = write(fd, &new_id, sizeof(new_id)); - close(fd); - if (n != sizeof(new_id)) { - NVSNAP_ERROR("NCCL restore: short write to %s", uid_path); - pthread_mutex_unlock(&g_nccl_mutex); - return; - } - - /* Signal ready */ - fd = open(ready_path, O_CREAT | O_WRONLY, 0666); - if (fd >= 0) - close(fd); - - NVSNAP_INFO("NCCL restore: rank 0 wrote unique ID to %s", uid_path); - } else { - /* Other ranks: poll for ready marker */ - int waited = 0; - while (access(ready_path, F_OK) != 0) { - usleep(NVSNAP_NCCL_RESTORE_POLL_US); - waited++; - if (waited * NVSNAP_NCCL_RESTORE_POLL_US > - NVSNAP_NCCL_RESTORE_TIMEOUT_S * 1000000) { - NVSNAP_ERROR("NCCL restore: timeout waiting for %s", - ready_path); - pthread_mutex_unlock(&g_nccl_mutex); - return; - } - } - - /* Read unique ID */ - int fd = open(uid_path, O_RDONLY); - if (fd < 0) { - NVSNAP_ERROR("NCCL restore: cannot open %s: %s", - uid_path, strerror(errno)); - pthread_mutex_unlock(&g_nccl_mutex); - return; - } - ssize_t n = read(fd, &new_id, sizeof(new_id)); - close(fd); - if (n != sizeof(new_id)) { - NVSNAP_ERROR("NCCL restore: short read from %s", uid_path); - pthread_mutex_unlock(&g_nccl_mutex); - return; - } - - NVSNAP_INFO("NCCL restore: rank %d read unique ID from %s", - entry->rank, uid_path); - } - - /* All ranks: create new communicator */ - ncclComm_t new_comm = NULL; - NVSNAP_INFO("NCCL restore: ncclCommInitRank(nranks=%d, rank=%d) for comm %d", - entry->nranks, entry->rank, i); - - ncclResult_t rc = real_ncclCommInitRank(&new_comm, entry->nranks, - new_id, entry->rank); - if (rc != ncclSuccess) { - NVSNAP_ERROR("NCCL restore: ncclCommInitRank failed rc=%d for comm %d", - (int)rc, i); - pthread_mutex_unlock(&g_nccl_mutex); - return; - } - - NVSNAP_INFO("NCCL restore: comm %d recreated: old=%p new=%p rank=%d/%d", - i, (void *)entry->real_comm, (void *)new_comm, - entry->rank, entry->nranks); - - /* Store remap: old_comm → new_comm */ - if (g_comm_remap_count < NVSNAP_NCCL_MAX_REMAP) { - g_comm_remap_old[g_comm_remap_count] = entry->real_comm; - g_comm_remap_new[g_comm_remap_count] = new_comm; - g_comm_remap_count++; - } - - /* Update tracking entry */ - entry->real_comm = new_comm; - entry->aborted = 0; - memcpy(&entry->unique_id, &new_id, sizeof(ncclUniqueId)); - - /* Cleanup shm (rank 0 only, after all ranks have read) */ - if (entry->rank == 0) { - /* Small delay to ensure other ranks have read the file */ - usleep(100000); /* 100ms */ - unlink(uid_path); - unlink(ready_path); - } - } - pthread_mutex_unlock(&g_nccl_mutex); - - NVSNAP_INFO("NCCL restore: done (%d communicators recreated, %d remaps)", - ncomms, g_comm_remap_count); -} - -/* ========================================================================= - * Collective hooks: transparently remap old comm pointers to new ones - * - * After restore, PyTorch/vLLM hold stale ncclComm_t pointers from before - * checkpoint. These hooks intercept collective calls, look up the old - * pointer in the remap table, and substitute the new communicator. - * ========================================================================= */ - -ncclResult_t ncclAllReduce(const void *sendbuff, void *recvbuff, size_t count, - ncclDataType_t datatype, ncclRedOp_t op, - ncclComm_t comm, cudaStream_t stream) -{ - pthread_once(&g_nccl_once, nvsnap_nccl_load_real); - if (!real_ncclAllReduce) { - real_ncclAllReduce = (ncclAllReduce_fn)nccl_real_dlsym(RTLD_NEXT, "ncclAllReduce"); - if (!real_ncclAllReduce) return ncclInternalError; - } - return real_ncclAllReduce(sendbuff, recvbuff, count, datatype, op, - nvsnap_nccl_remap_comm(comm), stream); -} - -ncclResult_t ncclBroadcast(const void *sendbuff, void *recvbuff, size_t count, - ncclDataType_t datatype, int root, - ncclComm_t comm, cudaStream_t stream) -{ - pthread_once(&g_nccl_once, nvsnap_nccl_load_real); - if (!real_ncclBroadcast) { - real_ncclBroadcast = (ncclBroadcast_fn)nccl_real_dlsym(RTLD_NEXT, "ncclBroadcast"); - if (!real_ncclBroadcast) return ncclInternalError; - } - return real_ncclBroadcast(sendbuff, recvbuff, count, datatype, root, - nvsnap_nccl_remap_comm(comm), stream); -} - -ncclResult_t ncclAllGather(const void *sendbuff, void *recvbuff, size_t sendcount, - ncclDataType_t datatype, - ncclComm_t comm, cudaStream_t stream) -{ - pthread_once(&g_nccl_once, nvsnap_nccl_load_real); - if (!real_ncclAllGather) { - real_ncclAllGather = (ncclAllGather_fn)nccl_real_dlsym(RTLD_NEXT, "ncclAllGather"); - if (!real_ncclAllGather) return ncclInternalError; - } - return real_ncclAllGather(sendbuff, recvbuff, sendcount, datatype, - nvsnap_nccl_remap_comm(comm), stream); -} - -ncclResult_t ncclReduceScatter(const void *sendbuff, void *recvbuff, - size_t recvcount, ncclDataType_t datatype, - ncclRedOp_t op, ncclComm_t comm, - cudaStream_t stream) -{ - pthread_once(&g_nccl_once, nvsnap_nccl_load_real); - if (!real_ncclReduceScatter) { - real_ncclReduceScatter = (ncclReduceScatter_fn)nccl_real_dlsym(RTLD_NEXT, "ncclReduceScatter"); - if (!real_ncclReduceScatter) return ncclInternalError; - } - return real_ncclReduceScatter(sendbuff, recvbuff, recvcount, datatype, - op, nvsnap_nccl_remap_comm(comm), stream); -} - -ncclResult_t ncclSend(const void *sendbuff, size_t count, - ncclDataType_t datatype, int peer, - ncclComm_t comm, cudaStream_t stream) -{ - pthread_once(&g_nccl_once, nvsnap_nccl_load_real); - if (!real_ncclSend) { - real_ncclSend = (ncclSend_fn)nccl_real_dlsym(RTLD_NEXT, "ncclSend"); - if (!real_ncclSend) return ncclInternalError; - } - return real_ncclSend(sendbuff, count, datatype, peer, - nvsnap_nccl_remap_comm(comm), stream); -} - -ncclResult_t ncclRecv(void *recvbuff, size_t count, - ncclDataType_t datatype, int peer, - ncclComm_t comm, cudaStream_t stream) -{ - pthread_once(&g_nccl_once, nvsnap_nccl_load_real); - if (!real_ncclRecv) { - real_ncclRecv = (ncclRecv_fn)nccl_real_dlsym(RTLD_NEXT, "ncclRecv"); - if (!real_ncclRecv) return ncclInternalError; - } - return real_ncclRecv(recvbuff, count, datatype, peer, - nvsnap_nccl_remap_comm(comm), stream); -} - -/* ========================================================================= - * dlsym override: redirect NCCL symbol lookups to our hooks - * - * Called from the dlsym override in zmq_intercept.c. - * Returns our wrapper function if the symbol matches, NULL otherwise. - * ========================================================================= */ - -void *nvsnap_nccl_symbol_override(const char *symbol) -{ - if (!nvsnap_nccl_is_enabled()) - return NULL; - /* ncclCommInitRank is handled by NvSnap's symbol table */ - if (strcmp(symbol, "ncclAllReduce") == 0) - return (void *)ncclAllReduce; - if (strcmp(symbol, "ncclBroadcast") == 0) - return (void *)ncclBroadcast; - if (strcmp(symbol, "ncclAllGather") == 0) - return (void *)ncclAllGather; - if (strcmp(symbol, "ncclReduceScatter") == 0) - return (void *)ncclReduceScatter; - if (strcmp(symbol, "ncclSend") == 0) - return (void *)ncclSend; - if (strcmp(symbol, "ncclRecv") == 0) - return (void *)ncclRecv; - return NULL; -} diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/quiesce.c b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/quiesce.c deleted file mode 100644 index b3c1c2a407..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/quiesce.c +++ /dev/null @@ -1,878 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * NVSNAP Quiescence Coordinator - * - * This module coordinates pre-checkpoint quiescence for io_uring and libuv. - * - * The problem: - * - io_uring with SQPOLL creates kernel threads that CRIU can't checkpoint - * - libuv/uvloop has C-level state that becomes stale after CRIU restore - * - We need a generic solution that works with ANY container, no app changes - * - * Solution: - * - LD_PRELOAD intercepts io_uring_setup(), libuv calls - * - On SIGUSR1 (pre-checkpoint): - * 1. Drain all io_uring rings (wait for pending I/O) - * 2. Stop SQPOLL kernel threads - * 3. Mark libuv loops for reinit after restore - * - On restore detection (NVSNAP_RESTORED=1): - * 1. Recreate io_uring instances (via transparent reinit in io_uring_intercept.c) - * 2. Reinit libuv loops with uv_loop_fork() - * 3. Handle invalidation for stale handles - * - * NOTE: io_uring tracking and reinit logic is in io_uring_intercept.c - * This file handles libuv tracking and overall quiescence coordination. - * - * Build: - * Part of libnvsnap_intercept.so - * - * Usage: - * LD_PRELOAD=/opt/nvsnap/lib/libnvsnap_intercept.so your_app - */ - -#define _GNU_SOURCE -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "nvsnap_intercept.h" - -/* io_uring syscall numbers */ -#ifndef __NR_io_uring_enter -#define __NR_io_uring_enter 426 -#endif - -/* io_uring constants (from linux/io_uring.h) */ -#define IORING_ENTER_GETEVENTS (1U << 0) - -/* Maximum tracked instances */ -#define MAX_LIBUV_LOOPS 64 - -/* - * ============================================================================= - * LIBUV TRACKING - * ============================================================================= - */ - -typedef struct { - void* loop; /* uv_loop_t pointer */ - int reinitialized; /* Has uv_loop_fork() been called? */ - pthread_t owner_thread; -} nvsnap_libuv_loop_t; - -static nvsnap_libuv_loop_t g_libuv_loops[MAX_LIBUV_LOOPS]; -static int g_libuv_count = 0; -static pthread_mutex_t g_libuv_mutex = PTHREAD_MUTEX_INITIALIZER; - -/* - * ============================================================================= - * QUIESCENCE STATE - * ============================================================================= - */ - -typedef enum { - NVSNAP_QUIESCE_NONE = 0, - NVSNAP_QUIESCE_REQUESTED, /* SIGUSR1 received, quiesce requested */ - NVSNAP_QUIESCE_IN_PROGRESS, /* One thread is performing quiescence */ - NVSNAP_QUIESCE_COMPLETE, /* All I/O drained, safe to checkpoint */ - NVSNAP_QUIESCE_RESUMED, /* SIGUSR2 received, resume normal operation */ -} nvsnap_quiesce_state_t; - -static atomic_int g_quiesce_state = NVSNAP_QUIESCE_NONE; -static _Atomic int g_is_restored = 0; -static int g_ack_pipe_fd = -1; -static int g_quiesce_pipe[2] = {-1, -1}; -static int g_quiesce_meta_thread_started = 0; -static int g_quiesce_worker_thread_started = 0; - - -/* Forward declarations */ -static void nvsnap_quiesce_libuv(void); -static void nvsnap_restore_libuv(void); -static void quiesce_signal_handler(int sig); -static void resume_signal_handler(int sig); -/* Metadata dump from io_uring_intercept.c */ -void nvsnap_dump_uvloop_metadata(void); - -/* - * ============================================================================= - * SIGACTION GUARD - * ============================================================================= - * - * Python's asyncio.loop.add_signal_handler(SIGUSR1, ...) (called by uvicorn - * during vLLM startup) replaces our C-level handler. We intercept sigaction() - * to prevent this: save their handler for chaining, keep ours installed. - * - * Guard only active when NVSNAP_QUIESCE_SIGNALS=1 is set. - * Restore pods do NOT set this, so CRIU is never affected. - */ -static int g_sigaction_guard_enabled = 0; -static struct sigaction g_saved_sigusr1 = { .sa_handler = SIG_DFL }; -static struct sigaction g_saved_sigusr2 = { .sa_handler = SIG_DFL }; -static int (*real_sigaction)(int, const struct sigaction *, struct sigaction *) = NULL; - -static void resolve_real_sigaction(void) { - /* Use __sigaction to get glibc's implementation directly. - * We can't use dlvsym(RTLD_NEXT, "sigaction", "GLIBC_2.2.5") because - * RTLD_NEXT resolves back to our own versioned symbol. */ - real_sigaction = (int (*)(int, const struct sigaction *, struct sigaction *)) - dlsym(RTLD_NEXT, "__sigaction"); - if (!real_sigaction) - real_sigaction = (int (*)(int, const struct sigaction *, struct sigaction *)) - dlsym(RTLD_NEXT, "sigaction"); -} - -/* - * Interpose sigaction with glibc version tag. - * - * Libraries linked against glibc (e.g. libtorch_cpu.so) reference the - * versioned symbol sigaction@GLIBC_2.2.5. An unversioned LD_PRELOAD - * override does NOT intercept versioned references — the dynamic linker - * resolves them directly to glibc. We must export our override with the - * same version tag to actually interpose these calls. - */ -__asm__(".symver nvsnap_sigaction,sigaction@@GLIBC_2.2.5"); -int nvsnap_sigaction(int signum, const struct sigaction *act, struct sigaction *oldact) { - if (!real_sigaction) { - resolve_real_sigaction(); - if (!real_sigaction) { errno = EINVAL; return -1; } - } - if (act != NULL && (signum == SIGUSR1 || signum == SIGUSR2)) { - if (g_sigaction_guard_enabled) { - NVSNAP_INFO("sigaction guard BLOCKED override of %s (handler=%p, guard=1)", - signum == SIGUSR1 ? "SIGUSR1" : "SIGUSR2", (void *)act->sa_handler); - if (signum == SIGUSR1) g_saved_sigusr1 = *act; - else g_saved_sigusr2 = *act; - if (oldact) return real_sigaction(signum, NULL, oldact); - return 0; - } else { - NVSNAP_INFO("sigaction ALLOWING override of %s (handler=%p, guard=0)", - signum == SIGUSR1 ? "SIGUSR1" : "SIGUSR2", (void *)act->sa_handler); - } - } - return real_sigaction(signum, act, oldact); -} - -/* - * Guard signal() the same way we guard sigaction(). - * Python's PyOS_setsig() may use signal() on some platforms, - * which would bypass our sigaction() guard. - */ -static void (*(*real_signal)(int, void (*)(int)))(int) = NULL; - -typedef void (*sighandler_t)(int); - -__asm__(".symver nvsnap_signal,signal@@GLIBC_2.2.5"); -sighandler_t nvsnap_signal(int signum, sighandler_t handler) { - if (!real_signal) { - /* Use bsd_signal to get glibc's implementation directly. - * Can't use RTLD_NEXT on "signal" — resolves back to us. */ - real_signal = (sighandler_t (*)(int, sighandler_t)) - dlsym(RTLD_NEXT, "bsd_signal"); - if (!real_signal) - real_signal = (sighandler_t (*)(int, sighandler_t)) - dlsym(RTLD_NEXT, "signal"); - } - if (!real_signal) { errno = EINVAL; return SIG_ERR; } - if (g_sigaction_guard_enabled && handler != SIG_ERR && - (signum == SIGUSR1 || signum == SIGUSR2)) { - /* Save their handler, keep ours installed */ - if (signum == SIGUSR1) g_saved_sigusr1.sa_handler = handler; - else g_saved_sigusr2.sa_handler = handler; - return SIG_DFL; /* Pretend old handler was default */ - } - return real_signal(signum, handler); -} - -static int install_signal_handler(int signum, void (*handler)(int)) { - if (!real_sigaction) return -1; - struct sigaction sa; - memset(&sa, 0, sizeof(sa)); - sa.sa_handler = handler; - /* SA_RESTART for SIGUSR1: don't disrupt app syscalls (NCCL proxy threads, - * Python I/O, etc). The quiesce worker thread polls every 50ms — that's - * fast enough. No SA_RESTART for SIGUSR2: we WANT EINTR to wake threads - * stuck in epoll_wait/poll after CRIU restore. */ - sa.sa_flags = (signum == SIGUSR1) ? SA_RESTART : 0; - sigemptyset(&sa.sa_mask); - return real_sigaction(signum, &sa, NULL); -} - -/* - * Quiesce worker thread: polls g_quiesce_state and performs quiescence. - * - * The SIGUSR1 signal handler only sets an atomic flag (async-signal-safe). - * Something must poll that flag and do the actual work. For io_uring-using - * processes, this happens inside io_uring_enter(). But NCCL worker processes - * sit idle between requests and never call io_uring_enter(), so the flag - * goes unnoticed. This thread ensures ALL processes perform quiescence - * regardless of their I/O pattern. - */ -static void* quiesce_worker_thread(void* arg) { - (void)arg; - char trigger_path[64]; - snprintf(trigger_path, sizeof(trigger_path), - "/dev/shm/nvsnap-quiesce-trigger-%d", getpid()); - - for (;;) { - /* File-based trigger: agent writes this file to request quiesce. - * Content doesn't matter — it's just a trigger. The checkpoint path - * is in a separate persistent file (/dev/shm/nvsnap-checkpoint-dir) - * read at save time, avoiding any race with SIGUSR1. */ - if (access(trigger_path, F_OK) == 0) { - NVSNAP_INFO("Quiesce trigger file detected: %s", trigger_path); - unlink(trigger_path); - - int expected = NVSNAP_QUIESCE_NONE; - atomic_compare_exchange_strong(&g_quiesce_state, - &expected, NVSNAP_QUIESCE_REQUESTED); - } - - if (atomic_load(&g_quiesce_state) == NVSNAP_QUIESCE_REQUESTED) { - nvsnap_perform_quiescence(); - } - - usleep(50000); /* 50ms poll — file check doesn't need 1ms */ - } - return NULL; -} - -static void* quiesce_meta_thread(void* arg) { - (void)arg; - char buf[16]; - for (;;) { - ssize_t n = read(g_quiesce_pipe[0], buf, sizeof(buf)); - if (n > 0) { - nvsnap_dump_uvloop_metadata(); - continue; - } - if (n < 0 && errno == EINTR) { - continue; - } - usleep(1000); - } - return NULL; -} - -static void wake_quiesce_meta_thread(void) { - if (g_quiesce_pipe[1] < 0) { - return; - } - const char sig = 'Q'; - (void)write(g_quiesce_pipe[1], &sig, 1); -} - -/* - * ============================================================================= - * SIGNAL HANDLERS - * ============================================================================= - */ - -/* - * Quiesce handler (SIGUSR1) - called before CRIU checkpoint - * - * We can't do the actual quiescence work here because signal handlers - * can only call async-signal-safe functions. Instead, we set a flag - * and let the interception code handle it. - */ -static void quiesce_signal_handler(int sig) { - (void)sig; - - /* Metadata-only mode: avoid full quiesce, just dump pointers. */ - const char* meta_only = getenv("NVSNAP_QUIESCE_METADATA_ONLY"); - if (meta_only && strcmp(meta_only, "1") == 0) { - wake_quiesce_meta_thread(); - const char msg[] = "[NVSNAP] SIGUSR1: Metadata-only quiesce requested\n"; - (void)write(STDERR_FILENO, msg, sizeof(msg) - 1); - return; - } - - /* Set quiesce requested flag */ - atomic_store(&g_quiesce_state, NVSNAP_QUIESCE_REQUESTED); - - /* Trigger metadata dump thread */ - wake_quiesce_meta_thread(); - - /* NOTE: Do NOT call nvsnap_zmq_handle_checkpoint() here. - * zmq_ctx_checkpoint() modifies ZMQ state — if CRIU captures this, - * the restored process has broken ZMQ IPC. ZMQ restore is handled - * independently via zmq_reinit_all_if_restored(). */ - - /* Notify via write (async-signal-safe) */ - const char msg[] = "[NVSNAP] SIGUSR1: Quiesce requested\n"; - (void)write(STDERR_FILENO, msg, sizeof(msg) - 1); - - /* Chain to saved handler (e.g., Python's asyncio handler) */ - if (g_saved_sigusr1.sa_handler != SIG_DFL && - g_saved_sigusr1.sa_handler != SIG_IGN && - g_saved_sigusr1.sa_handler != NULL) { - g_saved_sigusr1.sa_handler(sig); - } -} - -/* Resume handler (SIGUSR2) - called after checkpoint or on cancel. - * Must be minimal: set flag and return. No logging (can produce millions - * of messages if signal is re-delivered), no chaining to unknown handlers - * (can amplify signal delivery). */ -static void resume_signal_handler(int sig) { - (void)sig; - atomic_store(&g_quiesce_state, NVSNAP_QUIESCE_RESUMED); -} - -/* - * ============================================================================= - * EXPORTED: Check if we're in a restored process - * ============================================================================= - */ - -int nvsnap_is_restored(void) { - return g_is_restored; -} - -/* - * ============================================================================= - * LIBUV TRACKING FUNCTIONS - * ============================================================================= - */ - -/* Track a libuv loop */ -int nvsnap_track_libuv_loop(void* loop) { - if (!loop) return -1; - - pthread_mutex_lock(&g_libuv_mutex); - - /* Check if already tracked */ - for (int i = 0; i < g_libuv_count; i++) { - if (g_libuv_loops[i].loop == loop) { - pthread_mutex_unlock(&g_libuv_mutex); - return 0; /* Already tracked */ - } - } - - if (g_libuv_count >= MAX_LIBUV_LOOPS) { - NVSNAP_WARN("Too many libuv loops (%d)", g_libuv_count); - pthread_mutex_unlock(&g_libuv_mutex); - return -1; - } - - g_libuv_loops[g_libuv_count].loop = loop; - g_libuv_loops[g_libuv_count].reinitialized = 0; - g_libuv_loops[g_libuv_count].owner_thread = pthread_self(); - g_libuv_count++; - - NVSNAP_DEBUG("Tracked libuv loop %p (total=%d)", loop, g_libuv_count); - - pthread_mutex_unlock(&g_libuv_mutex); - return 0; -} - -/* Check if loop needs reinit and reinit if needed */ -int nvsnap_ensure_libuv_loop_ready(void* loop) { - if (!g_is_restored || !loop) { - return 0; /* Not restored, nothing to do */ - } - - pthread_mutex_lock(&g_libuv_mutex); - - /* Find the loop */ - for (int i = 0; i < g_libuv_count; i++) { - if (g_libuv_loops[i].loop == loop) { - if (g_libuv_loops[i].reinitialized) { - pthread_mutex_unlock(&g_libuv_mutex); - return 0; /* Already reinitialized */ - } - - /* Need to reinitialize */ - NVSNAP_INFO("Reinitializing libuv loop %p after CRIU restore", loop); - - /* Lookup real uv_loop_fork */ - static int (*real_uv_loop_fork)(void*) = NULL; - if (!real_uv_loop_fork) { - real_uv_loop_fork = dlsym(RTLD_DEFAULT, "uv_loop_fork"); - } - - if (real_uv_loop_fork) { - int err = real_uv_loop_fork(loop); - if (err != 0) { - NVSNAP_WARN("uv_loop_fork(%p) failed: %d", loop, err); - } else { - NVSNAP_INFO("uv_loop_fork(%p) succeeded", loop); - } - } else { - NVSNAP_WARN("uv_loop_fork not found, loop %p may be broken", loop); - } - - g_libuv_loops[i].reinitialized = 1; - pthread_mutex_unlock(&g_libuv_mutex); - return 0; - } - } - - /* Not tracked yet, track and reinit */ - pthread_mutex_unlock(&g_libuv_mutex); - nvsnap_track_libuv_loop(loop); - - /* Recursively ensure it's ready (will reinit on this call) */ - return nvsnap_ensure_libuv_loop_ready(loop); -} - -/* Quiesce libuv (nothing to drain, just mark for reinit) */ -static void nvsnap_quiesce_libuv(void) { - pthread_mutex_lock(&g_libuv_mutex); - - NVSNAP_INFO("libuv quiesce: %d loops tracked", g_libuv_count); - - /* Mark all loops as needing reinit after restore */ - for (int i = 0; i < g_libuv_count; i++) { - g_libuv_loops[i].reinitialized = 0; - } - - pthread_mutex_unlock(&g_libuv_mutex); -} - -/* - * External function from libuv_intercept.c (if enabled) - * Weak symbol so we don't fail if libuv interception is disabled. - */ -__attribute__((weak)) void nvsnap_libuv_on_restore(void) { - /* Stub - libuv interception disabled */ -} - -/* After restore: reset reinit flags so next use triggers uv_loop_fork */ -static void nvsnap_restore_libuv(void) { - g_is_restored = 1; - - pthread_mutex_lock(&g_libuv_mutex); - - NVSNAP_INFO("libuv restore: marking %d loops for reinit", g_libuv_count); - - for (int i = 0; i < g_libuv_count; i++) { - g_libuv_loops[i].reinitialized = 0; - } - - pthread_mutex_unlock(&g_libuv_mutex); - - /* Also mark all tracked handles for reinit */ - nvsnap_libuv_on_restore(); - - /* Trigger ZMQ restore */ - NVSNAP_INFO("Triggering ZMQ restore"); - nvsnap_zmq_handle_restore(); -} - -/* - * ============================================================================= - * MAIN QUIESCENCE FUNCTIONS - * ============================================================================= - */ - -/* - * Perform quiescence - called from interception points or polling thread - * - * Returns: 1 if quiescence was performed, 0 if not needed - */ -int nvsnap_perform_quiescence(void) { - /* Atomically transition REQUESTED → IN_PROGRESS. This ensures exactly one - * thread performs quiescence even if both the worker thread and an - * io_uring_enter() call race to check the flag. */ - int expected = NVSNAP_QUIESCE_REQUESTED; - if (!atomic_compare_exchange_strong(&g_quiesce_state, &expected, - NVSNAP_QUIESCE_IN_PROGRESS)) { - return 0; - } - - NVSNAP_INFO("=== Starting quiescence ==="); - - /* 1. io_uring draining is handled by io_uring_intercept.c on io_uring_enter */ - NVSNAP_INFO("io_uring draining: handled by intercept layer"); - - /* 2. Prepare libuv for restore */ - nvsnap_quiesce_libuv(); - - /* 2.1 Capture uvloop loop pointers for restore */ - nvsnap_dump_uvloop_metadata(); - - /* 2.2-2.3 Multi-GPU only: NCCL destroy + P2P disable + D2H save. - * Single-GPU uses pure cuda-checkpoint (CRIU plugin handles everything). - * The agent writes /dev/shm/nvsnap-multi-gpu for multi-GPU workloads. */ - if (access("/dev/shm/nvsnap-multi-gpu", F_OK) == 0) { - NVSNAP_INFO("Multi-GPU detected — running NCCL/P2P/D2H quiesce"); - - /* 2.2 NCCL destroy + P2P disable */ - { - static int (*nvsnap_quiesce)(void) = NULL; - static int quiesce_checked = 0; - if (!quiesce_checked) { - nvsnap_quiesce = (int (*)(void))dlsym(RTLD_DEFAULT, - "nvsnap_pre_checkpoint_quiesce"); - quiesce_checked = 1; - } - if (nvsnap_quiesce) { - NVSNAP_INFO("NvSnap pre-checkpoint quiesce (NCCL destroy + P2P disable)"); - nvsnap_quiesce(); - NVSNAP_INFO("NvSnap pre-checkpoint quiesce done"); - } - } - - /* 2.3 GPU D2H save */ - { - static int (*nvsnap_save)(const char *) = NULL; - static int nvsnap_checked = 0; - if (!nvsnap_checked) { - nvsnap_save = (int (*)(const char *))dlsym(RTLD_DEFAULT, - "nvsnap_checkpoint_save"); - nvsnap_checked = 1; - } - if (nvsnap_save) { - char ckpt_dir[512] = {0}; - int pfd = open("/dev/shm/nvsnap-checkpoint-dir", O_RDONLY); - if (pfd >= 0) { - ssize_t n = read(pfd, ckpt_dir, sizeof(ckpt_dir) - 1); - if (n > 0) ckpt_dir[n] = '\0'; - close(pfd); - } - if (ckpt_dir[0]) { - mkdir(ckpt_dir, 0755); - NVSNAP_INFO("NvSnap GPU save: saving to %s", ckpt_dir); - int ret = nvsnap_save(ckpt_dir); - if (ret != 0) - NVSNAP_WARN("NvSnap GPU save failed: %d", ret); - else - NVSNAP_INFO("NvSnap GPU save: done"); - } - } - } - } else { - NVSNAP_INFO("Single-GPU — skipping NCCL/P2P/D2H (pure cuda-checkpoint)"); - } - - /* 2.5 Write per-PID quiesce done marker for agent polling */ - { - char marker[256]; - snprintf(marker, sizeof(marker), "/dev/shm/nvsnap-quiesce-done-%d", getpid()); - int mfd = open(marker, O_CREAT | O_WRONLY, 0666); - if (mfd >= 0) { - (void)write(mfd, "done", 4); - close(mfd); - NVSNAP_INFO("Wrote quiesce done marker: %s", marker); - } else { - NVSNAP_WARN("Failed to write quiesce marker: %s", marker); - } - } - - /* 3. Mark complete */ - atomic_store(&g_quiesce_state, NVSNAP_QUIESCE_COMPLETE); - - /* 4. Send ACK to checkpoint agent */ - if (g_ack_pipe_fd >= 0) { - char ack = 1; - if (write(g_ack_pipe_fd, &ack, 1) != 1) { - NVSNAP_WARN("Failed to send quiesce ACK"); - } else { - NVSNAP_INFO("Quiesce ACK sent"); - } - } - - NVSNAP_INFO("=== Quiescence complete ==="); - - /* 5. Spin until checkpoint done or cancelled */ - NVSNAP_INFO("Waiting for checkpoint or resume signal..."); - while (atomic_load(&g_quiesce_state) == NVSNAP_QUIESCE_COMPLETE) { - usleep(1000); /* 1ms poll */ - } - - if (atomic_load(&g_quiesce_state) == NVSNAP_QUIESCE_RESUMED) { - NVSNAP_INFO("Resumed after quiesce"); - atomic_store(&g_quiesce_state, NVSNAP_QUIESCE_NONE); - } - - return 1; -} - -/* - * Perform post-restore reinitialization - * - * Called when NVSNAP_RESTORED=1 is detected - */ -void nvsnap_perform_restore_reinit(void) { - static atomic_int reinit_done = 0; - int expected = 0; - if (!atomic_compare_exchange_strong(&reinit_done, &expected, 1)) return; - - NVSNAP_INFO("=== Starting post-restore reinitialization ==="); - - /* 1. io_uring reinit is handled lazily in io_uring_intercept.c on first io_uring_enter */ - NVSNAP_INFO("io_uring restore: handled lazily by transparent reinit in intercept layer"); - - /* 2. Mark libuv for reinit on first use */ - nvsnap_restore_libuv(); - - /* 3-5. Multi-GPU only: NCCL restore + P2P re-enable + H2D restore. - * Single-GPU: CRIU plugin Restore+Unlock handles everything. */ - if (access("/dev/shm/nvsnap-multi-gpu", F_OK) == 0) { - NVSNAP_INFO("Multi-GPU restore path"); - - /* 3. NCCL restore */ - nvsnap_nccl_restore(); - - /* 4. GPU post-restore: re-enable P2P */ - { - extern int nvsnap_gpu_post_restore(void); - int p2p = nvsnap_gpu_post_restore(); - if (p2p > 0) - NVSNAP_INFO("GPU post-restore: %d P2P pairs re-enabled", p2p); - } - } else { - NVSNAP_INFO("Single-GPU — CRIU plugin handles GPU restore"); - } - - /* 5. GPU memory restore (multi-GPU only: reload D2H saved data). - * Single-GPU: cuda-checkpoint already restored GPU memory. */ - if (access("/dev/shm/nvsnap-multi-gpu", F_OK) == 0) { - static int (*nvsnap_restore)(const char *) = NULL; - if (!nvsnap_restore) - nvsnap_restore = (int (*)(const char *))dlsym(RTLD_DEFAULT, - "nvsnap_checkpoint_restore_self"); - if (nvsnap_restore) { - char ckpt_dir[512] = {0}; - int fd = open("/dev/shm/nvsnap-checkpoint-dir", O_RDONLY); - if (fd >= 0) { - ssize_t n = read(fd, ckpt_dir, sizeof(ckpt_dir) - 1); - if (n > 0) ckpt_dir[n] = '\0'; - close(fd); - } - if (ckpt_dir[0]) { - NVSNAP_INFO("GPU restore: reloading from %s (pid=%d)", ckpt_dir, getpid()); - int ret = nvsnap_restore(ckpt_dir); - if (ret == 0) - NVSNAP_INFO("GPU restore: success (pid=%d)", getpid()); - else - NVSNAP_WARN("GPU restore: error %d (pid=%d)", ret, getpid()); - } - } - } - - NVSNAP_INFO("=== Post-restore reinitialization complete ==="); -} - -/* - * ============================================================================= - * INITIALIZATION - * ============================================================================= - */ - -static int g_quiesce_initialized = 0; - -/* Start quiesce worker thread. Safe to call multiple times. */ -void nvsnap_start_quiesce_worker(void) { - if (!g_quiesce_worker_thread_started) { - pthread_t wtid; - if (pthread_create(&wtid, NULL, quiesce_worker_thread, NULL) == 0) { - pthread_detach(wtid); - g_quiesce_worker_thread_started = 1; - } - } -} - -/* Called after fork() in child process. Restarts the quiesce worker thread - * since threads don't survive fork(). Also resets quiesce state. */ -/* Set in the fork child; the worker is created later from a normal context. */ -static volatile sig_atomic_t g_quiesce_worker_needs_restart = 0; - -/* Runs in the child after fork(). POSIX allows only async-signal-safe calls - * here until exec(), so this may not create the worker thread directly even - * though the child needs one (threads do not survive fork, and the trigger - * file is keyed on the child's own pid). - * - * Creating it here deadlocks the child: pthread_create allocates TLS and takes - * the loader lock, which a thread that did not survive the fork may hold. - * That is not theoretical -- it wedges CRIU. CRIU forks during dump, and with - * this library force-loaded into it via /etc/ld.so.preload the child hung - * while the parent blocked in wait4 forever, stalling the dump before seize - * completed. Isolating each behaviour showed installing signal handlers and - * starting threads are both harmless; only this handler wedged it. - * - * So: record the need and let nvsnap_quiesce_worker_restart_if_needed() create - * the thread from a safe context. Most forks are followed by exec, where the - * constructors re-run and start the worker normally; this covers the - * fork-without-exec case. */ -static void nvsnap_quiesce_atfork_child(void) -{ - /* Reset state — child starts fresh */ - g_quiesce_state = NVSNAP_QUIESCE_NONE; - g_quiesce_worker_thread_started = 0; - g_quiesce_meta_thread_started = 0; - - g_quiesce_worker_needs_restart = 1; - - /* Reset NCCL tracking */ - nvsnap_nccl_atfork_child(); -} - -/* Create the worker deferred by the atfork handler. Safe to call from any - * normal (non-atfork, non-signal) context; a no-op unless a fork left the - * child without one. */ -void nvsnap_quiesce_worker_restart_if_needed(void) -{ - if (!g_quiesce_worker_needs_restart) - return; - - nvsnap_start_quiesce_worker(); - - /* Drop the request only once a worker actually exists. pthread_create can - * fail (EAGAIN under thread pressure, RLIMIT_NPROC), and clearing the flag - * first would discard the request permanently, leaving a forked child with - * no quiesce poller and no way to notice. Leaving it set means the next - * caller retries. */ - if (g_quiesce_worker_thread_started) - g_quiesce_worker_needs_restart = 0; -} - -__attribute__((constructor(103))) /* Run after main init (101) and NvSnap init (102) */ -static void nvsnap_quiesce_register_atfork(void) -{ - if (nvsnap_self_disabled()) - return; - - /* Register fork handler: - * child (after fork): reset quiesce state, restart worker thread, - * clear NCCL tracking so children build fresh state. */ - pthread_atfork(NULL, NULL, nvsnap_quiesce_atfork_child); -} - -void nvsnap_quiesce_init(void) { - if (g_quiesce_initialized) return; - - /* Check if quiescence is disabled */ - const char* disable_quiesce = getenv("NVSNAP_DISABLE_QUIESCE"); - if (disable_quiesce && strcmp(disable_quiesce, "1") == 0) { - g_quiesce_initialized = 1; - return; - } - - /* Check if lightweight mode - skip quiescence/io_uring */ - const char* lightweight = getenv("NVSNAP_LIGHTWEIGHT"); - if (lightweight && strcmp(lightweight, "1") == 0) { - NVSNAP_DEBUG("Quiesce module disabled (lightweight mode)"); - g_quiesce_initialized = 1; - return; - } - - /* Check if we're a restored process */ - g_is_restored = (getenv("NVSNAP_RESTORED") != NULL); - - /* Get ACK pipe FD if provided */ - const char* ack_fd_str = getenv("NVSNAP_ACK_FD"); - if (ack_fd_str) { - g_ack_pipe_fd = atoi(ack_fd_str); - } - - /* - * NOTE: We do NOT install signal handlers by default anymore. - * PyTorch's distributed module uses signals internally, and our handlers - * can conflict. Only enable if NVSNAP_QUIESCE_SIGNALS=1 is set. - */ - /* Resolve real sigaction for the guard pass-through */ - resolve_real_sigaction(); - - const char* enable_signals = getenv("NVSNAP_QUIESCE_SIGNALS"); - if (enable_signals && strcmp(enable_signals, "1") == 0) { - /* Install handlers via real_sigaction (bypasses our guard) */ - if (install_signal_handler(SIGUSR1, quiesce_signal_handler) == -1) - NVSNAP_WARN("Failed to install SIGUSR1 quiesce handler: %s", strerror(errno)); - if (install_signal_handler(SIGUSR2, resume_signal_handler) == -1) - NVSNAP_WARN("Failed to install SIGUSR2 resume handler: %s", strerror(errno)); - - /* Enable guard — prevents asyncio from overriding our handlers */ - if (real_sigaction) - g_sigaction_guard_enabled = 1; - - if (pipe(g_quiesce_pipe) == 0) { - int flags = fcntl(g_quiesce_pipe[1], F_GETFL, 0); - if (flags >= 0) { - (void)fcntl(g_quiesce_pipe[1], F_SETFL, flags | O_NONBLOCK); - } - if (!g_quiesce_meta_thread_started) { - pthread_t tid; - if (pthread_create(&tid, NULL, quiesce_meta_thread, NULL) == 0) { - pthread_detach(tid); - g_quiesce_meta_thread_started = 1; - NVSNAP_INFO("Quiesce metadata thread started"); - } else { - NVSNAP_WARN("Failed to start quiesce metadata thread"); - } - } - } else { - NVSNAP_WARN("Failed to create quiesce metadata pipe: %s", strerror(errno)); - } - - /* Start quiesce worker thread — ensures quiescence runs even in - * processes that don't call io_uring_enter() (e.g. NCCL workers). */ - if (!g_quiesce_worker_thread_started) { - pthread_t wtid; - if (pthread_create(&wtid, NULL, quiesce_worker_thread, NULL) == 0) { - pthread_detach(wtid); - g_quiesce_worker_thread_started = 1; - NVSNAP_INFO("Quiesce worker thread started"); - } else { - NVSNAP_WARN("Failed to start quiesce worker thread"); - } - } - - NVSNAP_INFO("Signal handlers installed for quiescence"); - } - - /* Register atfork handler to restart worker thread in children. - * After fork(), only the calling thread survives — our worker thread - * (which polls for quiesce triggers) is gone. Re-create it in child. */ - NVSNAP_INFO("Quiesce module initialized (restored=%d, ack_fd=%d)", - g_is_restored, g_ack_pipe_fd); - - /* If restored, trigger reinit */ - if (g_is_restored) { - NVSNAP_INFO("Detected CRIU restore, triggering reinitialization"); - nvsnap_perform_restore_reinit(); - } - - g_quiesce_initialized = 1; -} - -/* - * ============================================================================= - * DIAGNOSTIC FUNCTIONS - * ============================================================================= - */ - -void nvsnap_dump_quiesce_state(FILE* out) { - fprintf(out, "\n=== NVSNAP Quiesce State ===\n"); - fprintf(out, "Quiesce state: %d\n", atomic_load(&g_quiesce_state)); - fprintf(out, "Is restored: %d\n", g_is_restored); - fprintf(out, "ACK pipe fd: %d\n", g_ack_pipe_fd); - - fprintf(out, "\nlibuv loops: %d\n", g_libuv_count); - pthread_mutex_lock(&g_libuv_mutex); - for (int i = 0; i < g_libuv_count; i++) { - fprintf(out, " loop=%p reinitialized=%d\n", - g_libuv_loops[i].loop, - g_libuv_loops[i].reinitialized); - } - pthread_mutex_unlock(&g_libuv_mutex); - - /* io_uring state is managed by io_uring_intercept.c */ - fprintf(out, "\nio_uring instances: see io_uring_intercept.c\n"); - - fprintf(out, "===========================\n\n"); -} diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/seccomp_intercept.c b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/seccomp_intercept.c deleted file mode 100644 index f6731478bf..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/seccomp_intercept.c +++ /dev/null @@ -1,767 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * seccomp-bpf Interception for io_uring Syscalls - * - * This module uses seccomp-bpf with USER_NOTIF to intercept io_uring syscalls - * at the kernel boundary, regardless of whether the calling code is dynamically - * or statically linked. This is critical for uvloop which statically links libuv. - * - * How it works: - * 1. Install a seccomp filter that sends USER_NOTIF for io_uring syscalls - * 2. A supervisor thread receives notifications via the listener fd - * 3. The supervisor can: - * - Log the syscall and its arguments - * - Execute the real syscall on behalf of the process - * - Return a fake result if needed - * - * This is especially useful after CRIU restore when io_uring state needs healing. - */ - -#define _GNU_SOURCE -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "nvsnap_intercept.h" - -/* io_uring syscall numbers */ -#ifndef __NR_io_uring_setup -#define __NR_io_uring_setup 425 -#endif -#ifndef __NR_io_uring_enter -#define __NR_io_uring_enter 426 -#endif -#ifndef __NR_io_uring_register -#define __NR_io_uring_register 427 -#endif - -/* seccomp user notification (if not defined) */ -#ifndef SECCOMP_FILTER_FLAG_NEW_LISTENER -#define SECCOMP_FILTER_FLAG_NEW_LISTENER (1UL << 3) -#endif - -#ifndef SECCOMP_USER_NOTIF_FLAG_CONTINUE -#define SECCOMP_USER_NOTIF_FLAG_CONTINUE (1UL << 0) -#endif - -/* io_uring flags we care about */ -#define IORING_SETUP_SQPOLL (1U << 1) -#define IORING_ENTER_GETEVENTS (1U << 0) -#define IORING_ENTER_SQ_WAKEUP (1U << 1) - -/* Architecture for seccomp */ -#if defined(__x86_64__) -#define SECCOMP_AUDIT_ARCH AUDIT_ARCH_X86_64 -#elif defined(__aarch64__) -#define SECCOMP_AUDIT_ARCH AUDIT_ARCH_AARCH64 -#else -#error "Unsupported architecture for seccomp" -#endif - -/* - * ============================================================================= - * IO_URING RING TRACKING - * ============================================================================= - */ - -/* io_uring ring info (from io_uring_params) */ -typedef struct { - int fd; - uint32_t sq_entries; - uint32_t cq_entries; - uint32_t flags; - /* Offsets from io_uring_params.sq_off */ - uint32_t sq_head_off; - uint32_t sq_tail_off; - uint32_t sq_ring_mask_off; - uint32_t sq_flags_off; - /* Offsets from io_uring_params.cq_off */ - uint32_t cq_head_off; - uint32_t cq_tail_off; - /* Mmap addresses (captured from /proc/pid/maps after setup) */ - uint64_t sq_ring_addr; - uint64_t cq_ring_addr; - /* Cached state */ - bool valid; -} ring_info_t; - -#define MAX_RINGS 16 - -/* - * ============================================================================= - * SECCOMP INTERCEPT STATE - * ============================================================================= - */ - -typedef struct { - bool installed; /* Is seccomp filter installed? */ - bool post_restore; /* Are we in post-restore mode? */ - int listener_fd; /* Notification listener fd */ - pthread_t supervisor_thread; /* Supervisor thread */ - bool supervisor_running; /* Is supervisor thread running? */ - int io_uring_enter_count; /* Count of io_uring_enter calls */ - int heal_count; /* Number of "healed" calls */ - - /* Ring tracking */ - ring_info_t rings[MAX_RINGS]; - int ring_count; - pthread_mutex_t ring_mutex; -} seccomp_state_t; - -static seccomp_state_t g_seccomp_state = { - .installed = false, - .listener_fd = -1, - .supervisor_running = false, - .ring_mutex = PTHREAD_MUTEX_INITIALIZER, -}; - -/* - * ============================================================================= - * RING STATE INSPECTION - * ============================================================================= - */ - -/* io_uring_params structure for capturing setup info */ -struct io_uring_params_capture { - uint32_t sq_entries; - uint32_t cq_entries; - uint32_t flags; - uint32_t sq_thread_cpu; - uint32_t sq_thread_idle; - uint32_t features; - uint32_t wq_fd; - uint32_t resv[3]; - struct { - uint32_t head; - uint32_t tail; - uint32_t ring_mask; - uint32_t ring_entries; - uint32_t flags; - uint32_t dropped; - uint32_t array; - uint32_t resv1; - uint64_t resv2; - } sq_off; - struct { - uint32_t head; - uint32_t tail; - uint32_t ring_mask; - uint32_t ring_entries; - uint32_t overflow; - uint32_t cqes; - uint32_t flags; - uint32_t resv1; - uint64_t resv2; - } cq_off; -}; - -/* - * Read memory from target process - */ -static ssize_t read_proc_mem(pid_t pid, uint64_t addr, void* buf, size_t len) { - char path[64]; - snprintf(path, sizeof(path), "/proc/%d/mem", pid); - - int fd = open(path, O_RDONLY); - if (fd < 0) { - NVSNAP_DEBUG("Cannot open %s: %s", path, strerror(errno)); - return -1; - } - - ssize_t ret = pread(fd, buf, len, addr); - close(fd); - - if (ret < 0) { - NVSNAP_DEBUG("Cannot read from %s at 0x%lx: %s", path, addr, strerror(errno)); - } - - return ret; -} - -/* - * Write memory to target process - */ -static ssize_t write_proc_mem(pid_t pid, uint64_t addr, const void* buf, size_t len) { - char path[64]; - snprintf(path, sizeof(path), "/proc/%d/mem", pid); - - int fd = open(path, O_WRONLY); - if (fd < 0) { - NVSNAP_DEBUG("Cannot open %s for write: %s", path, strerror(errno)); - return -1; - } - - ssize_t ret = pwrite(fd, buf, len, addr); - close(fd); - - if (ret < 0) { - NVSNAP_DEBUG("Cannot write to %s at 0x%lx: %s", path, addr, strerror(errno)); - } - - return ret; -} - -/* - * Find io_uring mmap address from /proc/pid/maps - * Returns the address of the [io_uring] mapping - */ -static uint64_t find_io_uring_mmap(pid_t pid, int ring_fd) { - char path[64]; - snprintf(path, sizeof(path), "/proc/%d/maps", pid); - - FILE* f = fopen(path, "r"); - if (!f) { - NVSNAP_DEBUG("Cannot open %s: %s", path, strerror(errno)); - return 0; - } - - char line[512]; - char fdinfo[64]; - snprintf(fdinfo, sizeof(fdinfo), "anon_inode:[io_uring]"); - - uint64_t found_addr = 0; - - while (fgets(line, sizeof(line), f)) { - if (strstr(line, fdinfo) || strstr(line, "[io_uring]")) { - uint64_t start, end; - if (sscanf(line, "%lx-%lx", &start, &end) == 2) { - found_addr = start; - NVSNAP_DEBUG("Found io_uring mmap at 0x%lx-0x%lx (fd=%d)", start, end, ring_fd); - break; - } - } - } - - fclose(f); - return found_addr; -} - -/* - * Read ring info from /proc/pid/fdinfo/fd - * Returns: sq_head, sq_tail, cq_head, cq_tail, sq_mask, cq_mask - */ -static int read_ring_from_fdinfo(pid_t pid, int ring_fd, - uint32_t* sq_head, uint32_t* sq_tail, - uint32_t* cq_head, uint32_t* cq_tail, - uint32_t* sq_mask, uint32_t* cq_mask) { - char path[64]; - snprintf(path, sizeof(path), "/proc/%d/fdinfo/%d", pid, ring_fd); - - FILE* f = fopen(path, "r"); - if (!f) { - NVSNAP_DEBUG("Cannot open %s: %s", path, strerror(errno)); - return -1; - } - - char line[256]; - *sq_head = *sq_tail = *cq_head = *cq_tail = 0; - *sq_mask = *cq_mask = 0; - - while (fgets(line, sizeof(line), f)) { - /* fdinfo format for io_uring: - * SqHead: N - * SqTail: N - * CqHead: N - * CqTail: N - * SqMask: N - * CqMask: N - */ - unsigned int val; - if (sscanf(line, "SqHead:\t%u", &val) == 1) *sq_head = val; - else if (sscanf(line, "SqTail:\t%u", &val) == 1) *sq_tail = val; - else if (sscanf(line, "CqHead:\t%u", &val) == 1) *cq_head = val; - else if (sscanf(line, "CqTail:\t%u", &val) == 1) *cq_tail = val; - else if (sscanf(line, "SqMask:\t%u", &val) == 1) *sq_mask = val; - else if (sscanf(line, "CqMask:\t%u", &val) == 1) *cq_mask = val; - } - - fclose(f); - return 0; -} - -/* - * Read ring indices from target process - * Uses /proc/pid/fdinfo which has direct ring state - */ -static int read_ring_state(pid_t pid, int ring_fd, - uint32_t* sq_head, uint32_t* sq_tail, - uint32_t* cq_head, uint32_t* cq_tail) { - uint32_t sq_mask, cq_mask; - return read_ring_from_fdinfo(pid, ring_fd, sq_head, sq_tail, - cq_head, cq_tail, &sq_mask, &cq_mask); -} - -/* - * Track a ring from io_uring_setup params - */ -static void track_ring_from_setup(pid_t pid, int ring_fd, uint64_t params_addr) { - struct io_uring_params_capture params; - - if (read_proc_mem(pid, params_addr, ¶ms, sizeof(params)) < 0) { - NVSNAP_WARN("Cannot read io_uring_params from pid=%d addr=0x%lx", pid, params_addr); - return; - } - - pthread_mutex_lock(&g_seccomp_state.ring_mutex); - - if (g_seccomp_state.ring_count >= MAX_RINGS) { - pthread_mutex_unlock(&g_seccomp_state.ring_mutex); - NVSNAP_WARN("Too many io_uring rings, cannot track fd=%d", ring_fd); - return; - } - - ring_info_t* ring = &g_seccomp_state.rings[g_seccomp_state.ring_count++]; - ring->fd = ring_fd; - ring->sq_entries = params.sq_entries; - ring->cq_entries = params.cq_entries; - ring->flags = params.flags; - ring->sq_head_off = params.sq_off.head; - ring->sq_tail_off = params.sq_off.tail; - ring->sq_ring_mask_off = params.sq_off.ring_mask; - ring->sq_flags_off = params.sq_off.flags; - ring->cq_head_off = params.cq_off.head; - ring->cq_tail_off = params.cq_off.tail; - ring->sq_ring_addr = 0; /* Will be filled later from /proc/pid/maps */ - ring->valid = true; - - pthread_mutex_unlock(&g_seccomp_state.ring_mutex); - - NVSNAP_INFO("Tracked ring fd=%d: sq=%u cq=%u flags=0x%x offsets(sqh=%u sqt=%u cqh=%u cqt=%u)", - ring_fd, params.sq_entries, params.cq_entries, params.flags, - params.sq_off.head, params.sq_off.tail, - params.cq_off.head, params.cq_off.tail); -} - -/* - * ============================================================================= - * SUPERVISOR THREAD - * ============================================================================= - */ - -static void* supervisor_thread_func(void* arg) { - (void)arg; - - int listener_fd = g_seccomp_state.listener_fd; - - NVSNAP_INFO("Supervisor thread started, listening on fd=%d", listener_fd); - - while (g_seccomp_state.supervisor_running) { - struct seccomp_notif req; - struct seccomp_notif_resp resp; - - memset(&req, 0, sizeof(req)); - memset(&resp, 0, sizeof(resp)); - - /* Wait for notification */ - if (ioctl(listener_fd, SECCOMP_IOCTL_NOTIF_RECV, &req) < 0) { - if (errno == EINTR) continue; - if (errno == ENOENT) continue; /* Target may have exited */ - NVSNAP_ERROR("SECCOMP_IOCTL_NOTIF_RECV failed: %s", strerror(errno)); - break; - } - - int syscall_nr = req.data.nr; - __u64* args = req.data.args; - pid_t pid = req.pid; - - /* Handle different syscalls */ - if (syscall_nr == __NR_io_uring_setup) { - uint64_t params_addr = args[1]; - - /* - * SQPOLL FLAG STRIPPING - * - * Read the io_uring_params struct from target process memory, - * strip the SQPOLL flag, and write it back BEFORE the syscall executes. - * - * Why strip SQPOLL? - * - SQPOLL creates a kernel polling thread for the io_uring ring - * - CRIU cannot properly checkpoint/restore this kernel thread - * - After restore, the SQPOLL thread is in a broken state - * - Without SQPOLL, libuv/uvloop uses direct io_uring_enter calls - * - Direct calls are properly restored by CRIU - * - * libuv handles missing SQPOLL gracefully - it's just an optimization. - */ - struct io_uring_params_capture params; - bool sqpoll_stripped = false; - - if (read_proc_mem(pid, params_addr, ¶ms, sizeof(params)) == (ssize_t)sizeof(params)) { - uint32_t orig_flags = params.flags; - - if (params.flags & IORING_SETUP_SQPOLL) { - /* Strip SQPOLL flag */ - params.flags &= ~IORING_SETUP_SQPOLL; - - /* Write modified params back to target process */ - /* Only need to write the flags field, which is at offset 8 (after sq_entries and cq_entries) */ - uint64_t flags_addr = params_addr + offsetof(struct io_uring_params_capture, flags); - - if (write_proc_mem(pid, flags_addr, ¶ms.flags, sizeof(params.flags)) == sizeof(params.flags)) { - sqpoll_stripped = true; - NVSNAP_WARN("[SECCOMP] io_uring_setup: STRIPPED SQPOLL flag! 0x%x -> 0x%x (for checkpoint/restore compatibility)", - orig_flags, params.flags); - } else { - NVSNAP_ERROR("[SECCOMP] io_uring_setup: Failed to write modified flags to pid=%u addr=0x%llx", - pid, (unsigned long long)flags_addr); - } - } - - NVSNAP_INFO("[SECCOMP] io_uring_setup(entries=%u, params=0x%llx, flags=0x%x%s) from pid=%u", - (unsigned)args[0], (unsigned long long)args[1], - sqpoll_stripped ? params.flags : orig_flags, - sqpoll_stripped ? " [SQPOLL stripped]" : "", - pid); - } else { - NVSNAP_WARN("[SECCOMP] io_uring_setup(entries=%u, params=0x%llx) - could not read params from pid=%u", - (unsigned)args[0], (unsigned long long)args[1], pid); - } - - } else if (syscall_nr == __NR_io_uring_enter) { - g_seccomp_state.io_uring_enter_count++; - int ring_fd = (int)args[0]; - unsigned int to_submit = (unsigned int)args[1]; - unsigned int min_complete = (unsigned int)args[2]; - unsigned int flags = (unsigned int)args[3]; - - NVSNAP_INFO("[SECCOMP] io_uring_enter(fd=%d, submit=%u, complete=%u, flags=0x%x) #%d %s", - ring_fd, to_submit, min_complete, flags, - g_seccomp_state.io_uring_enter_count, - g_seccomp_state.post_restore ? "[POST-RESTORE]" : ""); - - /* In post-restore mode, inspect ring state and detect anomalies */ - if (g_seccomp_state.post_restore) { - uint32_t sq_head, sq_tail, cq_head, cq_tail; - if (read_ring_state(pid, ring_fd, &sq_head, &sq_tail, &cq_head, &cq_tail) == 0) { - uint32_t sq_pending = sq_tail - sq_head; - uint32_t cq_pending = cq_tail - cq_head; - - NVSNAP_INFO(" Ring state: sq_head=%u sq_tail=%u (pending=%u) cq_head=%u cq_tail=%u (pending=%u)", - sq_head, sq_tail, sq_pending, - cq_head, cq_tail, cq_pending); - - /* - * Detect stale state after restore: - * - First io_uring_enter call (#1) after restore - * - Ring indices suggest previous activity (sq_tail > 0) - * - But this is supposedly a fresh restore - * - * In a properly quiesced checkpoint, indices should be 0. - * If they're non-zero, the checkpoint wasn't clean. - */ - if (g_seccomp_state.io_uring_enter_count == 1) { - if (sq_head != 0 || cq_head != 0) { - NVSNAP_WARN(" POTENTIAL STALE STATE: First call after restore but indices non-zero!"); - NVSNAP_WARN(" This suggests checkpoint wasn't properly quiesced."); - NVSNAP_WARN(" Expected: sq_head=0 cq_head=0, Got: sq_head=%u cq_head=%u", - sq_head, cq_head); - - /* - * TODO: Implement healing here - * - * Option 1: Write to /proc/pid/mem to reset indices - * - Need mmap address from /proc/pid/maps - * - Risky: might corrupt ring state further - * - * Option 2: Return EAGAIN to force retry - * - Might help app recover - * - But doesn't fix the fundamental mismatch - * - * Option 3: Close the ring fd and hope app recreates - * - Too aggressive, likely to crash - * - * For now, just warn and continue. - * The real fix is proper quiescence before checkpoint. - */ - g_seccomp_state.heal_count++; - } else { - NVSNAP_INFO(" Ring state looks clean (indices at 0) - good!"); - } - } - } - } - - } else if (syscall_nr == __NR_io_uring_register) { - NVSNAP_INFO("[SECCOMP] io_uring_register(fd=%d, opcode=%u) from pid=%u", - (int)args[0], (unsigned)args[1], pid); - } - - /* Check if target is still valid before responding */ - if (ioctl(listener_fd, SECCOMP_IOCTL_NOTIF_ID_VALID, &req.id) < 0) { - NVSNAP_WARN("Target process exited before we could respond"); - continue; - } - - /* - * Let the syscall proceed normally by using CONTINUE flag. - * This tells the kernel to execute the original syscall. - */ - resp.id = req.id; - resp.flags = SECCOMP_USER_NOTIF_FLAG_CONTINUE; - resp.val = 0; - resp.error = 0; - - if (ioctl(listener_fd, SECCOMP_IOCTL_NOTIF_SEND, &resp) < 0) { - if (errno == ENOENT) { - NVSNAP_WARN("Target process exited before response"); - } else { - NVSNAP_ERROR("SECCOMP_IOCTL_NOTIF_SEND failed: %s", strerror(errno)); - } - } - } - - NVSNAP_INFO("Supervisor thread exiting"); - return NULL; -} - -/* - * ============================================================================= - * SECCOMP FILTER INSTALLATION - * ============================================================================= - */ - -int nvsnap_seccomp_install_filter(void) { - if (g_seccomp_state.installed) { - NVSNAP_WARN("seccomp filter already installed"); - return 0; - } - - /* - * BPF filter program: - * - Check architecture - * - Check syscall number - * - If io_uring syscall, USER_NOTIF - * - Otherwise, ALLOW - */ - struct sock_filter filter[] = { - /* Load architecture */ - BPF_STMT(BPF_LD | BPF_W | BPF_ABS, - (offsetof(struct seccomp_data, arch))), - /* Check architecture */ - BPF_JUMP(BPF_JMP | BPF_JEQ | BPF_K, SECCOMP_AUDIT_ARCH, 1, 0), - BPF_STMT(BPF_RET | BPF_K, SECCOMP_RET_ALLOW), /* Wrong arch, allow */ - - /* Load syscall number */ - BPF_STMT(BPF_LD | BPF_W | BPF_ABS, - (offsetof(struct seccomp_data, nr))), - - /* Check for io_uring_setup (425) */ - BPF_JUMP(BPF_JMP | BPF_JEQ | BPF_K, __NR_io_uring_setup, 0, 1), - BPF_STMT(BPF_RET | BPF_K, SECCOMP_RET_USER_NOTIF), - - /* Check for io_uring_enter (426) */ - BPF_JUMP(BPF_JMP | BPF_JEQ | BPF_K, __NR_io_uring_enter, 0, 1), - BPF_STMT(BPF_RET | BPF_K, SECCOMP_RET_USER_NOTIF), - - /* Check for io_uring_register (427) */ - BPF_JUMP(BPF_JMP | BPF_JEQ | BPF_K, __NR_io_uring_register, 0, 1), - BPF_STMT(BPF_RET | BPF_K, SECCOMP_RET_USER_NOTIF), - - /* Allow all other syscalls */ - BPF_STMT(BPF_RET | BPF_K, SECCOMP_RET_ALLOW), - }; - - struct sock_fprog prog = { - .len = sizeof(filter) / sizeof(filter[0]), - .filter = filter, - }; - - /* Allow setting seccomp filters */ - if (prctl(PR_SET_NO_NEW_PRIVS, 1, 0, 0, 0) < 0) { - NVSNAP_ERROR("prctl(NO_NEW_PRIVS) failed: %s", strerror(errno)); - return -1; - } - - /* Install the filter with NEW_LISTENER flag to get notification fd */ - int listener_fd = syscall(__NR_seccomp, SECCOMP_SET_MODE_FILTER, - SECCOMP_FILTER_FLAG_NEW_LISTENER, &prog); - - if (listener_fd < 0) { - NVSNAP_ERROR("seccomp(FILTER, NEW_LISTENER) failed: %s", strerror(errno)); - NVSNAP_ERROR("This requires kernel 5.0+ with USER_NOTIF support"); - return -1; - } - - g_seccomp_state.listener_fd = listener_fd; - g_seccomp_state.installed = true; - - NVSNAP_INFO("seccomp-bpf filter installed, listener fd=%d", listener_fd); - - /* Start supervisor thread */ - g_seccomp_state.supervisor_running = true; - if (pthread_create(&g_seccomp_state.supervisor_thread, NULL, - supervisor_thread_func, NULL) != 0) { - /* - * CRITICAL: The seccomp filter is already installed at kernel level - * and CANNOT be removed. If we can't start the supervisor thread, - * all io_uring syscalls will hang waiting for a response that never comes. - * - * We must NOT: - * - Close listener_fd (would cause ENOSYS for all io_uring calls) - * - Set installed = false (filter IS installed at kernel level) - * - * Options: - * 1. Retry thread creation - * 2. Abort the process (unrecoverable state) - * 3. Try to handle notifications in-line (complex) - * - * We try retry first, then abort if that fails. - */ - NVSNAP_ERROR("Failed to create supervisor thread: %s", strerror(errno)); - - /* Retry once */ - usleep(1000); /* Small delay */ - if (pthread_create(&g_seccomp_state.supervisor_thread, NULL, - supervisor_thread_func, NULL) != 0) { - NVSNAP_ERROR("FATAL: Supervisor thread creation failed on retry"); - NVSNAP_ERROR("FATAL: seccomp filter is installed but no listener!"); - NVSNAP_ERROR("FATAL: Process cannot use io_uring - aborting"); - /* - * We must abort because: - * - Filter is installed, can't be removed - * - Without supervisor, io_uring syscalls will block forever - * - Setting installed=false would lie about kernel state - * - Closing listener_fd would cause ENOSYS errors - */ - abort(); - } - NVSNAP_WARN("Supervisor thread started on retry"); - } - - NVSNAP_INFO("Supervisor thread started for io_uring interception"); - - return 0; -} - -/* - * Stop the seccomp supervisor - */ -void nvsnap_seccomp_stop(void) { - if (!g_seccomp_state.installed) { - return; - } - - g_seccomp_state.supervisor_running = false; - - if (g_seccomp_state.listener_fd >= 0) { - close(g_seccomp_state.listener_fd); - g_seccomp_state.listener_fd = -1; - } - - pthread_join(g_seccomp_state.supervisor_thread, NULL); - - NVSNAP_INFO("seccomp supervisor stopped"); -} - -/* - * Finalize seccomp interception - close the listener fd so CRIU can checkpoint. - * - * This should be called after all io_uring_setup calls have completed - * (e.g., after app startup). The SQPOLL stripping has already happened, - * so the seccomp filter is no longer needed. - * - * After this call: - * - io_uring syscalls will proceed normally (no interception) - * - The process can be checkpointed by CRIU - */ -void nvsnap_seccomp_finalize(void) { - if (!g_seccomp_state.installed) { - return; - } - - NVSNAP_INFO("seccomp: finalizing - closing listener fd for CRIU compatibility"); - - /* - * Stop the supervisor thread gracefully - */ - g_seccomp_state.supervisor_running = false; - - /* - * Close the listener fd. This has two effects: - * 1. The supervisor thread will exit (ioctl will fail with EBADF) - * 2. CRIU will not see the special seccomp notify fd - * - * After closing, io_uring syscalls will return ENOSYS because - * the filter is still installed but no one is listening. - * - * Actually, we need SECCOMP_USER_NOTIF_FLAG_CONTINUE behavior. - * When listener is closed and filter is still active: - * - Kernel returns -ENOSYS for the syscall - * - * This is a problem! We need a different approach... - * - * Solution: We can use dup2 to replace the listener fd with /dev/null, - * then close it. The supervisor thread will get EBADF on ioctl. - * But the filter is still active... - * - * Actually, the safest approach is: - * 1. Don't close the fd - * 2. Mark it as "external" for CRIU - * 3. On restore, recreate the listener - * - * But for now, let's try a simpler approach: - * Don't install seccomp at all during startup. - * Instead, modify CRIU to strip SQPOLL during restore. - * - * Wait - we already have that in CRIU's restorer.c! - * - * Let me try a different approach: - * Close the listener and hope the kernel allows the syscalls through. - */ - - if (g_seccomp_state.listener_fd >= 0) { - close(g_seccomp_state.listener_fd); - g_seccomp_state.listener_fd = -1; - } - - /* Wait for supervisor to exit */ - pthread_join(g_seccomp_state.supervisor_thread, NULL); - - NVSNAP_INFO("seccomp: finalized - listener closed, io_uring setup complete"); - NVSNAP_INFO("seccomp: Note: io_uring syscalls may now return ENOSYS until process restart"); -} - -/* - * Enable post-restore mode - subsequent io_uring calls may need healing - */ -void nvsnap_seccomp_set_post_restore(bool post_restore) { - g_seccomp_state.post_restore = post_restore; - if (post_restore) { - g_seccomp_state.io_uring_enter_count = 0; - g_seccomp_state.heal_count = 0; - NVSNAP_INFO("seccomp: post-restore mode enabled - monitoring io_uring calls"); - } -} - -/* - * Check if seccomp filter is installed - */ -bool nvsnap_seccomp_is_installed(void) { - return g_seccomp_state.installed; -} - -/* - * Get statistics - */ -void nvsnap_seccomp_get_stats(int* enter_count, int* heal_count) { - if (enter_count) *enter_count = g_seccomp_state.io_uring_enter_count; - if (heal_count) *heal_count = g_seccomp_state.heal_count; -} diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/self_disable.c b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/self_disable.c deleted file mode 100644 index aa30adb6f5..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/self_disable.c +++ /dev/null @@ -1,73 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * Stay out of our own tooling. - * - * Workloads enable interception by writing /etc/ld.so.preload. The loader - * honours that file for EVERY process that execs in the mount namespace, not - * just the workload -- so the CRIU that nsenters in to dump the container is - * force-loaded with this library too, as are the cuda-checkpoint and - * iptables-restore helpers CRIU execs from the bundle. - * - * Inside CRIU that is fatal: the dump wedges with CRIU blocked forever in - * wait4() before it finishes seizing the task tree. Reproduced with a plain - * `sleep` victim and no GPU -- preloaded into the victim only, the dump - * succeeds (813-line dump.log); preloaded into CRIU only, it hangs (36-line - * dump.log, zero images). It is the load into CRIU that breaks, not the load - * into the workload. - * - * No environment gate can undo this after the fact: our constructors run - * before any NVSNAP_* variable is consulted, and NVSNAP_LIGHTWEIGHT=1, - * NVSNAP_DISABLE_QUIESCE=1 and NVSNAP_LOG_LEVEL=0 were each measured to still - * hang. The check therefore has to happen before we touch anything, which is - * what this file provides -- every constructor calls it first and returns - * early, leaving the process completely untouched. - */ - -#define _GNU_SOURCE -#include -#include -#include -#include - -#include "nvsnap_intercept.h" - -/* Everything CRIU runs during dump/restore lives here: criu itself, the - * cuda-checkpoint it shells out to, iptables-restore for the network lock, - * and our own restore helpers. Matching on the directory covers them all - * without an executable-name list that drifts as the bundle changes. */ -#define NVSNAP_BUNDLE_PREFIX "/criu-bundle/" - -int nvsnap_self_disabled(void) -{ - /* Resolved once: the answer cannot change within a process image, and the - * constructors that call this run before threads exist. */ - static int cached = -1; - - if (cached >= 0) - return cached; - - /* Explicit override, for callers that stage the bundle somewhere else. */ - const char *env = getenv("NVSNAP_INTERCEPT_DISABLE"); - if (env && strcmp(env, "1") == 0) { - cached = 1; - return cached; - } - - char exe[PATH_MAX]; - ssize_t n = readlink("/proc/self/exe", exe, sizeof(exe) - 1); - if (n > 0) { - exe[n] = '\0'; - if (strncmp(exe, NVSNAP_BUNDLE_PREFIX, - sizeof(NVSNAP_BUNDLE_PREFIX) - 1) == 0) { - cached = 1; - return cached; - } - } - - cached = 0; - return cached; -} diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/zmq_intercept.c b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/zmq_intercept.c deleted file mode 100644 index 3e1edaeeed..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/src/zmq_intercept.c +++ /dev/null @@ -1,2618 +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 -#include -#include -#include -#include -#include -#include -#include -#include - -#include "nvsnap_intercept.h" - -#ifndef ZMQ_EVENTS -#define ZMQ_EVENTS 15 -#endif -#ifndef ZMQ_TYPE -#define ZMQ_TYPE 16 -#endif -#ifndef ZMQ_IO_THREADS -#define ZMQ_IO_THREADS 1 -#endif -#ifndef ZMQ_LAST_ENDPOINT -#define ZMQ_LAST_ENDPOINT 32 -#endif -#ifndef ZMQ_ROUTER -#define ZMQ_ROUTER 6 -#endif -#ifndef ZMQ_DEALER -#define ZMQ_DEALER 5 -#endif -#ifndef ZMQ_REQ -#define ZMQ_REQ 3 -#endif -#ifndef ZMQ_IDENTITY -#define ZMQ_IDENTITY 5 -#endif -#ifndef ZMQ_PROBE_ROUTER -#define ZMQ_PROBE_ROUTER 51 -#endif -#ifndef ZMQ_PAIR -#define ZMQ_PAIR 0 -#endif -#ifndef ZMQ_EVENT_ALL -#define ZMQ_EVENT_ALL 0xFFFF -#endif -#ifndef ZMQ_ROUTER_MANDATORY -#define ZMQ_ROUTER_MANDATORY 33 -#endif -#ifndef ZMQ_IMMEDIATE -#define ZMQ_IMMEDIATE 39 -#endif - -typedef struct zmq_msg_t zmq_msg_t; - -typedef struct zmq_pollitem_t { - void *socket; - int fd; - short events; - short revents; -} zmq_pollitem_t; - -void *zmq_ctx_new(void); -int zmq_ctx_term(void *); -int zmq_ctx_destroy(void *); -int zmq_ctx_shutdown(void *); -int zmq_ctx_set(void *, int, int); -int zmq_ctx_get(void *, int); -void *zmq_socket(void *, int); -int zmq_close(void *); -int zmq_bind(void *, const char *); -int zmq_unbind(void *, const char *); -int zmq_connect(void *, const char *); -int zmq_disconnect(void *, const char *); -int zmq_setsockopt(void *, int, const void *, size_t); -int zmq_getsockopt(void *, int, void *, size_t *); -int zmq_send(void *, const void *, size_t, int); -int zmq_recv(void *, void *, size_t, int); -int zmq_msg_send(zmq_msg_t *, void *, int); -int zmq_msg_recv(zmq_msg_t *, void *, int); -int zmq_msg_init(zmq_msg_t *); -int zmq_msg_init_size(zmq_msg_t *, size_t); -int zmq_msg_init_data(zmq_msg_t *, void *, size_t, void (*)(void *, void *), void *); -int zmq_msg_close(zmq_msg_t *); -void *zmq_msg_data(zmq_msg_t *); -size_t zmq_msg_size(zmq_msg_t *); -int zmq_poll(zmq_pollitem_t *, int, long); -int zmq_proxy(void *, void *, void *); -int zmq_errno(void); - -enum nvsnap_zmq_op_type { - NVSNAP_ZMQ_OP_SETSOCKOPT = 1, - NVSNAP_ZMQ_OP_BIND, - NVSNAP_ZMQ_OP_CONNECT, - NVSNAP_ZMQ_OP_UNBIND, - NVSNAP_ZMQ_OP_DISCONNECT, -}; - -struct nvsnap_zmq_op { - enum nvsnap_zmq_op_type type; - int option; - size_t len; - void *data; - char *endpoint; - struct nvsnap_zmq_op *next; -}; - -struct nvsnap_zmq_socket; - -typedef struct nvsnap_zmq_ctx { - uint64_t magic; - void *real; - int reinit_done; - int destroyed; - struct nvsnap_zmq_socket *sockets; - pthread_mutex_t lock; - struct nvsnap_zmq_ctx *next; -} nvsnap_zmq_ctx_t; - -typedef struct nvsnap_zmq_socket { - uint64_t magic; - nvsnap_zmq_ctx_t *ctx; - void *real; - int type; - int closed; - int replay_after_restore; - int replay_failed; /* Set if replay bind/connect failed */ - int rebuilding; - int rebuild_count; - int forced_connect_done; - int monitor_started; - void *monitor_sock; - pthread_t monitor_thread; - struct nvsnap_zmq_op *ops_head; - struct nvsnap_zmq_op *ops_tail; - struct nvsnap_zmq_socket *next; -} nvsnap_zmq_socket_t; - -#define NVSNAP_ZMQ_CTX_MAGIC 0x5a6d715f63747831ULL -#define NVSNAP_ZMQ_SOCK_MAGIC 0x5a6d715f736f636bULL - -static pthread_once_t g_zmq_once = PTHREAD_ONCE_INIT; -static int g_zmq_enabled = -1; -static atomic_int g_zmq_restore_seen = 0; -static int g_zmq_trace = -1; -static pthread_mutex_t g_zmq_ctx_list_lock = PTHREAD_MUTEX_INITIALIZER; -static nvsnap_zmq_ctx_t *g_zmq_ctx_list = NULL; -static pthread_mutex_t g_zmq_load_lock = PTHREAD_MUTEX_INITIALIZER; -static void *g_zmq_handle = NULL; -static void *g_zmq_newns_handle = NULL; -static int g_zmq_use_newns = 0; -static pthread_mutex_t g_zmq_recover_lock = PTHREAD_MUTEX_INITIALIZER; -static atomic_int g_zmq_recovering = 0; - -static void *(*real_zmq_ctx_new)(void); -static int (*real_zmq_ctx_term)(void *); -static int (*real_zmq_ctx_destroy)(void *); -static int (*real_zmq_ctx_shutdown)(void *); -static int (*real_zmq_ctx_set)(void *, int, int); -static int (*real_zmq_ctx_get)(void *, int); -static void *(*real_zmq_socket)(void *, int); -static int (*real_zmq_close)(void *); -static int (*real_zmq_bind)(void *, const char *); -static int (*real_zmq_unbind)(void *, const char *); -static int (*real_zmq_connect)(void *, const char *); -static int (*real_zmq_disconnect)(void *, const char *); -static int (*real_zmq_setsockopt)(void *, int, const void *, size_t); -static int (*real_zmq_getsockopt)(void *, int, void *, size_t *); -static int (*real_zmq_send)(void *, const void *, size_t, int); -static int (*real_zmq_recv)(void *, void *, size_t, int); -static int (*real_zmq_msg_send)(zmq_msg_t *, void *, int); -static int (*real_zmq_msg_recv)(zmq_msg_t *, void *, int); -static int (*real_zmq_msg_init)(zmq_msg_t *); -static int (*real_zmq_msg_init_size)(zmq_msg_t *, size_t); -static int (*real_zmq_msg_init_data)(zmq_msg_t *, void *, size_t, void (*)(void *, void *), void *); -static int (*real_zmq_msg_close)(zmq_msg_t *); -static void *(*real_zmq_msg_data)(zmq_msg_t *); -static size_t (*real_zmq_msg_size)(zmq_msg_t *); -static int (*real_zmq_poll)(zmq_pollitem_t *, int, long); -static int (*real_zmq_proxy)(void *, void *, void *); -static int (*real_zmq_errno)(void); -static int (*real_zmq_socket_monitor)(void *, const char *, int); -static const char *(*real_zmq_strerror)(int); -static void *(*real_dlsym)(void *, const char *); -static void *(*real_dlopen)(const char *, int); -static int (*real_pthread_setname_np)(pthread_t, const char *); -static int (*real_prctl)(int, ...); -static int (*old_zmq_ctx_shutdown)(void *); -static int (*old_zmq_ctx_term)(void *); -static int (*old_zmq_close)(void *); - -static void nvsnap_zmq_reinit_ctx(nvsnap_zmq_ctx_t *ctx); -static void nvsnap_zmq_register_ctx(nvsnap_zmq_ctx_t *ctx); -static void nvsnap_zmq_count_ops(nvsnap_zmq_socket_t *sock, int *connects, int *binds); -static void *nvsnap_zmq_symbol_override(const char *symbol); -static int nvsnap_zmq_restore_detected(void); -static int nvsnap_zmq_trace_enabled(void); -static void nvsnap_zmq_switch_to_newns_if_needed(void); -static void nvsnap_zmq_load_real(void); - -static void nvsnap_zmq_reset_real(void) { - real_zmq_ctx_new = NULL; - real_zmq_ctx_term = NULL; - real_zmq_ctx_destroy = NULL; - real_zmq_ctx_shutdown = NULL; - real_zmq_ctx_set = NULL; - real_zmq_ctx_get = NULL; - real_zmq_socket = NULL; - real_zmq_close = NULL; - real_zmq_bind = NULL; - real_zmq_unbind = NULL; - real_zmq_connect = NULL; - real_zmq_disconnect = NULL; - real_zmq_setsockopt = NULL; - real_zmq_getsockopt = NULL; - real_zmq_send = NULL; - real_zmq_recv = NULL; - real_zmq_msg_send = NULL; - real_zmq_msg_recv = NULL; - real_zmq_msg_init = NULL; - real_zmq_msg_init_size = NULL; - real_zmq_msg_init_data = NULL; - real_zmq_msg_close = NULL; - real_zmq_poll = NULL; - real_zmq_proxy = NULL; - real_zmq_errno = NULL; -} - -static void nvsnap_zmq_resolve_symbol(void **fn, const char *name, void *handle) { - if (*fn) { - return; - } - if (!real_dlsym) { - real_dlsym = dlvsym(RTLD_NEXT, "dlsym", "GLIBC_2.2.5"); - } - if (real_dlsym) { - *fn = real_dlsym(handle, name); - } else { - *fn = dlsym(handle, name); - } -} - -static int nvsnap_zmq_is_libzmq_path(const char *filename) { - if (!filename || !filename[0]) { - return 0; - } - if (strstr(filename, "libzmq") != NULL) { - return 1; - } - return 0; -} - -static const char *nvsnap_zmq_override_path(void) { - const char *path = getenv("NVSNAP_ZMQ_LIB_PATH"); - if (path && path[0]) { - return path; - } - return NULL; -} - -static int nvsnap_zmq_force_terminate_enabled(void) { - const char *env = getenv("NVSNAP_ZMQ_FORCE_TERMINATE"); - return env && env[0] == '1'; -} - -static int nvsnap_zmq_recovery_delay_ms(void) { - const char *env = getenv("NVSNAP_ZMQ_RECOVERY_DELAY_MS"); - if (env && *env) { - char *end = NULL; - long val = strtol(env, &end, 10); - if (end && *end == '\0' && val > 0) { - return (int)val; - } - } - return 0; -} - -static int nvsnap_zmq_ehostunreach_max_ms(void) { - const char *env = getenv("NVSNAP_ZMQ_EHOSTUNREACH_MAX_MS"); - if (env && *env) { - char *end = NULL; - long val = strtol(env, &end, 10); - if (end && *end == '\0' && val > 0) { - return (int)val; - } - } - return 0; -} - -static int nvsnap_zmq_ehostunreach_sleep_ms(void) { - const char *env = getenv("NVSNAP_ZMQ_EHOSTUNREACH_SLEEP_MS"); - if (env && *env) { - char *end = NULL; - long val = strtol(env, &end, 10); - if (end && *end == '\0' && val > 0) { - return (int)val; - } - } - return 100; -} - -static int nvsnap_zmq_gate_max_ms(void) { - const char *env = getenv("NVSNAP_ZMQ_GATE_MAX_MS"); - if (env && *env) { - char *end = NULL; - long val = strtol(env, &end, 10); - if (end && *end == '\0' && val > 0) { - return (int)val; - } - } - return 1000; -} - -static int nvsnap_zmq_gate_sleep_ms(void) { - const char *env = getenv("NVSNAP_ZMQ_GATE_SLEEP_MS"); - if (env && *env) { - char *end = NULL; - long val = strtol(env, &end, 10); - if (end && *end == '\0' && val > 0) { - return (int)val; - } - } - return 50; -} - -static int nvsnap_zmq_kill_bg_threads_enabled(void) { - const char *env = getenv("NVSNAP_ZMQ_KILL_BG_THREADS"); - return env && env[0] == '1'; -} - -static int nvsnap_zmq_skip_old_close_enabled(void) { - const char *env = getenv("NVSNAP_ZMQ_SKIP_OLD_CLOSE"); - return env && env[0] == '1'; -} - -static int nvsnap_zmq_rebuild_on_einval_enabled(void) { - const char *env = getenv("NVSNAP_ZMQ_REBUILD_ON_EINVAL"); - if (!env || !env[0]) { - return 0; - } - return env[0] == '1'; -} - -static int nvsnap_zmq_rebuild_on_first_send_enabled(void) { - const char *env = getenv("NVSNAP_ZMQ_REBUILD_ON_FIRST_SEND"); - if (!env || !env[0]) { - return 0; - } - return env[0] == '1'; -} - -static int nvsnap_zmq_is_bg_thread_name(const char *name) { - if (!name || !name[0]) { - return 0; - } - return strncmp(name, "ZMQbg", 5) == 0; -} - -static void nvsnap_zmq_maybe_exit_bg_thread(const char *name, const char *src) { - if (!nvsnap_zmq_kill_bg_threads_enabled()) { - return; - } - if (!nvsnap_zmq_restore_detected()) { - return; - } - if (!nvsnap_zmq_is_bg_thread_name(name)) { - return; - } - NVSNAP_WARN("ZMQ bg thread detected post-restore, exiting name=%s src=%s tid=%ld", - name ? name : "(null)", src ? src : "(unknown)", (long)syscall(SYS_gettid)); - pthread_exit(NULL); -} - -static int nvsnap_zmq_io_threads_override(void) { - const char *env = getenv("NVSNAP_ZMQ_IO_THREADS"); - if (env && *env) { - char *end = NULL; - long val = strtol(env, &end, 10); - if (end && *end == '\0' && val >= 0) { - return (int)val; - } - } - return -1; -} - -static void nvsnap_zmq_kill_bg_threads(void) { - if (!nvsnap_zmq_kill_bg_threads_enabled()) { - return; - } - DIR *dir = opendir("/proc/self/task"); - if (!dir) { - NVSNAP_WARN("ZMQ kill threads: failed to open /proc/self/task errno=%d", errno); - return; - } - pid_t self_tid = (pid_t)syscall(SYS_gettid); - int killed = 0; - struct dirent *ent = NULL; - while ((ent = readdir(dir)) != NULL) { - if (ent->d_name[0] < '0' || ent->d_name[0] > '9') { - continue; - } - pid_t tid = (pid_t)strtol(ent->d_name, NULL, 10); - if (tid <= 0 || tid == self_tid) { - continue; - } - char comm_path[256]; - snprintf(comm_path, sizeof(comm_path), "/proc/self/task/%s/comm", ent->d_name); - FILE *f = fopen(comm_path, "r"); - if (!f) { - continue; - } - char comm[64]; - if (fgets(comm, sizeof(comm), f)) { - comm[strcspn(comm, "\n")] = '\0'; - if (strncmp(comm, "ZMQbg", 5) == 0) { - NVSNAP_INFO("ZMQ kill thread tid=%d comm=%s", tid, comm); - kill(tid, SIGKILL); - killed++; - } - } - fclose(f); - } - closedir(dir); - NVSNAP_INFO("ZMQ kill threads completed killed=%d", killed); -} - -static void nvsnap_zmq_set_recovering(int recovering) { - pthread_mutex_lock(&g_zmq_recover_lock); - g_zmq_recovering = recovering; - pthread_mutex_unlock(&g_zmq_recover_lock); -} - -static int nvsnap_zmq_is_recovering(void) { - int recovering = 0; - pthread_mutex_lock(&g_zmq_recover_lock); - recovering = g_zmq_recovering; - pthread_mutex_unlock(&g_zmq_recover_lock); - return recovering; -} - -static void nvsnap_zmq_count_state(int *ctxs, int *socks) { - int ctx_count = 0; - int sock_count = 0; - pthread_mutex_lock(&g_zmq_ctx_list_lock); - nvsnap_zmq_ctx_t *ctx = g_zmq_ctx_list; - while (ctx) { - ctx_count++; - nvsnap_zmq_socket_t *sock = ctx->sockets; - while (sock) { - sock_count++; - sock = sock->next; - } - ctx = ctx->next; - } - pthread_mutex_unlock(&g_zmq_ctx_list_lock); - if (ctxs) { - *ctxs = ctx_count; - } - if (socks) { - *socks = sock_count; - } -} - -static void nvsnap_zmq_gate_if_recovering(const char *op) { - if (!nvsnap_zmq_restore_detected()) { - return; - } - if (!nvsnap_zmq_is_recovering()) { - return; - } - int max_ms = nvsnap_zmq_gate_max_ms(); - int sleep_ms = nvsnap_zmq_gate_sleep_ms(); - int waited = 0; - while (nvsnap_zmq_is_recovering() && waited < max_ms) { - usleep((useconds_t)sleep_ms * 1000); - waited += sleep_ms; - } - if (nvsnap_zmq_is_recovering()) { - NVSNAP_WARN("ZMQ gate timed out op=%s waited_ms=%d", op ? op : "unknown", waited); - } else if (waited > 0 && nvsnap_zmq_trace_enabled()) { - NVSNAP_INFO("ZMQ gate cleared op=%s waited_ms=%d", op ? op : "unknown", waited); - } -} - -static void nvsnap_zmq_load_old(void) { - if (old_zmq_ctx_shutdown || old_zmq_ctx_term || old_zmq_close) { - return; - } - if (!real_dlsym) { - real_dlsym = dlvsym(RTLD_NEXT, "dlsym", "GLIBC_2.2.5"); - } - if (real_dlsym) { - old_zmq_ctx_shutdown = real_dlsym(RTLD_NEXT, "zmq_ctx_shutdown"); - old_zmq_ctx_term = real_dlsym(RTLD_NEXT, "zmq_ctx_term"); - old_zmq_close = real_dlsym(RTLD_NEXT, "zmq_close"); - } else { - old_zmq_ctx_shutdown = dlsym(RTLD_NEXT, "zmq_ctx_shutdown"); - old_zmq_ctx_term = dlsym(RTLD_NEXT, "zmq_ctx_term"); - old_zmq_close = dlsym(RTLD_NEXT, "zmq_close"); - } - if (!old_zmq_ctx_shutdown || !old_zmq_ctx_term || !old_zmq_close) { - void *lib = dlopen("libzmq.so.5", RTLD_LAZY); - if (!lib) { - lib = dlopen("libzmq.so", RTLD_LAZY); - } - if (lib) { - if (!old_zmq_ctx_shutdown) { - old_zmq_ctx_shutdown = dlsym(lib, "zmq_ctx_shutdown"); - } - if (!old_zmq_ctx_term) { - old_zmq_ctx_term = dlsym(lib, "zmq_ctx_term"); - } - if (!old_zmq_close) { - old_zmq_close = dlsym(lib, "zmq_close"); - } - } - } -} - -static void nvsnap_zmq_shutdown_ctxs_before_switch(void) { - pthread_once(&g_zmq_once, nvsnap_zmq_load_real); - if (!real_zmq_ctx_shutdown) { - return; - } - pthread_mutex_lock(&g_zmq_ctx_list_lock); - nvsnap_zmq_ctx_t *ctx = g_zmq_ctx_list; - while (ctx) { - if (ctx->real && !ctx->destroyed) { - NVSNAP_INFO("ZMQ pre-switch shutdown ctx=%p real=%p", ctx, ctx->real); - real_zmq_ctx_shutdown(ctx->real); - } - ctx = ctx->next; - } - pthread_mutex_unlock(&g_zmq_ctx_list_lock); -} - -static int nvsnap_zmq_pre_shutdown_enabled(void) { - const char *env = getenv("NVSNAP_ZMQ_PRE_SHUTDOWN"); - return env && env[0] == '1'; -} - -void *dlopen(const char *filename, int flags) { - if (!real_dlsym) { - real_dlsym = dlvsym(RTLD_NEXT, "dlsym", "GLIBC_2.2.5"); - } - if (!real_dlopen) { - real_dlopen = real_dlsym ? real_dlsym(RTLD_NEXT, "dlopen") : dlsym(RTLD_NEXT, "dlopen"); - } - if (nvsnap_zmq_restore_detected() && nvsnap_zmq_is_libzmq_path(filename)) { - pthread_mutex_lock(&g_zmq_load_lock); - if (!g_zmq_newns_handle) { - const char *override = nvsnap_zmq_override_path(); - void *lib = NULL; - if (override) { - lib = dlmopen(LM_ID_NEWLM, override, flags); - } - if (!lib) { - lib = dlmopen(LM_ID_NEWLM, filename, flags); - } - if (!lib) { - lib = dlmopen(LM_ID_NEWLM, "libzmq.so.5", flags); - } - if (!lib) { - lib = dlmopen(LM_ID_NEWLM, "libzmq.so", flags); - } - if (lib) { - g_zmq_newns_handle = lib; - g_zmq_use_newns = 1; - nvsnap_zmq_reset_real(); - nvsnap_zmq_load_real(); - NVSNAP_INFO("ZMQ dlopen redirected to newns handle=%p path=%s", - lib, override ? override : (filename ? filename : "(null)")); - } - } - pthread_mutex_unlock(&g_zmq_load_lock); - if (g_zmq_newns_handle) { - return g_zmq_newns_handle; - } - } - return real_dlopen ? real_dlopen(filename, flags) : NULL; -} - -static void nvsnap_zmq_load_real(void) { - void *handle = g_zmq_use_newns && g_zmq_newns_handle ? g_zmq_newns_handle : RTLD_NEXT; - if (!real_dlsym) { - real_dlsym = dlvsym(RTLD_NEXT, "dlsym", "GLIBC_2.2.5"); - } - g_zmq_handle = handle; - - nvsnap_zmq_resolve_symbol((void **)&real_zmq_ctx_new, "zmq_ctx_new", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_ctx_term, "zmq_ctx_term", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_ctx_destroy, "zmq_ctx_destroy", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_ctx_shutdown, "zmq_ctx_shutdown", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_ctx_set, "zmq_ctx_set", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_ctx_get, "zmq_ctx_get", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_socket, "zmq_socket", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_close, "zmq_close", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_bind, "zmq_bind", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_unbind, "zmq_unbind", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_connect, "zmq_connect", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_disconnect, "zmq_disconnect", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_setsockopt, "zmq_setsockopt", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_getsockopt, "zmq_getsockopt", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_send, "zmq_send", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_recv, "zmq_recv", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_msg_send, "zmq_msg_send", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_msg_recv, "zmq_msg_recv", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_msg_init, "zmq_msg_init", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_msg_init_size, "zmq_msg_init_size", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_msg_init_data, "zmq_msg_init_data", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_msg_close, "zmq_msg_close", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_msg_data, "zmq_msg_data", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_msg_size, "zmq_msg_size", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_poll, "zmq_poll", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_proxy, "zmq_proxy", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_errno, "zmq_errno", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_strerror, "zmq_strerror", handle); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_socket_monitor, "zmq_socket_monitor", handle); - - if (!g_zmq_use_newns && (!real_zmq_ctx_new || !real_zmq_socket)) { - const char *override = nvsnap_zmq_override_path(); - void *lib = NULL; - if (override) { - lib = dlopen(override, RTLD_LAZY); - } - if (!lib) { - lib = dlopen("libzmq.so.5", RTLD_LAZY); - } - if (!lib) { - lib = dlopen("libzmq.so", RTLD_LAZY); - } - if (lib) { - g_zmq_handle = lib; - nvsnap_zmq_resolve_symbol((void **)&real_zmq_ctx_new, "zmq_ctx_new", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_ctx_term, "zmq_ctx_term", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_ctx_destroy, "zmq_ctx_destroy", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_ctx_shutdown, "zmq_ctx_shutdown", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_ctx_set, "zmq_ctx_set", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_ctx_get, "zmq_ctx_get", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_socket, "zmq_socket", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_close, "zmq_close", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_bind, "zmq_bind", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_unbind, "zmq_unbind", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_connect, "zmq_connect", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_disconnect, "zmq_disconnect", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_setsockopt, "zmq_setsockopt", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_getsockopt, "zmq_getsockopt", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_send, "zmq_send", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_recv, "zmq_recv", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_msg_send, "zmq_msg_send", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_msg_recv, "zmq_msg_recv", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_msg_init, "zmq_msg_init", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_msg_init_size, "zmq_msg_init_size", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_msg_init_data, "zmq_msg_init_data", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_msg_close, "zmq_msg_close", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_msg_data, "zmq_msg_data", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_msg_size, "zmq_msg_size", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_poll, "zmq_poll", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_proxy, "zmq_proxy", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_errno, "zmq_errno", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_strerror, "zmq_strerror", lib); - nvsnap_zmq_resolve_symbol((void **)&real_zmq_socket_monitor, "zmq_socket_monitor", lib); - } - } -} - -static void nvsnap_zmq_switch_to_newns_if_needed(void) { - if (!nvsnap_zmq_restore_detected()) { - return; - } - pthread_mutex_lock(&g_zmq_load_lock); - if (!g_zmq_use_newns) { - if (nvsnap_zmq_pre_shutdown_enabled()) { - nvsnap_zmq_shutdown_ctxs_before_switch(); - } - const char *override = nvsnap_zmq_override_path(); - void *lib = NULL; - if (override) { - lib = dlmopen(LM_ID_NEWLM, override, RTLD_LAZY | RTLD_LOCAL); - } - if (!lib) { - lib = dlmopen(LM_ID_NEWLM, "libzmq.so.5", RTLD_LAZY | RTLD_LOCAL); - } - if (!lib) { - lib = dlmopen(LM_ID_NEWLM, "libzmq.so", RTLD_LAZY | RTLD_LOCAL); - } - if (lib) { - g_zmq_newns_handle = lib; - g_zmq_use_newns = 1; - nvsnap_zmq_reset_real(); - nvsnap_zmq_load_real(); - NVSNAP_INFO("ZMQ switched to new link namespace handle=%p path=%s", - lib, override ? override : "libzmq.so.5"); - } else { - NVSNAP_WARN("ZMQ newns dlmopen failed: %s", dlerror()); - } - } - pthread_mutex_unlock(&g_zmq_load_lock); -} - -static void *nvsnap_zmq_symbol_override(const char *symbol) { - if (!symbol) { - return NULL; - } - if (strcmp(symbol, "zmq_ctx_new") == 0) return (void *)zmq_ctx_new; - if (strcmp(symbol, "zmq_ctx_term") == 0) return (void *)zmq_ctx_term; - if (strcmp(symbol, "zmq_ctx_destroy") == 0) return (void *)zmq_ctx_destroy; - if (strcmp(symbol, "zmq_ctx_shutdown") == 0) return (void *)zmq_ctx_shutdown; - if (strcmp(symbol, "zmq_ctx_set") == 0) return (void *)zmq_ctx_set; - if (strcmp(symbol, "zmq_ctx_get") == 0) return (void *)zmq_ctx_get; - if (strcmp(symbol, "zmq_socket") == 0) return (void *)zmq_socket; - if (strcmp(symbol, "zmq_close") == 0) return (void *)zmq_close; - if (strcmp(symbol, "zmq_bind") == 0) return (void *)zmq_bind; - if (strcmp(symbol, "zmq_unbind") == 0) return (void *)zmq_unbind; - if (strcmp(symbol, "zmq_connect") == 0) return (void *)zmq_connect; - if (strcmp(symbol, "zmq_disconnect") == 0) return (void *)zmq_disconnect; - if (strcmp(symbol, "zmq_setsockopt") == 0) return (void *)zmq_setsockopt; - if (strcmp(symbol, "zmq_getsockopt") == 0) return (void *)zmq_getsockopt; - if (strcmp(symbol, "zmq_send") == 0) return (void *)zmq_send; - if (strcmp(symbol, "zmq_recv") == 0) return (void *)zmq_recv; - if (strcmp(symbol, "zmq_msg_send") == 0) return (void *)zmq_msg_send; - if (strcmp(symbol, "zmq_msg_recv") == 0) return (void *)zmq_msg_recv; - if (strcmp(symbol, "zmq_msg_init") == 0) return (void *)zmq_msg_init; - if (strcmp(symbol, "zmq_msg_init_size") == 0) return (void *)zmq_msg_init_size; - if (strcmp(symbol, "zmq_msg_init_data") == 0) return (void *)zmq_msg_init_data; - if (strcmp(symbol, "zmq_msg_close") == 0) return (void *)zmq_msg_close; - if (strcmp(symbol, "zmq_poll") == 0) return (void *)zmq_poll; - if (strcmp(symbol, "zmq_proxy") == 0) return (void *)zmq_proxy; - if (strcmp(symbol, "zmq_errno") == 0) return (void *)zmq_errno; - return NULL; -} - -/* - * dlsym override — unversioned, same as the working 2-library setup. - * Unversioned dlsym override — routes symbol lookups to our wrappers. - * NvSnap's symbol table handles CUDA + NCCL routing. - * NvSnap handles ZMQ routing. - */ -extern void *nvsnap_lookup_symbol(const char *name); /* nvsnap_symbol_table.c */ - -void *dlsym(void *handle, const char *symbol) { - if (!real_dlsym) { - real_dlsym = dlvsym(RTLD_NEXT, "dlsym", "GLIBC_2.2.5"); - } - /* Only intercept global lookups (RTLD_DEFAULT, RTLD_NEXT). - * When a specific library handle is passed (e.g., nvsnap_resolve_real - * calling dlsym(libcudart_handle, "cudaSetDevice")), pass through directly. - * Without this, we return our own hooks instead of the real functions, - * causing NULL resolution → segfault during CUDA init. */ - if (handle == RTLD_DEFAULT || handle == RTLD_NEXT) { - void *override = nvsnap_zmq_symbol_override(symbol); - if (override) return override; - override = nvsnap_lookup_symbol(symbol); - if (override) return override; - } - return real_dlsym ? real_dlsym(handle, symbol) : NULL; -} - -int pthread_setname_np(pthread_t thread, const char *name) { - if (!real_pthread_setname_np) { - if (!real_dlsym) { - real_dlsym = dlvsym(RTLD_NEXT, "dlsym", "GLIBC_2.2.5"); - } - if (real_dlsym) - real_pthread_setname_np = real_dlsym(RTLD_NEXT, "pthread_setname_np"); - } - if (!real_pthread_setname_np) { - errno = ENOSYS; - return -1; - } - if (nvsnap_zmq_kill_bg_threads_enabled() && - nvsnap_zmq_restore_detected() && - nvsnap_zmq_is_bg_thread_name(name)) { - if (pthread_equal(thread, pthread_self())) { - nvsnap_zmq_maybe_exit_bg_thread(name, "pthread_setname_np"); - } else { - NVSNAP_WARN("ZMQ bg thread named in other thread post-restore name=%s", name); - } - } - return real_pthread_setname_np(thread, name); -} - -int prctl(int option, ...) { - if (!real_prctl) { - if (!real_dlsym) { - real_dlsym = dlvsym(RTLD_NEXT, "dlsym", "GLIBC_2.2.5"); - } - real_prctl = real_dlsym ? real_dlsym(RTLD_NEXT, "prctl") - : dlsym(RTLD_NEXT, "prctl"); - } - if (!real_prctl) { - errno = ENOSYS; - return -1; - } - va_list ap; - unsigned long arg2 = 0; - unsigned long arg3 = 0; - unsigned long arg4 = 0; - unsigned long arg5 = 0; - va_start(ap, option); - arg2 = va_arg(ap, unsigned long); - arg3 = va_arg(ap, unsigned long); - arg4 = va_arg(ap, unsigned long); - arg5 = va_arg(ap, unsigned long); - va_end(ap); - if (option == PR_SET_NAME && nvsnap_zmq_kill_bg_threads_enabled() && - nvsnap_zmq_restore_detected()) { - const char *name = (const char *)arg2; - if (nvsnap_zmq_is_bg_thread_name(name)) { - nvsnap_zmq_maybe_exit_bg_thread(name, "prctl"); - } - } - return real_prctl(option, arg2, arg3, arg4, arg5); -} - -static int nvsnap_zmq_is_enabled(void) { - if (g_zmq_enabled >= 0) { - return g_zmq_enabled; - } - const char *env = getenv("NVSNAP_ZMQ_INTERCEPT"); - if (!env) { - g_zmq_enabled = 1; - return g_zmq_enabled; - } - if (env[0] != '\0' && env[0] != '0' && env[0] != 'f' && env[0] != 'F') { - g_zmq_enabled = 1; - } else { - g_zmq_enabled = 0; - } - return g_zmq_enabled; -} - -static int nvsnap_zmq_trace_enabled(void) { - if (g_zmq_trace >= 0) { - return g_zmq_trace; - } - const char *env = getenv("NVSNAP_ZMQ_TRACE"); - if (env && env[0] != '\0' && env[0] != '0' && env[0] != 'f' && env[0] != 'F') { - g_zmq_trace = 1; - } else { - g_zmq_trace = 0; - } - return g_zmq_trace; -} - -static int nvsnap_zmq_restore_detected(void) { - if (g_zmq_restore_seen) { - return 1; - } - if (access("/nvsnap-lib/.restored", F_OK) == 0 || - access("/nvsnap/.restored", F_OK) == 0 || - access("/var/run/nvsnap/.restored", F_OK) == 0) { - g_zmq_restore_seen = 1; - } - return g_zmq_restore_seen; -} - -static nvsnap_zmq_ctx_t *nvsnap_zmq_ctx_from_ptr(void *ctx) { - nvsnap_zmq_ctx_t *proxy = (nvsnap_zmq_ctx_t *)ctx; - if (proxy && proxy->magic == NVSNAP_ZMQ_CTX_MAGIC) { - return proxy; - } - if (!ctx || !nvsnap_zmq_is_enabled()) { - return NULL; - } - pthread_mutex_lock(&g_zmq_ctx_list_lock); - nvsnap_zmq_ctx_t *iter = g_zmq_ctx_list; - while (iter) { - if (iter->real == ctx) { - pthread_mutex_unlock(&g_zmq_ctx_list_lock); - return iter; - } - iter = iter->next; - } - pthread_mutex_unlock(&g_zmq_ctx_list_lock); - nvsnap_zmq_ctx_t *wrapped = calloc(1, sizeof(*wrapped)); - if (!wrapped) { - return NULL; - } - wrapped->magic = NVSNAP_ZMQ_CTX_MAGIC; - wrapped->real = ctx; - pthread_mutex_init(&wrapped->lock, NULL); - nvsnap_zmq_register_ctx(wrapped); - NVSNAP_INFO("ZMQ ctx wrap pid=%d ctx=%p real=%p", getpid(), wrapped, wrapped->real); - return wrapped; -} - -static nvsnap_zmq_socket_t *nvsnap_zmq_sock_from_ptr(void *sock) { - nvsnap_zmq_socket_t *proxy = (nvsnap_zmq_socket_t *)sock; - if (proxy && proxy->magic == NVSNAP_ZMQ_SOCK_MAGIC) { - return proxy; - } - return NULL; -} - -static int nvsnap_zmq_errno(void) { - if (real_zmq_errno) { - return real_zmq_errno(); - } - return errno; -} - -static const char *nvsnap_zmq_strerror(int err) { - if (real_zmq_strerror) { - return real_zmq_strerror(err); - } - return strerror(err); -} - -int zmq_errno(void) { - pthread_once(&g_zmq_once, nvsnap_zmq_load_real); - if (!nvsnap_zmq_is_enabled()) { - return real_zmq_errno ? real_zmq_errno() : errno; - } - return nvsnap_zmq_errno(); -} - -static int nvsnap_zmq_ipc_path(const char *endpoint, const char **path_out) { - static const char prefix[] = "ipc://"; - size_t len; - const char *path; - - if (!endpoint) { - return 0; - } - len = strlen(prefix); - if (strncmp(endpoint, prefix, len) != 0) { - return 0; - } - path = endpoint + len; - if (!path[0] || path[0] == '@') { - return 0; - } - if (path_out) { - *path_out = path; - } - return 1; -} - -static void nvsnap_zmq_add_op(nvsnap_zmq_socket_t *sock, enum nvsnap_zmq_op_type type, - int option, const void *data, size_t len, const char *endpoint) { - if (nvsnap_zmq_trace_enabled() || nvsnap_zmq_restore_detected()) { - NVSNAP_INFO("ZMQ op add pid=%d sock=%p real=%p type=%d endpoint=%s", - getpid(), sock, sock ? sock->real : NULL, type, - endpoint ? endpoint : "(null)"); - } - struct nvsnap_zmq_op *op = calloc(1, sizeof(*op)); - if (!op) { - return; - } - op->type = type; - op->option = option; - if (data && len > 0) { - op->data = malloc(len); - if (op->data) { - memcpy(op->data, data, len); - op->len = len; - } - } - if (endpoint) { - op->endpoint = strdup(endpoint); - } - if (!sock->ops_head) { - sock->ops_head = op; - sock->ops_tail = op; - } else { - sock->ops_tail->next = op; - sock->ops_tail = op; - } -} - -static void nvsnap_zmq_replay_ops_ex(nvsnap_zmq_socket_t *sock, int allow_bind) { - struct nvsnap_zmq_op *op = sock->ops_head; - int op_count = 0; - if (nvsnap_zmq_restore_detected() || nvsnap_zmq_trace_enabled()) { - NVSNAP_INFO("ZMQ replay_ops_ex starting pid=%d sock=%p real=%p allow_bind=%d", - getpid(), sock, sock ? sock->real : NULL, allow_bind); - } - while (op) { - op_count++; - switch (op->type) { - case NVSNAP_ZMQ_OP_SETSOCKOPT: - if (op->data && op->len) { - if (real_zmq_setsockopt(sock->real, op->option, op->data, op->len) != 0) { - NVSNAP_WARN("ZMQ replay setsockopt failed opt=%d errno=%d", op->option, nvsnap_zmq_errno()); - } - } - break; - case NVSNAP_ZMQ_OP_BIND: - if (op->endpoint) { - if (!allow_bind && nvsnap_zmq_restore_detected()) { - if (nvsnap_zmq_trace_enabled()) { - NVSNAP_INFO("ZMQ replay bind deferred endpoint=%s", op->endpoint); - } - break; - } - const char *ipc_path = NULL; - if (nvsnap_zmq_ipc_path(op->endpoint, &ipc_path)) { - unlink(ipc_path); - } - int attempts = 0; - int rc = -1; - int err = 0; - int max_attempts = nvsnap_zmq_restore_detected() ? 10 : 1; - do { - rc = real_zmq_bind(sock->real, op->endpoint); - if (rc == 0) { - break; - } - err = nvsnap_zmq_errno(); - attempts++; - usleep(100 * 1000); - } while (attempts < max_attempts); - if (rc != 0) { - NVSNAP_WARN("ZMQ replay bind failed endpoint=%s errno=%d err=%s attempts=%d", - op->endpoint, err, nvsnap_zmq_strerror(err), attempts); - /* For IPC: try unlinking stale socket file and retry once more */ - const char *retry_ipc_path = NULL; - if (nvsnap_zmq_ipc_path(op->endpoint, &retry_ipc_path)) { - NVSNAP_INFO("ZMQ replay bind: unlinking stale IPC socket %s and retrying", - retry_ipc_path); - unlink(retry_ipc_path); - rc = real_zmq_bind(sock->real, op->endpoint); - if (rc == 0) { - NVSNAP_INFO("ZMQ replay bind succeeded after IPC unlink endpoint=%s", - op->endpoint); - } else { - NVSNAP_WARN("ZMQ replay bind still failed after IPC unlink endpoint=%s errno=%d", - op->endpoint, nvsnap_zmq_errno()); - sock->replay_failed = 1; - } - } else { - sock->replay_failed = 1; - } - } else if (nvsnap_zmq_trace_enabled()) { - NVSNAP_INFO("ZMQ replay bind ok endpoint=%s", op->endpoint); - } - } - break; - case NVSNAP_ZMQ_OP_CONNECT: - if (op->endpoint) { - int attempts = 0; - int rc = -1; - int err = 0; - const char *ipc_path = NULL; - if (nvsnap_zmq_ipc_path(op->endpoint, &ipc_path)) { - int wait_attempts = 0; - while (access(ipc_path, F_OK) != 0 && wait_attempts < 50) { - usleep(100 * 1000); - wait_attempts++; - } - } - do { - rc = real_zmq_connect(sock->real, op->endpoint); - if (rc == 0) { - break; - } - err = nvsnap_zmq_errno(); - if (err != ENOENT && err != ECONNREFUSED) { - break; - } - attempts++; - usleep(100 * 1000); - } while (attempts < 50); - if (rc != 0) { - NVSNAP_WARN("ZMQ replay connect failed endpoint=%s errno=%d attempts=%d", op->endpoint, err, attempts); - sock->replay_failed = 1; - } else if (nvsnap_zmq_trace_enabled()) { - NVSNAP_INFO("ZMQ replay connect ok endpoint=%s attempts=%d", op->endpoint, attempts); - } - } - break; - case NVSNAP_ZMQ_OP_UNBIND: - if (op->endpoint) { - if (real_zmq_unbind(sock->real, op->endpoint) != 0) { - NVSNAP_WARN("ZMQ replay unbind failed endpoint=%s errno=%d", op->endpoint, nvsnap_zmq_errno()); - } - } - break; - case NVSNAP_ZMQ_OP_DISCONNECT: - if (op->endpoint) { - if (real_zmq_disconnect(sock->real, op->endpoint) != 0) { - NVSNAP_WARN("ZMQ replay disconnect failed endpoint=%s errno=%d", op->endpoint, nvsnap_zmq_errno()); - } - } - break; - default: - break; - } - op = op->next; - } - if (nvsnap_zmq_restore_detected() || nvsnap_zmq_trace_enabled()) { - NVSNAP_INFO("ZMQ replay_ops_ex completed pid=%d sock=%p ops_processed=%d", - getpid(), sock, op_count); - } -} - -static void nvsnap_zmq_force_reconnect(nvsnap_zmq_socket_t *sock) { - if (!sock || !sock->real) { - return; - } - struct nvsnap_zmq_op *op = sock->ops_head; - while (op) { - if (op->type == NVSNAP_ZMQ_OP_CONNECT && op->endpoint) { - int rc = real_zmq_connect(sock->real, op->endpoint); - NVSNAP_WARN("ZMQ force reconnect sock=%p endpoint=%s rc=%d errno=%d", - sock, op->endpoint, rc, nvsnap_zmq_errno()); - } - op = op->next; - } -} - -static void nvsnap_zmq_force_connect_from_last_endpoint(nvsnap_zmq_socket_t *sock) { - if (!sock || !sock->real || sock->forced_connect_done) { - return; - } - if (!nvsnap_zmq_restore_detected()) { - return; - } - int connects = 0; - int binds = 0; - nvsnap_zmq_count_ops(sock, &connects, &binds); - if (connects > 0) { - return; - } - if (sock->type == ZMQ_ROUTER) { - return; - } - if (!real_zmq_getsockopt) { - return; - } - char endpoint[256]; - size_t endpoint_len = sizeof(endpoint); - endpoint[0] = '\0'; - if (real_zmq_getsockopt(sock->real, ZMQ_LAST_ENDPOINT, endpoint, &endpoint_len) != 0 || - !endpoint[0]) { - return; - } - if (strncmp(endpoint, "ipc://", 6) != 0) { - return; - } - const char *path = endpoint + 6; - if (!path[0]) { - return; - } - if (access(path, F_OK) != 0) { - NVSNAP_WARN("ZMQ forced connect skipped; ipc path missing endpoint=%s", endpoint); - return; - } - sock->forced_connect_done = 1; - nvsnap_zmq_add_op(sock, NVSNAP_ZMQ_OP_CONNECT, 0, NULL, 0, endpoint); - int rc = real_zmq_connect(sock->real, endpoint); - NVSNAP_WARN("ZMQ forced connect from last endpoint sock=%p endpoint=%s rc=%d errno=%d", - sock, endpoint, rc, nvsnap_zmq_errno()); -} - -static void nvsnap_zmq_dump_ctx_sockets(nvsnap_zmq_ctx_t *ctx, const char *tag) { - if (!ctx) { - return; - } - nvsnap_zmq_socket_t *sock = ctx->sockets; - while (sock) { - int connects = 0; - int binds = 0; - nvsnap_zmq_count_ops(sock, &connects, &binds); - int type = sock->type; - char endpoint[256]; - size_t endpoint_len = sizeof(endpoint); - endpoint[0] = '\0'; - if (real_zmq_getsockopt) { - if (real_zmq_getsockopt(sock->real, ZMQ_LAST_ENDPOINT, endpoint, &endpoint_len) != 0) { - endpoint[0] = '\0'; - } - } - NVSNAP_INFO("ZMQ ctx dump %s pid=%d sock=%p real=%p type=%d ops_connect=%d ops_bind=%d endpoint=%s", - tag ? tag : "restore", getpid(), sock, sock->real, type, connects, binds, - endpoint[0] ? endpoint : "(none)"); - sock = sock->next; - } -} - -static void nvsnap_zmq_dump_ctx_sockets_to_file(nvsnap_zmq_ctx_t *ctx, const char *tag) { - if (!ctx) { - return; - } - char path[256]; - snprintf(path, sizeof(path), "/tmp/nvsnap-zmq-%d.log", (int)getpid()); - FILE *fp = fopen(path, "a"); - if (!fp) { - return; - } - nvsnap_zmq_socket_t *sock = ctx->sockets; - while (sock) { - int connects = 0; - int binds = 0; - nvsnap_zmq_count_ops(sock, &connects, &binds); - int type = sock->type; - char endpoint[256]; - size_t endpoint_len = sizeof(endpoint); - endpoint[0] = '\0'; - if (real_zmq_getsockopt) { - if (real_zmq_getsockopt(sock->real, ZMQ_LAST_ENDPOINT, endpoint, &endpoint_len) != 0) { - endpoint[0] = '\0'; - } - } - fprintf(fp, - "ZMQ ctx dump %s pid=%d sock=%p real=%p type=%d ops_connect=%d ops_bind=%d endpoint=%s\n", - tag ? tag : "restore", (int)getpid(), sock, sock->real, type, connects, binds, - endpoint[0] ? endpoint : "(none)"); - sock = sock->next; - } - fclose(fp); -} - -static void nvsnap_zmq_log_first_call(const char *where) { - static __thread int logged = 0; - if (logged) { - return; - } - logged = 1; - char cmdline[512]; - cmdline[0] = '\0'; - FILE *cfp = fopen("/proc/self/cmdline", "r"); - if (cfp) { - size_t n = fread(cmdline, 1, sizeof(cmdline) - 1, cfp); - fclose(cfp); - cmdline[n] = '\0'; - for (size_t i = 0; i < n; i++) { - if (cmdline[i] == '\0') { - cmdline[i] = ' '; - } - } - } - const char *bases[] = {"/tmp", "/nvsnap-lib"}; - for (size_t i = 0; i < sizeof(bases) / sizeof(bases[0]); i++) { - char path[256]; - snprintf(path, sizeof(path), "%s/nvsnap-zmq-%d.log", bases[i], (int)getpid()); - FILE *fp = fopen(path, "a"); - if (!fp) { - continue; - } - fprintf(fp, "ZMQ first call pid=%d where=%s cmdline=%s\n", - (int)getpid(), where ? where : "unknown", cmdline[0] ? cmdline : "(none)"); - fclose(fp); - } -} - -static size_t nvsnap_zmq_capture_identity(void *sock, unsigned char *buf, size_t bufsize) { - if (!sock || !buf || bufsize == 0 || !real_zmq_getsockopt) { - return 0; - } - size_t size = bufsize; - if (real_zmq_getsockopt(sock, ZMQ_IDENTITY, buf, &size) != 0) { - return 0; - } - if (size > 0) { - NVSNAP_INFO("ZMQ identity captured pid=%d sock=%p len=%zu first=%02x%02x%02x%02x", - getpid(), sock, size, - buf[0], size > 1 ? buf[1] : 0, - size > 2 ? buf[2] : 0, size > 3 ? buf[3] : 0); - } - return size; -} - -static void nvsnap_zmq_apply_identity(void *sock, const unsigned char *buf, size_t size) { - if (!sock || !buf || size == 0 || !real_zmq_setsockopt) { - return; - } - NVSNAP_INFO("ZMQ identity apply pid=%d sock=%p len=%zu first=%02x%02x%02x%02x", - getpid(), sock, size, - buf[0], size > 1 ? buf[1] : 0, - size > 2 ? buf[2] : 0, size > 3 ? buf[3] : 0); - if (real_zmq_setsockopt(sock, ZMQ_IDENTITY, buf, size) != 0) { - NVSNAP_WARN("ZMQ restore sockopt IDENTITY failed errno=%d", nvsnap_zmq_errno()); - } -} - -typedef struct nvsnap_zmq_monitor_ctx { - void *mon; - void *target; -} nvsnap_zmq_monitor_ctx_t; - -static void *nvsnap_zmq_monitor_thread(void *arg) { - nvsnap_zmq_monitor_ctx_t *ctx = (nvsnap_zmq_monitor_ctx_t *)arg; - if (!ctx || !ctx->mon || !real_zmq_recv) { - free(ctx); - return NULL; - } - for (;;) { - unsigned char event_buf[64]; - char addr_buf[256]; - int rc = real_zmq_recv(ctx->mon, event_buf, sizeof(event_buf), 0); - if (rc == -1) { - break; - } - size_t size = (size_t)rc; - const unsigned char *data = event_buf; - unsigned int event = 0; - unsigned int value = 0; - if (size >= 6 && data) { - event = (unsigned int)(data[0] | (data[1] << 8)); - value = (unsigned int)(data[2] | (data[3] << 8) | (data[4] << 16) | (data[5] << 24)); - } - rc = real_zmq_recv(ctx->mon, addr_buf, sizeof(addr_buf) - 1, 0); - if (rc >= 0) { - size_t copy_len = (size_t)rc < sizeof(addr_buf) - 1 ? (size_t)rc : sizeof(addr_buf) - 1; - addr_buf[copy_len] = '\0'; - NVSNAP_INFO("ZMQ monitor event pid=%d sock=%p event=%u value=%u addr=%s", - getpid(), ctx->target, event, value, addr_buf); - } - } - free(ctx); - return NULL; -} - -static void nvsnap_zmq_start_monitor(nvsnap_zmq_socket_t *sock) { - if (!sock) { - return; - } - if (sock->monitor_started) { - NVSNAP_INFO("ZMQ monitor already started pid=%d sock=%p", getpid(), sock); - return; - } - NVSNAP_INFO("ZMQ monitor start pid=%d sock=%p real=%p type=%d", - getpid(), sock, sock->real, sock->type); - if (!real_zmq_socket_monitor) { - NVSNAP_INFO("ZMQ monitor unavailable (no zmq_socket_monitor) pid=%d sock=%p", getpid(), sock); - return; - } - if (!real_zmq_socket) { - NVSNAP_INFO("ZMQ monitor unavailable (no zmq_socket) pid=%d sock=%p", getpid(), sock); - return; - } - char endpoint[128]; - snprintf(endpoint, sizeof(endpoint), "inproc://nvsnap-mon-%d-%p", getpid(), (void *)sock); - if (real_zmq_socket_monitor(sock->real, endpoint, ZMQ_EVENT_ALL) != 0) { - NVSNAP_WARN("ZMQ monitor enable failed pid=%d sock=%p errno=%d", getpid(), sock, nvsnap_zmq_errno()); - return; - } - void *mon = real_zmq_socket(sock->ctx->real, ZMQ_PAIR); - if (!mon) { - NVSNAP_WARN("ZMQ monitor socket create failed pid=%d sock=%p errno=%d", getpid(), sock, nvsnap_zmq_errno()); - return; - } - if (real_zmq_connect && real_zmq_connect(mon, endpoint) != 0) { - NVSNAP_WARN("ZMQ monitor connect failed pid=%d sock=%p errno=%d", getpid(), sock, nvsnap_zmq_errno()); - real_zmq_close(mon); - return; - } - nvsnap_zmq_monitor_ctx_t *ctx = (nvsnap_zmq_monitor_ctx_t *)calloc(1, sizeof(*ctx)); - if (!ctx) { - real_zmq_close(mon); - return; - } - ctx->mon = mon; - ctx->target = sock; - sock->monitor_sock = mon; - sock->monitor_started = 1; - if (pthread_create(&sock->monitor_thread, NULL, nvsnap_zmq_monitor_thread, ctx) != 0) { - NVSNAP_WARN("ZMQ monitor thread failed pid=%d sock=%p", getpid(), sock); - real_zmq_close(mon); - sock->monitor_sock = NULL; - sock->monitor_started = 0; - free(ctx); - return; - } - pthread_detach(sock->monitor_thread); - NVSNAP_INFO("ZMQ monitor started pid=%d sock=%p endpoint=%s", getpid(), sock, endpoint); -} - -static void nvsnap_zmq_apply_restore_sockopts(nvsnap_zmq_socket_t *sock) { - if (!sock || !sock->real) { - return; - } - if (!nvsnap_zmq_restore_detected()) { - return; - } - if (!real_zmq_setsockopt) { - return; - } - int one = 1; - if (real_zmq_setsockopt(sock->real, ZMQ_IMMEDIATE, &one, sizeof(one)) != 0) { - NVSNAP_WARN("ZMQ restore sockopt IMMEDIATE failed errno=%d", nvsnap_zmq_errno()); - } - if (sock->type == ZMQ_ROUTER) { - if (real_zmq_setsockopt(sock->real, ZMQ_ROUTER_MANDATORY, &one, sizeof(one)) != 0) { - NVSNAP_WARN("ZMQ restore sockopt ROUTER_MANDATORY failed errno=%d", nvsnap_zmq_errno()); - } - } - if (sock->type == ZMQ_DEALER || sock->type == ZMQ_REQ) { - if (real_zmq_setsockopt(sock->real, ZMQ_PROBE_ROUTER, &one, sizeof(one)) != 0) { - NVSNAP_WARN("ZMQ restore sockopt PROBE_ROUTER failed errno=%d", nvsnap_zmq_errno()); - } - } -} - -static int nvsnap_zmq_rebuild_socket(nvsnap_zmq_socket_t *sock, const char *reason) { - if (!sock || !sock->ctx || sock->closed) { - return -1; - } - if (!real_zmq_socket) { - return -1; - } - unsigned char identity[256]; - size_t identity_len = nvsnap_zmq_capture_identity(sock->real, identity, sizeof(identity)); - pthread_mutex_lock(&sock->ctx->lock); - if (sock->rebuilding) { - pthread_mutex_unlock(&sock->ctx->lock); - return 0; - } - sock->rebuilding = 1; - void *old_sock = sock->real; - NVSNAP_WARN("ZMQ rebuild socket pid=%d sock=%p real=%p type=%d reason=%s", - getpid(), sock, old_sock, sock->type, reason ? reason : "unknown"); - sock->real = real_zmq_socket(sock->ctx->real, sock->type); - if (!sock->real) { - NVSNAP_WARN("ZMQ rebuild failed: socket create type=%d errno=%d", - sock->type, nvsnap_zmq_errno()); - sock->real = old_sock; - sock->rebuilding = 0; - pthread_mutex_unlock(&sock->ctx->lock); - return -1; - } - if (identity_len > 0) { - nvsnap_zmq_apply_identity(sock->real, identity, identity_len); - } - nvsnap_zmq_apply_restore_sockopts(sock); - nvsnap_zmq_start_monitor(sock); - nvsnap_zmq_replay_ops_ex(sock, 1); - nvsnap_zmq_force_connect_from_last_endpoint(sock); - sock->replay_after_restore = 1; - sock->rebuild_count++; - if (old_sock && old_sock != sock->real && !nvsnap_zmq_skip_old_close_enabled()) { - nvsnap_zmq_load_old(); - int (*close_fn)(void *) = old_zmq_close ? old_zmq_close : real_zmq_close; - if (close_fn && close_fn(old_sock) != 0) { - NVSNAP_WARN("ZMQ rebuild: closing old socket failed errno=%d", nvsnap_zmq_errno()); - } - } - sock->rebuilding = 0; - pthread_mutex_unlock(&sock->ctx->lock); - return 0; -} - -static void nvsnap_zmq_replay_ops(nvsnap_zmq_socket_t *sock) { - nvsnap_zmq_replay_ops_ex(sock, 1); -} - -static void nvsnap_zmq_count_ops(nvsnap_zmq_socket_t *sock, int *connects, int *binds) { - if (connects) { - *connects = 0; - } - if (binds) { - *binds = 0; - } - if (!sock) { - return; - } - struct nvsnap_zmq_op *op = sock->ops_head; - while (op) { - if (op->type == NVSNAP_ZMQ_OP_CONNECT) { - if (connects) { - (*connects)++; - } - if (nvsnap_zmq_restore_detected() || nvsnap_zmq_trace_enabled()) { - NVSNAP_INFO("ZMQ op count: connect endpoint=%s", op->endpoint ? op->endpoint : "(null)"); - } - } else if (op->type == NVSNAP_ZMQ_OP_BIND) { - if (binds) { - (*binds)++; - } - if (nvsnap_zmq_restore_detected() || nvsnap_zmq_trace_enabled()) { - NVSNAP_INFO("ZMQ op count: bind endpoint=%s", op->endpoint ? op->endpoint : "(null)"); - } - } - op = op->next; - } -} - -static void nvsnap_zmq_maybe_replay_after_restore(nvsnap_zmq_socket_t *sock) { - if (!sock || !sock->ctx) { - return; - } - if (!nvsnap_zmq_restore_detected()) { - return; - } - /* Socket replay disabled — CRIU restores socket FDs directly. - * Replaying (close + reopen + rebind) destroys the CRIU-restored - * connections and breaks IPC channels. */ - return; - if (sock->replay_after_restore && !sock->replay_failed) { - return; - } - /* If replay_failed, attempt re-replay once */ - if (sock->replay_failed) { - NVSNAP_WARN("ZMQ re-replay after previous failure sock=%p real=%p", sock, sock->real); - sock->replay_failed = 0; /* Reset to allow one retry */ - } - sock->replay_after_restore = 1; - int connects = 0; - int binds = 0; - nvsnap_zmq_count_ops(sock, &connects, &binds); - NVSNAP_INFO("ZMQ replay after restore sock=%p real=%p ops_connect=%d ops_bind=%d", - sock, sock->real, connects, binds); - pthread_mutex_lock(&sock->ctx->lock); - nvsnap_zmq_replay_ops(sock); - pthread_mutex_unlock(&sock->ctx->lock); -} - -/* - * nvsnap_zmq_reinit_ctx: Full destroy/recreate of ZMQ context after CRIU restore - * - * ZMQ contexts contain kernel thread state that cannot be safely restored by CRIU. - * This function implements a clean-slate approach: - * 1. Create NEW context with IO threads enabled (required for bind/connect) - * 2. Create NEW sockets (old sockets have stale kernel state) - * 3. Replay ALL tracked operations (setsockopt, bind, connect) IMMEDIATELY - * 4. Clean up old context and sockets (optional, controlled by env vars) - * - * CRITICAL FIX (2026-02-01): - * Previously, binds were deferred until first I/O operation (allow_bind=0). - * This caused failures because: - * - With IO_THREADS=0: No threads available for bind, returned EAGAIN - * - With IO_THREADS=1: Threads not ready yet, timing-dependent failures - * - * Solution: Set IO_THREADS=1 by default and replay binds IMMEDIATELY during reinit. - * This ensures fresh threads are available and endpoints are established before - * any application I/O operations. - */ -static void nvsnap_zmq_reinit_ctx(nvsnap_zmq_ctx_t *ctx) { - if (!ctx || ctx->reinit_done || ctx->destroyed) { - return; - } - if (!nvsnap_zmq_restore_detected()) { - return; - } - void *old_ctx = ctx->real; - void *new_ctx = real_zmq_ctx_new(); - if (!new_ctx) { - NVSNAP_WARN("ZMQ reinit failed: ctx_new error"); - return; - } - int io_threads = nvsnap_zmq_io_threads_override(); - if (io_threads < 0) { - // Default to 1 IO thread for restore (required for bind/connect) - io_threads = 1; - } - if (real_zmq_ctx_set) { - if (real_zmq_ctx_set(new_ctx, ZMQ_IO_THREADS, io_threads) == 0) { - NVSNAP_INFO("ZMQ reinit ctx set io_threads=%d", io_threads); - } else { - NVSNAP_WARN("ZMQ reinit ctx set io_threads=%d failed errno=%d", - io_threads, nvsnap_zmq_errno()); - } - } - NVSNAP_INFO("ZMQ reinit ctx pid=%d ctx=%p old_real=%p new_real=%p force_term=%d", - getpid(), ctx, old_ctx, new_ctx, nvsnap_zmq_force_terminate_enabled()); - ctx->real = new_ctx; - nvsnap_zmq_socket_t *sock = ctx->sockets; - if (!sock) { - NVSNAP_WARN("ZMQ reinit ctx pid=%d ctx=%p has NO sockets", getpid(), ctx); - } - /* Collect old socket pointers so we can close them AFTER context shutdown. - * zmq_ctx_shutdown() sends ETERM to threads blocked on sockets in the old - * context, but only if those sockets haven't been zmq_close()'d yet. - * Previous bug: closing old sockets first removed them from the context's - * socket list, so shutdown found nothing to signal → threads stayed blocked. */ - void *old_socks[64]; - int old_sock_count = 0; - while (sock) { - if (!sock->closed) { - int connects = 0; - int binds = 0; - void *old_sock = sock->real; - unsigned char identity[256]; - size_t identity_len = nvsnap_zmq_capture_identity(old_sock, identity, sizeof(identity)); - nvsnap_zmq_count_ops(sock, &connects, &binds); - NVSNAP_INFO("ZMQ reinit socket pid=%d sock=%p real=%p ops_connect=%d ops_bind=%d", - getpid(), sock, sock->real, connects, binds); - sock->real = real_zmq_socket(ctx->real, sock->type); - if (sock->real) { - if (identity_len > 0) { - nvsnap_zmq_apply_identity(sock->real, identity, identity_len); - } - nvsnap_zmq_apply_restore_sockopts(sock); - /* Monitor start is deferred to after reinit completes — - * see nvsnap_zmq_reinit_all_if_restored() below. Starting - * monitors here would create ZMQ sockets and threads while - * ctx->lock is held, risking deadlock if SIGUSR2 interrupts - * the internal ZMQ calls or if callbacks re-enter the lock. */ - // CRITICAL FIX: Enable binds during reinit (allow_bind=1) - // Previously allowed_bind=0 deferred binds until first I/O, - // but by then IO threads may not be ready, causing EAGAIN errors. - // With fresh context + IO threads, we can bind immediately. - nvsnap_zmq_replay_ops_ex(sock, 1); - nvsnap_zmq_force_connect_from_last_endpoint(sock); - sock->replay_after_restore = 1; // Mark as replayed - } else { - NVSNAP_WARN("ZMQ reinit failed: socket create type=%d errno=%d", sock->type, nvsnap_zmq_errno()); - } - /* Save old socket for deferred close (don't close yet!) */ - if (old_sock && old_sock != sock->real && old_sock_count < 64) { - old_socks[old_sock_count++] = old_sock; - } - } - sock = sock->next; - } - /* STEP 1: Shutdown old context FIRST to wake threads blocked in zmq_poll/recv/send. - * zmq_ctx_shutdown() sends a "stop" command to each socket's owner thread via - * the socket's mailbox signaler (eventfd). This causes poll() to return on the - * app thread, and zmq_recv/zmq_msg_recv returns -1 with errno=ETERM (156). - * CRITICAL: Old sockets must still be in the context's socket list for shutdown - * to find them and send ETERM. That's why we close them AFTER shutdown. */ - if (old_ctx && old_ctx != ctx->real) { - nvsnap_zmq_load_old(); - if (old_zmq_ctx_shutdown) { - NVSNAP_INFO("ZMQ shutdown old ctx=%p to wake blocked threads (ETERM)", old_ctx); - old_zmq_ctx_shutdown(old_ctx); - } - } - /* STEP 2: Brief delay to let ETERM propagate to blocked threads. - * After shutdown, threads wake from poll(), process the stop command, - * and return ETERM. Our intercepted zmq_recv/zmq_msg_recv catches ETERM - * and retries on the new socket (sock->real already updated above). - * The gate mechanism (g_zmq_recovering) holds them until we finish. */ - usleep(10000); /* 10ms for ETERM delivery */ - /* STEP 3: Now close old sockets (threads have exited them by now). */ - if (!nvsnap_zmq_skip_old_close_enabled()) { - nvsnap_zmq_load_old(); - int (*close_fn)(void *) = old_zmq_close ? old_zmq_close : real_zmq_close; - for (int i = 0; i < old_sock_count; i++) { - if (close_fn && close_fn(old_socks[i]) != 0) { - NVSNAP_WARN("ZMQ reinit: closing old socket %p failed errno=%d", - old_socks[i], nvsnap_zmq_errno()); - } - } - } - ctx->reinit_done = 1; - NVSNAP_INFO("ZMQ reinit completed (shutdown-before-close, %d old sockets closed)", - old_sock_count); -} - -void nvsnap_zmq_reinit_all_if_restored(void) { - if (!nvsnap_zmq_is_enabled()) { - return; - } - if (!nvsnap_zmq_restore_detected()) { - return; - } - /* Context teardown deadlocks: zmq_ctx_destroy waits for IO threads, but - * IO threads are stuck in epoll_wait on stale fds. Instead, rely on the - * patched libzmq's epoll reinit (checks /run/criu-restored) + SIGUSR2 - * to wake IO threads from epoll_wait. */ - NVSNAP_INFO("ZMQ restore detected — relying on patched libzmq epoll reinit + SIGUSR2 wake"); - return; - nvsnap_zmq_log_first_call("reinit_all"); - nvsnap_zmq_set_recovering(1); - nvsnap_zmq_kill_bg_threads(); - int delay_ms = nvsnap_zmq_recovery_delay_ms(); - if (delay_ms > 0) { - NVSNAP_INFO("ZMQ recovery delay %d ms", delay_ms); - usleep((useconds_t)delay_ms * 1000); - } - nvsnap_zmq_switch_to_newns_if_needed(); - int ctx_count = 0; - int sock_count = 0; - nvsnap_zmq_count_state(&ctx_count, &sock_count); - NVSNAP_INFO("ZMQ reinit all (pid=%d) starting ctxs=%d socks=%d handle=%p newns_handle=%p", - getpid(), ctx_count, sock_count, g_zmq_handle, g_zmq_newns_handle); - pthread_once(&g_zmq_once, nvsnap_zmq_load_real); - pthread_mutex_lock(&g_zmq_ctx_list_lock); - int count = 0; - nvsnap_zmq_ctx_t *ctx = g_zmq_ctx_list; - while (ctx) { - pthread_mutex_lock(&ctx->lock); - nvsnap_zmq_reinit_ctx(ctx); - nvsnap_zmq_dump_ctx_sockets(ctx, "post_reinit"); - nvsnap_zmq_dump_ctx_sockets_to_file(ctx, "post_reinit"); - pthread_mutex_unlock(&ctx->lock); - count++; - ctx = ctx->next; - } - pthread_mutex_unlock(&g_zmq_ctx_list_lock); - - /* Start socket monitors AFTER all locks are released. Monitor startup - * creates ZMQ inproc sockets and threads — doing this under ctx->lock - * risks deadlock if SIGUSR2 interrupts the internal ZMQ calls or if - * libzmq callbacks re-enter our wrapper code. */ - ctx = g_zmq_ctx_list; - while (ctx) { - nvsnap_zmq_socket_t *sock = ctx->sockets; - while (sock) { - if (!sock->closed && sock->real && !sock->monitor_started) { - nvsnap_zmq_start_monitor(sock); - } - sock = sock->next; - } - ctx = ctx->next; - } - - NVSNAP_INFO("ZMQ reinit all (pid=%d) completed ctxs=%d", getpid(), count); - nvsnap_zmq_set_recovering(0); -} - -static void nvsnap_zmq_register_ctx(nvsnap_zmq_ctx_t *ctx) { - pthread_mutex_lock(&g_zmq_ctx_list_lock); - ctx->next = g_zmq_ctx_list; - g_zmq_ctx_list = ctx; - pthread_mutex_unlock(&g_zmq_ctx_list_lock); - NVSNAP_INFO("ZMQ ctx register pid=%d ctx=%p real=%p", getpid(), ctx, ctx->real); -} - -static void nvsnap_zmq_unregister_ctx(nvsnap_zmq_ctx_t *ctx) { - pthread_mutex_lock(&g_zmq_ctx_list_lock); - nvsnap_zmq_ctx_t **iter = &g_zmq_ctx_list; - while (*iter) { - if (*iter == ctx) { - *iter = ctx->next; - break; - } - iter = &(*iter)->next; - } - pthread_mutex_unlock(&g_zmq_ctx_list_lock); - NVSNAP_INFO("ZMQ ctx unregister pid=%d ctx=%p real=%p", getpid(), ctx, ctx ? ctx->real : NULL); -} - -static void nvsnap_zmq_maybe_reinit(nvsnap_zmq_ctx_t *ctx) { - if (!ctx) { - return; - } - if (!nvsnap_zmq_restore_detected()) { - return; - } - /* Context teardown disabled — libzmq's epoll reinit is sufficient. - * Tearing down contexts breaks async ZMQ (zmq.asyncio) used by SGLang. */ - return; -} - -void *zmq_ctx_new(void) { - pthread_once(&g_zmq_once, nvsnap_zmq_load_real); - if (!real_zmq_ctx_new) { - errno = ENOSYS; - return NULL; - } - if (!nvsnap_zmq_is_enabled()) { - return real_zmq_ctx_new(); - } - void *real = real_zmq_ctx_new(); - if (!real) { - return NULL; - } - nvsnap_zmq_ctx_t *ctx = calloc(1, sizeof(*ctx)); - if (!ctx) { - real_zmq_ctx_term(real); - return NULL; - } - ctx->magic = NVSNAP_ZMQ_CTX_MAGIC; - ctx->real = real; - pthread_mutex_init(&ctx->lock, NULL); - nvsnap_zmq_register_ctx(ctx); - return ctx; -} - -int zmq_ctx_term(void *ctxp) { - pthread_once(&g_zmq_once, nvsnap_zmq_load_real); - if (!real_zmq_ctx_term) { - errno = ENOSYS; - return -1; - } - if (!nvsnap_zmq_is_enabled()) { - return real_zmq_ctx_term(ctxp); - } - nvsnap_zmq_ctx_t *ctx = nvsnap_zmq_ctx_from_ptr(ctxp); - if (!ctx) { - return real_zmq_ctx_term(ctxp); - } - ctx->destroyed = 1; - nvsnap_zmq_unregister_ctx(ctx); - return real_zmq_ctx_term(ctx->real); -} - -int zmq_ctx_destroy(void *ctxp) { - pthread_once(&g_zmq_once, nvsnap_zmq_load_real); - if (!real_zmq_ctx_destroy && real_zmq_ctx_term) { - real_zmq_ctx_destroy = real_zmq_ctx_term; - } - if (!real_zmq_ctx_destroy) { - errno = ENOSYS; - return -1; - } - if (!nvsnap_zmq_is_enabled()) { - return real_zmq_ctx_destroy(ctxp); - } - nvsnap_zmq_ctx_t *ctx = nvsnap_zmq_ctx_from_ptr(ctxp); - if (!ctx) { - return real_zmq_ctx_destroy(ctxp); - } - ctx->destroyed = 1; - nvsnap_zmq_unregister_ctx(ctx); - return real_zmq_ctx_destroy(ctx->real); -} - -int zmq_ctx_shutdown(void *ctxp) { - pthread_once(&g_zmq_once, nvsnap_zmq_load_real); - if (!real_zmq_ctx_shutdown) { - errno = ENOSYS; - return -1; - } - if (!nvsnap_zmq_is_enabled()) { - return real_zmq_ctx_shutdown(ctxp); - } - nvsnap_zmq_ctx_t *ctx = nvsnap_zmq_ctx_from_ptr(ctxp); - if (!ctx) { - return real_zmq_ctx_shutdown(ctxp); - } - return real_zmq_ctx_shutdown(ctx->real); -} - -int zmq_ctx_set(void *ctxp, int option, int optval) { - pthread_once(&g_zmq_once, nvsnap_zmq_load_real); - if (!nvsnap_zmq_is_enabled()) { - return real_zmq_ctx_set(ctxp, option, optval); - } - nvsnap_zmq_ctx_t *ctx = nvsnap_zmq_ctx_from_ptr(ctxp); - if (!ctx) { - return real_zmq_ctx_set(ctxp, option, optval); - } - return real_zmq_ctx_set(ctx->real, option, optval); -} - -int zmq_ctx_get(void *ctxp, int option) { - pthread_once(&g_zmq_once, nvsnap_zmq_load_real); - if (!nvsnap_zmq_is_enabled()) { - return real_zmq_ctx_get(ctxp, option); - } - nvsnap_zmq_ctx_t *ctx = nvsnap_zmq_ctx_from_ptr(ctxp); - if (!ctx) { - return real_zmq_ctx_get(ctxp, option); - } - return real_zmq_ctx_get(ctx->real, option); -} - -void *zmq_socket(void *ctxp, int type) { - nvsnap_zmq_log_first_call("zmq_socket"); - pthread_once(&g_zmq_once, nvsnap_zmq_load_real); - if (!nvsnap_zmq_is_enabled()) { - return real_zmq_socket(ctxp, type); - } - nvsnap_zmq_ctx_t *ctx = nvsnap_zmq_ctx_from_ptr(ctxp); - if (!ctx) { - return real_zmq_socket(ctxp, type); - } - nvsnap_zmq_maybe_reinit(ctx); - void *real = real_zmq_socket(ctx->real, type); - if (!real) { - return NULL; - } - nvsnap_zmq_socket_t *sock = calloc(1, sizeof(*sock)); - if (!sock) { - real_zmq_close(real); - return NULL; - } - sock->magic = NVSNAP_ZMQ_SOCK_MAGIC; - sock->ctx = ctx; - sock->real = real; - sock->type = type; - nvsnap_zmq_apply_restore_sockopts(sock); - if (nvsnap_zmq_restore_detected()) { - nvsnap_zmq_start_monitor(sock); - } - pthread_mutex_lock(&ctx->lock); - sock->next = ctx->sockets; - ctx->sockets = sock; - pthread_mutex_unlock(&ctx->lock); - nvsnap_zmq_dump_ctx_sockets_to_file(ctx, "socket_create"); - if (nvsnap_zmq_trace_enabled()) { - NVSNAP_INFO("ZMQ socket pid=%d sock=%p real=%p type=%d", getpid(), sock, sock->real, type); - } - return sock; -} - -int zmq_close(void *sockp) { - pthread_once(&g_zmq_once, nvsnap_zmq_load_real); - if (!nvsnap_zmq_is_enabled()) { - return real_zmq_close(sockp); - } - nvsnap_zmq_socket_t *sock = nvsnap_zmq_sock_from_ptr(sockp); - if (!sock) { - return real_zmq_close(sockp); - } - sock->closed = 1; - return real_zmq_close(sock->real); -} - -int zmq_bind(void *sockp, const char *endpoint) { - nvsnap_zmq_log_first_call("zmq_bind"); - pthread_once(&g_zmq_once, nvsnap_zmq_load_real); - if (!nvsnap_zmq_is_enabled()) { - return real_zmq_bind(sockp, endpoint); - } - nvsnap_zmq_socket_t *sock = nvsnap_zmq_sock_from_ptr(sockp); - if (!sock) { - return real_zmq_bind(sockp, endpoint); - } - nvsnap_zmq_maybe_reinit(sock->ctx); - nvsnap_zmq_add_op(sock, NVSNAP_ZMQ_OP_BIND, 0, NULL, 0, endpoint); - int rc = real_zmq_bind(sock->real, endpoint); - if (nvsnap_zmq_trace_enabled()) { - NVSNAP_INFO("ZMQ bind pid=%d sock=%p real=%p endpoint=%s rc=%d errno=%d", - getpid(), sockp, sock->real, endpoint ? endpoint : "(null)", rc, nvsnap_zmq_errno()); - } - return rc; -} - -int zmq_unbind(void *sockp, const char *endpoint) { - pthread_once(&g_zmq_once, nvsnap_zmq_load_real); - if (!nvsnap_zmq_is_enabled()) { - return real_zmq_unbind(sockp, endpoint); - } - nvsnap_zmq_socket_t *sock = nvsnap_zmq_sock_from_ptr(sockp); - if (!sock) { - return real_zmq_unbind(sockp, endpoint); - } - nvsnap_zmq_add_op(sock, NVSNAP_ZMQ_OP_UNBIND, 0, NULL, 0, endpoint); - return real_zmq_unbind(sock->real, endpoint); -} - -int zmq_connect(void *sockp, const char *endpoint) { - nvsnap_zmq_log_first_call("zmq_connect"); - pthread_once(&g_zmq_once, nvsnap_zmq_load_real); - if (!nvsnap_zmq_is_enabled()) { - return real_zmq_connect(sockp, endpoint); - } - nvsnap_zmq_socket_t *sock = nvsnap_zmq_sock_from_ptr(sockp); - if (!sock) { - return real_zmq_connect(sockp, endpoint); - } - nvsnap_zmq_maybe_reinit(sock->ctx); - nvsnap_zmq_add_op(sock, NVSNAP_ZMQ_OP_CONNECT, 0, NULL, 0, endpoint); - int rc = real_zmq_connect(sock->real, endpoint); - if (nvsnap_zmq_trace_enabled() || nvsnap_zmq_restore_detected()) { - NVSNAP_INFO("ZMQ connect pid=%d sock=%p real=%p endpoint=%s rc=%d errno=%d", - getpid(), sockp, sock->real, endpoint ? endpoint : "(null)", rc, nvsnap_zmq_errno()); - } - return rc; -} - -int zmq_disconnect(void *sockp, const char *endpoint) { - pthread_once(&g_zmq_once, nvsnap_zmq_load_real); - if (!nvsnap_zmq_is_enabled()) { - return real_zmq_disconnect(sockp, endpoint); - } - nvsnap_zmq_socket_t *sock = nvsnap_zmq_sock_from_ptr(sockp); - if (!sock) { - return real_zmq_disconnect(sockp, endpoint); - } - nvsnap_zmq_add_op(sock, NVSNAP_ZMQ_OP_DISCONNECT, 0, NULL, 0, endpoint); - return real_zmq_disconnect(sock->real, endpoint); -} - -int zmq_setsockopt(void *sockp, int option, const void *optval, size_t optvallen) { - pthread_once(&g_zmq_once, nvsnap_zmq_load_real); - if (!nvsnap_zmq_is_enabled()) { - return real_zmq_setsockopt(sockp, option, optval, optvallen); - } - nvsnap_zmq_socket_t *sock = nvsnap_zmq_sock_from_ptr(sockp); - if (!sock) { - return real_zmq_setsockopt(sockp, option, optval, optvallen); - } - nvsnap_zmq_add_op(sock, NVSNAP_ZMQ_OP_SETSOCKOPT, option, optval, optvallen, NULL); - return real_zmq_setsockopt(sock->real, option, optval, optvallen); -} - -int zmq_getsockopt(void *sockp, int option, void *optval, size_t *optvallen) { - pthread_once(&g_zmq_once, nvsnap_zmq_load_real); - if (!nvsnap_zmq_is_enabled()) { - return real_zmq_getsockopt(sockp, option, optval, optvallen); - } - nvsnap_zmq_socket_t *sock = nvsnap_zmq_sock_from_ptr(sockp); - if (!sock) { - return real_zmq_getsockopt(sockp, option, optval, optvallen); - } - return real_zmq_getsockopt(sock->real, option, optval, optvallen); -} - -static void nvsnap_zmq_log_sock_state(const char *tag, nvsnap_zmq_socket_t *sock) { - if (!sock || !sock->real) { - return; - } - int type = 0; - int events = 0; - size_t type_len = sizeof(type); - size_t events_len = sizeof(events); - char endpoint[256]; - endpoint[0] = '\0'; - size_t endpoint_len = sizeof(endpoint); - if (real_zmq_getsockopt) { - (void)real_zmq_getsockopt(sock->real, ZMQ_TYPE, &type, &type_len); - (void)real_zmq_getsockopt(sock->real, ZMQ_EVENTS, &events, &events_len); - if (real_zmq_getsockopt(sock->real, ZMQ_LAST_ENDPOINT, endpoint, &endpoint_len) == 0) { - endpoint[sizeof(endpoint) - 1] = '\0'; - } else { - endpoint[0] = '\0'; - } - } - NVSNAP_INFO("ZMQ state %s pid=%d sock=%p real=%p type=%d events=0x%x endpoint=%s", - tag ? tag : "unknown", getpid(), sock, sock->real, type, events, - endpoint[0] ? endpoint : "(none)"); -} - -int zmq_send(void *sockp, const void *buf, size_t len, int flags) { - /* ALWAYS LOG: Verify library is loaded */ - NVSNAP_TRACE("zmq_send entry pid=%d sock=%p len=%zu flags=0x%x", getpid(), sockp, len, flags); - - pthread_once(&g_zmq_once, nvsnap_zmq_load_real); - if (!nvsnap_zmq_is_enabled()) { - NVSNAP_TRACE("zmq_send disabled, passthrough pid=%d", getpid()); - return real_zmq_send(sockp, buf, len, flags); - } - if (nvsnap_zmq_trace_enabled()) { - NVSNAP_INFO("ZMQ send entry pid=%d sock=%p len=%zu flags=0x%x", - getpid(), sockp, len, flags); - } - nvsnap_zmq_socket_t *sock = nvsnap_zmq_sock_from_ptr(sockp); - if (!sock) { - NVSNAP_WARN("ZMQ send untracked sock pid=%d sock=%p len=%zu flags=0x%x", - getpid(), sockp, len, flags); - return real_zmq_send(sockp, buf, len, flags); - } - nvsnap_zmq_maybe_reinit(sock->ctx); - nvsnap_zmq_maybe_replay_after_restore(sock); - nvsnap_zmq_gate_if_recovering("send"); - int rc = real_zmq_send(sock->real, buf, len, flags); - int send_err = nvsnap_zmq_errno(); - /* After restore, SIGUSR2 causes EINTR or old context shutdown causes ETERM. - * Re-resolve sock->real (now points to new socket) and retry once. */ - if (rc < 0 && (send_err == EINTR || send_err == 156 /* ETERM */) - && nvsnap_zmq_restore_detected()) { - NVSNAP_INFO("ZMQ send got %s pid=%d sock=%p - switching to new socket", - send_err == EINTR ? "EINTR" : "ETERM", getpid(), sockp); - nvsnap_zmq_maybe_reinit(sock->ctx); - nvsnap_zmq_maybe_replay_after_restore(sock); - nvsnap_zmq_gate_if_recovering("send_restore_retry"); - rc = real_zmq_send(sock->real, buf, len, flags); - } - if (rc < 0 && nvsnap_zmq_errno() == EHOSTUNREACH && nvsnap_zmq_restore_detected()) { - NVSNAP_WARN("ZMQ send EHOSTUNREACH pid=%d sock=%p real=%p", getpid(), sock, sock->real); - nvsnap_zmq_force_reconnect(sock); - int max_ms = nvsnap_zmq_ehostunreach_max_ms(); - int sleep_ms = nvsnap_zmq_ehostunreach_sleep_ms(); - int waited = 0; - while (max_ms > 0 && waited < max_ms) { - usleep((useconds_t)sleep_ms * 1000); - waited += sleep_ms; - rc = real_zmq_send(sock->real, buf, len, flags); - if (rc >= 0 || nvsnap_zmq_errno() != EHOSTUNREACH) { - NVSNAP_WARN("ZMQ send retry rc=%d errno=%d waited_ms=%d", rc, nvsnap_zmq_errno(), waited); - break; - } - } - if (rc < 0 && nvsnap_zmq_errno() == EHOSTUNREACH) { - if (nvsnap_zmq_rebuild_socket(sock, "send_ehostunreach") == 0) { - nvsnap_zmq_force_reconnect(sock); - rc = real_zmq_send(sock->real, buf, len, flags); - NVSNAP_WARN("ZMQ send after rebuild rc=%d errno=%d", rc, nvsnap_zmq_errno()); - } - } - } - if (rc < 0 && nvsnap_zmq_errno() == EAGAIN) { - nvsnap_zmq_log_sock_state("send_eagain", sock); - nvsnap_zmq_maybe_replay_after_restore(sock); - rc = real_zmq_send(sock->real, buf, len, flags); - } - if (nvsnap_zmq_trace_enabled()) { - NVSNAP_INFO("ZMQ send sock=%p real=%p len=%zu flags=0x%x rc=%d errno=%d", - sockp, sock->real, len, flags, rc, nvsnap_zmq_errno()); - } - return rc; -} - -int zmq_recv(void *sockp, void *buf, size_t len, int flags) { - /* ALWAYS LOG: Verify library is loaded */ - NVSNAP_TRACE("zmq_recv entry pid=%d sock=%p len=%zu flags=0x%x", getpid(), sockp, len, flags); - - pthread_once(&g_zmq_once, nvsnap_zmq_load_real); - if (!nvsnap_zmq_is_enabled()) { - NVSNAP_TRACE("zmq_recv disabled, passthrough pid=%d", getpid()); - return real_zmq_recv(sockp, buf, len, flags); - } - if (nvsnap_zmq_trace_enabled()) { - NVSNAP_INFO("ZMQ recv entry pid=%d sock=%p len=%zu flags=0x%x", - getpid(), sockp, len, flags); - } - nvsnap_zmq_socket_t *sock = nvsnap_zmq_sock_from_ptr(sockp); - if (!sock) { - NVSNAP_WARN("ZMQ recv untracked sock pid=%d sock=%p len=%zu flags=0x%x", - getpid(), sockp, len, flags); - return real_zmq_recv(sockp, buf, len, flags); - } - nvsnap_zmq_maybe_reinit(sock->ctx); - nvsnap_zmq_maybe_replay_after_restore(sock); - nvsnap_zmq_gate_if_recovering("recv"); - int rc = real_zmq_recv(sock->real, buf, len, flags); - int recv_err = nvsnap_zmq_errno(); - /* After restore, SIGUSR2 interrupts blocked recv with EINTR, or old - * context shutdown delivers ETERM. Either way, re-resolve sock->real - * (now points to new socket after reinit) and retry on the new socket. */ - if (rc < 0 && (recv_err == EINTR || recv_err == 156 /* ETERM */) - && nvsnap_zmq_restore_detected()) { - NVSNAP_INFO("ZMQ recv got %s pid=%d sock=%p - switching to new socket", - recv_err == EINTR ? "EINTR" : "ETERM", getpid(), sockp); - nvsnap_zmq_maybe_reinit(sock->ctx); - nvsnap_zmq_maybe_replay_after_restore(sock); - nvsnap_zmq_gate_if_recovering("recv_restore_retry"); - rc = real_zmq_recv(sock->real, buf, len, flags); - recv_err = nvsnap_zmq_errno(); - } - if (rc < 0 && recv_err == EAGAIN) { - nvsnap_zmq_log_sock_state("recv_eagain", sock); - nvsnap_zmq_maybe_replay_after_restore(sock); - rc = real_zmq_recv(sock->real, buf, len, flags); - } - if (nvsnap_zmq_trace_enabled()) { - NVSNAP_INFO("ZMQ recv sock=%p real=%p len=%zu flags=0x%x rc=%d errno=%d", - sockp, sock->real, len, flags, rc, nvsnap_zmq_errno()); - } - return rc; -} - -int zmq_msg_send(zmq_msg_t *msg, void *sockp, int flags) { - /* ALWAYS LOG: Verify library is loaded */ - NVSNAP_TRACE("zmq_msg_send entry pid=%d sock=%p flags=0x%x", getpid(), sockp, flags); - - pthread_once(&g_zmq_once, nvsnap_zmq_load_real); - if (!nvsnap_zmq_is_enabled()) { - NVSNAP_TRACE("zmq_msg_send disabled, passthrough pid=%d", getpid()); - return real_zmq_msg_send(msg, sockp, flags); - } - if (nvsnap_zmq_trace_enabled()) { - NVSNAP_INFO("ZMQ msg_send entry pid=%d sock=%p flags=0x%x", - getpid(), sockp, flags); - } - nvsnap_zmq_socket_t *sock = nvsnap_zmq_sock_from_ptr(sockp); - if (!sock) { - NVSNAP_WARN("ZMQ msg_send untracked sock pid=%d sock=%p flags=0x%x", - getpid(), sockp, flags); - return real_zmq_msg_send(msg, sockp, flags); - } - if (nvsnap_zmq_restore_detected() && nvsnap_zmq_rebuild_on_first_send_enabled() && - sock->rebuild_count == 0) { - (void)nvsnap_zmq_rebuild_socket(sock, "first_send_after_restore"); - } - nvsnap_zmq_maybe_reinit(sock->ctx); - nvsnap_zmq_maybe_replay_after_restore(sock); - nvsnap_zmq_gate_if_recovering("msg_send"); - int rc = real_zmq_msg_send(msg, sock->real, flags); - int msg_send_err = nvsnap_zmq_errno(); - /* After restore, SIGUSR2 causes EINTR or old context shutdown causes ETERM. - * Re-resolve sock->real (now points to new socket) and retry once. */ - if (rc < 0 && (msg_send_err == EINTR || msg_send_err == 156 /* ETERM */) - && nvsnap_zmq_restore_detected()) { - NVSNAP_INFO("ZMQ msg_send got %s pid=%d sock=%p - switching to new socket", - msg_send_err == EINTR ? "EINTR" : "ETERM", getpid(), sockp); - nvsnap_zmq_maybe_reinit(sock->ctx); - nvsnap_zmq_maybe_replay_after_restore(sock); - nvsnap_zmq_gate_if_recovering("msg_send_restore_retry"); - rc = real_zmq_msg_send(msg, sock->real, flags); - } - - /* DETAILED EINVAL LOGGING */ - if (rc < 0 && nvsnap_zmq_errno() == 22) { /* EINVAL */ - NVSNAP_WARN("!!! zmq_msg_send EINVAL pid=%d sock=%p real=%p rc=%d", - getpid(), sockp, sock->real, rc); - nvsnap_zmq_log_sock_state("msg_send_EINVAL", sock); - - /* Check socket options */ - int type = 0, events = 0; - size_t sz = sizeof(type); - if (real_zmq_getsockopt && real_zmq_getsockopt(sock->real, ZMQ_TYPE, &type, &sz) == 0) { - NVSNAP_WARN("!!! Socket type=%d", type); - } - sz = sizeof(events); - if (real_zmq_getsockopt && real_zmq_getsockopt(sock->real, ZMQ_EVENTS, &events, &sz) == 0) { - NVSNAP_WARN("!!! Socket events=0x%x", events); - } - - /* Try to get more error context */ - NVSNAP_WARN("!!! Attempting replay_after_restore due to EINVAL"); - nvsnap_zmq_maybe_replay_after_restore(sock); - rc = real_zmq_msg_send(msg, sock->real, flags); - NVSNAP_WARN("!!! After replay: rc=%d errno=%d", rc, nvsnap_zmq_errno()); - - if (rc < 0 && nvsnap_zmq_restore_detected() && nvsnap_zmq_rebuild_on_einval_enabled()) { - if (nvsnap_zmq_rebuild_socket(sock, "msg_send_einval") == 0) { - rc = real_zmq_msg_send(msg, sock->real, flags); - NVSNAP_WARN("!!! After rebuild: rc=%d errno=%d", rc, nvsnap_zmq_errno()); - } - } - } else if (rc < 0 && nvsnap_zmq_errno() == EHOSTUNREACH && nvsnap_zmq_restore_detected()) { - NVSNAP_WARN("!!! zmq_msg_send EHOSTUNREACH pid=%d sock=%p real=%p", - getpid(), sockp, sock->real); - nvsnap_zmq_force_reconnect(sock); - int max_ms = nvsnap_zmq_ehostunreach_max_ms(); - int sleep_ms = nvsnap_zmq_ehostunreach_sleep_ms(); - int waited = 0; - while (max_ms > 0 && waited < max_ms) { - usleep((useconds_t)sleep_ms * 1000); - waited += sleep_ms; - rc = real_zmq_msg_send(msg, sock->real, flags); - if (rc >= 0 || nvsnap_zmq_errno() != EHOSTUNREACH) { - NVSNAP_WARN("!!! msg_send retry rc=%d errno=%d waited_ms=%d", - rc, nvsnap_zmq_errno(), waited); - break; - } - } - if (rc < 0 && nvsnap_zmq_errno() == EHOSTUNREACH) { - if (nvsnap_zmq_rebuild_socket(sock, "msg_send_ehostunreach") == 0) { - nvsnap_zmq_force_reconnect(sock); - rc = real_zmq_msg_send(msg, sock->real, flags); - NVSNAP_WARN("!!! After rebuild: rc=%d errno=%d", rc, nvsnap_zmq_errno()); - } - } - } else if (rc < 0 && nvsnap_zmq_errno() == EAGAIN) { - nvsnap_zmq_log_sock_state("msg_send_eagain", sock); - nvsnap_zmq_maybe_replay_after_restore(sock); - rc = real_zmq_msg_send(msg, sock->real, flags); - } - if (rc >= 0 && nvsnap_zmq_errno() == EAGAIN) { - nvsnap_zmq_log_sock_state("msg_send_errno_eagain", sock); - } - if (nvsnap_zmq_trace_enabled()) { - NVSNAP_INFO("ZMQ msg_send pid=%d sock=%p real=%p flags=0x%x rc=%d errno=%d", - getpid(), sockp, sock->real, flags, rc, nvsnap_zmq_errno()); - } - return rc; -} - -int zmq_msg_recv(zmq_msg_t *msg, void *sockp, int flags) { - /* ALWAYS LOG: Verify library is loaded */ - NVSNAP_TRACE("zmq_msg_recv entry pid=%d sock=%p flags=0x%x", getpid(), sockp, flags); - - pthread_once(&g_zmq_once, nvsnap_zmq_load_real); - if (!nvsnap_zmq_is_enabled()) { - NVSNAP_TRACE("zmq_msg_recv disabled, passthrough pid=%d", getpid()); - return real_zmq_msg_recv(msg, sockp, flags); - } - if (nvsnap_zmq_trace_enabled()) { - NVSNAP_INFO("ZMQ msg_recv entry pid=%d sock=%p flags=0x%x", - getpid(), sockp, flags); - } - nvsnap_zmq_socket_t *sock = nvsnap_zmq_sock_from_ptr(sockp); - if (!sock) { - NVSNAP_WARN("ZMQ msg_recv untracked sock pid=%d sock=%p flags=0x%x", - getpid(), sockp, flags); - return real_zmq_msg_recv(msg, sockp, flags); - } - nvsnap_zmq_maybe_reinit(sock->ctx); - nvsnap_zmq_maybe_replay_after_restore(sock); - nvsnap_zmq_gate_if_recovering("msg_recv"); - int rc = real_zmq_msg_recv(msg, sock->real, flags); - int msg_recv_err = nvsnap_zmq_errno(); - /* After restore, SIGUSR2 interrupts blocked msg_recv with EINTR, or old - * context shutdown delivers ETERM. Either way, re-resolve sock->real - * (now points to new socket after reinit) and retry on the new socket. */ - if (rc < 0 && (msg_recv_err == EINTR || msg_recv_err == 156 /* ETERM */) - && nvsnap_zmq_restore_detected()) { - NVSNAP_INFO("ZMQ msg_recv got %s pid=%d sock=%p - switching to new socket", - msg_recv_err == EINTR ? "EINTR" : "ETERM", getpid(), sockp); - nvsnap_zmq_maybe_reinit(sock->ctx); - nvsnap_zmq_maybe_replay_after_restore(sock); - nvsnap_zmq_gate_if_recovering("msg_recv_restore_retry"); - rc = real_zmq_msg_recv(msg, sock->real, flags); - msg_recv_err = nvsnap_zmq_errno(); - } - if (rc < 0 && msg_recv_err == EAGAIN) { - nvsnap_zmq_log_sock_state("msg_recv_eagain", sock); - nvsnap_zmq_maybe_replay_after_restore(sock); - rc = real_zmq_msg_recv(msg, sock->real, flags); - } - if (nvsnap_zmq_trace_enabled()) { - NVSNAP_INFO("ZMQ msg_recv pid=%d sock=%p real=%p flags=0x%x rc=%d errno=%d", - getpid(), sockp, sock->real, flags, rc, nvsnap_zmq_errno()); - } - return rc; -} - -int zmq_msg_init(zmq_msg_t *msg) { - pthread_once(&g_zmq_once, nvsnap_zmq_load_real); - return real_zmq_msg_init(msg); -} - -int zmq_msg_init_size(zmq_msg_t *msg, size_t size) { - pthread_once(&g_zmq_once, nvsnap_zmq_load_real); - return real_zmq_msg_init_size(msg, size); -} - -int zmq_msg_init_data(zmq_msg_t *msg, void *data, size_t size, void (*ffn)(void *, void *), void *hint) { - pthread_once(&g_zmq_once, nvsnap_zmq_load_real); - return real_zmq_msg_init_data(msg, data, size, ffn, hint); -} - -int zmq_msg_close(zmq_msg_t *msg) { - pthread_once(&g_zmq_once, nvsnap_zmq_load_real); - return real_zmq_msg_close(msg); -} - -static pthread_once_t g_zmq_poll_first_call = PTHREAD_ONCE_INIT; -static void nvsnap_zmq_poll_first_call_reinit(void) { - int restore_detected = nvsnap_zmq_restore_detected(); - NVSNAP_WARN("ZMQ poll first-call check pid=%d restore_detected=%d", getpid(), restore_detected); - if (restore_detected) { - NVSNAP_WARN("ZMQ poll on first call post-restore pid=%d - forcing global reinit", getpid()); - nvsnap_zmq_reinit_all_if_restored(); - } -} - -int zmq_poll(zmq_pollitem_t *items, int nitems, long timeout) { - NVSNAP_TRACE("zmq_poll entry pid=%d nitems=%d", getpid(), nitems); - - pthread_once(&g_zmq_once, nvsnap_zmq_load_real); - - if (!nvsnap_zmq_is_enabled()) { - NVSNAP_TRACE("zmq_poll disabled, passthrough pid=%d", getpid()); - return real_zmq_poll(items, nitems, timeout); - } - - NVSNAP_TRACE("zmq_poll enabled, first-call check pid=%d", getpid()); - /* CRITICAL: Check for restore on first poll call to catch early polls before library init */ - pthread_once(&g_zmq_poll_first_call, nvsnap_zmq_poll_first_call_reinit); - - nvsnap_zmq_socket_t **wrapped = NULL; - const int log_poll = (nvsnap_zmq_trace_enabled() || nvsnap_zmq_restore_detected()); - if (log_poll && nitems > 0) { - wrapped = calloc((size_t)nitems, sizeof(*wrapped)); - NVSNAP_INFO("ZMQ poll entry pid=%d nitems=%d timeout=%ld restore_detected=%d", - getpid(), nitems, timeout, nvsnap_zmq_restore_detected()); - } - int wrapped_count = 0; - int unwrapped_count = 0; - for (int i = 0; i < nitems; i++) { - nvsnap_zmq_socket_t *sock = nvsnap_zmq_sock_from_ptr(items[i].socket); - if (sock) { - wrapped_count++; - nvsnap_zmq_maybe_reinit(sock->ctx); - nvsnap_zmq_maybe_replay_after_restore(sock); - if (wrapped) { - wrapped[i] = sock; - } - items[i].socket = sock->real; - } else { - unwrapped_count++; - if (log_poll) { - NVSNAP_WARN("ZMQ poll pid=%d item[%d] socket=%p NOT WRAPPED (magic check failed)", - getpid(), i, items[i].socket); - } - } - } - if (log_poll && unwrapped_count > 0) { - NVSNAP_WARN("ZMQ poll pid=%d has %d unwrapped sockets out of %d total - triggering global reinit as fallback", - getpid(), unwrapped_count, nitems); - /* Fallback: if we have unwrapped sockets and we're post-restore, force global reinit */ - if (nvsnap_zmq_restore_detected()) { - nvsnap_zmq_reinit_all_if_restored(); - } - } - nvsnap_zmq_gate_if_recovering("poll"); - int rc = real_zmq_poll(items, nitems, timeout); - int err = nvsnap_zmq_errno(); - /* After restore, old context shutdown causes ETERM (156) on blocked poll. - * Re-resolve socket pointers from wrappers (now point to new sockets) and retry. */ - if (rc < 0 && nvsnap_zmq_restore_detected() && - (err == EINTR || err == EAGAIN || err == 156 /* ETERM */)) { - if (err == 156) { - NVSNAP_INFO("ZMQ poll got ETERM pid=%d - old ctx shutdown, re-resolving to new sockets", - getpid()); - /* Re-resolve socket pointers: wrappers now point to new (reinited) sockets */ - for (int i = 0; i < nitems; i++) { - nvsnap_zmq_socket_t *sock = wrapped ? wrapped[i] : nvsnap_zmq_sock_from_ptr(items[i].socket); - if (sock) { - nvsnap_zmq_maybe_reinit(sock->ctx); - nvsnap_zmq_maybe_replay_after_restore(sock); - items[i].socket = sock->real; /* Now points to NEW socket */ - if (wrapped) wrapped[i] = sock; - } - } - nvsnap_zmq_gate_if_recovering("poll_eterm_retry"); - } - for (int retry = 0; retry < 3 && rc < 0 && - (err == EINTR || err == EAGAIN || err == 156); retry++) { - if (log_poll) { - NVSNAP_INFO("ZMQ poll retry pid=%d attempt=%d errno=%d", - getpid(), retry + 1, err); - } - rc = real_zmq_poll(items, nitems, timeout); - err = nvsnap_zmq_errno(); - } - } - if (log_poll) { - NVSNAP_INFO("ZMQ poll pid=%d nitems=%d timeout=%ld rc=%d errno=%d", - getpid(), nitems, timeout, rc, err); - } - if (wrapped && rc <= 0) { - for (int i = 0; i < nitems; i++) { - if (!wrapped[i]) { - continue; - } - NVSNAP_INFO("ZMQ poll item pid=%d idx=%d events=0x%x revents=0x%x", - getpid(), i, items[i].events, items[i].revents); - nvsnap_zmq_log_sock_state("poll_rc_le_0", wrapped[i]); - } - } - if (wrapped) { - free(wrapped); - } - return rc; -} - -int zmq_proxy(void *frontend, void *backend, void *capture) { - pthread_once(&g_zmq_once, nvsnap_zmq_load_real); - if (!nvsnap_zmq_is_enabled()) { - return real_zmq_proxy(frontend, backend, capture); - } - nvsnap_zmq_socket_t *front = nvsnap_zmq_sock_from_ptr(frontend); - nvsnap_zmq_socket_t *back = nvsnap_zmq_sock_from_ptr(backend); - nvsnap_zmq_socket_t *cap = nvsnap_zmq_sock_from_ptr(capture); - return real_zmq_proxy(front ? front->real : frontend, - back ? back->real : backend, - cap ? cap->real : capture); -} - -/* - * ============================================================================= - * ZMQ CHECKPOINT/RESTORE SUPPORT - * ============================================================================= - */ - -/* New ZMQ checkpoint API function pointers */ -typedef int (*zmq_get_all_contexts_fn_t)(void ***contexts, int *count); -typedef int (*zmq_ctx_checkpoint_fn_t)(void *context, void **checkpoint, int flags); -typedef int (*zmq_ctx_restore_fn_t)(void *checkpoint, void **context, int flags); -typedef int (*zmq_checkpoint_destroy_fn_t)(void *checkpoint); - -static zmq_get_all_contexts_fn_t real_zmq_get_all_contexts = NULL; -static zmq_ctx_checkpoint_fn_t real_zmq_ctx_checkpoint = NULL; -static zmq_ctx_restore_fn_t real_zmq_ctx_restore = NULL; -static zmq_checkpoint_destroy_fn_t real_zmq_checkpoint_destroy = NULL; - -static pthread_once_t g_zmq_ckpt_once = PTHREAD_ONCE_INIT; -static void **g_saved_checkpoints = NULL; -static int g_num_checkpoints = 0; - -#define ZMQ_CKPT_FILE "/var/run/nvsnap/zmq-ckpt.dat" - -/** - * Load ZMQ checkpoint API symbols - */ -static void nvsnap_zmq_load_checkpoint_api(void) { - if (!g_zmq_handle) { - return; /* libzmq not loaded */ - } - - real_zmq_get_all_contexts = dlsym(g_zmq_handle, "zmq_get_all_contexts"); - real_zmq_ctx_checkpoint = dlsym(g_zmq_handle, "zmq_ctx_checkpoint"); - real_zmq_ctx_restore = dlsym(g_zmq_handle, "zmq_ctx_restore"); - real_zmq_checkpoint_destroy = dlsym(g_zmq_handle, "zmq_checkpoint_destroy"); - - if (real_zmq_ctx_checkpoint && real_zmq_ctx_restore) { - NVSNAP_INFO("ZMQ checkpoint API available"); - } else { - NVSNAP_WARN("ZMQ checkpoint API not available - using standard libzmq"); - } -} - -/** - * Checkpoint all ZMQ contexts (called on SIGUSR1) - */ -void nvsnap_zmq_checkpoint(void) { - pthread_once(&g_zmq_ckpt_once, nvsnap_zmq_load_checkpoint_api); - - if (!real_zmq_get_all_contexts || !real_zmq_ctx_checkpoint) { - NVSNAP_WARN("ZMQ checkpoint API not available"); - return; - } - - void **contexts = NULL; - int count = 0; - - /* Get all ZMQ contexts */ - int rc = real_zmq_get_all_contexts(&contexts, &count); - if (rc != 0) { - NVSNAP_ERROR("zmq_get_all_contexts failed: %d", rc); - return; - } - - NVSNAP_INFO("ZMQ checkpoint: found %d contexts", count); - - if (count == 0) { - free(contexts); - return; - } - - /* Allocate checkpoint array */ - g_saved_checkpoints = malloc(sizeof(void*) * count); - if (!g_saved_checkpoints) { - NVSNAP_ERROR("Out of memory for checkpoints"); - free(contexts); - return; - } - - /* Checkpoint each context */ - for (int i = 0; i < count; i++) { - void *checkpoint = NULL; - rc = real_zmq_ctx_checkpoint(contexts[i], &checkpoint, 0); - if (rc != 0) { - NVSNAP_ERROR("zmq_ctx_checkpoint failed for context %d: %d", i, rc); - continue; - } - - g_saved_checkpoints[i] = checkpoint; - g_num_checkpoints++; - - NVSNAP_INFO("ZMQ checkpoint: saved context %d/%d (ctx=%p ckpt=%p)", - i + 1, count, contexts[i], checkpoint); - } - - free(contexts); - - /* Save checkpoints to file for CRIU plugin */ - FILE *fp = fopen(ZMQ_CKPT_FILE, "w"); - if (fp) { - fprintf(fp, "%d\n", g_num_checkpoints); - for (int i = 0; i < g_num_checkpoints; i++) { - fprintf(fp, "%p\n", g_saved_checkpoints[i]); - } - fclose(fp); - NVSNAP_INFO("ZMQ checkpoint: saved %d checkpoints to %s", - g_num_checkpoints, ZMQ_CKPT_FILE); - } else { - NVSNAP_ERROR("Failed to save checkpoints to %s", ZMQ_CKPT_FILE); - } -} - -/** - * Restore ZMQ contexts (called on restore detection) - */ -void nvsnap_zmq_restore(void) { - pthread_once(&g_zmq_ckpt_once, nvsnap_zmq_load_checkpoint_api); - - if (!real_zmq_ctx_restore) { - NVSNAP_WARN("ZMQ restore API not available"); - return; - } - - /* Load checkpoints from file */ - FILE *fp = fopen(ZMQ_CKPT_FILE, "r"); - if (!fp) { - NVSNAP_DEBUG("No ZMQ checkpoint file found (normal for non-ZMQ apps)"); - return; - } - - int count = 0; - if (fscanf(fp, "%d\n", &count) != 1) { - NVSNAP_ERROR("Failed to read checkpoint count"); - fclose(fp); - return; - } - - NVSNAP_INFO("ZMQ restore: restoring %d contexts", count); - - for (int i = 0; i < count; i++) { - void *checkpoint = NULL; - if (fscanf(fp, "%p\n", &checkpoint) != 1) { - NVSNAP_ERROR("Failed to read checkpoint %d", i); - continue; - } - - void *new_ctx = NULL; - int rc = real_zmq_ctx_restore(checkpoint, &new_ctx, 0); - if (rc != 0) { - NVSNAP_ERROR("zmq_ctx_restore failed for context %d: %d", i, rc); - continue; - } - - NVSNAP_INFO("ZMQ restore: restored context %d/%d (ckpt=%p new_ctx=%p)", - i + 1, count, checkpoint, new_ctx); - - /* TODO: Update application's context pointers - * This requires tracking original context addresses - * For now, just create the contexts - app will use new ones - */ - } - - fclose(fp); - NVSNAP_INFO("ZMQ restore: completed restoring %d contexts", count); -} - -/** - * Public API for quiesce.c to call during checkpoint - */ -void nvsnap_zmq_handle_checkpoint(void) { - nvsnap_zmq_checkpoint(); -} - -/** - * Public API for init.c to call during restore detection - */ -void nvsnap_zmq_handle_restore(void) { - nvsnap_zmq_restore(); -} diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/run_cuda_tests.sh b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/run_cuda_tests.sh deleted file mode 100755 index 7a1b9a01e9..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/run_cuda_tests.sh +++ /dev/null @@ -1,54 +0,0 @@ -#!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# -# Build and run the CUDA interception test suite. -# Requires: GPU, CUDA runtime/driver available. -# -# Usage: -# ./tests/run_cuda_tests.sh # from lib/nvsnap_intercept/ -# make test-cuda # via Makefile target -# -set -euo pipefail - -SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" -ROOT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)" -TEST_BIN="$SCRIPT_DIR/test_cuda_intercept" -LIB="$ROOT_DIR/libnvsnap_intercept.so" - -echo "=== Building libnvsnap_intercept.so ===" -make -C "$ROOT_DIR" -j$(nproc) - -echo "" -echo "=== Compiling test_cuda_intercept ===" -gcc -g -O0 -Wall -Wextra \ - -o "$TEST_BIN" \ - "$SCRIPT_DIR/test_cuda_intercept.c" \ - -I"$ROOT_DIR/include" \ - -ldl -lpthread - -echo "" -echo "=== Running: CUDA interception tests (hooks ENABLED) ===" -NVSNAP_CUDA_INTERCEPT=1 \ -NVSNAP_NCCL_INTERCEPT=0 \ -NVSNAP_LOG_LEVEL=3 \ -LD_PRELOAD="$LIB" \ - "$TEST_BIN" -rc=$? - -echo "" -echo "=== Running: Disabled mode (hooks OFF, verify passthrough) ===" -NVSNAP_CUDA_INTERCEPT=0 \ -NVSNAP_NCCL_INTERCEPT=0 \ -NVSNAP_LOG_LEVEL=1 \ -LD_PRELOAD="$LIB" \ - "$TEST_BIN" || true - -echo "" -if [ $rc -eq 0 ]; then - echo "ALL TESTS PASSED" -else - echo "SOME TESTS FAILED (exit code $rc)" -fi -exit $rc diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/run_local_criu_test.sh b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/run_local_criu_test.sh deleted file mode 100755 index 54c0d7ad3f..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/run_local_criu_test.sh +++ /dev/null @@ -1,137 +0,0 @@ -#!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# -# Run this script from a terminal OUTSIDE Cursor to avoid inherited FDs -# -# Usage: ./run_local_criu_test.sh -# - -set -euo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -PROJECT_ROOT="$(dirname "$(dirname "$(dirname "$SCRIPT_DIR")")")" -CRIU="$PROJECT_ROOT/bin/criu" -LIB="$PROJECT_ROOT/lib/nvsnap_intercept/libnvsnap_intercept.so" - -CKPT_DIR="/tmp/nvsnap-local-test-$$" -mkdir -p "$CKPT_DIR" - -cleanup() { - pkill -f "local_test_app" 2>/dev/null || true - rm -rf "$CKPT_DIR" /tmp/local_test_app.py /tmp/local_test_pid 2>/dev/null || true -} -trap cleanup EXIT - -# Build intercept library if needed -if [ ! -f "$LIB" ]; then - echo "Building intercept library..." - cd "$PROJECT_ROOT/lib/nvsnap_intercept" && make -fi - -cat > /tmp/local_test_app.py << 'PYTHON' -#!/usr/bin/env python3 -import os, sys, time, signal - -counter = 0 -restored = os.environ.get("NVSNAP_RESTORED") == "1" - -def handler(sig, frame): - print(f"[APP] Signal {sig}, counter={counter}") - sys.stdout.flush() - -signal.signal(signal.SIGUSR1, handler) -signal.signal(signal.SIGUSR2, handler) - -with open("/tmp/local_test_pid", "w") as f: - f.write(str(os.getpid())) - -print(f"[APP] Started PID={os.getpid()} restored={restored}") -sys.stdout.flush() - -while counter < 100: - counter += 1 - if counter % 10 == 0: - print(f"[APP] counter={counter}") - sys.stdout.flush() - time.sleep(0.3) - -print("[APP] Finished") -PYTHON -chmod +x /tmp/local_test_app.py - -echo "========================================" -echo " LOCAL CHECKPOINT/RESTORE TEST" -echo "========================================" -echo "" -echo "CRIU: $CRIU" -echo "LIB: $LIB" -echo "CKPT: $CKPT_DIR" -echo "" - -echo "=== Step 1: Start application ===" -NVSNAP_LOG_LEVEL=3 LD_PRELOAD="$LIB" python3 /tmp/local_test_app.py & -APP_PID=$! -sleep 2 - -if [ -f /tmp/local_test_pid ]; then - APP_PID=$(cat /tmp/local_test_pid) -fi - -echo "App PID: $APP_PID" -echo "" - -echo "=== Step 2: Send quiesce signal (SIGUSR1) ===" -kill -USR1 $APP_PID 2>/dev/null || true -sleep 1 - -echo "" -echo "=== Step 3: CRIU Checkpoint ===" -sudo "$CRIU" dump \ - --tree $APP_PID \ - --images-dir "$CKPT_DIR" \ - --shell-job \ - -v2 - -if [ $? -eq 0 ]; then - echo "Checkpoint SUCCESS!" - ls "$CKPT_DIR"/*.img | wc -l -else - echo "Checkpoint FAILED!" - exit 1 -fi - -echo "" -echo "=== Step 4: CRIU Restore ===" -cd "$CKPT_DIR" -sudo NVSNAP_RESTORED=1 NVSNAP_LOG_LEVEL=3 LD_PRELOAD="$LIB" \ - "$CRIU" restore \ - --images-dir "$CKPT_DIR" \ - --shell-job \ - -d \ - -v2 - -sleep 2 - -echo "" -echo "=== Step 5: Verify restored process ===" -if pgrep -f local_test_app > /dev/null; then - RPID=$(pgrep -f local_test_app) - echo "Restored process running at PID=$RPID" - sleep 3 - - if ps -p $RPID > /dev/null 2>&1; then - echo "" - echo "========================================" - echo " SUCCESS: CHECKPOINT/RESTORE WORKS!" - echo "========================================" - sudo kill $RPID 2>/dev/null || true - else - echo "Process died after restore" - exit 1 - fi -else - echo "FAILED: Restored process not running" - exit 1 -fi diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_criu_integration.sh b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_criu_integration.sh deleted file mode 100755 index 50c0048cd7..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_criu_integration.sh +++ /dev/null @@ -1,149 +0,0 @@ -#!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# -# CRIU + NVSNAP Integration Test -# -# Prerequisites: -# sudo apt-get install criu -# # Or on RHEL/CentOS: sudo yum install criu -# -# This test: -# 1. Starts a CUDA program with our interception -# 2. Checkpoints it (CRIU for CPU, NVSNAP for GPU) -# 3. Kills it -# 4. Restores it -# 5. Verifies it continues correctly - -set -e - -SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" -LIB_DIR="$(dirname "$SCRIPT_DIR")" -CKPT_DIR="/tmp/nvsnap_criu_test_$$" - -echo "============================================================" -echo " CRIU + NVSNAP Integration Test" -echo "============================================================" -echo "" - -# Check prerequisites -if ! command -v criu &> /dev/null; then - echo "ERROR: criu not found. Install with:" - echo " sudo apt-get install criu" - exit 1 -fi - -if [ ! -f "$LIB_DIR/libnvsnap_intercept.so" ]; then - echo "ERROR: libnvsnap_intercept.so not found. Run 'make' first." - exit 1 -fi - -echo "Prerequisites OK" -echo " CRIU: $(criu --version 2>&1 | head -1)" -echo " Library: $LIB_DIR/libnvsnap_intercept.so" -echo "" - -# Create checkpoint directory -mkdir -p "$CKPT_DIR" -echo "Checkpoint directory: $CKPT_DIR" - -# Create a simple CUDA test program -cat > "$CKPT_DIR/cuda_counter.py" << 'PYEOF' -#!/usr/bin/env python3 -"""Simple CUDA counter that we'll checkpoint/restore""" -import os -import sys -import time -import signal - -# Signal handler to trigger GPU checkpoint -def checkpoint_handler(signum, frame): - print(f"[PID {os.getpid()}] Received checkpoint signal", flush=True) - import ctypes - try: - lib = ctypes.CDLL('./libnvsnap_intercept.so') - lib.nvsnap_checkpoint_to_dir.argtypes = [ctypes.c_char_p] - lib.nvsnap_checkpoint_to_dir.restype = ctypes.c_int - ret = lib.nvsnap_checkpoint_to_dir(os.environ.get('NVSNAP_CKPT_DIR', '/tmp/nvsnap_ckpt').encode()) - print(f"[PID {os.getpid()}] GPU checkpoint result: {ret}", flush=True) - except Exception as e: - print(f"[PID {os.getpid()}] GPU checkpoint error: {e}", flush=True) - -signal.signal(signal.SIGUSR1, checkpoint_handler) - -import torch - -print(f"[PID {os.getpid()}] Starting CUDA counter", flush=True) - -# Create a tensor on GPU -counter = torch.zeros(1, device='cuda') -print(f"[PID {os.getpid()}] Initial counter: {counter.item()}", flush=True) - -# Count forever (will be checkpointed/restored) -while True: - counter += 1 - val = counter.item() - if val % 10 == 0: - print(f"[PID {os.getpid()}] Counter: {val}", flush=True) - time.sleep(0.1) -PYEOF - -echo "" -echo "Step 1: Starting CUDA counter program..." -cd "$LIB_DIR" -NVSNAP_LOG_LEVEL=1 NVSNAP_CKPT_DIR="$CKPT_DIR" \ - LD_PRELOAD=./libnvsnap_intercept.so \ - python3 "$CKPT_DIR/cuda_counter.py" & -PID=$! -echo " Started PID: $PID" - -# Wait for it to count a bit -sleep 3 -echo "" -echo "Step 2: Triggering GPU checkpoint..." -kill -USR1 $PID -sleep 1 - -echo "" -echo "Step 3: Checkpointing with CRIU..." -# Note: CRIU checkpoint requires root for most operations -# and may not work with GPU processes without special handling -if sudo -n true 2>/dev/null; then - sudo criu dump -t $PID -D "$CKPT_DIR/criu" --shell-job -v4 2>"$CKPT_DIR/criu_dump.log" || { - echo " CRIU dump failed (expected - GPU processes need special handling)" - echo " See $CKPT_DIR/criu_dump.log for details" - } -else - echo " Skipping CRIU (needs sudo)" -fi - -echo "" -echo "Step 4: Killing process..." -kill $PID 2>/dev/null || true -wait $PID 2>/dev/null || true -echo " Process killed" - -echo "" -echo "Step 5: Checking checkpoint files..." -ls -la "$CKPT_DIR/" -echo "" -if [ -f "$CKPT_DIR/metadata.json" ]; then - echo "GPU checkpoint metadata:" - cat "$CKPT_DIR/metadata.json" -fi - -echo "" -echo "============================================================" -echo " Test Complete" -echo "============================================================" -echo "" -echo "GPU checkpoint saved to: $CKPT_DIR" -echo "" -echo "NOTE: Full CRIU restore of GPU processes requires:" -echo " 1. Kernel support for CRIU" -echo " 2. Special handling for GPU file descriptors" -echo " 3. Our GPU restore logic after CRIU restore" -echo "" -echo "This test demonstrates the GPU checkpoint mechanism works." -echo "Full integration with CRIU restore is a future milestone." diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_cuda_intercept.c b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_cuda_intercept.c deleted file mode 100644 index 557d0c8125..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_cuda_intercept.c +++ /dev/null @@ -1,637 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * Comprehensive GPU allocation interception test suite. - * - * Tests the entire interception chain: - * 1. PLT override (LD_PRELOAD symbol precedence) - * 2. dlsym override (versioned GLIBC_2.2.5 + GLIBC_2.34) - * 3. cuGetProcAddress override (CUDA 12+ PyTorch path) - * 4. Allocation tracking (cudaMalloc, cuMemAlloc_v2, cuMemMap VMM) - * 5. D2H save correctness (known pattern → save → verify) - * 6. Edge cases (table full, double free, disabled mode) - * - * Build (inside a CUDA container): - * gcc -g -o tests/test_cuda_intercept tests/test_cuda_intercept.c \ - * -I/usr/local/cuda/include -ldl -lpthread - * - * Run: - * NVSNAP_CUDA_INTERCEPT=1 LD_PRELOAD=./libnvsnap_intercept.so \ - * ./tests/test_cuda_intercept - * - * All CUDA functions resolved via dlsym at runtime — no link dependency - * on libcuda/libcudart (avoids interference with our LD_PRELOAD hooks). - */ -#define _GNU_SOURCE -#include -#include -#include -#include -#include -#include -#include -#include - -/* ─── Types ─────────────────────────────────────────────────────────── */ - -typedef int CUresult; -typedef int cudaError_t; -typedef unsigned long long CUdeviceptr; -typedef unsigned long long CUmemGenericAllocationHandle; -typedef unsigned long long cuuint64_t; - -/* VMM types */ -typedef enum { - CU_MEM_ALLOCATION_TYPE_PINNED = 1 -} CUmemAllocationType; - -typedef enum { - CU_MEM_LOCATION_TYPE_DEVICE = 1 -} CUmemLocationType; - -typedef enum { - CU_MEM_ACCESS_FLAGS_PROT_READWRITE = 3 -} CUmemAccess_flags; - -/* Minimal structs — sizes must match CUDA driver headers */ -typedef struct { - CUmemAllocationType type; - unsigned long long _pad1; - struct { CUmemLocationType type; int id; } location; - void *win32HandleMetaData; - struct { unsigned char compressionType; unsigned char gpuDirectRDMACapable; - unsigned short usage; unsigned char reserved[4]; } allocFlags; - unsigned long long _pad2; -} CUmemAllocationProp; - -typedef struct { - struct { CUmemLocationType type; int id; } location; - CUmemAccess_flags flags; -} CUmemAccessDesc; - -/* ─── Function pointers (resolved at runtime) ───────────────────────── */ - -static CUresult (*fn_cuInit)(unsigned) = NULL; -static cudaError_t (*fn_cudaSetDevice)(int) = NULL; -static cudaError_t (*fn_cudaMalloc)(void **, size_t) = NULL; -static cudaError_t (*fn_cudaFree)(void *) = NULL; -static cudaError_t (*fn_cudaMemcpy)(void *, const void *, size_t, int) = NULL; -static cudaError_t (*fn_cudaGetDeviceCount)(int *) = NULL; -static CUresult (*fn_cuMemAlloc_v2)(CUdeviceptr *, size_t) = NULL; -static CUresult (*fn_cuMemFree_v2)(CUdeviceptr) = NULL; - -/* VMM */ -static CUresult (*fn_cuMemAddressReserve)(CUdeviceptr *, size_t, size_t, CUdeviceptr, unsigned long long) = NULL; -static CUresult (*fn_cuMemCreate)(CUmemGenericAllocationHandle *, size_t, const CUmemAllocationProp *, unsigned long long) = NULL; -static CUresult (*fn_cuMemMap)(CUdeviceptr, size_t, size_t, CUmemGenericAllocationHandle, unsigned long long) = NULL; -static CUresult (*fn_cuMemSetAccess)(CUdeviceptr, size_t, const CUmemAccessDesc *, size_t) = NULL; -static CUresult (*fn_cuMemUnmap)(CUdeviceptr, size_t) = NULL; -static CUresult (*fn_cuMemRelease)(CUmemGenericAllocationHandle) = NULL; -static CUresult (*fn_cuMemAddressFree)(CUdeviceptr, size_t) = NULL; - -/* cuGetProcAddress */ -typedef CUresult (*cuGetProcAddress_v2_fn)(const char *, void **, int, cuuint64_t, int *); -static cuGetProcAddress_v2_fn fn_cuGetProcAddress_v2 = NULL; - -/* nvsnap_cuda_save — from our library */ -static int (*fn_nvsnap_cuda_save)(const char *) = NULL; - -/* ─── Helpers ───────────────────────────────────────────────────────── */ - -static int g_pass = 0, g_fail = 0; - -#define CHECK(cond, name) do { \ - if (cond) { printf(" PASS: %s\n", name); g_pass++; } \ - else { printf(" FAIL: %s\n", name); g_fail++; } \ -} while(0) - -/* Get the REAL dlsym (bypass our override) */ -static void *(*real_dlsym_fn)(void *, const char *) = NULL; -static void *get_real_sym(const char *lib, const char *sym) { - if (!real_dlsym_fn) - real_dlsym_fn = dlvsym(RTLD_NEXT, "dlsym", "GLIBC_2.2.5"); - void *h = dlopen(lib, RTLD_LAZY | RTLD_NOLOAD); - if (!h) h = dlopen(lib, RTLD_LAZY); - if (!h) return NULL; - return real_dlsym_fn ? real_dlsym_fn(h, sym) : NULL; -} - -static int resolve_all(void) { - /* Use dlsym (goes through our override) for functions we want hooked */ - fn_cuInit = dlsym(RTLD_DEFAULT, "cuInit"); - fn_cudaSetDevice = dlsym(RTLD_DEFAULT, "cudaSetDevice"); - fn_cudaMalloc = dlsym(RTLD_DEFAULT, "cudaMalloc"); - fn_cudaFree = dlsym(RTLD_DEFAULT, "cudaFree"); - fn_cudaMemcpy = dlsym(RTLD_DEFAULT, "cudaMemcpy"); - fn_cudaGetDeviceCount = dlsym(RTLD_DEFAULT, "cudaGetDeviceCount"); - fn_cuMemAlloc_v2 = dlsym(RTLD_DEFAULT, "cuMemAlloc_v2"); - fn_cuMemFree_v2 = dlsym(RTLD_DEFAULT, "cuMemFree_v2"); - - /* VMM */ - fn_cuMemAddressReserve = dlsym(RTLD_DEFAULT, "cuMemAddressReserve"); - fn_cuMemCreate = dlsym(RTLD_DEFAULT, "cuMemCreate"); - fn_cuMemMap = dlsym(RTLD_DEFAULT, "cuMemMap"); - fn_cuMemSetAccess = dlsym(RTLD_DEFAULT, "cuMemSetAccess"); - fn_cuMemUnmap = dlsym(RTLD_DEFAULT, "cuMemUnmap"); - fn_cuMemRelease = dlsym(RTLD_DEFAULT, "cuMemRelease"); - fn_cuMemAddressFree = dlsym(RTLD_DEFAULT, "cuMemAddressFree"); - - /* cuGetProcAddress */ - fn_cuGetProcAddress_v2 = dlsym(RTLD_DEFAULT, "cuGetProcAddress_v2"); - - /* nvsnap_cuda_save — resolve from our preloaded library */ - fn_nvsnap_cuda_save = dlsym(RTLD_DEFAULT, "nvsnap_cuda_save"); - - if (!fn_cuInit || !fn_cudaMalloc || !fn_cudaFree) { - fprintf(stderr, "Failed to resolve basic CUDA functions\n"); - return -1; - } - return 0; -} - -static void rmrf(const char *path) { - char cmd[512]; - snprintf(cmd, sizeof(cmd), "rm -rf %s", path); - system(cmd); -} - -static long file_size(const char *path) { - struct stat st; - if (stat(path, &st) != 0) return -1; - return st.st_size; -} - -/* Count entries in manifest JSON (simple — just count "addr" occurrences) */ -static int manifest_count(const char *dir) { - char path[512]; - snprintf(path, sizeof(path), "%s/gpu-manifest.json", dir); - FILE *f = fopen(path, "r"); - if (!f) return -1; - char buf[65536]; - size_t n = fread(buf, 1, sizeof(buf)-1, f); - buf[n] = '\0'; - fclose(f); - int count = 0; - char *p = buf; - while ((p = strstr(p, "\"addr\"")) != NULL) { count++; p++; } - return count; -} - -/* ─── Test 1: dlsym Override ────────────────────────────────────────── */ - -static void test_dlsym_override(void) { - printf("\n=== Test 1: dlsym Override ===\n"); - - /* Get pointers via dlsym (our override) and via real dlsym from libcuda/libcudart */ - void *our_cudaMalloc = dlsym(RTLD_DEFAULT, "cudaMalloc"); - void *real_cudaMalloc = get_real_sym("libcudart.so", "cudaMalloc"); - CHECK(our_cudaMalloc != NULL, "dlsym returns non-NULL for cudaMalloc"); - CHECK(real_cudaMalloc != NULL, "real cudaMalloc found in libcudart.so"); - CHECK(our_cudaMalloc != real_cudaMalloc, "dlsym(cudaMalloc) returns OUR hook, not libcudart's"); - - void *our_cuMemMap = dlsym(RTLD_DEFAULT, "cuMemMap"); - void *real_cuMemMap = get_real_sym("libcuda.so.1", "cuMemMap"); - CHECK(our_cuMemMap != NULL, "dlsym returns non-NULL for cuMemMap"); - if (real_cuMemMap) { - CHECK(our_cuMemMap != real_cuMemMap, "dlsym(cuMemMap) returns OUR hook, not libcuda's"); - } -} - -/* ─── Test 2: cuGetProcAddress Override ──────────────────────────────── */ - -static void test_cugetprocaddress(void) { - printf("\n=== Test 2: cuGetProcAddress Override ===\n"); - - if (!fn_cuGetProcAddress_v2) { - printf(" SKIP: cuGetProcAddress_v2 not available\n"); - return; - } - - /* Resolve cuMemMap via cuGetProcAddress — should get our hook */ - void *resolved_cuMemMap = NULL; - CUresult r = fn_cuGetProcAddress_v2("cuMemMap", &resolved_cuMemMap, 10020, 0, NULL); - CHECK(r == 0, "cuGetProcAddress_v2(cuMemMap) succeeds"); - CHECK(resolved_cuMemMap != NULL, "cuGetProcAddress_v2(cuMemMap) returns non-NULL"); - - /* Compare with real libcuda version */ - void *real_cuMemMap = get_real_sym("libcuda.so.1", "cuMemMap"); - if (real_cuMemMap && resolved_cuMemMap) { - CHECK(resolved_cuMemMap != real_cuMemMap, - "cuGetProcAddress(cuMemMap) returns OUR hook, not libcuda's"); - } - - /* Resolve something we DON'T hook — should pass through to real */ - void *resolved_cuCtxCreate = NULL; - r = fn_cuGetProcAddress_v2("cuCtxCreate", &resolved_cuCtxCreate, 2000, 0, NULL); - CHECK(r == 0, "cuGetProcAddress_v2(cuCtxCreate) succeeds (passthrough)"); - CHECK(resolved_cuCtxCreate != NULL, "cuGetProcAddress_v2(cuCtxCreate) returns non-NULL"); -} - -/* ─── Test 3: cudaMalloc Tracking ───────────────────────────────────── */ - -static void test_cudamalloc_tracking(void) { - printf("\n=== Test 3: cudaMalloc Tracking ===\n"); - - const char *dir = "/tmp/nvsnap_test_cudamalloc"; - rmrf(dir); - mkdir(dir, 0755); - - void *ptr = NULL; - cudaError_t err = fn_cudaMalloc(&ptr, 4096); - CHECK(err == 0, "cudaMalloc(4096) succeeds"); - CHECK(ptr != NULL, "cudaMalloc returns non-NULL pointer"); - - if (fn_nvsnap_cuda_save) { - int rc = fn_nvsnap_cuda_save(dir); - CHECK(rc == 0, "nvsnap_cuda_save succeeds"); - int count = manifest_count(dir); - CHECK(count >= 1, "manifest has >= 1 allocation after cudaMalloc"); - - /* Free and verify it's gone */ - fn_cudaFree(ptr); - rmrf(dir); - mkdir(dir, 0755); - fn_nvsnap_cuda_save(dir); - int count2 = manifest_count(dir); - CHECK(count2 == count - 1, "manifest has one fewer allocation after cudaFree"); - } else { - printf(" SKIP: nvsnap_cuda_save not found (not preloaded?)\n"); - fn_cudaFree(ptr); - } - rmrf(dir); -} - -/* ─── Test 4: cuMemAlloc_v2 Tracking ────────────────────────────────── */ - -static void test_cumemalloc_tracking(void) { - printf("\n=== Test 4: cuMemAlloc_v2 Tracking ===\n"); - - if (!fn_cuMemAlloc_v2 || !fn_cuMemFree_v2) { - printf(" SKIP: cuMemAlloc_v2/cuMemFree_v2 not available\n"); - return; - } - - const char *dir = "/tmp/nvsnap_test_cumemalloc"; - rmrf(dir); - mkdir(dir, 0755); - - CUdeviceptr dptr = 0; - CUresult r = fn_cuMemAlloc_v2(&dptr, 8192); - CHECK(r == 0, "cuMemAlloc_v2(8192) succeeds"); - CHECK(dptr != 0, "cuMemAlloc_v2 returns non-zero address"); - - if (fn_nvsnap_cuda_save) { - int rc = fn_nvsnap_cuda_save(dir); - CHECK(rc == 0, "nvsnap_cuda_save succeeds"); - int count = manifest_count(dir); - CHECK(count >= 1, "manifest has >= 1 allocation after cuMemAlloc_v2"); - } - - fn_cuMemFree_v2(dptr); - rmrf(dir); -} - -/* ─── Test 5: VMM (cuMemMap) Tracking ───────────────────────────────── */ - -static void test_vmm_tracking(void) { - printf("\n=== Test 5: VMM cuMemMap Tracking (PyTorch path) ===\n"); - - if (!fn_cuMemAddressReserve || !fn_cuMemCreate || !fn_cuMemMap || - !fn_cuMemSetAccess || !fn_cuMemUnmap || !fn_cuMemRelease || - !fn_cuMemAddressFree) { - printf(" SKIP: VMM APIs not available\n"); - return; - } - - const size_t SIZE = 2 * 1024 * 1024; /* 2 MB — CUDA VMM granularity */ - CUresult r; - - /* Reserve VA */ - CUdeviceptr va = 0; - r = fn_cuMemAddressReserve(&va, SIZE, 0, 0, 0); - CHECK(r == 0, "cuMemAddressReserve succeeds"); - - /* Create physical allocation */ - CUmemAllocationProp prop; - memset(&prop, 0, sizeof(prop)); - prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; - prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; - prop.location.id = 0; - - CUmemGenericAllocationHandle handle = 0; - r = fn_cuMemCreate(&handle, SIZE, &prop, 0); - CHECK(r == 0, "cuMemCreate succeeds"); - - /* Map physical to VA — THIS is where our hook should fire */ - r = fn_cuMemMap(va, SIZE, 0, handle, 0); - CHECK(r == 0, "cuMemMap succeeds"); - - /* Set access */ - CUmemAccessDesc access; - access.location.type = CU_MEM_LOCATION_TYPE_DEVICE; - access.location.id = 0; - access.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; - r = fn_cuMemSetAccess(va, SIZE, &access, 1); - CHECK(r == 0, "cuMemSetAccess succeeds"); - - /* Verify tracking */ - const char *dir = "/tmp/nvsnap_test_vmm"; - rmrf(dir); - mkdir(dir, 0755); - - if (fn_nvsnap_cuda_save) { - int rc = fn_nvsnap_cuda_save(dir); - CHECK(rc == 0, "nvsnap_cuda_save after cuMemMap succeeds"); - int count = manifest_count(dir); - CHECK(count >= 1, "manifest has >= 1 allocation after cuMemMap"); - - /* Check the saved file has the right size */ - char fpath[512]; - snprintf(fpath, sizeof(fpath), "%s/gpu-alloc-0.bin", dir); - long fsize = file_size(fpath); - if (fsize >= 0) { - CHECK((size_t)fsize == SIZE, "saved file size matches allocation size (2 MB)"); - } - } - - /* Cleanup */ - fn_cuMemUnmap(va, SIZE); - fn_cuMemRelease(handle); - fn_cuMemAddressFree(va, SIZE); - - if (fn_nvsnap_cuda_save) { - rmrf(dir); - mkdir(dir, 0755); - fn_nvsnap_cuda_save(dir); - int count2 = manifest_count(dir); - CHECK(count2 >= 0, "manifest valid after cuMemUnmap"); - } - rmrf(dir); -} - -/* ─── Test 6: D2H Save Data Integrity ───────────────────────────────── */ - -static void test_d2h_integrity(void) { - printf("\n=== Test 6: D2H Save Data Integrity ===\n"); - - if (!fn_cudaMemcpy || !fn_nvsnap_cuda_save) { - printf(" SKIP: cudaMemcpy or nvsnap_cuda_save not available\n"); - return; - } - - const size_t SIZE = 1024 * 1024; /* 1 MB */ - const char *dir = "/tmp/nvsnap_test_d2h"; - rmrf(dir); - mkdir(dir, 0755); - - /* Allocate and fill with known pattern */ - void *dptr = NULL; - fn_cudaMalloc(&dptr, SIZE); - CHECK(dptr != NULL, "cudaMalloc for D2H test succeeds"); - - unsigned int *host_pattern = malloc(SIZE); - for (size_t i = 0; i < SIZE / sizeof(unsigned int); i++) - host_pattern[i] = 0xDEADBEEF ^ (unsigned int)i; - - fn_cudaMemcpy(dptr, host_pattern, SIZE, 1 /* H2D */); - - /* Save */ - int rc = fn_nvsnap_cuda_save(dir); - CHECK(rc == 0, "nvsnap_cuda_save succeeds"); - - /* Find the allocation file — look for any gpu-alloc-*.bin */ - int verified = 0; - for (int i = 0; i < 100; i++) { - char fpath[512]; - snprintf(fpath, sizeof(fpath), "%s/gpu-alloc-%d.bin", dir, i); - long fsize = file_size(fpath); - if (fsize == (long)SIZE) { - /* Read and compare */ - FILE *f = fopen(fpath, "rb"); - if (f) { - unsigned int *readback = malloc(SIZE); - fread(readback, 1, SIZE, f); - fclose(f); - int match = memcmp(host_pattern, readback, SIZE) == 0; - CHECK(match, "D2H saved data matches original pattern (1 MB)"); - verified = 1; - free(readback); - } - break; - } - } - if (!verified) { - CHECK(0, "D2H saved file with correct size found"); - } - - free(host_pattern); - fn_cudaFree(dptr); - rmrf(dir); -} - -/* ─── Test 7: cuGetProcAddress → Allocate → Track ───────────────────── */ - -static void test_cugetprocaddr_allocate(void) { - printf("\n=== Test 7: cuGetProcAddress → Allocate → Track ===\n"); - - if (!fn_cuGetProcAddress_v2 || !fn_nvsnap_cuda_save) { - printf(" SKIP: cuGetProcAddress_v2 or nvsnap_cuda_save not available\n"); - return; - } - - /* Resolve cudaMalloc-equivalent via cuGetProcAddress (the PyTorch path) */ - CUresult (*resolved_cuMemAlloc)(CUdeviceptr *, size_t) = NULL; - CUresult r = fn_cuGetProcAddress_v2("cuMemAlloc", (void **)&resolved_cuMemAlloc, 2000, 0, NULL); - CHECK(r == 0 && resolved_cuMemAlloc != NULL, "cuGetProcAddress resolves cuMemAlloc"); - - CUresult (*resolved_cuMemFree)(CUdeviceptr) = NULL; - fn_cuGetProcAddress_v2("cuMemFree", (void **)&resolved_cuMemFree, 2000, 0, NULL); - - if (!resolved_cuMemAlloc) return; - - const char *dir = "/tmp/nvsnap_test_getproc"; - rmrf(dir); - mkdir(dir, 0755); - - /* Allocate via the cuGetProcAddress-resolved function */ - CUdeviceptr dptr = 0; - r = resolved_cuMemAlloc(&dptr, 16384); - CHECK(r == 0, "cuGetProcAddress-resolved cuMemAlloc(16384) succeeds"); - - /* Verify it was tracked */ - int rc = fn_nvsnap_cuda_save(dir); - CHECK(rc == 0, "nvsnap_cuda_save succeeds"); - int count = manifest_count(dir); - CHECK(count >= 1, "allocation via cuGetProcAddress path is tracked in manifest"); - - if (resolved_cuMemFree) resolved_cuMemFree(dptr); - rmrf(dir); -} - -/* ─── Test 8: Empty Save ────────────────────────────────────────────── */ - -static void test_empty_save(void) { - printf("\n=== Test 8: Empty Save (no allocations) ===\n"); - - if (!fn_nvsnap_cuda_save) { - printf(" SKIP: nvsnap_cuda_save not available\n"); - return; - } - - /* Note: there may be leftover allocations from previous tests. - * We just verify save doesn't crash with any count. */ - const char *dir = "/tmp/nvsnap_test_empty"; - rmrf(dir); - mkdir(dir, 0755); - int rc = fn_nvsnap_cuda_save(dir); - CHECK(rc == 0 || rc == -1, "nvsnap_cuda_save returns without crash"); - rmrf(dir); -} - -/* ─── Test 9: Save to Invalid Path ──────────────────────────────────── */ - -static void test_save_invalid_path(void) { - printf("\n=== Test 9: Save to Invalid Path ===\n"); - - if (!fn_nvsnap_cuda_save) { - printf(" SKIP: nvsnap_cuda_save not available\n"); - return; - } - - /* Allocate something so there's data to save */ - void *dptr = NULL; - fn_cudaMalloc(&dptr, 1024); - - int rc = fn_nvsnap_cuda_save("/nonexistent/nvsnap_test_bad"); - CHECK(rc == -1, "nvsnap_cuda_save returns -1 for invalid path"); - /* Process must still be alive */ - CHECK(1, "process survived save to invalid path"); - - fn_cudaFree(dptr); -} - -/* ─── Test 10: Multi-Device Tracking ────────────────────────────────── */ - -static void test_multidevice(void) { - printf("\n=== Test 10: Multi-Device Tracking ===\n"); - - int ndev = 0; - fn_cudaGetDeviceCount(&ndev); - if (ndev < 2) { - printf(" SKIP: need >= 2 GPUs, have %d\n", ndev); - return; - } - - /* Allocate on device 0 */ - fn_cudaSetDevice(0); - void *p0 = NULL; - fn_cudaMalloc(&p0, 4096); - - /* Allocate on device 1 */ - fn_cudaSetDevice(1); - void *p1 = NULL; - fn_cudaMalloc(&p1, 4096); - - CHECK(p0 != NULL && p1 != NULL, "allocations on both devices succeed"); - CHECK(p0 != p1, "allocations on different devices have different addresses"); - - /* Save and check manifest has both devices */ - const char *dir = "/tmp/nvsnap_test_multidev"; - rmrf(dir); - mkdir(dir, 0755); - if (fn_nvsnap_cuda_save) { - fn_nvsnap_cuda_save(dir); - /* Read manifest and check for device 0 and device 1 */ - char path[512]; - snprintf(path, sizeof(path), "%s/gpu-manifest.json", dir); - FILE *f = fopen(path, "r"); - if (f) { - char buf[65536]; - size_t n = fread(buf, 1, sizeof(buf)-1, f); - buf[n] = '\0'; - fclose(f); - CHECK(strstr(buf, "\"device\": 0") != NULL, "manifest contains device 0"); - CHECK(strstr(buf, "\"device\": 1") != NULL, "manifest contains device 1"); - } - } - - fn_cudaSetDevice(0); fn_cudaFree(p0); - fn_cudaSetDevice(1); fn_cudaFree(p1); - fn_cudaSetDevice(0); - rmrf(dir); -} - -/* ─── Test 11: Thread Safety ────────────────────────────────────────── */ - -#define THREAD_COUNT 4 -#define ALLOCS_PER_THREAD 50 - -static void *thread_alloc_free(void *arg) { - (void)arg; - for (int i = 0; i < ALLOCS_PER_THREAD; i++) { - void *p = NULL; - if (fn_cudaMalloc(&p, 1024) == 0 && p) { - fn_cudaFree(p); - } - } - return NULL; -} - -static void test_thread_safety(void) { - printf("\n=== Test 11: Thread Safety ===\n"); - - pthread_t threads[THREAD_COUNT]; - for (int i = 0; i < THREAD_COUNT; i++) - pthread_create(&threads[i], NULL, thread_alloc_free, NULL); - for (int i = 0; i < THREAD_COUNT; i++) - pthread_join(threads[i], NULL); - - CHECK(1, "concurrent alloc/free completed without crash"); - - /* Save should work after concurrent ops */ - if (fn_nvsnap_cuda_save) { - const char *dir = "/tmp/nvsnap_test_threads"; - rmrf(dir); - mkdir(dir, 0755); - int rc = fn_nvsnap_cuda_save(dir); - CHECK(rc == 0, "nvsnap_cuda_save succeeds after concurrent alloc/free"); - rmrf(dir); - } -} - -/* ─── Main ──────────────────────────────────────────────────────────── */ - -int main(void) { - printf("=== NVSNAP CUDA Interception Test Suite ===\n"); - printf("NVSNAP_CUDA_INTERCEPT=%s\n", getenv("NVSNAP_CUDA_INTERCEPT") ?: "(unset)"); - printf("LD_PRELOAD=%s\n\n", getenv("LD_PRELOAD") ?: "(unset)"); - - if (resolve_all() < 0) { - fprintf(stderr, "FATAL: Cannot resolve CUDA functions\n"); - return 1; - } - - /* Initialize CUDA */ - if (fn_cuInit) fn_cuInit(0); - fn_cudaSetDevice(0); - - /* Run tests */ - test_dlsym_override(); - test_cugetprocaddress(); - test_cudamalloc_tracking(); - test_cumemalloc_tracking(); - test_vmm_tracking(); - test_d2h_integrity(); - test_cugetprocaddr_allocate(); - test_empty_save(); - test_save_invalid_path(); - test_multidevice(); - test_thread_safety(); - - /* Summary */ - printf("\n=== Results: %d passed, %d failed ===\n", g_pass, g_fail); - return g_fail > 0 ? 1 : 0; -} diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_dlsym_recursion.c b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_dlsym_recursion.c deleted file mode 100644 index 2955223780..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_dlsym_recursion.c +++ /dev/null @@ -1,144 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * Test: dlsym override does NOT cause infinite recursion. - * - * In a single-library build, our dlsym override intercepts ALL dlsym calls. - * If NVSNAP_LOAD_REAL uses dlsym(RTLD_NEXT, "cudaMalloc") and our override - * returns our own hook, we get infinite recursion → hang/segfault. - * - * This test catches that bug by: - * 1. Calling dlsym(RTLD_DEFAULT, "cudaMalloc") — should return our hook - * 2. Calling the returned function — should NOT hang (must reach real libcudart) - * 3. Verifying dlsym for non-hooked symbols still works (e.g., "printf") - * - * Build: - * gcc -g -o tests/test_dlsym_recursion tests/test_dlsym_recursion.c -ldl - * - * Run: - * LD_PRELOAD=./libnvsnap_intercept.so ./tests/test_dlsym_recursion - */ -#define _GNU_SOURCE -#include -#include -#include -#include -#include -#include - -static volatile int g_alarm_fired = 0; - -static void alarm_handler(int sig) { - (void)sig; - g_alarm_fired = 1; - const char msg[] = "FAIL: dlsym recursion detected (alarm timeout)\n"; - write(STDERR_FILENO, msg, sizeof(msg) - 1); - _exit(1); -} - -int main(void) { - printf("=== Test: dlsym override recursion safety ===\n"); - - /* Set 5-second alarm — if we hang, it fires and fails the test */ - signal(SIGALRM, alarm_handler); - alarm(5); - - int pass = 0, fail = 0; - - /* Test 1: dlsym for a hooked symbol returns non-NULL */ - void *sym = dlsym(RTLD_DEFAULT, "cudaMalloc"); - if (sym) { - printf(" PASS: dlsym(cudaMalloc) = %p\n", sym); - pass++; - } else { - printf(" SKIP: cudaMalloc not found (no CUDA runtime)\n"); - } - - /* Test 2: dlsym for cuMemMap */ - sym = dlsym(RTLD_DEFAULT, "cuMemMap"); - if (sym) { - printf(" PASS: dlsym(cuMemMap) = %p\n", sym); - pass++; - } else { - printf(" SKIP: cuMemMap not found\n"); - } - - /* Test 3: dlsym for non-hooked symbol still works */ - sym = dlsym(RTLD_DEFAULT, "printf"); - if (sym) { - printf(" PASS: dlsym(printf) = %p (non-hooked passthrough works)\n", sym); - pass++; - } else { - printf(" FAIL: dlsym(printf) returned NULL — override is broken\n"); - fail++; - } - - /* Test 4: dlsym(RTLD_DEFAULT, "dlsym") — meta-test */ - sym = dlsym(RTLD_DEFAULT, "dlsym"); - if (sym) { - printf(" PASS: dlsym(dlsym) = %p (no infinite recursion)\n", sym); - pass++; - } else { - printf(" FAIL: dlsym(dlsym) returned NULL\n"); - fail++; - } - - /* Test 5: Verify nvsnap_resolve_real exists (merged build) */ - void *(*resolve_fn)(const char *, const char *) = dlsym(RTLD_DEFAULT, "nvsnap_resolve_real"); - if (resolve_fn) { - printf(" PASS: nvsnap_resolve_real found (merged build confirmed)\n"); - pass++; - - /* Test 6: nvsnap_resolve_real returns real cudaMalloc, not our hook */ - void *real = resolve_fn("cudaMalloc", "libcudart.so"); - void *hook = dlsym(RTLD_DEFAULT, "cudaMalloc"); - if (real && hook && real != hook) { - printf(" PASS: resolve_real(cudaMalloc) != dlsym(cudaMalloc) — no self-resolution\n"); - pass++; - } else if (real && hook && real == hook) { - printf(" FAIL: resolve_real returns our own hook — dladdr self-check broken\n"); - fail++; - } else { - printf(" SKIP: cudaMalloc not available for comparison\n"); - } - } else { - printf(" SKIP: nvsnap_resolve_real not found (not merged build?)\n"); - } - - /* Test 7: dlsym with specific library handle passes through (no hook) */ - { - void *libc = dlopen("libc.so.6", RTLD_LAZY | RTLD_NOLOAD); - if (libc) { - void *real_printf = dlsym(libc, "printf"); - void *hook_printf = dlsym(RTLD_DEFAULT, "printf"); - /* With specific handle, should get the REAL function. - * Our override should NOT intercept specific handles. */ - if (real_printf) { - printf(" PASS: dlsym(libc_handle, printf) = %p (specific handle passthrough)\n", - real_printf); - pass++; - } else { - printf(" FAIL: dlsym(libc_handle, printf) returned NULL\n"); - fail++; - } - } - } - - /* Test 8: Rapid dlsym calls don't hang (stress test) */ - for (int i = 0; i < 1000; i++) { - dlsym(RTLD_DEFAULT, "cudaMalloc"); - dlsym(RTLD_DEFAULT, "cuMemMap"); - dlsym(RTLD_DEFAULT, "ncclCommInitRank"); - dlsym(RTLD_DEFAULT, "printf"); - } - printf(" PASS: 4000 dlsym calls completed without hang\n"); - pass++; - - alarm(0); /* Cancel alarm */ - - printf("\n=== Results: %d passed, %d failed ===\n", pass, fail); - return fail > 0 ? 1 : 0; -} diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_full_checkpoint_restore.sh b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_full_checkpoint_restore.sh deleted file mode 100755 index 3b505bdbe8..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_full_checkpoint_restore.sh +++ /dev/null @@ -1,303 +0,0 @@ -#!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# GPU Checkpoint/Restore Test using cuda-checkpoint + CRIU -# -# This script demonstrates full GPU checkpoint/restore: -# 1. Start GPU process with LD_PRELOAD for interception -# 2. Lock CUDA state with cuda-checkpoint -# 3. Checkpoint CUDA state with cuda-checkpoint -# 4. CRIU dump (CPU state + file descriptors) -# 5. CRIU restore -# 6. Restore CUDA state with cuda-checkpoint -# 7. Unlock and resume -# -# Requirements: -# - NVIDIA driver 555+ (for cuda-checkpoint support) -# - cuda-checkpoint binary -# - CRIU with CUDA plugin -# - libnvsnap_intercept.so built - -set -e - -# ============================================================================ -# Configuration -# ============================================================================ - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -LIB_DIR="$(dirname "$SCRIPT_DIR")" - -_repo_root="$(cd "$LIB_DIR/../.." && pwd)" - -# cuda-checkpoint location (check env var, then sibling checkout, then PATH) -if [ -n "${CUDA_CHECKPOINT:-}" ]; then - CUDA_CKPT="$CUDA_CHECKPOINT" -elif [ -x "$_repo_root/../cuda-checkpoint/bin/x86_64_Linux/cuda-checkpoint" ]; then - CUDA_CKPT="$_repo_root/../cuda-checkpoint/bin/x86_64_Linux/cuda-checkpoint" -elif command -v cuda-checkpoint &>/dev/null; then - CUDA_CKPT="$(command -v cuda-checkpoint)" -else - echo "ERROR: cuda-checkpoint not found" - echo "Set CUDA_CHECKPOINT env var or add to PATH" - exit 1 -fi - -# CRIU binary (check env var, then sibling-built fork, then PATH) -if [ -n "${NVSNAP_CRIU:-}" ]; then - CRIU="$NVSNAP_CRIU" -elif [ -x "$_repo_root/../criu/criu/criu" ]; then - CRIU="$_repo_root/../criu/criu/criu" -elif command -v criu &>/dev/null; then - CRIU="$(command -v criu)" -else - echo "ERROR: criu not found" - echo "Set NVSNAP_CRIU env var or build the CRIU fork as a sibling of this repo" - exit 1 -fi - -# CRIU CUDA plugin directory -if [ -n "${NVSNAP_CRIU_PLUGIN_DIR:-}" ]; then - PLUGIN_DIR="$NVSNAP_CRIU_PLUGIN_DIR" -elif [ -d "$_repo_root/../criu/plugins/cuda" ]; then - PLUGIN_DIR="$_repo_root/../criu/plugins/cuda" -elif [ -d "/usr/lib/criu" ]; then - PLUGIN_DIR="/usr/lib/criu" -else - PLUGIN_DIR="" -fi - -# Test type (simple or pytorch) -TEST_TYPE="${1:-simple}" -CHECKPOINT_DIR="/tmp/nvsnap_checkpoint" -CRIU_IMG_DIR="/tmp/criu_img" -TEST_OUTPUT="/tmp/gpu_test_output.txt" - -# ============================================================================ -# Helper functions -# ============================================================================ - -cleanup() { - echo "[Cleanup] Stopping any previous test processes..." - pkill -9 -f test_simple_checkpoint 2>/dev/null || true - pkill -9 -f "python.*test_pytorch" 2>/dev/null || true - rm -rf "$CRIU_IMG_DIR" "$CHECKPOINT_DIR" "$TEST_OUTPUT" - mkdir -p "$CRIU_IMG_DIR" -} - -check_requirements() { - echo "[Check] Verifying requirements..." - - # Check driver version - DRIVER_VER=$(nvidia-smi --query-gpu=driver_version --format=csv,noheader 2>/dev/null | head -1) - DRIVER_MAJOR=$(echo "$DRIVER_VER" | cut -d. -f1) - if [ -z "$DRIVER_MAJOR" ] || [ "$DRIVER_MAJOR" -lt 555 ]; then - echo "ERROR: NVIDIA driver 555+ required (found: $DRIVER_VER)" - exit 1 - fi - echo " NVIDIA driver: $DRIVER_VER (OK)" - - # Check cuda-checkpoint - if ! "$CUDA_CKPT" --help &>/dev/null; then - echo "ERROR: cuda-checkpoint not working" - exit 1 - fi - echo " cuda-checkpoint: $CUDA_CKPT (OK)" - - # Check CRIU - if ! "$CRIU" --version &>/dev/null; then - echo "ERROR: CRIU not working" - exit 1 - fi - echo " CRIU: $CRIU (OK)" - - # Check library - if [ ! -f "$LIB_DIR/libnvsnap_intercept.so" ]; then - echo "ERROR: libnvsnap_intercept.so not found" - exit 1 - fi - echo " Library: $LIB_DIR/libnvsnap_intercept.so (OK)" - - echo "" -} - -start_simple_test() { - echo "[Start] Launching simple CUDA test..." - cd "$LIB_DIR" - # Note: LD_PRELOAD not needed - cuda-checkpoint handles GPU state directly - ./tests/test_simple_checkpoint > "$TEST_OUTPUT" 2>&1 & - TEST_PID=$! - echo " PID: $TEST_PID" - sleep 4 -} - -start_pytorch_test() { - echo "[Start] Launching PyTorch test..." - cd "$LIB_DIR" - # Note: LD_PRELOAD not needed - cuda-checkpoint handles GPU state directly - - python3 -c " -import torch -import time -import os - -print(f'PyTorch PID: {os.getpid()}') -print(f'CUDA available: {torch.cuda.is_available()}') -print(f'Device: {torch.cuda.get_device_name(0)}') - -# Create model and data -model = torch.nn.Linear(1024, 1024).cuda() -data = torch.randn(64, 1024).cuda() - -# Forward pass to ensure GPU is active -output = model(data) -print(f'Model output shape: {output.shape}') -print(f'Output sum: {output.sum().item():.4f}') - -# Keep alive for checkpoint -print('Waiting for checkpoint...') -while True: - time.sleep(1) -" > "$TEST_OUTPUT" 2>&1 & - TEST_PID=$! - echo " PID: $TEST_PID" - sleep 5 -} - -verify_running() { - if ! ps -p $TEST_PID > /dev/null 2>&1; then - echo "ERROR: Test process died" - cat "$TEST_OUTPUT" - exit 1 - fi - - STATE=$("$CUDA_CKPT" --get-state --pid $TEST_PID 2>/dev/null | tail -1) - echo " CUDA state: $STATE" - - if [ "$STATE" != "running" ]; then - echo "ERROR: CUDA not in running state" - exit 1 - fi -} - -do_checkpoint() { - echo "" - echo "[Lock] Locking CUDA state..." - "$CUDA_CKPT" --action lock --pid $TEST_PID --timeout 10000 2>/dev/null - echo " Lock: OK" - - echo "" - echo "[Checkpoint] Checkpointing CUDA state..." - "$CUDA_CKPT" --action checkpoint --pid $TEST_PID 2>/dev/null - - STATE=$("$CUDA_CKPT" --get-state --pid $TEST_PID 2>/dev/null | tail -1) - echo " CUDA state: $STATE" - - echo "" - echo "[CRIU Dump] Dumping CPU state..." - if [ -n "$PLUGIN_DIR" ]; then - "$CRIU" dump -t $TEST_PID -D "$CRIU_IMG_DIR" --shell-job -L "$PLUGIN_DIR" 2>/dev/null - else - "$CRIU" dump -t $TEST_PID -D "$CRIU_IMG_DIR" --shell-job 2>/dev/null - fi - echo " Created $(ls "$CRIU_IMG_DIR"/*.img 2>/dev/null | wc -l) image files" -} - -do_restore() { - echo "" - echo "[CRIU Restore] Restoring CPU state..." - if [ -n "$PLUGIN_DIR" ]; then - "$CRIU" restore -d -D "$CRIU_IMG_DIR" --shell-job -L "$PLUGIN_DIR" 2>/dev/null - else - "$CRIU" restore -d -D "$CRIU_IMG_DIR" --shell-job 2>/dev/null - fi - sleep 2 - - # Get restored PID (should be same as original) - if [ "$TEST_TYPE" = "simple" ]; then - NEW_PID=$(pgrep -f test_simple_checkpoint 2>/dev/null | head -1) - else - NEW_PID=$(pgrep -f "python.*test_pytorch" 2>/dev/null | head -1) - fi - - if [ -z "$NEW_PID" ]; then - echo "ERROR: No restored process found" - exit 1 - fi - echo " Restored PID: $NEW_PID" - TEST_PID=$NEW_PID - - echo "" - echo "[Restore CUDA] Restoring CUDA state..." - "$CUDA_CKPT" --action restore --pid $TEST_PID 2>/dev/null - echo " Restore: OK" - - echo "" - echo "[Unlock] Unlocking CUDA..." - "$CUDA_CKPT" --action unlock --pid $TEST_PID 2>/dev/null - - STATE=$("$CUDA_CKPT" --get-state --pid $TEST_PID 2>/dev/null | tail -1) - echo " CUDA state: $STATE" -} - -verify_restored() { - echo "" - echo "[Verify] Checking restored process..." - sleep 3 - - if ps -p $TEST_PID > /dev/null 2>&1; then - echo " Process: ALIVE" - else - echo " Process: DEAD" - cat "$TEST_OUTPUT" - exit 1 - fi - - if grep -q "Verification PASSED" "$TEST_OUTPUT" 2>/dev/null; then - echo " GPU data: VERIFIED" - elif grep -q "output" "$TEST_OUTPUT" 2>/dev/null; then - echo " GPU computation: OK" - fi -} - -final_cleanup() { - kill -9 $TEST_PID 2>/dev/null || true -} - -# ============================================================================ -# Main -# ============================================================================ - -echo "==========================================" -echo " GPU Checkpoint/Restore Test" -echo " Test type: $TEST_TYPE" -echo "==========================================" -echo "" - -cleanup -check_requirements - -if [ "$TEST_TYPE" = "pytorch" ]; then - start_pytorch_test -else - start_simple_test -fi - -verify_running -do_checkpoint -do_restore -verify_restored -final_cleanup - -echo "" -echo "==========================================" -echo " TEST PASSED!" -echo "==========================================" -echo "" -echo "Summary:" -echo " - Process started with GPU" -echo " - cuda-checkpoint lock/checkpoint: OK" -echo " - CRIU dump: OK" -echo " - CRIU restore: OK" -echo " - cuda-checkpoint restore/unlock: OK" -echo " - Process alive after restore: YES" diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_in_docker.sh b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_in_docker.sh deleted file mode 100755 index 4bff0e100d..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_in_docker.sh +++ /dev/null @@ -1,124 +0,0 @@ -#!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# -# Test intercept library inside Docker container -# -# This creates a reproducible test environment with PyTorch/vLLM. -# -# Usage: -# ./tests/test_in_docker.sh # Run all tests -# ./tests/test_in_docker.sh dlsym # Just dlsym test -# ./tests/test_in_docker.sh torch # PyTorch distributed -# ./tests/test_in_docker.sh torch-light # Lightweight mode -# ./tests/test_in_docker.sh vllm # Full vLLM - -set -euo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -LIB_DIR="$(dirname "$SCRIPT_DIR")" - -# Build library locally first -echo "=== Building intercept library ===" -cd "$LIB_DIR" -make clean && make -echo "" - -# Get test type -TEST_TYPE="${1:-all}" - -# Docker image with PyTorch and CUDA -DOCKER_IMAGE="docker.io/pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime" - -# Check for GPU -if ! command -v nvidia-smi &>/dev/null; then - echo "WARNING: nvidia-smi not found, tests may fail without GPU" - GPU_FLAG="" -else - GPU_FLAG="--gpus all" -fi - -run_in_docker() { - local cmd="$1" - docker run --rm -it \ - $GPU_FLAG \ - -v "$LIB_DIR":/nvsnap_intercept:ro \ - -v "$SCRIPT_DIR":/tests:ro \ - -e NVSNAP_LOG_LEVEL=4 \ - -e PYTHONFAULTHANDLER=1 \ - -w /tests \ - "$DOCKER_IMAGE" \ - bash -c "$cmd" -} - -case "$TEST_TYPE" in - dlsym) - echo "=== Testing dlsym in Docker ===" - run_in_docker " - apt-get update -qq && apt-get install -qq -y gcc > /dev/null - gcc -o /tmp/test_dlsym /tests/test_dlsym_basic.c -ldl -lpthread - echo '--- Without intercept ---' - /tmp/test_dlsym - echo '' - echo '--- With intercept ---' - LD_PRELOAD=/nvsnap_intercept/libnvsnap_intercept.so /tmp/test_dlsym - " - ;; - - torch) - echo "=== Testing PyTorch distributed in Docker ===" - run_in_docker " - echo '--- Without intercept ---' - python3 /tests/test_pytorch_distributed.py - echo '' - echo '--- With intercept ---' - LD_PRELOAD=/nvsnap_intercept/libnvsnap_intercept.so python3 /tests/test_pytorch_distributed.py - " - ;; - - torch-light) - echo "=== Testing PyTorch distributed (lightweight mode) ===" - run_in_docker " - NVSNAP_LIGHTWEIGHT=1 LD_PRELOAD=/nvsnap_intercept/libnvsnap_intercept.so \ - python3 /tests/test_pytorch_distributed.py - " - ;; - - vllm) - echo "=== Testing with vLLM ===" - docker run --rm -it \ - $GPU_FLAG \ - -v "$LIB_DIR":/nvsnap_intercept:ro \ - -e NVSNAP_LOG_LEVEL=4 \ - -e PYTHONFAULTHANDLER=1 \ - docker.io/vllm/vllm-openai:v0.6.6.post1 \ - bash -c " - echo '--- Without intercept ---' - python3 -c 'import vllm; print(\"vLLM:\", vllm.__version__)' || echo 'vLLM import OK' - echo '' - echo '--- With intercept ---' - LD_PRELOAD=/nvsnap_intercept/libnvsnap_intercept.so \ - python3 -c 'import vllm; print(\"vLLM:\", vllm.__version__)' || echo 'vLLM import FAILED' - " - ;; - - all) - echo "=== Running all tests in Docker ===" - echo "" - "$0" dlsym - echo "" - "$0" torch - echo "" - "$0" torch-light - ;; - - *) - echo "Unknown test: $TEST_TYPE" - echo "Usage: $0 [dlsym|torch|torch-light|vllm|all]" - exit 1 - ;; -esac - -echo "" -echo "=== Done ===" diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_library_safety.c b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_library_safety.c deleted file mode 100644 index 4972ba563d..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_library_safety.c +++ /dev/null @@ -1,375 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * Comprehensive library safety test suite. - * - * Tests everything our merged library does that could break an application, - * WITHOUT requiring a GPU. Runs locally in <2 seconds. - * - * Tests: - * 1. Library loads and unloads cleanly - * 2. dlsym override: no recursion, passthrough works, hooked symbols returned - * 3. sigaction guard: SIGUSR1/SIGUSR2 handlers installed, guard blocks overrides - * 4. Signal delivery: SIGUSR1 sets quiesce flag, SIGUSR2 sets resume flag - * 5. Quiesce worker thread: starts, polls trigger files - * 6. Fork safety: child process gets fresh state, worker thread restarts - * 7. Constructor ordering: NvSnap (101) before NvSnap (102) before atfork (103) - * 8. ZMQ interception: dlsym("zmq_ctx_new") returns hook if ZMQ loaded - * 9. NCCL symbol routing: dlsym("ncclCommInitRank") returns hook - * 10. Thread safety: concurrent dlsym calls don't crash - * 11. Trigger file mechanism: write file, detect, quiesce fires - * 12. Checkpoint path file: /dev/shm/nvsnap-checkpoint-dir readable - * - * Build: - * gcc -g -o tests/test_library_safety tests/test_library_safety.c -ldl -lpthread - * - * Run: - * LD_PRELOAD=./libnvsnap_intercept.so NVSNAP_QUIESCE_SIGNALS=1 \ - * NVSNAP_NCCL_INTERCEPT=1 NVSNAP_CUDA_INTERCEPT=1 \ - * ./tests/test_library_safety - */ -#define _GNU_SOURCE -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -static int g_pass = 0, g_fail = 0, g_skip = 0; - -#define CHECK(cond, name) do { \ - if (cond) { printf(" PASS: %s\n", name); g_pass++; } \ - else { printf(" FAIL: %s\n", name); g_fail++; } \ -} while(0) - -#define SKIP(name) do { printf(" SKIP: %s\n", name); g_skip++; } while(0) - -/* ─── Test 1: dlsym override safety ─────────────────────────────────── */ - -static void test_dlsym_override(void) { - printf("\n=== Test 1: dlsym Override Safety ===\n"); - - /* Set alarm — any hang = test failure */ - alarm(5); - - /* Hooked CUDA symbols return non-NULL */ - void *p = dlsym(RTLD_DEFAULT, "cudaMalloc"); - CHECK(p != NULL, "dlsym(cudaMalloc) returns non-NULL"); - - p = dlsym(RTLD_DEFAULT, "cuMemMap"); - CHECK(p != NULL, "dlsym(cuMemMap) returns non-NULL"); - - p = dlsym(RTLD_DEFAULT, "ncclCommInitRank"); - CHECK(p != NULL, "dlsym(ncclCommInitRank) returns non-NULL"); - - /* Non-hooked symbols pass through */ - p = dlsym(RTLD_DEFAULT, "printf"); - CHECK(p != NULL, "dlsym(printf) passthrough works"); - - p = dlsym(RTLD_DEFAULT, "nonexistent_symbol_xyz"); - CHECK(p == NULL, "dlsym(nonexistent) returns NULL"); - - /* dlsym(dlsym) doesn't recurse */ - p = dlsym(RTLD_DEFAULT, "dlsym"); - CHECK(p != NULL, "dlsym(dlsym) no infinite recursion"); - - /* Stress test */ - for (int i = 0; i < 10000; i++) { - dlsym(RTLD_DEFAULT, "cudaMalloc"); - dlsym(RTLD_DEFAULT, "printf"); - } - CHECK(1, "20000 rapid dlsym calls without hang"); - - alarm(0); -} - -/* ─── Test 2: nvsnap_resolve_real self-check ──────────────────────── */ - -static void test_resolve_real(void) { - printf("\n=== Test 2: nvsnap_resolve_real Self-Check ===\n"); - - void *(*resolve)(const char *, const char *) = - dlsym(RTLD_DEFAULT, "nvsnap_resolve_real"); - if (!resolve) { - SKIP("nvsnap_resolve_real not found (not merged build?)"); - return; - } - - CHECK(resolve != NULL, "nvsnap_resolve_real exists in merged library"); - - /* resolve_real should find real functions from system libraries */ - void *real_printf = resolve("printf", "libc.so.6"); - CHECK(real_printf != NULL, "resolve_real(printf, libc.so.6) finds real printf"); - - /* resolve_real for CUDA symbols: may be NULL if no GPU libs loaded */ - void *real_malloc = resolve("cudaMalloc", "libcudart.so"); - if (real_malloc) { - /* Verify it's NOT our hook */ - void *our_hook = dlsym(RTLD_DEFAULT, "cudaMalloc"); - CHECK(real_malloc != our_hook, - "resolve_real(cudaMalloc) != dlsym(cudaMalloc) — self-check works"); - } else { - SKIP("cudaMalloc not in libcudart.so (no CUDA runtime loaded)"); - } -} - -/* ─── Test 3: sigaction guard ───────────────────────────────────────── */ - -static volatile sig_atomic_t g_sigusr1_count = 0; -static volatile sig_atomic_t g_sigusr2_count = 0; - -static void test_sigusr1_handler(int sig) { (void)sig; g_sigusr1_count++; } -static void test_sigusr2_handler(int sig) { (void)sig; g_sigusr2_count++; } - -static void test_sigaction_guard(void) { - printf("\n=== Test 3: sigaction Guard ===\n"); - - /* Our library's sigaction interpose should be active */ - void *sa = dlsym(RTLD_DEFAULT, "sigaction"); - CHECK(sa != NULL, "sigaction is available"); - - /* Try to override SIGUSR1 — guard should block if NVSNAP_QUIESCE_SIGNALS=1 */ - struct sigaction act, old; - memset(&act, 0, sizeof(act)); - act.sa_handler = test_sigusr1_handler; - int rc = sigaction(SIGUSR1, &act, &old); - CHECK(rc == 0, "sigaction(SIGUSR1) call succeeds (may be blocked by guard)"); - - /* Send SIGUSR1 to self — should be handled by NVSNAP's handler (if guard active) - * or our test handler (if guard not active) */ - g_sigusr1_count = 0; - kill(getpid(), SIGUSR1); - usleep(10000); /* 10ms for delivery */ - - /* We can't assert which handler ran (depends on NVSNAP_QUIESCE_SIGNALS), - * but the process must survive */ - CHECK(1, "SIGUSR1 delivery did not crash process"); -} - -/* ─── Test 4: Signal delivery and quiesce state ─────────────────────── */ - -static void test_signal_delivery(void) { - printf("\n=== Test 4: Signal Delivery ===\n"); - - /* SIGUSR2 should be handled (noop or resume handler) */ - kill(getpid(), SIGUSR2); - usleep(10000); - CHECK(1, "SIGUSR2 delivery did not crash process"); - - /* Multiple rapid signals */ - for (int i = 0; i < 100; i++) { - kill(getpid(), SIGUSR2); - } - usleep(50000); - CHECK(1, "100 rapid SIGUSR2 signals handled without crash"); -} - -/* ─── Test 5: Fork safety ───────────────────────────────────────────── */ - -static void test_fork_safety(void) { - printf("\n=== Test 5: Fork Safety ===\n"); - - pid_t child = fork(); - if (child == 0) { - /* Child: verify we can dlsym without hanging */ - alarm(3); - void *p = dlsym(RTLD_DEFAULT, "cudaMalloc"); - if (!p) _exit(1); - - /* Verify we can send signals */ - kill(getpid(), SIGUSR2); - usleep(10000); - - _exit(0); - } - - int status; - waitpid(child, &status, 0); - CHECK(WIFEXITED(status) && WEXITSTATUS(status) == 0, - "child process: dlsym + signal works after fork"); -} - -/* ─── Test 6: Trigger file mechanism ────────────────────────────────── */ - -static void test_trigger_file(void) { - printf("\n=== Test 6: Trigger File Mechanism ===\n"); - - /* Write a trigger file for our PID */ - char path[64]; - snprintf(path, sizeof(path), "/dev/shm/nvsnap-quiesce-trigger-%d", getpid()); - int fd = open(path, O_CREAT | O_WRONLY, 0644); - if (fd < 0) { - SKIP("cannot write to /dev/shm (not available?)"); - return; - } - write(fd, "test", 4); - close(fd); - - CHECK(access(path, F_OK) == 0, "trigger file created"); - - /* Wait for worker thread to detect and unlink it (up to 200ms) */ - int detected = 0; - for (int i = 0; i < 20; i++) { - usleep(10000); /* 10ms */ - if (access(path, F_OK) != 0) { - detected = 1; - break; - } - } - CHECK(detected, "worker thread detected and consumed trigger file within 200ms"); - - /* Clean up if not consumed */ - unlink(path); -} - -/* ─── Test 7: Checkpoint path file ──────────────────────────────────── */ - -static void test_checkpoint_path(void) { - printf("\n=== Test 7: Checkpoint Path File ===\n"); - - const char *path = "/dev/shm/nvsnap-checkpoint-dir"; - const char *test_dir = "/tmp/nvsnap-test-checkpoint-12345"; - - /* Write checkpoint path */ - int fd = open(path, O_CREAT | O_WRONLY | O_TRUNC, 0644); - if (fd < 0) { - SKIP("cannot write to /dev/shm"); - return; - } - write(fd, test_dir, strlen(test_dir)); - close(fd); - - /* Read it back */ - char buf[256] = {0}; - fd = open(path, O_RDONLY); - CHECK(fd >= 0, "checkpoint path file readable"); - if (fd >= 0) { - ssize_t n = read(fd, buf, sizeof(buf) - 1); - close(fd); - CHECK(n > 0 && strcmp(buf, test_dir) == 0, - "checkpoint path content matches what was written"); - } - - unlink(path); -} - -/* ─── Test 8: Thread safety ─────────────────────────────────────────── */ - -#define THREAD_COUNT 8 -#define ITERS_PER_THREAD 1000 - -static void *thread_dlsym_stress(void *arg) { - (void)arg; - for (int i = 0; i < ITERS_PER_THREAD; i++) { - dlsym(RTLD_DEFAULT, "cudaMalloc"); - dlsym(RTLD_DEFAULT, "cuMemMap"); - dlsym(RTLD_DEFAULT, "ncclCommInitRank"); - dlsym(RTLD_DEFAULT, "printf"); - dlsym(RTLD_DEFAULT, "nonexistent"); - } - return NULL; -} - -static void test_thread_safety(void) { - printf("\n=== Test 8: Thread Safety ===\n"); - - pthread_t threads[THREAD_COUNT]; - for (int i = 0; i < THREAD_COUNT; i++) - pthread_create(&threads[i], NULL, thread_dlsym_stress, NULL); - for (int i = 0; i < THREAD_COUNT; i++) - pthread_join(threads[i], NULL); - - CHECK(1, "8 threads × 5000 dlsym calls completed without crash"); -} - -/* ─── Test 9: Constructor exports ───────────────────────────────────── */ - -static void test_exports(void) { - printf("\n=== Test 9: Critical Symbol Exports ===\n"); - - /* NvSnap symbols */ - CHECK(dlsym(RTLD_DEFAULT, "nvsnap_nccl_quiesce") != NULL, - "nvsnap_nccl_quiesce exported"); - CHECK(dlsym(RTLD_DEFAULT, "nvsnap_perform_quiescence") != NULL, - "nvsnap_perform_quiescence exported"); - - /* NvSnap symbols */ - CHECK(dlsym(RTLD_DEFAULT, "nvsnap_checkpoint_save") != NULL, - "nvsnap_checkpoint_save exported"); - CHECK(dlsym(RTLD_DEFAULT, "nvsnap_pre_checkpoint_quiesce") != NULL, - "nvsnap_pre_checkpoint_quiesce exported"); - CHECK(dlsym(RTLD_DEFAULT, "nvsnap_post_restore_resume") != NULL, - "nvsnap_post_restore_resume exported"); - CHECK(dlsym(RTLD_DEFAULT, "nvsnap_resolve_real") != NULL, - "nvsnap_resolve_real exported"); - - /* CUDA hooks */ - CHECK(dlsym(RTLD_DEFAULT, "cudaMalloc") != NULL, "cudaMalloc hook exported"); - CHECK(dlsym(RTLD_DEFAULT, "cudaFree") != NULL, "cudaFree hook exported"); - CHECK(dlsym(RTLD_DEFAULT, "cuMemMap") != NULL, "cuMemMap hook exported"); - CHECK(dlsym(RTLD_DEFAULT, "cuMemUnmap") != NULL, "cuMemUnmap hook exported"); - CHECK(dlsym(RTLD_DEFAULT, "cuMemAlloc_v2") != NULL, "cuMemAlloc_v2 hook exported"); - - /* NCCL hooks */ - CHECK(dlsym(RTLD_DEFAULT, "ncclCommInitRank") != NULL, - "ncclCommInitRank hook exported"); - CHECK(dlsym(RTLD_DEFAULT, "ncclCommInitRankConfig") != NULL, - "ncclCommInitRankConfig hook exported"); - - /* sigaction interpose */ - CHECK(dlsym(RTLD_DEFAULT, "sigaction") != NULL, "sigaction interpose exported"); -} - -/* ─── Test 10: No duplicate symbols (single hook per function) ──────── */ - -static void test_no_duplicate_hooks(void) { - printf("\n=== Test 10: No Duplicate Hooks ===\n"); - - /* In a merged build, there should be exactly ONE ncclCommInitRank. - * If both NvSnap and NvSnap export it, the linker picks one but - * the other is dead code — this is OK but we verify which one won. */ - void *nccl_hook = dlsym(RTLD_DEFAULT, "ncclCommInitRank"); - CHECK(nccl_hook != NULL, "ncclCommInitRank has exactly one hook"); - - /* The symbol table should route to the same one */ - void *(*lookup)(const char *) = dlsym(RTLD_DEFAULT, "nvsnap_lookup_symbol"); - if (lookup) { - void *table_nccl = lookup("ncclCommInitRank"); - CHECK(table_nccl == nccl_hook, - "nvsnap_lookup_symbol(ncclCommInitRank) matches PLT symbol"); - } -} - -/* ─── Main ──────────────────────────────────────────────────────────── */ - -int main(void) { - printf("=== Merged Library Safety Test Suite ===\n"); - printf("PID: %d\n", getpid()); - printf("LD_PRELOAD: %s\n", getenv("LD_PRELOAD") ?: "(unset)"); - printf("NVSNAP_QUIESCE_SIGNALS: %s\n", getenv("NVSNAP_QUIESCE_SIGNALS") ?: "(unset)"); - printf("NVSNAP_NCCL_INTERCEPT: %s\n", getenv("NVSNAP_NCCL_INTERCEPT") ?: "(unset)"); - - test_dlsym_override(); - test_resolve_real(); - test_sigaction_guard(); - test_signal_delivery(); - test_fork_safety(); - test_trigger_file(); - test_checkpoint_path(); - test_thread_safety(); - test_exports(); - test_no_duplicate_hooks(); - - printf("\n=== Results: %d passed, %d failed, %d skipped ===\n", - g_pass, g_fail, g_skip); - return g_fail > 0 ? 1 : 0; -} diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_libuv_intercept.c b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_libuv_intercept.c deleted file mode 100644 index 544cc914b7..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_libuv_intercept.c +++ /dev/null @@ -1,94 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * Test libuv interception with explicit loop initialization. - * - * Compile: gcc -o test_libuv_intercept test_libuv_intercept.c -luv - * Run: NVSNAP_LOG_LEVEL=3 LD_PRELOAD=../libnvsnap_intercept.so ./test_libuv_intercept - */ - -#include -#include -#include - -static int timer_count = 0; - -static void timer_callback(uv_timer_t* handle) { - timer_count++; - printf("Timer fired! count=%d\n", timer_count); - if (timer_count >= 3) { - uv_timer_stop(handle); - uv_close((uv_handle_t*)handle, NULL); - } -} - -int main() { - int r; - - printf("=== Libuv Interception Test ===\n"); - fflush(stdout); - - /* Allocate and explicitly init a loop (tests uv_loop_init interception) */ - uv_loop_t* loop = malloc(sizeof(uv_loop_t)); - if (!loop) { - fprintf(stderr, "Failed to allocate loop\n"); - return 1; - } - - printf("Calling uv_loop_init()...\n"); - fflush(stdout); - r = uv_loop_init(loop); - if (r != 0) { - fprintf(stderr, "uv_loop_init failed: %s (code %d)\n", uv_strerror(r), r); - free(loop); - return 1; - } - printf("Loop initialized: %p\n", (void*)loop); - fflush(stdout); - - /* Create a timer (tests uv_timer_init interception) */ - uv_timer_t* timer = malloc(sizeof(uv_timer_t)); - printf("Calling uv_timer_init()...\n"); - fflush(stdout); - r = uv_timer_init(loop, timer); - if (r != 0) { - fprintf(stderr, "uv_timer_init failed: %s\n", uv_strerror(r)); - return 1; - } - printf("Timer initialized\n"); - fflush(stdout); - - /* Start timer - fire every 50ms, 3 times */ - printf("Calling uv_timer_start()...\n"); - fflush(stdout); - r = uv_timer_start(timer, timer_callback, 50, 50); - if (r != 0) { - fprintf(stderr, "uv_timer_start failed: %s\n", uv_strerror(r)); - return 1; - } - printf("Timer started, running loop...\n"); - fflush(stdout); - - /* Run the event loop (tests uv_run interception) */ - printf("Calling uv_run()...\n"); - fflush(stdout); - r = uv_run(loop, UV_RUN_DEFAULT); - printf("Loop exited with: %d\n", r); - fflush(stdout); - - /* Cleanup */ - printf("Calling uv_loop_close()...\n"); - fflush(stdout); - r = uv_loop_close(loop); - if (r != 0) { - fprintf(stderr, "uv_loop_close failed: %s\n", uv_strerror(r)); - } - free(timer); - free(loop); - - printf("=== Test passed ===\n"); - return 0; -} diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_libuv_simple.c b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_libuv_simple.c deleted file mode 100644 index 497a215a2b..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_libuv_simple.c +++ /dev/null @@ -1,60 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * Simple libuv test to verify interception works with dynamically linked libuv. - * - * Compile: gcc -o test_libuv_simple test_libuv_simple.c -luv - * Run: LD_PRELOAD=../libnvsnap_intercept.so ./test_libuv_simple - */ - -#include -#include -#include - -static void timer_callback(uv_timer_t* handle) { - printf("Timer fired!\n"); - uv_timer_stop(handle); - uv_stop(uv_default_loop()); -} - -int main() { - printf("=== Simple libuv test ===\n"); - - /* Initialize the default loop */ - uv_loop_t* loop = uv_default_loop(); - if (!loop) { - fprintf(stderr, "Failed to get default loop\n"); - return 1; - } - printf("Loop initialized: %p\n", (void*)loop); - - /* Create a timer */ - uv_timer_t timer; - int r = uv_timer_init(loop, &timer); - if (r != 0) { - fprintf(stderr, "uv_timer_init failed: %s\n", uv_strerror(r)); - return 1; - } - printf("Timer initialized\n"); - - /* Start timer - fire after 100ms */ - r = uv_timer_start(&timer, timer_callback, 100, 0); - if (r != 0) { - fprintf(stderr, "uv_timer_start failed: %s\n", uv_strerror(r)); - return 1; - } - printf("Timer started, running loop...\n"); - - /* Run the event loop */ - r = uv_run(loop, UV_RUN_DEFAULT); - printf("Loop exited with: %d\n", r); - - /* Cleanup */ - uv_loop_close(loop); - - printf("=== Test passed ===\n"); - return 0; -} diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_local_checkpoint.sh b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_local_checkpoint.sh deleted file mode 100755 index baaad0e672..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_local_checkpoint.sh +++ /dev/null @@ -1,242 +0,0 @@ -#!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# -# Local checkpoint/restore test WITHOUT Kubernetes -# -# This tests our intercept library with CRIU directly. -# Much faster iteration than deploying to k8s. -# -# Prerequisites: -# - CRIU installed (apt install criu or from our fork) -# - cuda-checkpoint tool -# - GPU available -# - Root privileges (for CRIU) -# -# Usage: -# sudo ./test_local_checkpoint.sh # Full test -# sudo ./test_local_checkpoint.sh --no-gpu # CPU-only test (no cuda-checkpoint) - -set -euo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -LIB_DIR="$(dirname "$SCRIPT_DIR")" -PROJECT_ROOT="$(dirname "$(dirname "$LIB_DIR")")" -LIB_PATH="${LIB_DIR}/libnvsnap_intercept.so" - -# Checkpoint directory -CKPT_DIR="/tmp/nvsnap-local-test-$$" - -# Options -USE_GPU=1 -for arg in "$@"; do - case $arg in - --no-gpu) USE_GPU=0 ;; - esac -done - -cleanup() { - echo "=== Cleanup ===" - # Kill any remaining test processes - pkill -f "test_checkpoint_app" 2>/dev/null || true - rm -rf "$CKPT_DIR" -} -trap cleanup EXIT - -echo "========================================" -echo " Local Checkpoint/Restore Test" -echo "========================================" -echo "" -echo "Library: $LIB_PATH" -echo "Checkpoint dir: $CKPT_DIR" -echo "GPU mode: $USE_GPU" -echo "" - -# Check prerequisites -if [ ! -f "$LIB_PATH" ]; then - echo "Building intercept library..." - cd "$LIB_DIR" && make -fi - -if ! command -v criu &>/dev/null; then - echo "ERROR: criu not found. Install with: apt install criu" - exit 1 -fi - -if [ "$USE_GPU" -eq 1 ]; then - if ! command -v nvidia-smi &>/dev/null; then - echo "WARNING: nvidia-smi not found, disabling GPU mode" - USE_GPU=0 - fi -fi - -mkdir -p "$CKPT_DIR" - -# Create a simple test application -cat > /tmp/test_checkpoint_app.py << 'PYTHON' -#!/usr/bin/env python3 -"""Simple app for checkpoint/restore testing""" -import os -import sys -import time -import signal - -# Track state -counter = 0 -restored = os.environ.get("NVSNAP_RESTORED") == "1" - -def signal_handler(sig, frame): - print(f"[APP] Received signal {sig}, counter={counter}") - if sig == signal.SIGUSR1: - print("[APP] Preparing for checkpoint...") - sys.stdout.flush() - -signal.signal(signal.SIGUSR1, signal_handler) -signal.signal(signal.SIGUSR2, signal_handler) - -print(f"[APP] Started, PID={os.getpid()}, restored={restored}") -sys.stdout.flush() - -# If GPU available, do some CUDA work -try: - import torch - if torch.cuda.is_available(): - print(f"[APP] CUDA available: {torch.cuda.get_device_name(0)}") - x = torch.zeros(1000, device='cuda') - print(f"[APP] Created CUDA tensor: {x.shape}") -except ImportError: - print("[APP] PyTorch not available, CPU-only mode") -except Exception as e: - print(f"[APP] CUDA init error: {e}") - -# Main loop -print("[APP] Entering main loop...") -while True: - counter += 1 - if counter % 10 == 0: - print(f"[APP] Counter: {counter}") - sys.stdout.flush() - time.sleep(0.5) -PYTHON -chmod +x /tmp/test_checkpoint_app.py - -echo "=== Step 1: Starting test application ===" -cd /tmp - -# Start the app with our intercept library -NVSNAP_LOG_LEVEL=3 LD_PRELOAD="$LIB_PATH" \ - python3 /tmp/test_checkpoint_app.py & -APP_PID=$! - -echo "App PID: $APP_PID" -sleep 3 # Let it run a bit - -# Verify it's running -if ! kill -0 $APP_PID 2>/dev/null; then - echo "ERROR: App failed to start" - exit 1 -fi -echo "App is running" - -echo "" -echo "=== Step 2: Sending quiesce signal (SIGUSR1) ===" -kill -USR1 $APP_PID -sleep 1 - -echo "" -echo "=== Step 3: Creating checkpoint ===" - -# GPU checkpoint (if enabled) -if [ "$USE_GPU" -eq 1 ]; then - CUDA_CKPT="${PROJECT_ROOT}/bin/criu-bundle/cuda-checkpoint" - if [ -f "$CUDA_CKPT" ]; then - echo "Locking GPU with cuda-checkpoint..." - $CUDA_CKPT --lock --pid $APP_PID || echo "cuda-checkpoint lock failed (may be OK if no GPU context)" - else - echo "WARNING: cuda-checkpoint not found at $CUDA_CKPT" - fi -fi - -# CRIU dump -echo "Running CRIU dump..." -criu dump \ - --tree $APP_PID \ - --images-dir "$CKPT_DIR" \ - --leave-running \ - --shell-job \ - --tcp-established \ - -v4 \ - -o "$CKPT_DIR/dump.log" || { - echo "CRIU dump failed! Log:" - cat "$CKPT_DIR/dump.log" | tail -50 - exit 1 - } - -echo "Checkpoint created successfully!" -ls -la "$CKPT_DIR" - -# GPU unlock (if enabled) -if [ "$USE_GPU" -eq 1 ] && [ -f "$CUDA_CKPT" ]; then - echo "Unlocking GPU..." - $CUDA_CKPT --unlock --pid $APP_PID || true -fi - -# Send resume signal -echo "" -echo "=== Step 4: Sending resume signal (SIGUSR2) ===" -kill -USR2 $APP_PID -sleep 2 - -echo "" -echo "=== Step 5: Stopping original process ===" -kill $APP_PID 2>/dev/null || true -wait $APP_PID 2>/dev/null || true -sleep 1 - -echo "" -echo "=== Step 6: Restoring from checkpoint ===" - -# Set restored flag -export NVSNAP_RESTORED=1 -export NVSNAP_LOG_LEVEL=3 -export LD_PRELOAD="$LIB_PATH" - -# CRIU restore -echo "Running CRIU restore..." -criu restore \ - --images-dir "$CKPT_DIR" \ - --shell-job \ - --tcp-established \ - -v4 \ - -o "$CKPT_DIR/restore.log" & -RESTORE_PID=$! - -sleep 3 - -# Check if restored process is running -if kill -0 $RESTORE_PID 2>/dev/null; then - echo "Restored process is running!" - - # Let it run a bit - sleep 5 - - # Check it's still alive - if kill -0 $RESTORE_PID 2>/dev/null; then - echo "" - echo "========================================" - echo " SUCCESS: Checkpoint/Restore works!" - echo "========================================" - - # Cleanup - kill $RESTORE_PID 2>/dev/null || true - else - echo "ERROR: Restored process died!" - cat "$CKPT_DIR/restore.log" | tail -30 - exit 1 - fi -else - echo "ERROR: Restore failed!" - cat "$CKPT_DIR/restore.log" | tail -50 - exit 1 -fi diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_local_criu_uvloop.py b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_local_criu_uvloop.py deleted file mode 100644 index a4e76ab49a..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_local_criu_uvloop.py +++ /dev/null @@ -1,273 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -Local CRIU checkpoint/restore test with uvloop. - -This test: -1. Starts a uvloop-based async server -2. Checkpoints it with CRIU -3. Restores it with seccomp interception enabled -4. Verifies io_uring calls are intercepted - -Run this script with sudo (CRIU requires root for some operations). - -Usage: - sudo python3 tests/test_local_criu_uvloop.py -""" - -import os -import sys -import time -import signal -import asyncio -import tempfile -import subprocess -import shutil - -# Check if we're running as a restored process -IS_RESTORED = os.environ.get('GPUCR_POST_RESTORE') == '1' -CHECKPOINT_READY = '/tmp/gpucr_test_ready' -CHECKPOINT_DIR = '/tmp/gpucr_test_checkpoint' - -def log(msg): - """Simple logging with timestamp""" - import datetime - ts = datetime.datetime.now().strftime('%H:%M:%S.%f')[:-3] - print(f"[{ts}] [TEST] {msg}", flush=True) - -async def run_server(): - """Simple async server that uses io_uring via uvloop""" - - log(f"Server starting (PID={os.getpid()}, restored={IS_RESTORED})") - - counter = 0 - - while True: - counter += 1 - log(f"Server tick #{counter}") - - # Signal ready for checkpoint after first tick - if counter == 1 and not IS_RESTORED: - log("Creating ready marker for checkpoint") - with open(CHECKPOINT_READY, 'w') as f: - f.write(str(os.getpid())) - - # Do some async I/O (triggers io_uring) - await asyncio.sleep(1.0) - - # If restored, exit after a few ticks to verify it worked - if IS_RESTORED and counter >= 3: - log("Restored server completed successfully!") - return - -def run_target_process(): - """Run the target process that will be checkpointed""" - try: - import uvloop - log(f"uvloop version: {uvloop.__version__}") - uvloop.install() - except ImportError: - log("WARNING: uvloop not available, using default asyncio") - - try: - asyncio.run(run_server()) - except KeyboardInterrupt: - log("Server interrupted") - -def do_checkpoint(pid, checkpoint_dir): - """Checkpoint the target process using CRIU""" - - criu_path = os.path.join(os.path.dirname(__file__), '..', '..', '..', 'bin', 'criu') - if not os.path.exists(criu_path): - criu_path = 'criu' # Try system criu - - log(f"Checkpointing PID {pid} to {checkpoint_dir}") - - cmd = [ - criu_path, 'dump', - '-t', str(pid), - '-D', checkpoint_dir, - '--shell-job', - '--tcp-close', - '-v4', - '-o', os.path.join(checkpoint_dir, 'dump.log'), - ] - - log(f"Running: {' '.join(cmd)}") - - result = subprocess.run(cmd, capture_output=True, text=True) - - if result.returncode != 0: - log(f"CRIU dump failed: {result.stderr}") - return False - - log("Checkpoint successful!") - return True - -def do_restore(checkpoint_dir): - """Restore the process with seccomp interception""" - - criu_path = os.path.join(os.path.dirname(__file__), '..', '..', '..', 'bin', 'criu') - if not os.path.exists(criu_path): - criu_path = 'criu' - - lib_path = os.path.join(os.path.dirname(__file__), '..', 'libgpucr_intercept.so') - - log(f"Restoring from {checkpoint_dir} with seccomp interception") - - # Set environment for restored process - env = os.environ.copy() - env['GPUCR_POST_RESTORE'] = '1' - env['GPUCR_SECCOMP_ENABLED'] = '1' - env['GPUCR_LOG_LEVEL'] = '3' - env['LD_PRELOAD'] = lib_path - - cmd = [ - criu_path, 'restore', - '-D', checkpoint_dir, - '--shell-job', - '-v4', - '-o', os.path.join(checkpoint_dir, 'restore.log'), - ] - - log(f"Running: {' '.join(cmd)}") - log(f"With LD_PRELOAD={lib_path}") - log(f"With GPUCR_SECCOMP_ENABLED=1 GPUCR_POST_RESTORE=1") - - result = subprocess.run(cmd, env=env, capture_output=True, text=True) - - if result.returncode != 0: - log(f"CRIU restore failed: {result.stderr}") - # Print restore log - restore_log = os.path.join(checkpoint_dir, 'restore.log') - if os.path.exists(restore_log): - log("=== Restore log (last 50 lines) ===") - with open(restore_log) as f: - lines = f.readlines() - for line in lines[-50:]: - print(line.rstrip()) - return False - - log("Restore completed!") - return True - -def main(): - """Main test driver""" - - # Check if we're the restored process - if IS_RESTORED: - log("Running as restored process") - run_target_process() - return 0 - - log("=" * 60) - log("Local CRIU + uvloop + seccomp interception test") - log("=" * 60) - - # Check if running as root - if os.geteuid() != 0: - log("ERROR: This test requires root for CRIU operations") - log("Run with: sudo python3 tests/test_local_criu_uvloop.py") - return 1 - - # Clean up previous state - if os.path.exists(CHECKPOINT_READY): - os.remove(CHECKPOINT_READY) - if os.path.exists(CHECKPOINT_DIR): - shutil.rmtree(CHECKPOINT_DIR) - os.makedirs(CHECKPOINT_DIR, exist_ok=True) - - # Start target process - log("Starting target process...") - - lib_path = os.path.join(os.path.dirname(__file__), '..', 'libgpucr_intercept.so') - venv_python = os.path.join(os.path.dirname(__file__), '..', 'venv', 'bin', 'python3') - - if not os.path.exists(venv_python): - venv_python = sys.executable - - env = os.environ.copy() - env['GPUCR_SECCOMP_ENABLED'] = '1' - env['GPUCR_LOG_LEVEL'] = '3' - env['LD_PRELOAD'] = lib_path - - target = subprocess.Popen( - [venv_python, __file__], - env=env, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - ) - - log(f"Target process started: PID={target.pid}") - - # Wait for ready marker - log("Waiting for target to be ready...") - for i in range(30): - if os.path.exists(CHECKPOINT_READY): - with open(CHECKPOINT_READY) as f: - actual_pid = int(f.read().strip()) - log(f"Target ready (PID from marker: {actual_pid})") - break - time.sleep(0.5) - else: - log("ERROR: Timeout waiting for target to be ready") - target.kill() - return 1 - - # Give it a moment to settle - time.sleep(0.5) - - # Show target output so far - log("=== Target output before checkpoint ===") - # Non-blocking read - import select - while select.select([target.stdout], [], [], 0)[0]: - line = target.stdout.readline() - if line: - print(line.decode().rstrip()) - else: - break - log("=== End target output ===") - - # Checkpoint - if not do_checkpoint(actual_pid, CHECKPOINT_DIR): - target.kill() - return 1 - - # Target should have been killed by checkpoint - target.wait() - log("Target process terminated by checkpoint") - - # Wait a moment - time.sleep(1) - - # Restore - log("=" * 40) - log("Restoring process with seccomp interception...") - log("=" * 40) - - if not do_restore(CHECKPOINT_DIR): - return 1 - - log("=" * 60) - log("TEST PASSED!") - log("=" * 60) - - return 0 - -if __name__ == '__main__': - sys.exit(main()) diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_local_vllm.sh b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_local_vllm.sh deleted file mode 100755 index 915961cfaf..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_local_vllm.sh +++ /dev/null @@ -1,110 +0,0 @@ -#!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# -# Local vLLM test with libnvsnap_intercept.so -# -# This script tests the intercept library locally without Kubernetes. -# Much faster iteration for debugging segfaults and other issues. -# -# Prerequisites: -# - CUDA installed locally -# - vLLM installed: pip install vllm -# - GPU available -# -# Usage: -# ./test_local_vllm.sh # Run with intercept library -# ./test_local_vllm.sh --no-preload # Run without intercept (baseline) -# ./test_local_vllm.sh --gdb # Run with gdb for debugging -# ./test_local_vllm.sh --strace # Run with strace -# ./test_local_vllm.sh --lightweight # Run with NVSNAP_LIGHTWEIGHT=1 - -set -euo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -LIB_DIR="$(dirname "$SCRIPT_DIR")" -LIB_PATH="${LIB_DIR}/libnvsnap_intercept.so" - -# Parse arguments -USE_PRELOAD=1 -USE_GDB=0 -USE_STRACE=0 -LIGHTWEIGHT=0 - -for arg in "$@"; do - case $arg in - --no-preload) USE_PRELOAD=0 ;; - --gdb) USE_GDB=1 ;; - --strace) USE_STRACE=1 ;; - --lightweight) LIGHTWEIGHT=1 ;; - --help|-h) - echo "Usage: $0 [OPTIONS]" - echo "" - echo "Options:" - echo " --no-preload Run without LD_PRELOAD (baseline test)" - echo " --gdb Run under gdb for debugging" - echo " --strace Run with strace" - echo " --lightweight Disable io_uring/libuv interception" - exit 0 - ;; - esac -done - -# Build library if needed -echo "=== Building intercept library ===" -cd "$LIB_DIR" -make clean && make -echo "" - -# Check if library exists -if [ ! -f "$LIB_PATH" ]; then - echo "ERROR: Library not found at $LIB_PATH" - exit 1 -fi - -echo "=== Library built: $LIB_PATH ===" -ls -la "$LIB_PATH" -echo "" - -# Setup environment -export CUDA_VISIBLE_DEVICES=0 -export NVSNAP_LOG_LEVEL=4 # Debug level -export PYTHONFAULTHANDLER=1 -export PYTHONUNBUFFERED=1 - -if [ $LIGHTWEIGHT -eq 1 ]; then - export NVSNAP_LIGHTWEIGHT=1 - echo "=== Lightweight mode enabled (no io_uring/libuv interception) ===" -fi - -# Build the command -VLLM_CMD="python3 -m vllm.entrypoints.openai.api_server \ - --model TinyLlama/TinyLlama-1.1B-Chat-v1.0 \ - --host 127.0.0.1 \ - --port 8000 \ - --max-model-len 512 \ - --gpu-memory-utilization 0.3" - -if [ $USE_PRELOAD -eq 1 ]; then - export LD_PRELOAD="$LIB_PATH" - echo "=== Running with LD_PRELOAD=$LIB_PATH ===" -else - echo "=== Running WITHOUT LD_PRELOAD (baseline) ===" -fi - -echo "" -echo "=== Starting vLLM ===" -echo "Command: $VLLM_CMD" -echo "" - -if [ $USE_GDB -eq 1 ]; then - echo "=== Running under GDB ===" - echo "Type 'run' to start, 'bt' for backtrace on crash" - gdb -ex "set follow-fork-mode child" -ex "run" --args $VLLM_CMD -elif [ $USE_STRACE -eq 1 ]; then - echo "=== Running with strace ===" - strace -f -e trace=openat,socket,connect,bind $VLLM_CMD -else - $VLLM_CMD -fi diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_pytorch.py b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_pytorch.py deleted file mode 100644 index 2c56642709..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_pytorch.py +++ /dev/null @@ -1,287 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -Test: PyTorch Interception - -This is the CRITICAL test - if interception works here, it works for real workloads. - -This test progressively verifies: -1. Basic tensor creation (cudaMalloc) -2. Tensor operations (kernel launches) -3. cuBLAS operations (torch.mm) -4. cuDNN operations (torch.nn.Conv2d) -5. Multiple allocations and frees -6. Autograd backward pass - -Run with: - GPUCR_LOG_LEVEL=3 LD_PRELOAD=./libgpucr_intercept.so python3 test_pytorch.py -""" - -import sys -import os - -def check_gpucr_loaded(): - """Check if GPUCR interception library is loaded.""" - preload = os.environ.get('LD_PRELOAD', '') - if 'libgpucr_intercept' not in preload: - print("WARNING: GPUCR interception library not loaded!") - print("Run with: LD_PRELOAD=./libgpucr_intercept.so python3 test_pytorch.py") - return False - return True - -def test_basic_tensor(): - """Test 1: Basic tensor allocation.""" - print("\n=== Test 1: Basic Tensor Allocation ===") - import torch - - print(f"PyTorch version: {torch.__version__}") - print(f"CUDA available: {torch.cuda.is_available()}") - - if not torch.cuda.is_available(): - print("CUDA not available, skipping GPU tests") - return False - - print(f"CUDA device: {torch.cuda.get_device_name(0)}") - - # Allocate tensor - x = torch.zeros(1024, 1024, device='cuda') - print(f"Allocated tensor: shape={x.shape}, dtype={x.dtype}, device={x.device}") - print(f"Memory: {x.numel() * x.element_size()} bytes") - - # Verify it works - x.fill_(42.0) - assert x.mean().item() == 42.0, "Tensor verification failed" - print("Tensor verification: PASSED") - - del x - torch.cuda.synchronize() - print("Freed tensor") - - return True - -def test_tensor_operations(): - """Test 2: Tensor operations (kernel launches).""" - print("\n=== Test 2: Tensor Operations ===") - import torch - - a = torch.randn(1000, 1000, device='cuda') - b = torch.randn(1000, 1000, device='cuda') - - # Element-wise operations - c = a + b - d = a * b - e = torch.relu(c) - - print(f"Element-wise ops: shape={e.shape}") - - # Reduction - mean = e.mean() - print(f"Mean: {mean.item():.4f}") - - del a, b, c, d, e - torch.cuda.synchronize() - print("Tensor operations: PASSED") - - return True - -def test_cublas(): - """Test 3: cuBLAS operations (matrix multiply).""" - print("\n=== Test 3: cuBLAS Operations (torch.mm) ===") - import torch - - # Matrix multiply uses cuBLAS - a = torch.randn(512, 512, device='cuda') - b = torch.randn(512, 512, device='cuda') - - c = torch.mm(a, b) - print(f"Matrix multiply: {a.shape} x {b.shape} = {c.shape}") - - # Batched matmul - a_batch = torch.randn(8, 64, 128, device='cuda') - b_batch = torch.randn(8, 128, 64, device='cuda') - c_batch = torch.bmm(a_batch, b_batch) - print(f"Batched matmul: {a_batch.shape} x {b_batch.shape} = {c_batch.shape}") - - del a, b, c, a_batch, b_batch, c_batch - torch.cuda.synchronize() - print("cuBLAS operations: PASSED") - - return True - -def test_cudnn(): - """Test 4: cuDNN operations (convolution).""" - print("\n=== Test 4: cuDNN Operations (Conv2d) ===") - import torch - import torch.nn as nn - - # Convolution uses cuDNN - conv = nn.Conv2d(3, 64, kernel_size=3, padding=1).cuda() - x = torch.randn(1, 3, 224, 224, device='cuda') - - y = conv(x) - print(f"Conv2d: {x.shape} -> {y.shape}") - - # BatchNorm also uses cuDNN - bn = nn.BatchNorm2d(64).cuda() - y = bn(y) - print(f"BatchNorm2d: {y.shape}") - - # MaxPool - pool = nn.MaxPool2d(2).cuda() - y = pool(y) - print(f"MaxPool2d: {y.shape}") - - del conv, bn, pool, x, y - torch.cuda.synchronize() - print("cuDNN operations: PASSED") - - return True - -def test_autograd(): - """Test 5: Autograd backward pass.""" - print("\n=== Test 5: Autograd Backward Pass ===") - import torch - import torch.nn as nn - - # Simple model - model = nn.Sequential( - nn.Linear(100, 50), - nn.ReLU(), - nn.Linear(50, 10), - ).cuda() - - x = torch.randn(32, 100, device='cuda', requires_grad=True) - y = torch.randint(0, 10, (32,), device='cuda') - - # Forward - output = model(x) - loss = nn.functional.cross_entropy(output, y) - print(f"Forward pass: loss={loss.item():.4f}") - - # Backward - loss.backward() - print(f"Backward pass completed") - print(f"Input gradient shape: {x.grad.shape}") - - del model, x, y, output, loss - torch.cuda.synchronize() - print("Autograd: PASSED") - - return True - -def test_memory_stress(): - """Test 6: Memory allocation stress test.""" - print("\n=== Test 6: Memory Allocation Stress ===") - import torch - - tensors = [] - total_bytes = 0 - - # Allocate many tensors - for i in range(100): - size = (256 + i * 10, 256 + i * 10) - t = torch.randn(*size, device='cuda') - tensors.append(t) - total_bytes += t.numel() * t.element_size() - - print(f"Allocated {len(tensors)} tensors, total {total_bytes / 1024 / 1024:.1f} MB") - - # Free half - for t in tensors[:50]: - del t - tensors = tensors[50:] - torch.cuda.synchronize() - print(f"Freed 50 tensors, {len(tensors)} remaining") - - # Allocate more - for i in range(50): - t = torch.randn(512, 512, device='cuda') - tensors.append(t) - print(f"Allocated 50 more, {len(tensors)} total") - - # Free all - del tensors - torch.cuda.synchronize() - torch.cuda.empty_cache() - print("Freed all tensors") - - print("Memory stress test: PASSED") - return True - -def print_memory_stats(): - """Print CUDA memory statistics.""" - import torch - if torch.cuda.is_available(): - print("\n=== CUDA Memory Stats ===") - print(f"Allocated: {torch.cuda.memory_allocated() / 1024 / 1024:.1f} MB") - print(f"Cached: {torch.cuda.memory_reserved() / 1024 / 1024:.1f} MB") - print(f"Max allocated: {torch.cuda.max_memory_allocated() / 1024 / 1024:.1f} MB") - -def main(): - print("=" * 60) - print("GPUCR PyTorch Interception Test") - print("=" * 60) - - gpucr_loaded = check_gpucr_loaded() - - tests = [ - ("Basic Tensor", test_basic_tensor), - ("Tensor Operations", test_tensor_operations), - ("cuBLAS (torch.mm)", test_cublas), - ("cuDNN (Conv2d)", test_cudnn), - ("Autograd", test_autograd), - ("Memory Stress", test_memory_stress), - ] - - results = {} - - for name, test_fn in tests: - try: - if test_fn(): - results[name] = "PASSED" - else: - results[name] = "SKIPPED" - except Exception as e: - results[name] = f"FAILED: {e}" - import traceback - traceback.print_exc() - - print_memory_stats() - - print("\n" + "=" * 60) - print("SUMMARY") - print("=" * 60) - for name, result in results.items(): - status = "✓" if result == "PASSED" else ("○" if result == "SKIPPED" else "✗") - print(f" {status} {name}: {result}") - - failed = sum(1 for r in results.values() if r.startswith("FAILED")) - if failed > 0: - print(f"\n{failed} test(s) FAILED") - return 1 - - if gpucr_loaded: - print("\n✓ All tests passed with GPUCR interception!") - print(" Check the logs above to verify allocations were tracked.") - else: - print("\n○ Tests passed but GPUCR was not loaded.") - print(" Run with LD_PRELOAD to verify interception.") - - return 0 - -if __name__ == "__main__": - sys.exit(main()) diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_pytorch_distributed.py b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_pytorch_distributed.py deleted file mode 100644 index 15664a4330..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_pytorch_distributed.py +++ /dev/null @@ -1,104 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -Test PyTorch distributed init with libgpucr_intercept.so - -This isolates the exact code path that segfaults in vLLM: -- torch.distributed.init_process_group() -- _create_c10d_store() for TCP store - -Run with: - # Without intercept (baseline) - python3 test_pytorch_distributed.py - - # With intercept library - LD_PRELOAD=../libgpucr_intercept.so python3 test_pytorch_distributed.py - - # With intercept + debug - GPUCR_LOG_LEVEL=4 LD_PRELOAD=../libgpucr_intercept.so python3 test_pytorch_distributed.py - - # Under gdb - GPUCR_LOG_LEVEL=4 gdb -ex run --args python3 test_pytorch_distributed.py -""" - -import os -import sys -import socket - -print(f"PID: {os.getpid()}") -print(f"LD_PRELOAD: {os.environ.get('LD_PRELOAD', 'not set')}") -print() - -# Check if CUDA is available -try: - import torch - print(f"PyTorch version: {torch.__version__}") - print(f"CUDA available: {torch.cuda.is_available()}") - if torch.cuda.is_available(): - print(f"CUDA device: {torch.cuda.get_device_name(0)}") -except ImportError: - print("ERROR: PyTorch not installed") - sys.exit(1) - -print() - -# Find a free port -def find_free_port(): - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(('127.0.0.1', 0)) - return s.getsockname()[1] - -port = find_free_port() -print(f"Using port: {port}") - -# This is the exact code path that segfaults -print() -print("=== Testing torch.distributed.init_process_group() ===") -print("This creates a TCP store - the exact point where vLLM segfaults") -print() - -try: - # Set up environment for single-process distributed - os.environ['MASTER_ADDR'] = '127.0.0.1' - os.environ['MASTER_PORT'] = str(port) - os.environ['RANK'] = '0' - os.environ['WORLD_SIZE'] = '1' - - print("Calling torch.distributed.init_process_group('gloo')...") - - # This is where vLLM segfaults - torch.distributed.init_process_group( - backend='gloo', - init_method=f'tcp://127.0.0.1:{port}', - rank=0, - world_size=1 - ) - - print("SUCCESS: init_process_group completed!") - - # Clean up - torch.distributed.destroy_process_group() - print("Cleaned up process group") - -except Exception as e: - print(f"FAILED: {type(e).__name__}: {e}") - import traceback - traceback.print_exc() - sys.exit(1) - -print() -print("=== All tests passed ===") diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_quiesce.py b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_quiesce.py deleted file mode 100644 index 37f8819b8b..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_quiesce.py +++ /dev/null @@ -1,197 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -Test quiescence functionality of libgpucr_intercept.so - -This test: -1. Creates io_uring instances (via uvloop or direct) -2. Sends SIGUSR1 to trigger quiesce -3. Verifies that I/O is drained -4. Sends SIGUSR2 to resume - -Usage: - GPUCR_LOG_LEVEL=3 LD_PRELOAD=./libgpucr_intercept.so python3 tests/test_quiesce.py -""" - -import os -import sys -import time -import signal -import threading - -def test_asyncio_basic(): - """Test basic asyncio quiescence""" - import asyncio - - print("[TEST] Basic asyncio quiescence") - - # Create some async work - completed = [] - - async def background_work(n): - for i in range(5): - await asyncio.sleep(0.1) - completed.append(f"task-{n}-{i}") - return f"task-{n}-done" - - async def main(): - # Start background tasks - tasks = [asyncio.create_task(background_work(i)) for i in range(3)] - - # Let them run a bit - await asyncio.sleep(0.2) - - print(f"[TEST] Before quiesce: {len(completed)} items completed") - - # Trigger quiesce via SIGUSR1 (from another thread) - def send_quiesce(): - time.sleep(0.1) - print("[TEST] Sending SIGUSR1 (quiesce)") - os.kill(os.getpid(), signal.SIGUSR1) - time.sleep(0.5) # Wait for quiesce to complete - print("[TEST] Sending SIGUSR2 (resume)") - os.kill(os.getpid(), signal.SIGUSR2) - - t = threading.Thread(target=send_quiesce) - t.start() - - # Wait for tasks (should resume after SIGUSR2) - results = await asyncio.gather(*tasks) - t.join() - - print(f"[TEST] After resume: {len(completed)} items completed") - print(f"[TEST] Results: {results}") - - return True - - return asyncio.run(main()) - -def test_uvloop_quiesce(): - """Test uvloop quiescence (if available)""" - try: - import uvloop - except ImportError: - print("[TEST] uvloop not installed, skipping uvloop test") - return True - - import asyncio - - print("[TEST] uvloop quiescence") - - # Install uvloop - uvloop.install() - - async def uvloop_work(): - results = [] - for i in range(10): - await asyncio.sleep(0.05) - results.append(i) - return results - - async def main(): - # Create some work - task = asyncio.create_task(uvloop_work()) - - # Trigger quiesce - def send_quiesce(): - time.sleep(0.1) - print("[TEST] Sending SIGUSR1 to uvloop") - os.kill(os.getpid(), signal.SIGUSR1) - time.sleep(0.3) - print("[TEST] Sending SIGUSR2 to resume") - os.kill(os.getpid(), signal.SIGUSR2) - - t = threading.Thread(target=send_quiesce) - t.start() - - result = await task - t.join() - - print(f"[TEST] uvloop completed: {len(result)} items") - return len(result) == 10 - - return asyncio.run(main()) - -def test_io_uring_tracking(): - """Test that io_uring instances are tracked""" - print("[TEST] io_uring tracking") - - # Try to use io_uring via aiofiles or direct syscall - try: - import asyncio - - async def io_work(): - # This should create io_uring if uvloop is installed - await asyncio.sleep(0.1) - return True - - result = asyncio.run(io_work()) - print(f"[TEST] io_uring tracking test: {'PASS' if result else 'FAIL'}") - return result - except Exception as e: - print(f"[TEST] io_uring tracking error: {e}") - return False - -def main(): - print("=" * 60) - print("GPUCR Quiescence Test Suite") - print("=" * 60) - print(f"PID: {os.getpid()}") - print(f"LD_PRELOAD: {os.environ.get('LD_PRELOAD', 'NOT SET')}") - print() - - # Check if library is loaded - if 'libgpucr_intercept.so' not in os.environ.get('LD_PRELOAD', ''): - print("WARNING: libgpucr_intercept.so not in LD_PRELOAD") - print(" Run with: LD_PRELOAD=./libgpucr_intercept.so python3 " + __file__) - print() - - tests = [ - ("Basic asyncio", test_asyncio_basic), - ("io_uring tracking", test_io_uring_tracking), - ("uvloop", test_uvloop_quiesce), - ] - - results = {} - for name, test_fn in tests: - print(f"\n{'='*60}") - print(f"Running: {name}") - print('=' * 60) - try: - results[name] = test_fn() - except Exception as e: - print(f"[TEST] {name} EXCEPTION: {e}") - import traceback - traceback.print_exc() - results[name] = False - - # Summary - print(f"\n{'='*60}") - print("Test Results:") - print('=' * 60) - for name, passed in results.items(): - status = "PASS" if passed else "FAIL" - print(f" {name}: {status}") - - all_passed = all(results.values()) - print() - print(f"Overall: {'PASS' if all_passed else 'FAIL'}") - - return 0 if all_passed else 1 - -if __name__ == "__main__": - sys.exit(main()) diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_seccomp.py b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_seccomp.py deleted file mode 100644 index 155f407ffc..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_seccomp.py +++ /dev/null @@ -1,82 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -Test seccomp-bpf interception of io_uring syscalls. - -This test verifies that the seccomp filter intercepts io_uring_enter -even when called from statically-linked code (uvloop). - -Run with: - GPUCR_SECCOMP_ENABLED=1 GPUCR_LOG_LEVEL=4 LD_PRELOAD=./libgpucr_intercept.so python3 tests/test_seccomp.py -""" - -import os -import sys -import asyncio - -def test_with_uvloop(): - """Test io_uring interception via uvloop (statically links libuv)""" - try: - import uvloop - print("uvloop available, testing with uvloop event loop") - uvloop.install() - except ImportError: - print("uvloop not available, testing with default asyncio") - - async def simple_io(): - """Do some async I/O to trigger io_uring operations""" - print("Starting async I/O operations...") - - # These operations may use io_uring internally - await asyncio.sleep(0.1) - print("After sleep 1") - - await asyncio.sleep(0.1) - print("After sleep 2") - - # Try some actual I/O - proc = await asyncio.create_subprocess_exec( - 'echo', 'hello', - stdout=asyncio.subprocess.PIPE - ) - stdout, _ = await proc.communicate() - print(f"Subprocess output: {stdout.decode().strip()}") - - print("Async I/O completed") - - asyncio.run(simple_io()) - -def main(): - print("=" * 60) - print("seccomp-bpf io_uring interception test") - print("=" * 60) - - print(f"PID: {os.getpid()}") - print(f"GPUCR_SECCOMP_ENABLED: {os.environ.get('GPUCR_SECCOMP_ENABLED', 'not set')}") - print(f"GPUCR_POST_RESTORE: {os.environ.get('GPUCR_POST_RESTORE', 'not set')}") - print(f"GPUCR_LOG_LEVEL: {os.environ.get('GPUCR_LOG_LEVEL', 'not set')}") - print() - - test_with_uvloop() - - print() - print("=" * 60) - print("Test completed - check output for [SECCOMP] log lines") - print("=" * 60) - -if __name__ == "__main__": - main() diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_seccomp_direct.c b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_seccomp_direct.c deleted file mode 100644 index 06fe4a45e3..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_seccomp_direct.c +++ /dev/null @@ -1,107 +0,0 @@ -/* -SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 -*/ - -/* - * Direct io_uring test for seccomp interception - * - * This uses raw syscalls to test that seccomp intercepts io_uring, - * independent of any library linking. - * - * Compile: gcc -o test_seccomp_direct tests/test_seccomp_direct.c - * Run: NVSNAP_SECCOMP_ENABLED=1 NVSNAP_LOG_LEVEL=3 LD_PRELOAD=./libnvsnap_intercept.so ./test_seccomp_direct - */ - -#define _GNU_SOURCE -#include -#include -#include -#include -#include -#include -#include -#include - -/* io_uring syscall numbers */ -#ifndef __NR_io_uring_setup -#define __NR_io_uring_setup 425 -#endif -#ifndef __NR_io_uring_enter -#define __NR_io_uring_enter 426 -#endif - -/* io_uring structures */ -struct io_uring_params { - uint32_t sq_entries; - uint32_t cq_entries; - uint32_t flags; - uint32_t sq_thread_cpu; - uint32_t sq_thread_idle; - uint32_t features; - uint32_t wq_fd; - uint32_t resv[3]; - struct { - uint32_t head, tail, ring_mask, ring_entries; - uint32_t flags, dropped, array, resv1; - uint64_t resv2; - } sq_off; - struct { - uint32_t head, tail, ring_mask, ring_entries; - uint32_t overflow, cqes, flags, resv1; - uint64_t resv2; - } cq_off; -}; - -int main(void) { - printf("Direct io_uring syscall test\n"); - printf("PID: %d\n\n", getpid()); - - /* Test io_uring_setup */ - printf("Calling io_uring_setup(32, params)...\n"); - - struct io_uring_params params; - memset(¶ms, 0, sizeof(params)); - - int fd = syscall(__NR_io_uring_setup, 32, ¶ms); - - if (fd < 0) { - printf("io_uring_setup failed: %s (errno=%d)\n", strerror(errno), errno); - if (errno == ENOSYS) { - printf("Kernel does not support io_uring\n"); - } - return 1; - } - - printf("io_uring_setup returned fd=%d\n", fd); - printf("sq_entries=%u, cq_entries=%u, features=0x%x\n", - params.sq_entries, params.cq_entries, params.features); - - /* Test io_uring_enter (with nothing to do) */ - printf("\nCalling io_uring_enter(fd=%d, to_submit=0, min_complete=0, flags=0)...\n", fd); - - int ret = syscall(__NR_io_uring_enter, fd, 0, 0, 0, NULL); - - if (ret < 0) { - printf("io_uring_enter failed: %s (errno=%d)\n", strerror(errno), errno); - } else { - printf("io_uring_enter returned %d\n", ret); - } - - /* Test io_uring_enter again */ - printf("\nCalling io_uring_enter(fd=%d, to_submit=0, min_complete=0, flags=1 GETEVENTS)...\n", fd); - - ret = syscall(__NR_io_uring_enter, fd, 0, 0, 1, NULL); /* 1 = IORING_ENTER_GETEVENTS */ - - if (ret < 0) { - printf("io_uring_enter failed: %s (errno=%d)\n", strerror(errno), errno); - } else { - printf("io_uring_enter returned %d\n", ret); - } - - /* Cleanup */ - close(fd); - - printf("\nTest completed!\n"); - return 0; -} diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_seccomp_uvloop.py b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_seccomp_uvloop.py deleted file mode 100644 index 337f9de606..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_seccomp_uvloop.py +++ /dev/null @@ -1,136 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -Test seccomp interception with uvloop (statically linked libuv). - -This test verifies that: -1. io_uring syscalls from uvloop are intercepted via seccomp -2. We can see SQPOLL mode being used -3. The intercepted calls proceed normally - -Run with: - source venv/bin/activate - GPUCR_SECCOMP_ENABLED=1 GPUCR_LOG_LEVEL=3 LD_PRELOAD=./libgpucr_intercept.so python3 tests/test_seccomp_uvloop.py - -For post-restore simulation: - GPUCR_SECCOMP_ENABLED=1 GPUCR_POST_RESTORE=1 GPUCR_LOG_LEVEL=3 LD_PRELOAD=./libgpucr_intercept.so python3 tests/test_seccomp_uvloop.py -""" - -import os -import sys -import asyncio - -IS_POST_RESTORE = os.environ.get('GPUCR_POST_RESTORE') == '1' - -def log(msg): - import datetime - ts = datetime.datetime.now().strftime('%H:%M:%S.%f')[:-3] - mode = "POST-RESTORE" if IS_POST_RESTORE else "NORMAL" - print(f"[{ts}] [{mode}] {msg}", flush=True) - -async def do_io_operations(): - """Perform various async I/O operations that trigger io_uring""" - - log("Starting I/O operations...") - - # Operation 1: Simple sleep (timer) - log("Op 1: asyncio.sleep(0.1)") - await asyncio.sleep(0.1) - log("Op 1 complete") - - # Operation 2: Another sleep - log("Op 2: asyncio.sleep(0.1)") - await asyncio.sleep(0.1) - log("Op 2 complete") - - # Operation 3: File I/O (may use io_uring depending on loop impl) - log("Op 3: Reading /proc/self/stat") - loop = asyncio.get_event_loop() - with open('/proc/self/stat', 'r') as f: - content = f.read() - log(f"Op 3 complete (read {len(content)} bytes)") - - # Operation 4: Subprocess (creates pipes, may use io_uring) - log("Op 4: Running subprocess 'echo hello'") - proc = await asyncio.create_subprocess_exec( - 'echo', 'hello', - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - stdout, stderr = await proc.communicate() - log(f"Op 4 complete: {stdout.decode().strip()}") - - # Operation 5: Socket operation (network I/O) - log("Op 5: Opening socket to localhost (expected to fail)") - try: - reader, writer = await asyncio.wait_for( - asyncio.open_connection('127.0.0.1', 12345), - timeout=0.1 - ) - writer.close() - await writer.wait_closed() - except (ConnectionRefusedError, asyncio.TimeoutError): - log("Op 5 complete (connection refused/timeout as expected)") - - # Operation 6: Concurrent operations - log("Op 6: Running 5 concurrent sleeps") - await asyncio.gather( - asyncio.sleep(0.05), - asyncio.sleep(0.05), - asyncio.sleep(0.05), - asyncio.sleep(0.05), - asyncio.sleep(0.05), - ) - log("Op 6 complete") - - log("All I/O operations completed successfully!") - -def main(): - print("=" * 70) - print("seccomp + uvloop io_uring interception test") - print("=" * 70) - print(f"PID: {os.getpid()}") - print(f"GPUCR_SECCOMP_ENABLED: {os.environ.get('GPUCR_SECCOMP_ENABLED', 'not set')}") - print(f"GPUCR_POST_RESTORE: {os.environ.get('GPUCR_POST_RESTORE', 'not set')}") - print(f"GPUCR_LOG_LEVEL: {os.environ.get('GPUCR_LOG_LEVEL', 'not set')}") - print() - - # Try to use uvloop - try: - import uvloop - print(f"uvloop version: {uvloop.__version__}") - uvloop.install() - print("Using uvloop event loop (statically links libuv)") - except ImportError: - print("WARNING: uvloop not available, using default asyncio") - print(" io_uring interception may not be exercised") - - print() - print("-" * 70) - - # Run the async operations - asyncio.run(do_io_operations()) - - print("-" * 70) - print() - print("=" * 70) - print("TEST COMPLETED") - print("Check output above for [SECCOMP] log lines showing intercepted calls") - print("=" * 70) - -if __name__ == '__main__': - main() diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_sglang_local.py b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_sglang_local.py deleted file mode 100644 index 834c1db020..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_sglang_local.py +++ /dev/null @@ -1,102 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -Local SGLang test for checkpoint/restore debugging. - -SGLang has different characteristics from vLLM: -- Different scheduler -- Different memory management -- May use different async patterns - -Run with: - LD_PRELOAD=/path/to/libgpucr_intercept.so python3 test_sglang_local.py -""" -import os -import sys -import time -import signal - -# NOTE: Don't close FDs here - multiprocessing spawn needs them -# FD closing should only be done for simple single-process tests - -# Write PID -with open("/tmp/sglang_local_pid", "w") as f: - f.write(str(os.getpid())) - -print(f"[SGLANG-TEST] PID={os.getpid()}", flush=True) -print(f"[SGLANG-TEST] GPUCR_RESTORED={os.environ.get('GPUCR_RESTORED', 'not set')}", flush=True) -print(f"[SGLANG-TEST] LD_PRELOAD={os.environ.get('LD_PRELOAD', 'not set')}", flush=True) - -# Signal handlers -def signal_handler(sig, frame): - print(f"[SGLANG-TEST] Received signal {sig}", flush=True) - -signal.signal(signal.SIGUSR1, signal_handler) -signal.signal(signal.SIGUSR2, signal_handler) - -print("[SGLANG-TEST] Importing sglang...", flush=True) - -try: - import sglang as sgl - from sglang import RuntimeEndpoint - - print(f"[SGLANG-TEST] SGLang version: {sgl.__version__}", flush=True) - - # Check if we can use the simple generation API - # SGLang has different API patterns than vLLM - print("[SGLANG-TEST] Testing SGLang Engine...", flush=True) - - # Use the Engine class directly for offline inference - from sglang import Engine - - print("[SGLANG-TEST] Creating Engine with TinyLlama...", flush=True) - print("[SGLANG-TEST] This may take a minute to load...", flush=True) - - # SGLang Engine for offline inference - engine = Engine( - model_path="TinyLlama/TinyLlama-1.1B-Chat-v1.0", - mem_fraction_static=0.5, # Use less GPU memory - ) - - print("[SGLANG-TEST] Engine created successfully!", flush=True) - - # Do a simple generation - print("[SGLANG-TEST] Running inference...", flush=True) - prompts = ["Hello, how are you?"] - - outputs = engine.generate(prompts, max_new_tokens=50) - - for i, output in enumerate(outputs): - print(f"[SGLANG-TEST] Prompt {i}: {prompts[i]!r}", flush=True) - print(f"[SGLANG-TEST] Output {i}: {output!r}", flush=True) - - print("[SGLANG-TEST] Inference complete!", flush=True) - print("[SGLANG-TEST] Entering idle loop (ready for checkpoint)...", flush=True) - - # Keep running for checkpoint - counter = 0 - while True: - counter += 1 - if counter % 30 == 0: - print(f"[SGLANG-TEST] Idle, counter={counter}", flush=True) - time.sleep(1) - -except Exception as e: - print(f"[SGLANG-TEST] ERROR: {type(e).__name__}: {e}", flush=True) - import traceback - traceback.print_exc() - sys.exit(1) diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_sglang_simple.py b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_sglang_simple.py deleted file mode 100644 index 3339b45755..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_sglang_simple.py +++ /dev/null @@ -1,48 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -Simple SGLang import test - check if it loads with our intercept library. -""" -import os -import sys - -print(f"PID: {os.getpid()}") -print(f"LD_PRELOAD: {os.environ.get('LD_PRELOAD', 'not set')}") -print() - -print("Importing sglang...") -try: - import sglang as sgl - print(f"SGLang version: {sgl.__version__}") - print("SUCCESS: SGLang imported!") -except Exception as e: - print(f"FAILED: {type(e).__name__}: {e}") - sys.exit(1) - -print() -print("Checking torch...") -try: - import torch - print(f"PyTorch version: {torch.__version__}") - print(f"CUDA available: {torch.cuda.is_available()}") - if torch.cuda.is_available(): - print(f"GPU: {torch.cuda.get_device_name(0)}") -except Exception as e: - print(f"PyTorch check failed: {e}") - -print() -print("All imports successful!") diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_uring_simple.py b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_uring_simple.py deleted file mode 100644 index bd2748d1a4..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_uring_simple.py +++ /dev/null @@ -1,24 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Simple io_uring/asyncio test""" -import asyncio - -async def test(): - await asyncio.sleep(0.1) - print('asyncio test passed') - -asyncio.run(test()) diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_uvloop_simple.py b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_uvloop_simple.py deleted file mode 100644 index 0784a2012e..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_uvloop_simple.py +++ /dev/null @@ -1,31 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Simple uvloop test""" -import asyncio - -try: - import uvloop - uvloop.install() - print('uvloop installed') -except ImportError: - print('uvloop not installed, using default event loop') - -async def test(): - await asyncio.sleep(0.1) - print('event loop test passed') - -asyncio.run(test()) diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_vllm.py b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_vllm.py deleted file mode 100644 index 904c1eac5a..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_vllm.py +++ /dev/null @@ -1,217 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -vLLM Interception Test - -This is the ULTIMATE test - vLLM uses: -- Multiple processes for tensor parallelism -- NCCL for inter-GPU communication -- Custom PagedAttention memory management -- Complex KV cache allocation strategies -- CUDA graphs for optimization - -If our interception works with vLLM, it works with anything. -""" - -import os -import sys -import time - -# Check if running with interception -INTERCEPTION_ENABLED = 'LD_PRELOAD' in os.environ and 'gpucr' in os.environ.get('LD_PRELOAD', '') - -def print_banner(msg): - print(f"\n{'='*60}") - print(f" {msg}") - print(f"{'='*60}\n") - -def test_vllm_import(): - """Test 1: Basic vLLM import - this triggers CUDA initialization""" - print_banner("Test 1: vLLM Import") - - try: - import vllm - print(f"✓ vLLM version: {vllm.__version__}") - return True - except ImportError as e: - print(f"✗ vLLM not installed: {e}") - return False - except Exception as e: - print(f"✗ vLLM import failed: {e}") - return False - -def test_vllm_offline_engine(): - """Test 2: Initialize vLLM LLM engine (offline mode)""" - print_banner("Test 2: vLLM Offline Engine (TinyLlama)") - - try: - from vllm import LLM, SamplingParams - - # Use a tiny model that fits in 6GB - # TinyLlama-1.1B is about 2.2GB in float16 - model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" - - print(f"Loading model: {model_name}") - print("This will download the model on first run (~2GB)...") - - # NOTE: vLLM 0.13+ uses multiprocessing by default. - # Each subprocess has its own interception state. - # Set GPUCR_LOG_LEVEL=1+ to see cudaMalloc calls in subprocess logs. - # For checkpoint/restore, each process is checkpointed individually. - - # Initialize with conservative settings for 6GB GPU - llm = LLM( - model=model_name, - tensor_parallel_size=1, # Single GPU - gpu_memory_utilization=0.7, # Leave some headroom - max_model_len=1024, # Reduce context length - trust_remote_code=True, - enforce_eager=True, # Disable CUDA graphs for simpler debugging - disable_log_stats=True, # Less noise - ) - - print("✓ vLLM LLM engine initialized successfully!") - - # Test inference - print("\nTesting inference...") - sampling_params = SamplingParams( - temperature=0.7, - max_tokens=50, - ) - - prompts = ["What is GPU checkpointing?"] - - start = time.time() - outputs = llm.generate(prompts, sampling_params) - elapsed = time.time() - start - - for output in outputs: - print(f"\nPrompt: {output.prompt}") - print(f"Generated: {output.outputs[0].text}") - - print(f"\n✓ Inference completed in {elapsed:.2f}s") - - # Clean up - del llm - - return True - - except Exception as e: - import traceback - print(f"✗ vLLM offline engine test failed: {e}") - traceback.print_exc() - return False - -def test_vllm_api_server(): - """Test 3: Start vLLM API server (tests multi-process)""" - print_banner("Test 3: vLLM API Server (Multi-Process)") - - # This would start the actual vLLM server - # For now, just report that this needs a separate test - print("Note: Full API server test requires running vLLM as a server") - print("This can be tested with:") - print(" LD_PRELOAD=./libgpucr_intercept.so python -m vllm.entrypoints.openai.api_server \\") - print(" --model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --port 8000") - print("\nSkipping automated API server test for now.") - return True - -def test_tensor_parallel(): - """Test 4: Tensor parallelism (requires multiple GPUs)""" - print_banner("Test 4: Tensor Parallelism") - - import torch - gpu_count = torch.cuda.device_count() - - if gpu_count < 2: - print(f"Only {gpu_count} GPU available - tensor parallelism requires 2+ GPUs") - print("Skipping tensor parallel test") - return True # Not a failure, just skipped - - try: - from vllm import LLM - - print(f"Testing with {gpu_count} GPUs") - llm = LLM( - model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", - tensor_parallel_size=gpu_count, - gpu_memory_utilization=0.5, - max_model_len=512, - ) - - print(f"✓ Tensor parallel initialization with {gpu_count} GPUs succeeded!") - del llm - return True - - except Exception as e: - print(f"✗ Tensor parallel test failed: {e}") - return False - -def main(): - print_banner("vLLM INTERCEPTION TEST SUITE") - - if INTERCEPTION_ENABLED: - print("✓ Running WITH gpucr interception") - else: - print("⚠ Running WITHOUT gpucr interception (for baseline)") - - results = {} - - # Test 1: Basic import - results['import'] = test_vllm_import() - if not results['import']: - print("\n\n⚠ vLLM not installed. Install with:") - print(" pip install vllm") - return 1 - - # Test 2: Offline engine - results['offline_engine'] = test_vllm_offline_engine() - - # Test 3: API server info - results['api_server'] = test_vllm_api_server() - - # Test 4: Tensor parallel - results['tensor_parallel'] = test_tensor_parallel() - - # Summary - print_banner("TEST SUMMARY") - - passed = sum(1 for v in results.values() if v) - total = len(results) - - for name, result in results.items(): - status = "✓ PASS" if result else "✗ FAIL" - print(f" {name}: {status}") - - print(f"\nTotal: {passed}/{total} tests passed") - - if INTERCEPTION_ENABLED: - print("\n" + "="*60) - print("INTERCEPTION NOTE:") - print("="*60) - print("vLLM uses multiprocessing - GPU work happens in child processes.") - print("Each process has its own interception stats.") - print("") - print("To verify interception, look for '[INFO] [cudaMalloc]' lines above") - print("from the EngineCore subprocess - those show interception working!") - print("") - print("For checkpoint/restore, each process is handled individually.") - print("="*60) - - return 0 if passed == total else 1 - -if __name__ == "__main__": - sys.exit(main()) diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_vllm_checkpoint.sh b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_vllm_checkpoint.sh deleted file mode 100755 index 8fc530a158..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_vllm_checkpoint.sh +++ /dev/null @@ -1,194 +0,0 @@ -#!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# vLLM GPU Checkpoint/Restore Test -# -# This script tests full checkpoint/restore of a vLLM inference server. -# It uses cuda-checkpoint (NVIDIA's tool) for GPU state and CRIU for CPU state. -# -# Requirements: -# - NVIDIA driver 555+ (for cuda-checkpoint support) -# - cuda-checkpoint binary -# - CRIU with CUDA plugin -# - vLLM installed in Python environment -# -# Usage: -# sudo ./test_vllm_checkpoint.sh - -set -e - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -LIB_DIR="$(dirname "$SCRIPT_DIR")" -VENV_PYTHON="${VENV_PYTHON:?VENV_PYTHON must point at a Python venv with vLLM installed}" - -# Sibling-relative paths to other forks; override individually with env vars. -_repo_root="$(cd "$LIB_DIR/../.." && pwd)" - -# cuda-checkpoint location -if [ -n "${CUDA_CHECKPOINT:-}" ]; then - CUDA_CKPT="$CUDA_CHECKPOINT" -elif [ -x "$_repo_root/../cuda-checkpoint/bin/x86_64_Linux/cuda-checkpoint" ]; then - CUDA_CKPT="$_repo_root/../cuda-checkpoint/bin/x86_64_Linux/cuda-checkpoint" -elif command -v cuda-checkpoint &>/dev/null; then - CUDA_CKPT="$(command -v cuda-checkpoint)" -else - echo "ERROR: cuda-checkpoint not found (set CUDA_CHECKPOINT or install in PATH)" - exit 1 -fi - -# CRIU binary (built from our CRIU fork) -if [ -n "${NVSNAP_CRIU:-}" ]; then - CRIU="$NVSNAP_CRIU" -elif [ -x "$_repo_root/../criu/criu/criu" ]; then - CRIU="$_repo_root/../criu/criu/criu" -elif command -v criu &>/dev/null; then - CRIU="$(command -v criu)" -else - echo "ERROR: CRIU not found (set NVSNAP_CRIU or clone the CRIU fork as a sibling and build it)" - exit 1 -fi - -# CRIU CUDA plugin source (lives inside the CRIU fork tree) -PLUGIN="${NVSNAP_CRIU_PLUGIN_DIR:-$_repo_root/../criu/plugins/cuda}" - -echo "==========================================" -echo " vLLM Checkpoint/Restore Test" -echo "==========================================" - -# Cleanup -pkill -9 -f vllm_test 2>/dev/null || true -rm -rf /tmp/criu_img /tmp/vllm_output.txt -mkdir -p /tmp/criu_img - -# Create test script using a tiny model -cat > /tmp/vllm_test.py << 'EOF' -import os -import sys -import time - -os.environ["VLLM_USE_CUDA_GRAPH"] = "0" - -print(f"PID: {os.getpid()}", flush=True) - -from vllm import LLM, SamplingParams - -print("Loading TinyLlama model...", flush=True) -llm = LLM( - model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", - tensor_parallel_size=1, - gpu_memory_utilization=0.7, - max_model_len=256, - trust_remote_code=True, - enforce_eager=True, -) -print("Model loaded!", flush=True) - -params = SamplingParams(max_tokens=20, temperature=0.8) -outputs = llm.generate(["Hello, world!"], params) -print(f"Inference output: {outputs[0].outputs[0].text[:50]}...", flush=True) - -print("READY - waiting for checkpoint...", flush=True) -sys.stdout.flush() - -while True: - time.sleep(1) -EOF - -echo "" -echo "[1] Starting vLLM process..." -# Note: LD_PRELOAD not needed - cuda-checkpoint handles GPU state directly -$VENV_PYTHON /tmp/vllm_test.py > /tmp/vllm_output.txt 2>&1 & -PID=$! -echo " PID: $PID" - -echo " Waiting for model to load..." -for i in {1..180}; do - if grep -q "READY" /tmp/vllm_output.txt 2>/dev/null; then - echo " Model loaded after ${i}s" - break - fi - if ! ps -p $PID > /dev/null 2>&1; then - echo " FAILED: Process died during loading" - cat /tmp/vllm_output.txt - exit 1 - fi - sleep 1 -done - -if ! grep -q "READY" /tmp/vllm_output.txt 2>/dev/null; then - echo " TIMEOUT waiting for model" - cat /tmp/vllm_output.txt - kill -9 $PID 2>/dev/null || true - exit 1 -fi - -echo "" -grep -E "(PID:|Model loaded|Inference|READY)" /tmp/vllm_output.txt || true - -echo "" -echo "[2] CUDA state: $(${CUDA_CKPT} --get-state --pid $PID 2>/dev/null | tail -1)" - -echo "" -echo "[3] Locking CUDA..." -${CUDA_CKPT} --action lock --pid $PID --timeout 30000 2>/dev/null -echo " Lock: OK" - -echo "" -echo "[4] Checkpointing CUDA..." -${CUDA_CKPT} --action checkpoint --pid $PID 2>/dev/null -echo " State: $(${CUDA_CKPT} --get-state --pid $PID 2>/dev/null | tail -1)" - -echo "" -echo "[5] CRIU dump..." -$CRIU dump -t $PID -D /tmp/criu_img --shell-job --tcp-established -L $PLUGIN 2>&1 | grep -v "NVSNAP\|Calls\|Driver\|Runtime\|cuBLAS\|cuDNN\|NCCL\|Tracked\|Allocations\|Streams\|Events\|Modules\|Contexts\|Comms\|===" || true -echo " Created $(ls /tmp/criu_img/*.img 2>/dev/null | wc -l) images" - -echo "" -echo "[6] CRIU restore..." -$CRIU restore -d -D /tmp/criu_img --shell-job --tcp-established -L $PLUGIN 2>&1 | grep -v "NVSNAP\|Calls\|Driver\|Runtime\|cuBLAS\|cuDNN\|NCCL\|Tracked\|Allocations\|Streams\|Events\|Modules\|Contexts\|Comms\|===" | tail -5 || true -sleep 3 - -NEW_PID=$(pgrep -f vllm_test.py 2>/dev/null | head -1) -if [ -z "$NEW_PID" ]; then - echo " FAILED: No restored process" - exit 1 -fi -echo " Restored PID: $NEW_PID" - -echo "" -echo "[7] Restoring CUDA..." -${CUDA_CKPT} --action restore --pid $NEW_PID 2>/dev/null -echo " Restore: OK" - -echo "" -echo "[8] Unlocking CUDA..." -${CUDA_CKPT} --action unlock --pid $NEW_PID 2>/dev/null -echo " State: $(${CUDA_CKPT} --get-state --pid $NEW_PID 2>/dev/null | tail -1)" - -echo "" -echo "[9] Verifying..." -sleep 5 -if ps -p $NEW_PID > /dev/null 2>&1; then - echo " Process: ALIVE" -else - echo " Process: DEAD (may be io_uring issue)" -fi - -kill -9 $NEW_PID 2>/dev/null || true -pkill -9 -f EngineCore 2>/dev/null || true - -echo "" -echo "==========================================" -echo " vLLM CHECKPOINT/RESTORE TEST COMPLETE" -echo "==========================================" -echo "" -echo "Summary:" -echo " - cuda-checkpoint lock/checkpoint: OK" -echo " - CRIU dump: OK" -echo " - CRIU restore: OK" -echo " - cuda-checkpoint restore/unlock: OK" -echo "" -echo "Note: vLLM uses io_uring which has CRIU compatibility issues." -echo "The checkpoint/restore cycle works, but process may crash after" -echo "due to io_uring ring address remapping." diff --git a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_vllm_local.py b/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_vllm_local.py deleted file mode 100644 index 1bd9c8d9d7..0000000000 --- a/src/compute-plane-services/nvsnap/lib/nvsnap_intercept/tests/test_vllm_local.py +++ /dev/null @@ -1,93 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -Local vLLM test for checkpoint/restore debugging. - -This script starts vLLM with a small model, suitable for local testing. -Run with: - LD_PRELOAD=/path/to/libgpucr_intercept.so python3 test_vllm_local.py -""" -import os -import sys -import time -import signal - -# Close inherited FDs (important for CRIU) -for fd in range(3, 256): - try: - os.close(fd) - except: - pass - -# Write PID -with open("/tmp/vllm_local_pid", "w") as f: - f.write(str(os.getpid())) - -print(f"[vLLM-TEST] PID={os.getpid()}", flush=True) -print(f"[vLLM-TEST] GPUCR_RESTORED={os.environ.get('GPUCR_RESTORED', 'not set')}", flush=True) -print(f"[vLLM-TEST] LD_PRELOAD={os.environ.get('LD_PRELOAD', 'not set')}", flush=True) - -# Signal handlers -def signal_handler(sig, frame): - print(f"[vLLM-TEST] Received signal {sig}", flush=True) - -signal.signal(signal.SIGUSR1, signal_handler) -signal.signal(signal.SIGUSR2, signal_handler) - -print("[vLLM-TEST] Importing vLLM...", flush=True) -from vllm import LLM, SamplingParams - -print("[vLLM-TEST] Initializing model (TinyLlama)...", flush=True) -print("[vLLM-TEST] This may take a minute...", flush=True) - -try: - # Use a small model for testing - llm = LLM( - model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", - max_model_len=512, - gpu_memory_utilization=0.5, - trust_remote_code=True, - ) - print("[vLLM-TEST] Model loaded successfully!", flush=True) - - # Do a simple inference - sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=50) - prompts = ["Hello, how are you?"] - - print("[vLLM-TEST] Running inference...", flush=True) - outputs = llm.generate(prompts, sampling_params) - - for output in outputs: - print(f"[vLLM-TEST] Prompt: {output.prompt!r}", flush=True) - print(f"[vLLM-TEST] Output: {output.outputs[0].text!r}", flush=True) - - print("[vLLM-TEST] Inference complete!", flush=True) - print("[vLLM-TEST] Entering idle loop (ready for checkpoint)...", flush=True) - - # Keep running for checkpoint - counter = 0 - while True: - counter += 1 - if counter % 30 == 0: - print(f"[vLLM-TEST] Idle, counter={counter}", flush=True) - time.sleep(1) - -except Exception as e: - print(f"[vLLM-TEST] ERROR: {type(e).__name__}: {e}", flush=True) - import traceback - traceback.print_exc() - sys.exit(1) diff --git a/src/compute-plane-services/nvsnap/lib/sitecustomize/sitecustomize.py b/src/compute-plane-services/nvsnap/lib/sitecustomize/sitecustomize.py deleted file mode 100644 index 7223a3dade..0000000000 --- a/src/compute-plane-services/nvsnap/lib/sitecustomize/sitecustomize.py +++ /dev/null @@ -1,44 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""NvSnap runtime sitecustomize — prepends version-matched site-packages. - -Imported automatically by CPython's site.py during interpreter startup -because the directory containing this file is on PYTHONPATH (set by every -NvSnap-enabled pod's env). Detects the running interpreter's ABI tag and -prepends /nvsnap-lib/site-packages-cpXY to sys.path so `import uvloop` -(and any future patched packages) resolve to the NvSnap-patched copies -before any vendored or venv-installed stock copy. - -The patched native libraries (libzmq, libuv) are not Python-version- -keyed; they live directly under /nvsnap-lib and are routed via -LD_LIBRARY_PATH — see docs/GENERIC-PYTHON-INJECTION-DESIGN.md. - -This file is intentionally minimal: - - No imports beyond os/sys (no risk of breaking interpreter startup). - - No-op if /nvsnap-lib/site-packages-cpXY is absent (graceful when the - pod runs without NvSnap init containers). - - No chain-load of a customer-provided sitecustomize. If a future BYOC - workload ships its own sitecustomize.py, add an importlib chain here. -""" - -import os -import sys - -_PYVER = "cp{}{}".format(sys.version_info.major, sys.version_info.minor) -_PKG_DIR = "/nvsnap-lib/site-packages-" + _PYVER - -if os.path.isdir(_PKG_DIR) and _PKG_DIR not in sys.path: - sys.path.insert(0, _PKG_DIR) diff --git a/src/compute-plane-services/nvsnap/scripts/auto-inject-init.sh b/src/compute-plane-services/nvsnap/scripts/auto-inject-init.sh deleted file mode 100644 index 3c3b682626..0000000000 --- a/src/compute-plane-services/nvsnap/scripts/auto-inject-init.sh +++ /dev/null @@ -1,53 +0,0 @@ -#!/bin/sh -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# Auto-inject init payload — replaces the 4-init-container fan-out -# (get-uvloop, get-libuv, get-libzmq, get-nvsnap). This script runs as -# a single init container using the nvsnap-agent image; it copies all -# four payloads into the workload pod's /nvsnap-lib emptyDir volume. -# -# The agent image must have everything it needs already bundled: -# /criu-bundle/payload/wheels/uvloop-*.whl (multi-python wheels) -# /criu-bundle/payload/lib/libuv.so* -# /criu-bundle/payload/lib/libzmq.so* -# /criu-bundle/lib/libnvsnap_intercept.so -# /criu-bundle/sitecustomize/ (sitecustomize.py) -# -# Caller stamps PYTHONPATH=/nvsnap-lib/sitecustomize, LD_LIBRARY_PATH -# starts with /nvsnap-lib, and LD_PRELOAD points at the intercept lib — -# all from the webhook auto-inject patch in internal/webhook/auto_inject.go. - -set -eu - -DST=/nvsnap-lib - -# 1. uvloop wheels per Python ABI tag. sitecustomize picks the right -# one at runtime based on sys.version_info. -for whl in /criu-bundle/payload/wheels/uvloop-*.whl; do - if [ ! -e "$whl" ]; then - echo "auto-inject: no uvloop wheels found at /criu-bundle/payload/wheels/" >&2 - exit 1 - fi - tag=$(echo "$whl" | grep -oE 'cp3[0-9]+' | head -1) - if [ -z "$tag" ]; then - echo "auto-inject: cannot extract ABI tag from $whl" >&2 - exit 1 - fi - mkdir -p "$DST/site-packages-${tag}" - python3 -m zipfile -e "$whl" "$DST/site-packages-${tag}/" -done - -# 2. Patched native libs (libuv, libzmq). LD_LIBRARY_PATH picks them -# up before the system versions. -cp /criu-bundle/payload/lib/libuv.so* "$DST/" -cp /criu-bundle/payload/lib/libzmq.so* "$DST/" - -# 3. Intercept lib + sitecustomize. LD_PRELOAD'd into every workload -# process at startup; PYTHONPATH'd so the interpreter imports our -# sitecustomize.py before user code runs. -cp /criu-bundle/lib/libnvsnap_intercept.so "$DST/" -cp -r /criu-bundle/sitecustomize "$DST/" - -echo "auto-inject: payload installed under $DST" -ls -la "$DST" diff --git a/src/compute-plane-services/nvsnap/scripts/build-agent-app.sh b/src/compute-plane-services/nvsnap/scripts/build-agent-app.sh deleted file mode 100755 index 4e5714e957..0000000000 --- a/src/compute-plane-services/nvsnap/scripts/build-agent-app.sh +++ /dev/null @@ -1,88 +0,0 @@ -#!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# Build NVSNAP agent APP image (Go binaries + intercept library) -# This builds on top of the base image and is FAST (~30 seconds) -set -euo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -PROJECT_ROOT="$(dirname "$SCRIPT_DIR")" - -# Configuration -REGISTRY="${REGISTRY:-nvcr.io/0651155215864979/ncp-dev}" -BASE_IMAGE="${BASE_IMAGE:-${REGISTRY}/nvsnap-agent-base:v45}" -IMAGE_NAME="nvsnap-agent" -VERSION="${VERSION:-v0.7.119}" -CONTAINER_TOOL="${CONTAINER_TOOL:-}" -ENABLE_LIBUV_INTERCEPT="${ENABLE_LIBUV_INTERCEPT:-0}" - -if [ -z "$CONTAINER_TOOL" ]; then - if command -v nerdctl >/dev/null 2>&1; then - CONTAINER_TOOL="nerdctl" - elif command -v docker >/dev/null 2>&1; then - CONTAINER_TOOL="docker" - else - echo "No container tool found (install docker or nerdctl)" - exit 1 - fi -fi - -if [ "$CONTAINER_TOOL" = "docker" ] && [ -z "${DOCKER_HOST:-}" ]; then - if [ -S /var/run/docker.sock ]; then - export DOCKER_HOST="unix:///var/run/docker.sock" - fi -fi - -echo "=== Building NVSNAP Agent APP Image ===" -echo "Base: ${BASE_IMAGE}" -echo "Version: ${VERSION}" -echo "Enable libuv intercept: ${ENABLE_LIBUV_INTERCEPT}" -echo "" - -# Create build context from project root -BUILD_CONTEXT=$(mktemp -d) -trap "rm -rf $BUILD_CONTEXT" EXIT - -echo "Build context: $BUILD_CONTEXT" - -# Copy Dockerfile -cp "${PROJECT_ROOT}/docker/agent/Dockerfile.app" "${BUILD_CONTEXT}/Dockerfile" - -# Copy Go project (exclude large/unnecessary directories) -echo "Copying Go project..." -rsync -a \ - --exclude='.git' \ - --exclude='vendor' \ - --exclude='bin' \ - --exclude='docker' \ - --exclude='placeholder' \ - --exclude='*.tar.gz' \ - --exclude='node_modules' \ - "${PROJECT_ROOT}/" "${BUILD_CONTEXT}/" - -# Copy intercept library source -mkdir -p "${BUILD_CONTEXT}/lib" -cp -r "${PROJECT_ROOT}/lib/nvsnap_intercept" "${BUILD_CONTEXT}/lib/" - -# Build image -echo "" -echo "Building Docker image..." -"${CONTAINER_TOOL}" build \ - --platform linux/amd64 \ - --build-arg BASE_IMAGE="${BASE_IMAGE}" \ - --build-arg ENABLE_LIBUV_INTERCEPT="${ENABLE_LIBUV_INTERCEPT}" \ - -t "${REGISTRY}/${IMAGE_NAME}:${VERSION}" \ - "${BUILD_CONTEXT}" - -echo "" -echo "=== App image build complete ===" -echo "Image: ${REGISTRY}/${IMAGE_NAME}:${VERSION}" -echo "" - -# Auto-push option -if [ "${PUSH:-}" = "1" ] || [ "${PUSH:-}" = "true" ]; then - echo "Pushing..." - "${CONTAINER_TOOL}" push "${REGISTRY}/${IMAGE_NAME}:${VERSION}" - echo "Pushed: ${REGISTRY}/${IMAGE_NAME}:${VERSION}" -fi diff --git a/src/compute-plane-services/nvsnap/scripts/build-agent.sh b/src/compute-plane-services/nvsnap/scripts/build-agent.sh index 42d127fa4a..78a197fcc2 100755 --- a/src/compute-plane-services/nvsnap/scripts/build-agent.sh +++ b/src/compute-plane-services/nvsnap/scripts/build-agent.sh @@ -209,20 +209,11 @@ build_base() { } build_app() { - echo "=== Building APP image (Go binaries, intercept lib) ===" + echo "=== Building APP image (Go binaries + C helpers) ===" echo "Base: ${BASE_IMAGE}" echo "Image: ${APP_IMAGE}" echo "" - # Warn about untracked .c files in intercept library - local untracked - untracked=$(cd "$PROJECT_ROOT" && git ls-files --others --exclude-standard lib/nvsnap_intercept/src/*.c 2>/dev/null) - if [ -n "$untracked" ]; then - echo "WARNING: Untracked .c files in intercept library: $untracked" - echo " These will be missing from the Docker build context!" - echo " Run: git add $untracked" - fi - # Verify base image CRIU binary echo "Verifying base image CRIU binary..." local verify_id @@ -262,9 +253,6 @@ build_app() { --exclude='*.tar.gz' \ "${PROJECT_ROOT}/" "${BUILD_CTX}/" - # Copy intercept library source - cp -r "${PROJECT_ROOT}/lib/nvsnap_intercept" "${BUILD_CTX}/lib/" - # cuda-checkpoint wrapper only: the real binary is built from source in # the BASE image (Dockerfile.base cuda-cli-builder stage); Dockerfile.app # just re-overlays the wrapper at /criu-bundle/cuda-checkpoint. @@ -281,9 +269,6 @@ build_app() { $cache_flag \ --platform linux/amd64 \ --build-arg BASE_IMAGE="${BASE_IMAGE}" \ - --build-arg UVLOOP_IMAGE="${REGISTRY}/uvloop-builder:${NVSNAP_UVLOOP_VERSION}" \ - --build-arg LIBUV_IMAGE="${REGISTRY}/libuv-builder:${NVSNAP_LIBUV_VERSION}" \ - --build-arg LIBZMQ_IMAGE="${REGISTRY}/libzmq-builder:${NVSNAP_LIBZMQ_VERSION}" \ -t "${APP_IMAGE}" \ "${BUILD_CTX}" @@ -318,28 +303,6 @@ push_app() { echo "Pushing app image: ${APP_IMAGE}" docker push "${APP_IMAGE}" - # Auto-build and push nvsnap-init with matching agent version. - # nvsnap-init bundles the agent's intercept lib + patched dependencies. - # Using the same tag as the agent prevents build-ID mismatch on restore. - local INIT_IMAGE="${REGISTRY}/nvsnap-init:${APP_VERSION}" - echo "" - echo "Building nvsnap-init: ${INIT_IMAGE}" - # --network=host so the build container's apt-get can reach - # archive.ubuntu.com via the host's DNS + proxy config. Without - # this, on a corp VPN the build container ends up in an isolated - # bridge network with no working DNS for ubuntu.com mirrors. - docker build --no-cache --network=host \ - -t "${INIT_IMAGE}" \ - --build-arg UVLOOP_IMAGE="${REGISTRY}/uvloop-builder:${NVSNAP_UVLOOP_VERSION}" \ - --build-arg LIBUV_IMAGE="${REGISTRY}/libuv-builder:${NVSNAP_LIBUV_VERSION}" \ - --build-arg LIBZMQ_IMAGE="${REGISTRY}/libzmq-builder:${NVSNAP_LIBZMQ_VERSION}" \ - --build-arg PYZMQ_IMAGE="${REGISTRY}/pyzmq-builder:${NVSNAP_PYZMQ_VERSION}" \ - --build-arg AGENT_IMAGE="${APP_IMAGE}" \ - -f "${PROJECT_ROOT}/docker/init/Dockerfile" \ - "${PROJECT_ROOT}" - - echo "Pushing nvsnap-init: ${INIT_IMAGE}" - docker push "${INIT_IMAGE}" echo "Done." } @@ -418,9 +381,6 @@ sync_versions() { echo "=== Syncing versions from versions.sh to manifests ===" echo " APP_VERSION: ${APP_VERSION}" echo " BASE_VERSION: ${BASE_VERSION}" - echo " UVLOOP: ${NVSNAP_UVLOOP_VERSION}" - echo " LIBZMQ: ${NVSNAP_LIBZMQ_VERSION}" - echo " PYZMQ: ${NVSNAP_PYZMQ_VERSION}" echo " VLLM: ${NVSNAP_VLLM_VERSION}" echo "" @@ -452,11 +412,7 @@ sync_versions() { sync_workload_manifest() { local manifest="$1" sed -i \ - -e "s|image: nvcr.io/0651155215864979/ncp-dev/uvloop-builder:[^ ]*|image: ${REGISTRY}/uvloop-builder:${NVSNAP_UVLOOP_VERSION}|" \ - -e "s|image: nvcr.io/0651155215864979/ncp-dev/libzmq-builder:[^ ]*|image: ${REGISTRY}/libzmq-builder:${NVSNAP_LIBZMQ_VERSION}|" \ - -e "s|image: nvcr.io/0651155215864979/ncp-dev/pyzmq-builder:[^ ]*|image: ${REGISTRY}/pyzmq-builder:${NVSNAP_PYZMQ_VERSION}|" \ -e "s|image: nvcr.io/0651155215864979/ncp-dev/nvsnap-agent:[^ ]*|image: ${REGISTRY}/nvsnap-agent:${APP_VERSION}|" \ - -e "s|image: nvcr.io/0651155215864979/ncp-dev/nvsnap-init:[^ ]*|image: ${REGISTRY}/nvsnap-init:${NVSNAP_INIT_VERSION}|" \ -e "s|image: vllm/vllm-openai:[^ ]*|image: vllm/vllm-openai:${NVSNAP_VLLM_VERSION}|" \ "$manifest" # Restore manifests reference the agent's checkpoint hostPath. @@ -506,9 +462,6 @@ show_versions() { echo " App image: ${REGISTRY}/nvsnap-agent:${APP_VERSION}" echo "" echo " Dependencies:" - echo " uvloop: ${REGISTRY}/uvloop-builder:${NVSNAP_UVLOOP_VERSION}" - echo " libzmq: ${REGISTRY}/libzmq-builder:${NVSNAP_LIBZMQ_VERSION}" - echo " pyzmq: ${REGISTRY}/pyzmq-builder:${NVSNAP_PYZMQ_VERSION}" echo " vllm: vllm/vllm-openai:${NVSNAP_VLLM_VERSION}" echo "" echo " DOCKER_HOST: ${DOCKER_HOST:-}" diff --git a/src/compute-plane-services/nvsnap/scripts/build-deps.sh b/src/compute-plane-services/nvsnap/scripts/build-deps.sh deleted file mode 100755 index deb33c39d5..0000000000 --- a/src/compute-plane-services/nvsnap/scripts/build-deps.sh +++ /dev/null @@ -1,264 +0,0 @@ -#!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# Build NVSNAP dependency images reproducibly -# -# Usage: -# ./scripts/build-deps.sh libzmq # Build libzmq builder image -# ./scripts/build-deps.sh uvloop # Build uvloop wheel image -# ./scripts/build-deps.sh pyzmq # Build pyzmq wheel image -# ./scripts/build-deps.sh all # Build all -# ./scripts/build-deps.sh push # Push all to registry -# -# Environment: -# REGISTRY - Image registry (default: from versions.sh) -# PUSH=1 - Push after build -# NO_CACHE=1 - Force rebuild without cache - -set -euo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -PROJECT_ROOT="$(dirname "$SCRIPT_DIR")" - -source "${SCRIPT_DIR}/versions.sh" - -REGISTRY="${REGISTRY:-$NVSNAP_REGISTRY}" - -# Source repo locations (auto-discovered from sibling layout — see CONTRIBUTING.md) -source "$(dirname "${BASH_SOURCE[0]}")/_deps.sh" -nvsnap_resolve_sibling LIBZMQ_SRC libzmq -nvsnap_resolve_sibling UVLOOP_SRC uvloop -nvsnap_resolve_sibling PYZMQ_SRC pyzmq -nvsnap_resolve_sibling LIBUV_SRC libuv - -# Auto-detect DOCKER_HOST -if [ -z "${DOCKER_HOST:-}" ]; then - if [ -S /var/run/docker.sock ]; then - export DOCKER_HOST="unix:///var/run/docker.sock" - fi -fi - -CACHE_FLAG="" -if [ "${NO_CACHE:-0}" = "1" ]; then - CACHE_FLAG="--no-cache" -fi - -usage() { - echo "Usage: $0 [libzmq|uvloop|pyzmq|libuv|all|push]" - echo "" - echo "Builds dependency images from version-controlled Dockerfiles." - echo "Source repos must be cloned locally (see environment variables)." - echo "" - echo "Current versions:" - echo " libzmq: ${NVSNAP_LIBZMQ_VERSION}" - echo " uvloop: ${NVSNAP_UVLOOP_VERSION}" - echo " pyzmq: ${NVSNAP_PYZMQ_VERSION}" - echo " libuv: ${NVSNAP_LIBUV_VERSION}" -} - -build_libzmq() { - local IMAGE="${REGISTRY}/libzmq-builder:${NVSNAP_LIBZMQ_VERSION}" - echo "=== Building libzmq: ${IMAGE} ===" - - if [ ! -d "$LIBZMQ_SRC/src" ]; then - echo "ERROR: libzmq source not found at $LIBZMQ_SRC" - echo "Clone: git clone $LIBZMQ_SRC" - exit 1 - fi - - # Create clean build context (exclude .git, build artifacts) - local BUILD_CTX - BUILD_CTX=$(mktemp -d) - trap "rm -rf $BUILD_CTX" RETURN - - rsync -a \ - --exclude='.git' \ - --exclude='build/' \ - --exclude='*.o' \ - --exclude='*.so' \ - --exclude='*.so.*' \ - "$LIBZMQ_SRC/" "$BUILD_CTX/" - - docker build \ - $CACHE_FLAG \ - --platform linux/amd64 \ - -t "$IMAGE" \ - -f "${PROJECT_ROOT}/docker/libzmq/Dockerfile" \ - "$BUILD_CTX" - - echo "" - echo "Built: $IMAGE" - - if [ "${PUSH:-0}" = "1" ]; then - echo "Pushing $IMAGE..." - docker push "$IMAGE" - fi -} - -build_uvloop() { - local IMAGE="${REGISTRY}/uvloop-builder:${NVSNAP_UVLOOP_VERSION}" - echo "=== Building uvloop: ${IMAGE} ===" - - if [ ! -f "$UVLOOP_SRC/pyproject.toml" ]; then - echo "ERROR: uvloop source not found at $UVLOOP_SRC" - echo "Clone: git clone $UVLOOP_SRC" - exit 1 - fi - - local BUILD_CTX - BUILD_CTX=$(mktemp -d) - trap "rm -rf $BUILD_CTX" RETURN - - rsync -a \ - --exclude='.git' \ - --exclude='build/' \ - --exclude='dist/' \ - --exclude='*.egg-info' \ - --exclude='.eggs/' \ - --exclude='*.o' \ - --exclude='*.so' \ - --exclude='uvloop/loop.c' \ - "$UVLOOP_SRC/" "$BUILD_CTX/" - - docker build \ - $CACHE_FLAG \ - --platform linux/amd64 \ - -t "$IMAGE" \ - -f "${PROJECT_ROOT}/docker/uvloop/Dockerfile" \ - "$BUILD_CTX" - - echo "" - echo "Built: $IMAGE" - - if [ "${PUSH:-0}" = "1" ]; then - echo "Pushing $IMAGE..." - docker push "$IMAGE" - fi -} - -build_pyzmq() { - local IMAGE="${REGISTRY}/pyzmq-builder:${NVSNAP_PYZMQ_VERSION}" - echo "=== Building pyzmq: ${IMAGE} ===" - - if [ ! -f "$PYZMQ_SRC/pyproject.toml" ]; then - echo "ERROR: pyzmq source not found at $PYZMQ_SRC" - echo "Clone: git clone $PYZMQ_SRC" - exit 1 - fi - if [ ! -d "$LIBZMQ_SRC/src" ]; then - echo "ERROR: libzmq source not found at $LIBZMQ_SRC (needed to build pyzmq)" - exit 1 - fi - - # pyzmq needs both pyzmq-src and libzmq-src in the build context - local BUILD_CTX - BUILD_CTX=$(mktemp -d) - trap "rm -rf $BUILD_CTX" RETURN - - rsync -a \ - --exclude='.git' \ - --exclude='build/' \ - --exclude='dist/' \ - --exclude='*.egg-info' \ - --exclude='*.o' \ - --exclude='*.so' \ - --exclude='*.so.*' \ - "$PYZMQ_SRC/" "$BUILD_CTX/pyzmq-src/" - - rsync -a \ - --exclude='.git' \ - --exclude='build/' \ - --exclude='*.o' \ - --exclude='*.so' \ - --exclude='*.so.*' \ - "$LIBZMQ_SRC/" "$BUILD_CTX/libzmq-src/" - - docker build \ - $CACHE_FLAG \ - --platform linux/amd64 \ - -t "$IMAGE" \ - -f "${PROJECT_ROOT}/docker/pyzmq/Dockerfile" \ - "$BUILD_CTX" - - echo "" - echo "Built: $IMAGE" - - if [ "${PUSH:-0}" = "1" ]; then - echo "Pushing $IMAGE..." - docker push "$IMAGE" - fi -} - -build_libuv() { - local IMAGE="${REGISTRY}/libuv-builder:${NVSNAP_LIBUV_VERSION}" - echo "=== Building libuv: ${IMAGE} ===" - - if [ ! -f "$LIBUV_SRC/CMakeLists.txt" ]; then - echo "ERROR: libuv source not found at $LIBUV_SRC" - echo "Clone: git clone $LIBUV_SRC" - exit 1 - fi - - local BUILD_CTX - BUILD_CTX=$(mktemp -d) - trap "rm -rf $BUILD_CTX" RETURN - - git -C "$LIBUV_SRC" archive HEAD | tar -x -C "$BUILD_CTX" - - docker build \ - $CACHE_FLAG \ - --platform linux/amd64 \ - -t "$IMAGE" \ - -f "${PROJECT_ROOT}/docker/libuv/Dockerfile" \ - "$BUILD_CTX" - - echo "" - echo "Built: $IMAGE" - - if [ "${PUSH:-0}" = "1" ]; then - echo "Pushing $IMAGE..." - docker push "$IMAGE" - fi -} - -push_all() { - echo "=== Pushing all dependency images ===" - docker push "${REGISTRY}/libzmq-builder:${NVSNAP_LIBZMQ_VERSION}" - docker push "${REGISTRY}/uvloop-builder:${NVSNAP_UVLOOP_VERSION}" - docker push "${REGISTRY}/pyzmq-builder:${NVSNAP_PYZMQ_VERSION}" - docker push "${REGISTRY}/libuv-builder:${NVSNAP_LIBUV_VERSION}" - echo "Done." -} - -case "${1:-help}" in - libzmq) - build_libzmq - ;; - uvloop) - build_uvloop - ;; - pyzmq) - build_pyzmq - ;; - libuv) - build_libuv - ;; - all) - build_libzmq - build_uvloop - build_pyzmq - build_libuv - ;; - push) - push_all - ;; - help|--help|-h) - usage - ;; - *) - echo "Unknown command: $1" - usage - exit 1 - ;; -esac diff --git a/src/compute-plane-services/nvsnap/scripts/build-intercept-lib-local.sh b/src/compute-plane-services/nvsnap/scripts/build-intercept-lib-local.sh deleted file mode 100755 index 82dfd23ae2..0000000000 --- a/src/compute-plane-services/nvsnap/scripts/build-intercept-lib-local.sh +++ /dev/null @@ -1,28 +0,0 @@ -#!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -set -euo pipefail - -REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" -OUT_DIR="${1:-/tmp/nvsnap-local/nvsnap-lib}" -DOCKER_HOST_DEFAULT="unix:///var/run/docker.sock" -export DOCKER_HOST="${NVSNAP_DOCKER_HOST:-$DOCKER_HOST_DEFAULT}" -BUILD_IMAGE="${NVSNAP_INTERCEPT_BUILD_IMAGE:-nvsnap-intercept-build:local}" -ENABLE_LIBUV_INTERCEPT="${NVSNAP_ENABLE_LIBUV_INTERCEPT:-0}" - -mkdir -p "${OUT_DIR}" - -if ! docker image inspect "${BUILD_IMAGE}" >/dev/null 2>&1; then - echo "Building intercept toolchain image ${BUILD_IMAGE}..." - docker build -f "${REPO_ROOT}/scripts/Dockerfile.intercept-build" -t "${BUILD_IMAGE}" "${REPO_ROOT}/scripts" -fi - -echo "Building libnvsnap_intercept.so in ${BUILD_IMAGE}..." -docker run --rm \ - -v "${REPO_ROOT}/lib/nvsnap_intercept:/src" \ - -v "${OUT_DIR}:/out" \ - "${BUILD_IMAGE}" bash -lc \ - "make -C /src clean && make -C /src ENABLE_LIBUV_INTERCEPT=${ENABLE_LIBUV_INTERCEPT} && cp /src/libnvsnap_intercept.so /out/" - -echo "Built: ${OUT_DIR}/libnvsnap_intercept.so" diff --git a/src/compute-plane-services/nvsnap/scripts/build-libzmq-image.sh b/src/compute-plane-services/nvsnap/scripts/build-libzmq-image.sh deleted file mode 100755 index 745c00dd52..0000000000 --- a/src/compute-plane-services/nvsnap/scripts/build-libzmq-image.sh +++ /dev/null @@ -1,116 +0,0 @@ -#!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# Build libzmq builder image with checkpoint/restore support -# This image contains libzmq.so with checkpoint APIs for injection into workloads - -set -euo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -PROJECT_ROOT="$(dirname "$SCRIPT_DIR")" - -# Logging -log_info() { echo "[INFO] $*"; } -log_error() { echo "[ERROR] $*" >&2; } - -# Configuration -REGISTRY="${REGISTRY:-nvcr.io/0651155215864979/ncp-dev}" -IMAGE_NAME="libzmq-builder" -VERSION="${VERSION:-v4.3.6-checkpoint-v1}" -source "$(dirname "${BASH_SOURCE[0]}")/_deps.sh" -nvsnap_resolve_sibling LIBZMQ_SRC libzmq -LIBZMQ_BRANCH="${LIBZMQ_BRANCH:-checkpoint-restore-v1}" -CONTAINER_TOOL="${CONTAINER_TOOL:-docker}" - -# Validate source -if [ ! -d "$LIBZMQ_SRC" ]; then - log_error "libzmq source not found at $LIBZMQ_SRC" - exit 1 -fi - -log_info "Building libzmq builder image" -log_info "Version: $VERSION" -log_info "Source: $LIBZMQ_SRC" -log_info "Branch: $LIBZMQ_BRANCH" - -# Create build context -BUILD_CONTEXT=$(mktemp -d) -trap "rm -rf $BUILD_CONTEXT" EXIT - -log_info "Preparing build context..." - -# Copy libzmq source (include builds/ directory for CMake helpers) -rsync -a --exclude='.git' --exclude='build' \ - "$LIBZMQ_SRC/" "$BUILD_CONTEXT/libzmq-src/" - -# Create Dockerfile -cat > "$BUILD_CONTEXT/Dockerfile" <<'EOF' -FROM ubuntu:22.04 - -RUN apt-get update && apt-get install -y --no-install-recommends \ - build-essential \ - cmake \ - git \ - pkg-config \ - libsodium-dev \ - && rm -rf /var/lib/apt/lists/* - -# Copy libzmq source with checkpoint support -COPY libzmq-src/ /libzmq-src/ -WORKDIR /libzmq-src - -# Build and install -RUN mkdir -p build && cd build && \ - cmake -DCMAKE_INSTALL_PREFIX=/usr/local \ - -DCMAKE_BUILD_TYPE=Release \ - -DBUILD_TESTS=OFF \ - .. && \ - make -j$(nproc) && \ - make install && \ - ldconfig - -# Verify checkpoint API is present -RUN nm -D /usr/local/lib/libzmq.so | grep -E "zmq_ctx_checkpoint|zmq_ctx_restore|zmq_get_all_contexts" || \ - (echo "ERROR: Checkpoint API not found in libzmq.so"; exit 1) - -# Show version info -RUN echo "libzmq version:" && \ - strings /usr/local/lib/libzmq.so | grep -E "^4\.[0-9]\.[0-9]" | head -1 && \ - echo "Checkpoint API symbols:" && \ - nm -D /usr/local/lib/libzmq.so | grep checkpoint - -# Default command for testing -CMD ["/bin/bash"] -EOF - -# Build image -log_info "Building Docker image..." -FULL_IMAGE="${REGISTRY}/${IMAGE_NAME}:${VERSION}" - -cd "$BUILD_CONTEXT" - -"${CONTAINER_TOOL}" build \ - -t "$FULL_IMAGE" \ - -t "${REGISTRY}/${IMAGE_NAME}:latest" \ - . - -log_info "Built: $FULL_IMAGE" - -# Push to registry -if [ "${PUSH:-true}" = "true" ]; then - log_info "Pushing to registry..." - "${CONTAINER_TOOL}" push "$FULL_IMAGE" - "${CONTAINER_TOOL}" push "${REGISTRY}/${IMAGE_NAME}:latest" - log_info "Pushed: $FULL_IMAGE" -fi - -echo "" -log_info "==========================================" -log_info "libzmq builder image ready: $FULL_IMAGE" -log_info "==========================================" -echo "" -echo "Usage in vLLM pod:" -echo " Copy from image: COPY --from=$FULL_IMAGE /usr/local/lib/libzmq.so* /usr/local/lib/" -echo " Or init container to inject at runtime" -echo "" diff --git a/src/compute-plane-services/nvsnap/scripts/build-pyzmq-wheel.sh b/src/compute-plane-services/nvsnap/scripts/build-pyzmq-wheel.sh deleted file mode 100755 index d5e553dc68..0000000000 --- a/src/compute-plane-services/nvsnap/scripts/build-pyzmq-wheel.sh +++ /dev/null @@ -1,161 +0,0 @@ -#!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# Build pyzmq wheel linked against our patched libzmq (CRIU restore support) -# -# Stock pyzmq wheels bundle their own libzmq inside pyzmq.libs/. -# This means our patched libzmq (with epoll reinit for CRIU restore) -# is completely bypassed — pyzmq calls the bundled stock copy. -# -# This script builds pyzmq from source with ZMQ_PREFIX pointing to our -# patched libzmq installation. The resulting wheel links against the -# system libzmq.so.5 (no bundled copy), so at runtime it uses whatever -# libzmq.so.5 is on LD_LIBRARY_PATH — our patched version. -# -# Same pattern as build-uvloop-wheel.sh. - -set -euo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -PROJECT_ROOT="$(dirname "$SCRIPT_DIR")" - -# Logging -log_info() { echo "[INFO] $*"; } -log_error() { echo "[ERROR] $*" >&2; } - -# Configuration -REGISTRY="${REGISTRY:-nvcr.io/0651155215864979/ncp-dev}" -IMAGE_NAME="pyzmq-builder" -VERSION="${VERSION:-v27.2.0-gpucr1}" -source "$(dirname "${BASH_SOURCE[0]}")/_deps.sh" -nvsnap_resolve_sibling PYZMQ_SRC pyzmq -nvsnap_resolve_sibling LIBZMQ_SRC libzmq -CONTAINER_TOOL="${CONTAINER_TOOL:-docker}" -# Build against vLLM base to match GLIBC version (2.35, Ubuntu 22.04) -BUILDER_BASE="${BUILDER_BASE:-vllm/vllm-openai:v0.11.2}" -PUSH="${PUSH:-true}" - -# Validate sources -if [ ! -d "$PYZMQ_SRC" ]; then - log_error "pyzmq source not found at $PYZMQ_SRC" - exit 1 -fi -if [ ! -d "$LIBZMQ_SRC" ]; then - log_error "libzmq source not found at $LIBZMQ_SRC" - exit 1 -fi - -log_info "Building pyzmq builder image" -log_info "Version: $VERSION" -log_info "pyzmq source: $PYZMQ_SRC" -log_info "libzmq source: $LIBZMQ_SRC" - -# Create build context -BUILD_CONTEXT=$(mktemp -d) -trap "rm -rf $BUILD_CONTEXT" EXIT - -log_info "Preparing build context..." - -# Copy pyzmq source -rsync -a \ - --exclude='.git' \ - --exclude='build' \ - --exclude='dist' \ - --exclude='*.egg-info' \ - --exclude='__pycache__' \ - --exclude='*.pyc' \ - --exclude='*.so' \ - --exclude='buildutils/bundled' \ - "$PYZMQ_SRC/" "$BUILD_CONTEXT/pyzmq-src/" - -# Copy libzmq source (for building from source inside Docker) -rsync -a \ - --exclude='.git' \ - --exclude='build' \ - "$LIBZMQ_SRC/" "$BUILD_CONTEXT/libzmq-src/" - -# Create Dockerfile -cat > "$BUILD_CONTEXT/Dockerfile" <=3.0.0' 'packaging' 'scikit-build-core>=0.10' - -# Build pyzmq with ZMQ_PREFIX=/usr/local (our patched libzmq) -# PYZMQ_NO_BUNDLE=1 prevents fallback to bundled libzmq -RUN ZMQ_PREFIX=/usr/local PYZMQ_NO_BUNDLE=1 \\ - pip wheel . -w /wheels/ --no-build-isolation -v - -# Verify: the wheel should NOT contain bundled libzmq -RUN echo "=== Wheel contents ===" && \\ - unzip -l /wheels/pyzmq-*.whl | grep -i "libzmq" && \\ - echo "(should show NO bundled libzmq-*.so files)" || true - -# Verify: pyzmq loads and uses our patched libzmq -# Run from /tmp to avoid importing source tree instead of installed wheel -RUN pip install --force-reinstall /wheels/pyzmq-*.whl && \\ - cd /tmp && python3 -c "import zmq; print('pyzmq version:', zmq.__version__); print('zmq.zmq_version():', zmq.zmq_version()); print('Using system libzmq (CRIU restore support)')" - -# Minimal output image with just the wheel -FROM python:3.12-slim -COPY --from=builder /wheels/ /wheels/ -CMD ["ls", "-la", "/wheels/"] -DOCKERFILE - -# Build image -log_info "Building Docker image..." -FULL_IMAGE="${REGISTRY}/${IMAGE_NAME}:${VERSION}" - -cd "$BUILD_CONTEXT" - -"${CONTAINER_TOOL}" build \ - -t "$FULL_IMAGE" \ - -t "${REGISTRY}/${IMAGE_NAME}:latest" \ - . - -log_info "Built: $FULL_IMAGE" - -# Push to registry -if [ "$PUSH" = "true" ]; then - log_info "Pushing to registry..." - "${CONTAINER_TOOL}" push "$FULL_IMAGE" - "${CONTAINER_TOOL}" push "${REGISTRY}/${IMAGE_NAME}:latest" - log_info "Pushed: $FULL_IMAGE" -fi - -echo "" -log_info "==========================================" -log_info "pyzmq builder image ready: $FULL_IMAGE" -log_info "==========================================" -echo "" -echo "Usage in K8s init container:" -echo " image: $FULL_IMAGE" -echo " command: cp /wheels/pyzmq-*.whl /nvsnap-lib/" -echo "" -echo "Then in vLLM startup:" -echo " pip install --force-reinstall --no-deps /nvsnap-lib/pyzmq-*.whl" -echo " (LD_LIBRARY_PATH must include /nvsnap-lib/ where libzmq.so.5 lives)" diff --git a/src/compute-plane-services/nvsnap/scripts/build-uvloop-wheel.sh b/src/compute-plane-services/nvsnap/scripts/build-uvloop-wheel.sh deleted file mode 100755 index 515b04bbcc..0000000000 --- a/src/compute-plane-services/nvsnap/scripts/build-uvloop-wheel.sh +++ /dev/null @@ -1,102 +0,0 @@ -#!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# Build uvloop wheels (cp310/311/312/313) with CRIU checkpoint/restore -# support. See docs/GENERIC-PYTHON-INJECTION-DESIGN.md for the design. -# -# The patch adds ~15 lines of Cython to uvloop/loop.pyx: -# - detects CRIU restore via /run/criu-restored marker -# - calls uv_loop_fork() on first _run() after restore -# - reinitializes libuv kernel state (epoll, signal pipes, io_uring) -# -# Builder image: docker/uvloop/Dockerfile (manylinux_2_28_x86_64 base). -# Output image carries 4 wheels under /wheels/, one per Python ABI tag. - -set -euo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -PROJECT_ROOT="$(dirname "$SCRIPT_DIR")" - -log_info() { echo "[INFO] $*"; } -log_error() { echo "[ERROR] $*" >&2; } - -REGISTRY="${REGISTRY:-nvcr.io/0651155215864979/ncp-dev}" -IMAGE_NAME="uvloop-builder" -VERSION="${VERSION:-v0.22.1-multipy1}" -source "$(dirname "${BASH_SOURCE[0]}")/_deps.sh" -nvsnap_resolve_sibling UVLOOP_SRC uvloop -UVLOOP_BRANCH="${UVLOOP_BRANCH:-checkpoint-restore-v1}" -CONTAINER_TOOL="${CONTAINER_TOOL:-docker}" -PUSH="${PUSH:-true}" - -if [ ! -d "$UVLOOP_SRC" ]; then - log_error "uvloop source not found at $UVLOOP_SRC" - exit 1 -fi - -log_info "uvloop builder" -log_info " version: $VERSION" -log_info " source : $UVLOOP_SRC (branch $UVLOOP_BRANCH)" - -# Prepare build context: uvloop source + checked-in Dockerfile. -# We rsync the source rather than mounting so the Docker daemon can use it -# even when the daemon runs remotely. -BUILD_CONTEXT="$(mktemp -d)" -trap 'rm -rf "$BUILD_CONTEXT"' EXIT - -log_info "Preparing build context at $BUILD_CONTEXT" -rsync -a \ - --exclude='.git' \ - --exclude='build' \ - --exclude='dist' \ - --exclude='*.egg-info' \ - --exclude='__pycache__' \ - --exclude='*.pyc' \ - --exclude='*.so' \ - "$UVLOOP_SRC/" "$BUILD_CONTEXT/" - -# Use the checked-in Dockerfile (no inline heredoc — see CLAUDE.md rule #13). -cp "$PROJECT_ROOT/docker/uvloop/Dockerfile" "$BUILD_CONTEXT/Dockerfile" - -FULL_IMAGE="${REGISTRY}/${IMAGE_NAME}:${VERSION}" -log_info "Building $FULL_IMAGE" - -"${CONTAINER_TOOL}" build \ - -t "$FULL_IMAGE" \ - -t "${REGISTRY}/${IMAGE_NAME}:latest" \ - "$BUILD_CONTEXT" - -log_info "Built $FULL_IMAGE" - -if [ "$PUSH" = "true" ]; then - log_info "Pushing..." - "${CONTAINER_TOOL}" push "$FULL_IMAGE" - "${CONTAINER_TOOL}" push "${REGISTRY}/${IMAGE_NAME}:latest" - log_info "Pushed $FULL_IMAGE" -fi - -echo -log_info "==================================================" -log_info "uvloop multi-Python builder ready: $FULL_IMAGE" -log_info " /wheels/uvloop-*-cp310-cp310-*.whl" -log_info " /wheels/uvloop-*-cp311-cp311-*.whl" -log_info " /wheels/uvloop-*-cp312-cp312-*.whl" -log_info " /wheels/uvloop-*-cp313-cp313-*.whl" -log_info "==================================================" -echo -echo "Init container snippet (per docs/GENERIC-PYTHON-INJECTION-DESIGN.md):" -cat <<'EOF' - - name: get-uvloop - image: nvcr.io/0651155215864979/ncp-dev/uvloop-builder:VERSION - command: ["/bin/sh", "-c"] - args: - - | - for whl in /wheels/uvloop-*.whl; do - tag=$(echo "$whl" | grep -oE 'cp3[0-9]+' | head -1) - mkdir -p "/nvsnap-lib/site-packages-${tag}" - python3 -m zipfile -e "$whl" "/nvsnap-lib/site-packages-${tag}/" - done - volumeMounts: - - { name: nvsnap-lib, mountPath: /nvsnap-lib } -EOF diff --git a/src/compute-plane-services/nvsnap/scripts/local-uvloop-test.sh b/src/compute-plane-services/nvsnap/scripts/local-uvloop-test.sh deleted file mode 100755 index 9edc228afd..0000000000 --- a/src/compute-plane-services/nvsnap/scripts/local-uvloop-test.sh +++ /dev/null @@ -1,235 +0,0 @@ -#!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -set -euo pipefail - -REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" - -ROOT="${NVSNAP_LOCAL_ROOT:-/tmp/nvsnap-local}" -LIB_DIR="${ROOT}/nvsnap-lib" -RUN_DIR="${ROOT}/run" -CHECKPOINT_DIR="${ROOT}/checkpoints" -BUNDLE_DIR="${ROOT}/criu-bundle" -AGENT_BIN="${ROOT}/nvsnap-agent" - -IMAGE="${NVSNAP_LOCAL_IMAGE:-nvsnap-uvloop-test:local}" -SRC_CONTAINER="${NVSNAP_LOCAL_CONTAINER:-nvsnap-uvloop-src}" -RESTORE_CONTAINER="${NVSNAP_LOCAL_RESTORE_CONTAINER:-nvsnap-uvloop-restore}" - -AGENT_PORT="${NVSNAP_AGENT_PORT:-8081}" -CONTAINERD_SOCKET="${NVSNAP_CONTAINERD_SOCKET:-/run/containerd/containerd.sock}" -DOCKER_HOST_DEFAULT="unix:///var/run/docker.sock" -export DOCKER_HOST="${NVSNAP_DOCKER_HOST:-$DOCKER_HOST_DEFAULT}" - -usage() { - cat < - -Commands: - build Build uvloop test image + intercept lib - start Start uvloop test container - agent Build + start local agent - checkpoint Trigger checkpoint (cleans checkpoints first) - restore Restore latest checkpoint into new container - logs Show container logs - cleanup Stop container + agent + remove local dirs - -Env overrides: - NVSNAP_LOCAL_ROOT (default: /tmp/nvsnap-local) - NVSNAP_LOCAL_IMAGE (default: nvsnap-uvloop-test:local) - NVSNAP_LOCAL_CONTAINER (default: nvsnap-uvloop-src) - NVSNAP_AGENT_PORT (default: 8081) - NVSNAP_CONTAINERD_SOCKET (default: /run/containerd/containerd.sock) - NVSNAP_CONTAINERD_NAMESPACE (default: auto-detect) -EOF - exit 1 -} - -ensure_dirs() { - mkdir -p "${LIB_DIR}" "${RUN_DIR}" "${CHECKPOINT_DIR}" -} - -detect_namespace() { - local ns="${NVSNAP_CONTAINERD_NAMESPACE:-}" - if [[ -n "${ns}" ]]; then - echo "${ns}" - return - fi - if command -v ctr >/dev/null 2>&1; then - if ctr namespaces list 2>/dev/null | awk '{print $1}' | grep -q "^moby$"; then - echo "moby" - return - fi - if ctr namespaces list 2>/dev/null | awk '{print $1}' | grep -q "^default$"; then - echo "default" - return - fi - fi - echo "k8s.io" -} - -build_intercept() { - echo "Building intercept library (ubuntu:22.04 toolchain)..." - "${REPO_ROOT}/scripts/build-intercept-lib-local.sh" "${LIB_DIR}" -} - -build_image() { - echo "Building uvloop test image..." - docker build -f "${REPO_ROOT}/tests/uvloop-mp-test/Dockerfile.pip" -t "${IMAGE}" "${REPO_ROOT}" -} - -cmd_build() { - ensure_dirs - build_intercept - build_image - echo "Build complete." -} - -cmd_start() { - ensure_dirs - rm -f "${LIB_DIR}/.restored" "${LIB_DIR}/.force_uvloop_fork" "${LIB_DIR}/.debug_io_uring" "${LIB_DIR}/uvloop_loops."*".json" "${LIB_DIR}/uvloop_loops.json" 2>/dev/null || true - rm -f "${RUN_DIR}/.restored" "${RUN_DIR}/.force_uvloop_fork" "${RUN_DIR}/.debug_io_uring" "${RUN_DIR}/uvloop_loops."*".json" "${RUN_DIR}/uvloop_loops.json" 2>/dev/null || true - docker rm -f "${SRC_CONTAINER}" >/dev/null 2>&1 || true - echo "Starting uvloop test container..." - docker run -d --name "${SRC_CONTAINER}" --privileged \ - -p 8000:8000 \ - -e LD_PRELOAD=/nvsnap-lib/libnvsnap_intercept.so \ - -e NVSNAP_QUIESCE_SIGNALS=1 \ - -e NVSNAP_QUIESCE_METADATA_ONLY=1 \ - -v "${LIB_DIR}:/nvsnap-lib" \ - -v "${RUN_DIR}:/var/run/nvsnap" \ - "${IMAGE}" >/dev/null - echo "Container started: ${SRC_CONTAINER}" -} - -cmd_agent() { - ensure_dirs - if [[ ! -S "${CONTAINERD_SOCKET}" ]]; then - echo "ERROR: containerd socket not found at ${CONTAINERD_SOCKET}" - exit 1 - fi - - local ns - ns="$(detect_namespace)" - echo "Using containerd namespace: ${ns}" - - echo "Building agent..." - (cd "${REPO_ROOT}" && go build -o "${AGENT_BIN}" ./cmd/agent) - - if [[ -f "${ROOT}/agent.pid" ]]; then - kill "$(cat "${ROOT}/agent.pid")" >/dev/null 2>&1 || true - fi - - echo "Starting agent on :${AGENT_PORT}..." - "${AGENT_BIN}" \ - --listen ":${AGENT_PORT}" \ - --log-level debug \ - --containerd-socket "${CONTAINERD_SOCKET}" \ - --containerd-namespace "${ns}" \ - --checkpoint-dir "${CHECKPOINT_DIR}" \ - --criu-path "${CRIU_PATH:-/usr/local/sbin/criu}" \ - --cuda-checkpoint-path "${CUDA_CHECKPOINT_PATH:-/usr/local/bin/cuda-checkpoint}" \ - > "${ROOT}/agent.log" 2>&1 & - - echo $! > "${ROOT}/agent.pid" - sleep 1 - if ! kill -0 "$(cat "${ROOT}/agent.pid")" >/dev/null 2>&1; then - echo "Agent failed to start. Log:" - tail -n 20 "${ROOT}/agent.log" || true - exit 1 - fi - echo "Agent started (pid $(cat "${ROOT}/agent.pid"))" -} - -cmd_checkpoint() { - ensure_dirs - echo "Cleaning checkpoints in ${CHECKPOINT_DIR}..." - if ! rm -rf "${CHECKPOINT_DIR:?}/"* 2>/dev/null; then - printf '%s\n' "0mMurug@" | sudo -S rm -rf "${CHECKPOINT_DIR:?}/"* || true - fi - - local container_id - container_id="$(docker inspect -f '{{.Id}}' "${SRC_CONTAINER}")" - - local container_name - container_name="$(docker inspect -f '{{.Name}}' "${SRC_CONTAINER}" | sed 's,^/,,')" - echo "Triggering checkpoint for container ${container_id:0:12} (${container_name})..." - curl -s -X POST "http://127.0.0.1:${AGENT_PORT}/v1/checkpoint" \ - -H "Content-Type: application/json" \ - -d "{\"namespace\":\"local\",\"containerName\":\"${container_name}\",\"containerId\":\"${container_id}\"}" - echo "" -} - -cmd_logs() { - docker logs --tail 200 "${SRC_CONTAINER}" -} - -latest_checkpoint_id() { - if [[ -n "${NVSNAP_CHECKPOINT_ID:-}" ]]; then - echo "${NVSNAP_CHECKPOINT_ID}" - return - fi - local latest - latest="$(ls -1dt "${CHECKPOINT_DIR}/"* 2>/dev/null | head -n1 || true)" - if [[ -z "${latest}" ]]; then - echo "" - return - fi - basename "${latest}" -} - -cmd_restore() { - ensure_dirs - local checkpoint_id - checkpoint_id="$(latest_checkpoint_id)" - if [[ -z "${checkpoint_id}" ]]; then - echo "ERROR: no checkpoints found in ${CHECKPOINT_DIR}" - exit 1 - fi - - docker rm -f "${RESTORE_CONTAINER}" >/dev/null 2>&1 || true - # Stop source container to free port 8000. - docker rm -f "${SRC_CONTAINER}" >/dev/null 2>&1 || true - - echo "Restoring checkpoint ${checkpoint_id} into ${RESTORE_CONTAINER}..." - local force_fork="${NVSNAP_FORCE_UVLOOP_FORK:-1}" - docker run -d --name "${RESTORE_CONTAINER}" --privileged \ - -p 8000:8000 \ - -e CRIU_BUNDLE_PATH=/nvsnap \ - -e CHECKPOINT_PATH=/checkpoints \ - -e CHECKPOINT_ID="${checkpoint_id}" \ - -e NVSNAP_FORCE_UVLOOP_FORK="${force_fork}" \ - -e NVSNAP_LOG_FILE=/var/run/nvsnap/nvsnap.log \ - -e NVSNAP_LOG_LEVEL=1 \ - -e NVSNAP_DEBUG_IO_URING=1 \ - -v "${REPO_ROOT}/bin/criu-bundle:/nvsnap:ro" \ - -v "${CHECKPOINT_DIR}:/checkpoints" \ - -v "${LIB_DIR}:/nvsnap-lib" \ - -v "${RUN_DIR}:/var/run/nvsnap" \ - "${IMAGE}" /nvsnap/restore-entrypoint >/dev/null - - echo "Restore container started: ${RESTORE_CONTAINER}" -} - -cmd_cleanup() { - if [[ -f "${ROOT}/agent.pid" ]]; then - kill "$(cat "${ROOT}/agent.pid")" >/dev/null 2>&1 || true - rm -f "${ROOT}/agent.pid" - fi - docker rm -f "${SRC_CONTAINER}" >/dev/null 2>&1 || true - docker rm -f "${RESTORE_CONTAINER}" >/dev/null 2>&1 || true - rm -rf "${ROOT}" - echo "Cleaned ${ROOT}" -} - -case "${1:-}" in - build) cmd_build ;; - start) cmd_start ;; - agent) cmd_agent ;; - checkpoint) cmd_checkpoint ;; - restore) cmd_restore ;; - logs) cmd_logs ;; - cleanup) cmd_cleanup ;; - *) usage ;; -esac diff --git a/src/compute-plane-services/nvsnap/scripts/restore-bundle-init.sh b/src/compute-plane-services/nvsnap/scripts/restore-bundle-init.sh index 5dc5b765ea..1918d60008 100755 --- a/src/compute-plane-services/nvsnap/scripts/restore-bundle-init.sh +++ b/src/compute-plane-services/nvsnap/scripts/restore-bundle-init.sh @@ -18,7 +18,7 @@ # The intercept payload (libnvsnap_intercept.so, patched uvloop/libuv/ # libzmq, sitecustomize) is no longer staged: criu-v2 dumps and restores # in-namespace, so no userspace interception is injected into workloads. -# lib/nvsnap_intercept/ stays in-tree for future multi-GPU work. +# the interception stack was removed entirely. set -euo pipefail diff --git a/src/compute-plane-services/nvsnap/scripts/sync-versions.sh b/src/compute-plane-services/nvsnap/scripts/sync-versions.sh index f5ebdff9d6..66d4a81339 100755 --- a/src/compute-plane-services/nvsnap/scripts/sync-versions.sh +++ b/src/compute-plane-services/nvsnap/scripts/sync-versions.sh @@ -63,12 +63,7 @@ set_chart_tag() { declare -A IMAGES=( [nvsnap-agent]="$NVSNAP_APP_VERSION" [nvsnap-server]="$NVSNAP_SERVER_VERSION" - [nvsnap-init]="$NVSNAP_INIT_VERSION" [nvsnap-blobstore]="$NVSNAP_BLOBSTORE_VERSION" - [uvloop-builder]="$NVSNAP_UVLOOP_VERSION" - [libuv-builder]="$NVSNAP_LIBUV_VERSION" - [libzmq-builder]="$NVSNAP_LIBZMQ_VERSION" - [pyzmq-builder]="$NVSNAP_PYZMQ_VERSION" ) for name in "${!IMAGES[@]}"; do diff --git a/src/compute-plane-services/nvsnap/scripts/test-gpu-checkpoint.sh b/src/compute-plane-services/nvsnap/scripts/test-gpu-checkpoint.sh deleted file mode 100755 index 0def223bbc..0000000000 --- a/src/compute-plane-services/nvsnap/scripts/test-gpu-checkpoint.sh +++ /dev/null @@ -1,253 +0,0 @@ -#!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# NVSNAP GPU Checkpoint Test Script -# -# Tests GPU checkpoint/restore using cuda-checkpoint on a running GPU process. -# Can be run standalone or as part of the NVSNAP test suite. -# -# Usage: -# ./test-gpu-checkpoint.sh [--pid ] [--simple] -# -# Options: -# --pid Checkpoint an existing process -# --simple Run a simple CUDA test program instead of vLLM - -set -euo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -NVSNAP_ROOT="$(dirname "$SCRIPT_DIR")" -CHECKPOINT_DIR="${NVSNAP_CHECKPOINT_DIR:-/tmp/nvsnap-checkpoint}" - -RED='\033[0;31m' -GREEN='\033[0;32m' -YELLOW='\033[1;33m' -NC='\033[0m' - -log() { echo -e "${GREEN}[INFO]${NC} $1"; } -warn() { echo -e "${YELLOW}[WARN]${NC} $1"; } -error() { echo -e "${RED}[ERROR]${NC} $1"; exit 1; } - -# Check prerequisites -check_prereqs() { - log "Checking prerequisites..." - - if ! command -v cuda-checkpoint &>/dev/null; then - error "cuda-checkpoint not found. Install with: scripts/install-node.sh cuda-checkpoint" - fi - - if ! nvidia-smi &>/dev/null; then - error "nvidia-smi not available. Is the NVIDIA driver installed?" - fi - - log "Prerequisites OK" -} - -# Get GPU state for a process -get_gpu_state() { - local pid=$1 - cuda-checkpoint --get-state --pid "$pid" 2>/dev/null || echo "unknown" -} - -# Full checkpoint cycle -checkpoint_process() { - local pid=$1 - local timeout_ms=${2:-30000} - - log "Checkpointing process $pid..." - - # Get initial state - local state=$(get_gpu_state "$pid") - log "Initial GPU state: $state" - - if [[ "$state" != "running" ]]; then - error "Process $pid is not in running state (state: $state)" - fi - - # Lock - log "Locking GPU state..." - if ! sudo cuda-checkpoint --action lock --pid "$pid" --timeout "$timeout_ms"; then - error "Failed to lock GPU state" - fi - - state=$(get_gpu_state "$pid") - log "State after lock: $state" - - # Checkpoint - log "Checkpointing GPU state..." - if ! sudo cuda-checkpoint --action checkpoint --pid "$pid"; then - warn "Checkpoint failed, attempting restore..." - sudo cuda-checkpoint --action restore --pid "$pid" || true - sudo cuda-checkpoint --action unlock --pid "$pid" || true - error "Failed to checkpoint GPU state" - fi - - state=$(get_gpu_state "$pid") - log "State after checkpoint: $state" - - log "GPU checkpoint complete!" - return 0 -} - -# Restore GPU state -restore_process() { - local pid=$1 - - log "Restoring process $pid..." - - local state=$(get_gpu_state "$pid") - log "Current GPU state: $state" - - # Restore - log "Restoring GPU state..." - if ! sudo cuda-checkpoint --action restore --pid "$pid"; then - error "Failed to restore GPU state" - fi - - state=$(get_gpu_state "$pid") - log "State after restore: $state" - - # Unlock - log "Unlocking GPU state..." - if ! sudo cuda-checkpoint --action unlock --pid "$pid"; then - error "Failed to unlock GPU state" - fi - - state=$(get_gpu_state "$pid") - log "Final state: $state" - - log "GPU restore complete!" - return 0 -} - -# Run simple CUDA test -run_simple_test() { - log "Running simple CUDA checkpoint test..." - - # Build test program if needed - local test_prog="$NVSNAP_ROOT/lib/nvsnap_intercept/tests/test_simple_checkpoint" - if [[ ! -f "$test_prog" ]]; then - log "Building test program..." - make -C "$NVSNAP_ROOT/lib/nvsnap_intercept" tests/test_simple_checkpoint - fi - - # Start test program in background - log "Starting test program..." - "$test_prog" & - local pid=$! - - # Wait for it to initialize - sleep 2 - - if ! kill -0 "$pid" 2>/dev/null; then - error "Test program exited prematurely" - fi - - log "Test program running with PID $pid" - - # Do checkpoint cycle - checkpoint_process "$pid" - - log "Waiting 2 seconds..." - sleep 2 - - restore_process "$pid" - - # Verify process is still running - sleep 1 - if kill -0 "$pid" 2>/dev/null; then - log "✅ Test PASSED - Process still running after checkpoint/restore" - kill "$pid" 2>/dev/null || true - return 0 - else - error "❌ Test FAILED - Process died after restore" - fi -} - -# Full checkpoint/restore with CRIU -run_full_test() { - local pid=$1 - - log "Running full checkpoint/restore test with CRIU..." - - if ! command -v criu &>/dev/null; then - error "CRIU not found. Install with: scripts/install-node.sh criu" - fi - - mkdir -p "$CHECKPOINT_DIR" - - # GPU checkpoint - checkpoint_process "$pid" - - # CRIU dump - log "Running CRIU dump..." - if ! sudo criu dump -t "$pid" -D "$CHECKPOINT_DIR" --shell-job -v2; then - warn "CRIU dump failed, restoring GPU..." - restore_process "$pid" - error "CRIU dump failed" - fi - - log "Process checkpointed to $CHECKPOINT_DIR" - - # CRIU restore - log "Running CRIU restore..." - if ! sudo criu restore -D "$CHECKPOINT_DIR" --shell-job -d -v2; then - error "CRIU restore failed" - fi - - # Find new PID (CRIU may assign new PID) - local new_pid=$(pgrep -f "test_simple_checkpoint" | head -1) - if [[ -z "$new_pid" ]]; then - error "Could not find restored process" - fi - - log "Process restored with PID $new_pid" - - # GPU restore - restore_process "$new_pid" - - log "✅ Full checkpoint/restore complete!" -} - -# Main -main() { - local mode="simple" - local target_pid="" - - while [[ $# -gt 0 ]]; do - case "$1" in - --pid) - target_pid="$2" - shift 2 - ;; - --simple) - mode="simple" - shift - ;; - --full) - mode="full" - shift - ;; - *) - echo "Usage: $0 [--pid ] [--simple|--full]" - exit 1 - ;; - esac - done - - check_prereqs - - if [[ -n "$target_pid" ]]; then - checkpoint_process "$target_pid" - sleep 2 - restore_process "$target_pid" - elif [[ "$mode" == "simple" ]]; then - run_simple_test - else - run_simple_test - # run_full_test uses the simple test process - fi -} - -main "$@" diff --git a/src/compute-plane-services/nvsnap/scripts/test-vllm-zmq.sh b/src/compute-plane-services/nvsnap/scripts/test-vllm-zmq.sh deleted file mode 100755 index 27179cfdba..0000000000 --- a/src/compute-plane-services/nvsnap/scripts/test-vllm-zmq.sh +++ /dev/null @@ -1,294 +0,0 @@ -#!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# Automated vLLM Checkpoint/Restore Test -# Tests the complete flow: deploy → wait → checkpoint → restore → verify inference -# -# Uses kubectl exec + curl inside the pod for API access (reliable, no port-forward). -# Prints a structured timing summary at the end. -# -# Exit codes: 0 = PASS, 1 = FAIL (with which step failed) - -set -euo pipefail - -SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" -PROJECT_ROOT="$(dirname "$SCRIPT_DIR")" - -# Colors -GREEN='\033[0;32m' -RED='\033[0;31m' -YELLOW='\033[1;33m' -CYAN='\033[0;36m' -BOLD='\033[1m' -NC='\033[0m' - -log_info() { echo -e "${GREEN}[INFO]${NC} $*"; } -log_error() { echo -e "${RED}[ERROR]${NC} $*"; } -log_warn() { echo -e "${YELLOW}[WARN]${NC} $*"; } - -# Configuration -NAMESPACE="nvsnap-system" -POD_NAME="vllm-small" -CONTAINER_NAME="vllm" -RESTORE_POD_NAME="vllm-small-restored" -RESTORE_CONTAINER_NAME="restore" -VLLM_PORT=8000 - -# Timeouts (seconds) -POD_READY_TIMEOUT=600 # 10min: image pull + model download + torch compile -MODELS_POLL_TIMEOUT=600 # 10min: extra safety for /v1/models -MODELS_POLL_INTERVAL=30 -INFERENCE_POLL_TIMEOUT=300 # 5min: inference warmup -INFERENCE_POLL_INTERVAL=30 -RESTORE_READY_TIMEOUT=600 # 10min: CRIU restore + GPU restore -POST_MODELS_TIMEOUT=120 # 2min: post-restore /v1/models -POST_MODELS_INTERVAL=10 -POST_INFER_TIMEOUT=120 # 2min: post-restore /v1/completions -POST_INFER_INTERVAL=10 - -# ─── Timing infrastructure ────────────────────────────────────────────────── -declare -a STEP_NAMES=() -declare -a STEP_DURATIONS=() -declare -a STEP_RESULTS=() -TEST_START=$(date +%s) - -step_start() { - CURRENT_STEP_START=$(date +%s) -} - -step_done() { - local name="$1" result="$2" - local elapsed=$(( $(date +%s) - CURRENT_STEP_START )) - STEP_NAMES+=("$name") - STEP_DURATIONS+=("$elapsed") - STEP_RESULTS+=("$result") -} - -fmt_duration() { - local secs="$1" - printf "%dm %02ds" $((secs / 60)) $((secs % 60)) -} - -print_summary() { - local total_elapsed=$(( $(date +%s) - TEST_START )) - local overall="PASS" - - echo "" - echo -e "${BOLD}${CYAN}Step Duration Result${NC}" - echo -e "${CYAN}──────────────────────────────────────────────────${NC}" - for i in "${!STEP_NAMES[@]}"; do - local name="${STEP_NAMES[$i]}" - local dur=$(fmt_duration "${STEP_DURATIONS[$i]}") - local res="${STEP_RESULTS[$i]}" - local color="$GREEN" - if [ "$res" = "FAIL" ]; then - color="$RED" - overall="FAIL" - elif [ "$res" = "SKIP" ]; then - color="$YELLOW" - fi - printf "%-28s %-11s ${color}%s${NC}\n" "$name" "$dur" "$res" - done - echo -e "${CYAN}──────────────────────────────────────────────────${NC}" - local total_dur=$(fmt_duration "$total_elapsed") - local total_color="$GREEN" - if [ "$overall" = "FAIL" ]; then total_color="$RED"; fi - printf "%-28s %-11s ${total_color}%s${NC}\n" "Total" "$total_dur" "$overall" - echo "" - - if [ "$overall" = "FAIL" ]; then - return 1 - fi - return 0 -} - -# ─── Helper: call API via kubectl exec ──────────────────────────────────────── -# Runs curl inside the pod — no port-forward needed, always reliable. -# Usage: pod_curl [data] [timeout] -pod_curl() { - local pod="$1" container="$2" method="$3" path="$4" data="${5:-}" timeout="${6:-10}" - - # Clear LD_PRELOAD so the intercept library doesn't load into curl - if [ -n "$data" ]; then - kubectl exec -n "$NAMESPACE" "$pod" -c "$container" -- \ - env LD_PRELOAD= curl -s -m "$timeout" -X "$method" "http://localhost:${VLLM_PORT}${path}" \ - -H "Content-Type: application/json" -d "$data" 2>/dev/null - else - kubectl exec -n "$NAMESPACE" "$pod" -c "$container" -- \ - env LD_PRELOAD= curl -s -m "$timeout" -X "$method" "http://localhost:${VLLM_PORT}${path}" 2>/dev/null - fi -} - -# ─── Helper: poll until API responds ────────────────────────────────────────── -# Usage: poll_api -# Returns: 0 on success (sets POLL_RESULT), 1 on timeout -POLL_RESULT="" -poll_api() { - local pod="$1" container="$2" method="$3" path="$4" data="$5" - local pattern="$6" timeout_secs="$7" interval="$8" desc="$9" - - local deadline=$(( $(date +%s) + timeout_secs )) - local attempt=0 - while [ "$(date +%s)" -lt "$deadline" ]; do - attempt=$((attempt + 1)) - POLL_RESULT=$(pod_curl "$pod" "$container" "$method" "$path" "$data" 30 || true) - if echo "$POLL_RESULT" | grep -q "$pattern"; then - log_info " $desc OK (attempt $attempt)" - return 0 - fi - local remaining=$(( deadline - $(date +%s) )) - if [ "$remaining" -le 0 ]; then break; fi - log_info " attempt $attempt: $desc not ready, retrying in ${interval}s (${remaining}s left)..." - sleep "$interval" - done - log_error " $desc not responding after ${timeout_secs}s" - return 1 -} - -# ─── Fail handler ──────────────────────────────────────────────────────────── -fail() { - local step="$1" - step_done "$step" "FAIL" - log_error "FAILED at: $step" - print_summary || true - exit 1 -} - -log_info "==========================================" -log_info "vLLM Checkpoint/Restore Test" -log_info "==========================================" -echo "" - -# ─── Step 1: Clean up ──────────────────────────────────────────────────────── -log_info "Step 1: Cleaning up existing pods..." -kubectl delete pod $POD_NAME $RESTORE_POD_NAME -n $NAMESPACE --ignore-not-found -sleep 3 - -# ─── Step 2: Deploy vLLM ───────────────────────────────────────────────────── -log_info "Step 2: Deploying vLLM pod..." -kubectl apply -f "$PROJECT_ROOT/deploy/k8s/vllm-small.yaml" - -# ─── Step 3: Wait for pod ready (readiness probe checks /v1/models) ────────── -step_start -log_info "Step 3: Waiting for pod ready (up to ${POD_READY_TIMEOUT}s)..." -if kubectl wait --for=condition=ready pod/$POD_NAME -n $NAMESPACE --timeout=${POD_READY_TIMEOUT}s; then - step_done "Pod ready" "OK" -else - kubectl logs $POD_NAME -n $NAMESPACE -c $CONTAINER_NAME --tail=20 || true - fail "Pod ready" -fi - -# ─── Step 4: Verify /v1/models ─────────────────────────────────────────────── -step_start -log_info "Step 4: Verifying /v1/models responds..." -if poll_api "$POD_NAME" "$CONTAINER_NAME" GET /v1/models "" "TinyLlama" \ - "$MODELS_POLL_TIMEOUT" "$MODELS_POLL_INTERVAL" "/v1/models"; then - step_done "Models API ready" "OK" -else - kubectl logs $POD_NAME -n $NAMESPACE -c $CONTAINER_NAME --tail=30 || true - fail "Models API ready" -fi - -# ─── Step 5: Verify inference ──────────────────────────────────────────────── -step_start -log_info "Step 5: Verifying inference works before checkpoint..." -if poll_api "$POD_NAME" "$CONTAINER_NAME" POST /v1/completions \ - '{"model":"TinyLlama/TinyLlama-1.1B-Chat-v1.0","prompt":"Hello","max_tokens":5}' \ - '"choices"' "$INFERENCE_POLL_TIMEOUT" "$INFERENCE_POLL_INTERVAL" "/v1/completions"; then - echo "$POLL_RESULT" | python3 -m json.tool 2>/dev/null || echo "$POLL_RESULT" - step_done "Pre-checkpoint infer" "OK" -else - kubectl logs $POD_NAME -n $NAMESPACE -c $CONTAINER_NAME --tail=30 || true - fail "Pre-checkpoint infer" -fi - -# ─── Step 6: Get node ──────────────────────────────────────────────────────── -POD_NODE=$(kubectl get pod $POD_NAME -n $NAMESPACE -o jsonpath='{.spec.nodeName}') -log_info "Pod on node: $POD_NODE" - -# ─── Step 7: Create checkpoint ─────────────────────────────────────────────── -step_start -log_info "Step 7: Creating checkpoint..." -CHECKPOINT_OUTPUT=$(${SCRIPT_DIR}/checkpoint.sh create $POD_NAME $CONTAINER_NAME $NAMESPACE 2>&1) || true -CHECKPOINT_ID=$(echo "$CHECKPOINT_OUTPUT" | grep "Checkpoint ID:" | awk '{print $NF}') - -if [ -z "$CHECKPOINT_ID" ]; then - log_error "Failed to create checkpoint" - echo "$CHECKPOINT_OUTPUT" - fail "Checkpoint" -fi -log_info "Checkpoint: $CHECKPOINT_ID" -step_done "Checkpoint" "OK" - -# ─── Step 8: Delete original pod ───────────────────────────────────────────── -log_info "Step 8: Deleting original pod..." -kubectl delete pod $POD_NAME -n $NAMESPACE --wait=false -sleep 5 - -# ─── Step 9: Restore from checkpoint ───────────────────────────────────────── -step_start -log_info "Step 9: Restoring on node $POD_NODE..." -RESTORE_MANIFEST=$(mktemp) -trap "rm -f $RESTORE_MANIFEST" EXIT - -sed -e "s|value: \"vllm-small__nvsnap-system__[0-9-]*\"|value: \"$CHECKPOINT_ID\"|" \ - -e "s|nodeName: .*|nodeName: $POD_NODE|" \ - "$PROJECT_ROOT/deploy/k8s/vllm-small-restore.yaml" > "$RESTORE_MANIFEST" - -kubectl apply -f "$RESTORE_MANIFEST" - -log_info "Waiting for restore pod ready (up to ${RESTORE_READY_TIMEOUT}s)..." -log_info " (readiness probe polls /v1/models — succeeds only when vLLM is serving)" -if kubectl wait --for=condition=ready pod/$RESTORE_POD_NAME -n $NAMESPACE --timeout=${RESTORE_READY_TIMEOUT}s; then - step_done "Restore pod ready" "OK" -else - log_warn "Restore pod not ready, checking status..." - kubectl get pod $RESTORE_POD_NAME -n $NAMESPACE -o wide || true - kubectl logs $RESTORE_POD_NAME -n $NAMESPACE -c $RESTORE_CONTAINER_NAME --tail=30 || true - fail "Restore pod ready" -fi - -# ─── Step 10: Post-restore /v1/models ──────────────────────────────────────── -step_start -log_info "Step 10: Verifying /v1/models after restore..." -if poll_api "$RESTORE_POD_NAME" "$RESTORE_CONTAINER_NAME" GET /v1/models "" "TinyLlama" \ - "$POST_MODELS_TIMEOUT" "$POST_MODELS_INTERVAL" "post-restore /v1/models"; then - step_done "Post-restore models" "OK" -else - kubectl logs $RESTORE_POD_NAME -n $NAMESPACE -c $RESTORE_CONTAINER_NAME --tail=30 || true - fail "Post-restore models" -fi - -# ─── Step 11: Post-restore /v1/completions ─────────────────────────────────── -step_start -log_info "Step 11: Verifying /v1/completions after restore..." -if poll_api "$RESTORE_POD_NAME" "$RESTORE_CONTAINER_NAME" POST /v1/completions \ - '{"model":"TinyLlama/TinyLlama-1.1B-Chat-v1.0","prompt":"The meaning of life is","max_tokens":10}' \ - '"choices"' "$POST_INFER_TIMEOUT" "$POST_INFER_INTERVAL" "post-restore /v1/completions"; then - echo "$POLL_RESULT" | python3 -m json.tool 2>/dev/null || echo "$POLL_RESULT" - step_done "Post-restore infer" "OK" -else - log_warn "Post-restore /v1/completions not working" - step_done "Post-restore infer" "FAIL" -fi - -# ─── Diagnostics ───────────────────────────────────────────────────────────── -echo "" -log_info "Key restore events:" -kubectl logs $RESTORE_POD_NAME -n $NAMESPACE -c $RESTORE_CONTAINER_NAME 2>&1 | \ - grep -E "wakeRestoredThreads|uv_loop_fork|libzmq.*CRIU|ETERM|reinit completed|RESTORE_COMPLETE" | head -10 || true - -# ─── Summary ───────────────────────────────────────────────────────────────── -echo "" -log_info "==========================================" -if print_summary; then - log_info "TEST PASSED" -else - log_error "TEST FAILED" -fi -log_info "==========================================" -echo "" -log_info "Checkpoint: ${CHECKPOINT_ID:-}" -log_info "Restored pod: $RESTORE_POD_NAME" -log_info "Cleanup: kubectl delete pod $RESTORE_POD_NAME -n $NAMESPACE" diff --git a/src/compute-plane-services/nvsnap/scripts/validate-libzmq-fork.sh b/src/compute-plane-services/nvsnap/scripts/validate-libzmq-fork.sh deleted file mode 100755 index f11906d18d..0000000000 --- a/src/compute-plane-services/nvsnap/scripts/validate-libzmq-fork.sh +++ /dev/null @@ -1,230 +0,0 @@ -#!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# Validate libzmq fork with checkpoint/restore API -# This script builds libzmq from source and runs checkpoint tests - -set -euo pipefail - -# Colors -RED='\033[0;31m' -GREEN='\033[0;32m' -YELLOW='\033[1;33m' -BLUE='\033[0;34m' -NC='\033[0m' - -log_info() { echo -e "${GREEN}[INFO]${NC} $*"; } -log_warn() { echo -e "${YELLOW}[WARN]${NC} $*"; } -log_error() { echo -e "${RED}[ERROR]${NC} $*"; } -log_step() { echo -e "${BLUE}[STEP]${NC} $*"; } - -# Configuration -source "$(dirname "${BASH_SOURCE[0]}")/_deps.sh" -nvsnap_resolve_sibling LIBZMQ_SRC libzmq -BUILD_DIR="${LIBZMQ_SRC}/build" -INSTALL_PREFIX="${INSTALL_PREFIX:-/usr/local}" - -# Check if libzmq source exists -if [ ! -d "$LIBZMQ_SRC" ]; then - log_error "libzmq source not found at $LIBZMQ_SRC" - log_info "Clone from: https://github.com/balajinvda/libzmq.git" - log_info "Branch: checkpoint-restore-v1" - exit 1 -fi - -cd "$LIBZMQ_SRC" - -# Verify branch -current_branch=$(git branch --show-current 2>/dev/null || echo "unknown") -if [ "$current_branch" != "checkpoint-restore-v1" ]; then - log_warn "Current branch is '$current_branch', expected 'checkpoint-restore-v1'" - log_info "Switch with: git checkout checkpoint-restore-v1" -fi - -log_info "Validating libzmq fork with checkpoint/restore API" -log_info "Source: $LIBZMQ_SRC" -log_info "Branch: $current_branch" -echo "" - -# Step 1: Clean build directory -log_step "Step 1/6: Cleaning build directory..." -if [ -d "$BUILD_DIR" ]; then - log_info "Removing existing build directory" - rm -rf "$BUILD_DIR" -fi -mkdir -p "$BUILD_DIR" -cd "$BUILD_DIR" - -# Step 2: Check dependencies -log_step "Step 2/6: Checking build dependencies..." -missing_deps=() - -if ! command -v cmake >/dev/null 2>&1; then - missing_deps+=("cmake") -fi - -if ! command -v g++ >/dev/null 2>&1; then - missing_deps+=("g++") -fi - -if ! dpkg -l libsodium-dev 2>/dev/null | grep -q "^ii"; then - missing_deps+=("libsodium-dev") -fi - -if [ ${#missing_deps[@]} -gt 0 ]; then - log_error "Missing dependencies: ${missing_deps[*]}" - log_info "Install with: sudo apt-get install -y ${missing_deps[*]}" - exit 1 -fi - -log_info "All dependencies found ✓" -echo "" - -# Step 3: Configure with CMake -log_step "Step 3/6: Configuring with CMake..." -log_info "Install prefix: $INSTALL_PREFIX" - -if ! cmake -DCMAKE_INSTALL_PREFIX="$INSTALL_PREFIX" \ - -DCMAKE_BUILD_TYPE=Release \ - -DBUILD_TESTS=ON \ - ..; then - log_error "CMake configuration failed" - exit 1 -fi - -log_info "Configuration successful ✓" -echo "" - -# Step 4: Build libzmq -log_step "Step 4/6: Building libzmq..." -NPROC=$(nproc) -log_info "Building with $NPROC parallel jobs" - -if ! make -j"$NPROC"; then - log_error "Build failed" - exit 1 -fi - -log_info "Build successful ✓" -echo "" - -# Step 5: Verify checkpoint.cpp was built -log_step "Step 5/6: Verifying checkpoint API..." -if [ ! -f "lib/libzmq.so" ]; then - log_error "libzmq.so not found in build directory" - exit 1 -fi - -# Check if checkpoint symbols are present -if nm lib/libzmq.so | grep -q "zmq_ctx_checkpoint"; then - log_info "Found checkpoint API symbols ✓" -else - log_error "Checkpoint API symbols not found in libzmq.so" - log_info "Expected symbols: zmq_ctx_checkpoint, zmq_ctx_restore, zmq_get_all_contexts" - nm lib/libzmq.so | grep zmq_ctx || true - exit 1 -fi - -# List all checkpoint-related symbols -log_info "Checkpoint API symbols found:" -nm lib/libzmq.so | grep -E "zmq_(ctx_checkpoint|ctx_restore|get_all_contexts|checkpoint_destroy|ctx_checkpoint_resume)" | while read -r line; do - echo " $line" -done -echo "" - -# Step 6: Run tests -log_step "Step 6/6: Running checkpoint/restore tests..." - -# Find test binaries -test_binaries=( - "bin/test_checkpoint_basic" - "bin/test_checkpoint_socket" - "bin/test_checkpoint_endpoints" -) - -# Check if tests were built -tests_found=0 -for test_bin in "${test_binaries[@]}"; do - if [ -f "$test_bin" ]; then - tests_found=$((tests_found + 1)) - fi -done - -if [ $tests_found -eq 0 ]; then - log_error "No test binaries found in $BUILD_DIR/bin/" - log_info "Tests may not have been built. Check CMakeLists.txt" - ls -la bin/ 2>/dev/null || log_info "bin/ directory not found" - exit 1 -fi - -log_info "Found $tests_found checkpoint tests" -echo "" - -# Run each test -failed_tests=() -passed_tests=() - -for test_bin in "${test_binaries[@]}"; do - test_name=$(basename "$test_bin") - - if [ ! -f "$test_bin" ]; then - log_warn "Test not found: $test_name (skipping)" - continue - fi - - echo "----------------------------------------" - log_info "Running: $test_name" - echo "" - - # Set LD_LIBRARY_PATH to use built library - export LD_LIBRARY_PATH="$BUILD_DIR/lib:${LD_LIBRARY_PATH:-}" - - if ./"$test_bin"; then - echo "" - log_info "✓ $test_name PASSED" - passed_tests+=("$test_name") - else - echo "" - log_error "✗ $test_name FAILED" - failed_tests+=("$test_name") - fi - echo "" -done - -# Summary -echo "========================================" -echo "" -log_info "VALIDATION SUMMARY" -echo "" -log_info "Tests passed: ${#passed_tests[@]}" -for test in "${passed_tests[@]}"; do - echo " ✓ $test" -done - -if [ ${#failed_tests[@]} -gt 0 ]; then - echo "" - log_error "Tests failed: ${#failed_tests[@]}" - for test in "${failed_tests[@]}"; do - echo " ✗ $test" - done - echo "" - echo "========================================" - log_error "VALIDATION FAILED" - exit 1 -fi - -echo "" -echo "========================================" -log_info "✓ ALL TESTS PASSED" -echo "========================================" -echo "" - -# Show next steps -log_info "Next steps:" -echo " 1. Install libzmq: sudo make install && sudo ldconfig" -echo " 2. Build libzmq image: ./scripts/build-libzmq-image.sh" -echo " 3. Verify system install: ldconfig -p | grep libzmq" -echo "" - -exit 0 diff --git a/src/compute-plane-services/nvsnap/scripts/validate-test-env.sh b/src/compute-plane-services/nvsnap/scripts/validate-test-env.sh index 7ea31a2f40..de0b9f8740 100755 --- a/src/compute-plane-services/nvsnap/scripts/validate-test-env.sh +++ b/src/compute-plane-services/nvsnap/scripts/validate-test-env.sh @@ -30,20 +30,6 @@ echo "==========================================" ERRORS=0 -# 1. Check local library was built recently -echo -e "\n${YELLOW}[1/6] Checking local library build...${NC}" -LOCAL_LIB="lib/nvsnap_intercept/src/io_uring_intercept.c" -if [ -f "$LOCAL_LIB" ]; then - MODIFIED=$(stat -c %Y "$LOCAL_LIB" 2>/dev/null || stat -f %m "$LOCAL_LIB") - NOW=$(date +%s) - AGE_HOURS=$(( (NOW - MODIFIED) / 3600 )) - if [ $AGE_HOURS -gt 2 ]; then - echo -e "${YELLOW} WARNING: $LOCAL_LIB last modified $AGE_HOURS hours ago${NC}" - else - echo -e "${GREEN} OK: Source modified recently ($AGE_HOURS hours ago)${NC}" - fi -fi - # 2. Check Docker image exists locally echo -e "\n${YELLOW}[2/6] Checking Docker image...${NC}" IMAGE="nvcr.io/0651155215864979/ncp-dev/nvsnap-agent:$EXPECTED_VERSION" @@ -52,7 +38,7 @@ if docker image inspect "$IMAGE" >/dev/null 2>&1; then echo -e "${GREEN} OK: Image exists locally (created: $CREATED)${NC}" else echo -e "${RED} ERROR: Image $IMAGE not found locally${NC}" - echo " Run: ./scripts/build-agent-app.sh" + echo " Run: ./scripts/build-agent.sh app" ERRORS=$((ERRORS + 1)) fi diff --git a/src/compute-plane-services/nvsnap/scripts/versions.sh b/src/compute-plane-services/nvsnap/scripts/versions.sh index 74ad49419e..7fb4faef31 100755 --- a/src/compute-plane-services/nvsnap/scripts/versions.sh +++ b/src/compute-plane-services/nvsnap/scripts/versions.sh @@ -48,13 +48,6 @@ NVSNAP_SERVER_VERSION="${NVSNAP_SERVER_VERSION:-v0.0.31}" NVSNAP_BLOBSTORE_VERSION="${NVSNAP_BLOBSTORE_VERSION:-v0.0.1}" NVSNAP_L2WAIT_VERSION="${NVSNAP_L2WAIT_VERSION:-v0.0.1}" -# Dependency builder images — also reset to v0.0.1. They retain their -# CRIU/NCCL patches; the version is just a new tag on the new registry. -NVSNAP_UVLOOP_VERSION="${NVSNAP_UVLOOP_VERSION:-v0.0.1}" -NVSNAP_LIBZMQ_VERSION="${NVSNAP_LIBZMQ_VERSION:-v0.0.1}" -NVSNAP_PYZMQ_VERSION="${NVSNAP_PYZMQ_VERSION:-v0.0.1}" -NVSNAP_LIBUV_VERSION="${NVSNAP_LIBUV_VERSION:-v0.0.1}" - # Dependency / CRIU fork repos + refs. All forks now live under # github.com/balajinvda. These are consumed by ci/build-image.sh when a # dep-builder or base image actually has to be (re)built; day to day the @@ -75,17 +68,6 @@ NVSNAP_CRIU_REPO="${NVSNAP_CRIU_REPO:-https://github.com/balajinvda/criu.git}" # Only the clean-checkout path fetches, so a developer with a local ../criu # checkout never sees this; it breaks OSS clone-and-build only. NVSNAP_CRIU_REF="${NVSNAP_CRIU_REF:-169595fd8ff115690c35d70c1fae90a8d03a7321}" -NVSNAP_LIBZMQ_REPO="${NVSNAP_LIBZMQ_REPO:-https://github.com/balajinvda/libzmq.git}" -NVSNAP_LIBZMQ_REF="${NVSNAP_LIBZMQ_REF:-checkpoint-restore-v1}" -NVSNAP_LIBUV_REPO="${NVSNAP_LIBUV_REPO:-https://github.com/balajinvda/libuv.git}" -NVSNAP_LIBUV_REF="${NVSNAP_LIBUV_REF:-fix/issue-41-no-sqarray}" -NVSNAP_UVLOOP_REPO="${NVSNAP_UVLOOP_REPO:-https://github.com/balajinvda/uvloop.git}" -NVSNAP_UVLOOP_REF="${NVSNAP_UVLOOP_REF:-checkpoint-restore-v1}" -NVSNAP_PYZMQ_REPO="${NVSNAP_PYZMQ_REPO:-https://github.com/balajinvda/pyzmq.git}" -NVSNAP_PYZMQ_REF="${NVSNAP_PYZMQ_REF:-main}" - -# Combined init container — always matches agent version to prevent build-ID mismatch -NVSNAP_INIT_VERSION="${NVSNAP_INIT_VERSION:-${NVSNAP_APP_VERSION}}" # vLLM base image NVSNAP_VLLM_VERSION="${NVSNAP_VLLM_VERSION:-v0.20.0}" From c49254b26e16da96f63c039f25f27213577156b9 Mon Sep 17 00:00:00 2001 From: Balaji Ganesan Date: Thu, 24 Sep 2026 08:50:26 -0700 Subject: [PATCH 2/3] refactor(nvsnap): finish removing the interception stack from build and docs Follow-up to the code removal: the pieces that were off the agent build path and so did not block testing. docker/agent/Dockerfile was a legacy monolith referenced by nothing -- earlier greps matched it only as a prefix of Dockerfile.base and Dockerfile.app -- and it still built the interceptor. Deleted. scripts/build-vllm-image.sh existed solely to bake the patched uvloop and libzmq into a vLLM image, so it goes with them. Dockerfile.local keeps the local dev path but loses its intercept stage. ci/build-image.sh still built all four dependency-builder images from Dockerfiles this branch deleted, and still passed their refs as build args to the agent and init images. That would have broken CI, not just left dead code. Docs described the stack as current architecture: the README diagram placed libnvsnap_intercept.so inside every GPU pod, lib/README.md documented directories that no longer exist, and CONTRIBUTING.md listed nvsnap_cr.so as a build output. GENERIC-PYTHON-INJECTION-DESIGN.md moves to docs/archive alongside the other superseded design notes, with the docs index updated to say why. Also corrected the stale comments in mutate.go that still called the cachedir capture patches "auto-inject". Not included, and deliberately: the quiesce helpers and the /run/criu-restored marker in restore.go, wake_threads.go and quiesce_test.go still name the interceptor. On main both isMultiGPU and useCUDAInterposition are hardcoded false, so that code is unreachable, but removing it means touching the live restore path and deserves its own change rather than a rushed one at the end of this one. Relates to #2094 Co-Authored-By: Balaji Ganesan --- .../nvsnap/CONTRIBUTING.md | 7 +- src/compute-plane-services/nvsnap/README.md | 9 +- .../nvsnap/ci/build-image.sh | 56 ----- .../nvsnap/docker/agent/Dockerfile | 215 ------------------ .../nvsnap/docker/agent/Dockerfile.local | 16 -- .../nvsnap/docs/README.md | 2 +- .../GENERIC-PYTHON-INJECTION-DESIGN.md | 0 .../nvsnap/internal/agent/checkpoint_v2.go | 2 +- .../nvsnap/internal/webhook/mutate.go | 14 +- .../nvsnap/lib/README.md | 23 +- .../nvsnap/scripts/build-vllm-image.sh | 98 -------- .../scripts/retag-and-push-to-ncp-dev.sh | 5 - .../nvsnap/scripts/test-agent-driven-e2e.sh | 2 +- 13 files changed, 26 insertions(+), 423 deletions(-) delete mode 100644 src/compute-plane-services/nvsnap/docker/agent/Dockerfile rename src/compute-plane-services/nvsnap/docs/{ => archive}/GENERIC-PYTHON-INJECTION-DESIGN.md (100%) delete mode 100755 src/compute-plane-services/nvsnap/scripts/build-vllm-image.sh diff --git a/src/compute-plane-services/nvsnap/CONTRIBUTING.md b/src/compute-plane-services/nvsnap/CONTRIBUTING.md index 31ba78282e..9fc585ab40 100644 --- a/src/compute-plane-services/nvsnap/CONTRIBUTING.md +++ b/src/compute-plane-services/nvsnap/CONTRIBUTING.md @@ -18,10 +18,11 @@ The criu-v2 engine builds from public inputs only: in the base image's `cuda-cli-builder` stage, on the public CUDA driver checkpoint API (`cuCheckpointProcess*`, driver 570+ at runtime). Works for x86-64 and arm64 - no committed binaries. -- The Go agent and `nvsnap_cr.so`: this repository. +- The Go agent: this repository. -The legacy LD_PRELOAD injection stack (patched uvloop/libuv/libzmq) is not used -by criu-v2 and is not required to build. +The legacy LD_PRELOAD injection stack (patched uvloop/libuv/libzmq plus the +interception library) was removed once criu-v2 replaced it. The implementation +is preserved at the tag `archive/nvsnap-injection-stack`. ## Build diff --git a/src/compute-plane-services/nvsnap/README.md b/src/compute-plane-services/nvsnap/README.md index 400faeac10..ed5e9328f7 100644 --- a/src/compute-plane-services/nvsnap/README.md +++ b/src/compute-plane-services/nvsnap/README.md @@ -44,7 +44,7 @@ workload: A node-level **agent** (DaemonSet) performs capture/restore using `/proc` + cgroup inspection — runtime-agnostic, not tied to any container runtime API. A **server** provides a REST API, web UI, and checkpoint -catalog. An optional **admission webhook** auto-injects restore plumbing +catalog. An optional **admission webhook** injects restore plumbing into workload pods. Deeper design docs: @@ -172,8 +172,7 @@ optional webhook). Full guide, options, and troubleshooting: │ └─ peer-cascade HTTP server │ │ │ │ GPU Pod (unmodified image) │ - │ ├─ NIM / vLLM / SGLang / TRT-LLM │ - │ └─ libnvsnap_intercept.so (PRELOAD) │ + │ └─ NIM / vLLM / SGLang / TRT-LLM │ └───────────────────────────────────────┘ │ L1 same-node ─► L2 shared PVC ─► L3 object store @@ -190,9 +189,6 @@ tiers are present. See - **go-criu RPC, not CLI CRIU** — engines like vLLM run 900+ threads; CLI CRIU's per-thread ptrace detach takes minutes and often hangs, while RPC is seconds regardless of thread count. -- **`libnvsnap_intercept.so` (LD_PRELOAD)** — `io_uring` (via uvloop/libuv) - and `libzmq` epoll don't survive CRIU restore cleanly. Rather than fork - every engine, the library reinitializes them on a restore marker. - **Forked CRIU** — 26 patches for Kubernetes container support, io_uring, and the CUDA plugin. See [docs/THIRD-PARTY-FORKS.md](docs/THIRD-PARTY-FORKS.md). @@ -257,7 +253,6 @@ The full endpoint list (including agent-side cascade endpoints) is in ```text cmd/ binary entry points (agent, server, restore-entrypoint, gpu-restore, CLI) internal/ agent, server, webhook, CRIU, checkpointstore (Go) -lib/nvsnap_intercept/ LD_PRELOAD interception library (C) deploy/helm/nvsnap/ Helm chart deploy/k8s/ manifests + sample workloads docker/ Dockerfiles diff --git a/src/compute-plane-services/nvsnap/ci/build-image.sh b/src/compute-plane-services/nvsnap/ci/build-image.sh index 03a0747fee..8d33383a7f 100755 --- a/src/compute-plane-services/nvsnap/ci/build-image.sh +++ b/src/compute-plane-services/nvsnap/ci/build-image.sh @@ -8,7 +8,6 @@ # # is one of: # base server blobstore l2-wait # self-contained / CRIU -# uvloop libzmq libuv pyzmq # dependency builders (clone forks) # agent init # layer on the base + dep images # # Tags come from scripts/versions.sh (the single source of truth). The @@ -107,47 +106,6 @@ case "$COMPONENT" in push_img "$base_img" ;; - # ── dependency builders (context = the cloned fork) ───────────────── - uvloop) - img="$REG/uvloop-builder:$NVSNAP_UVLOOP_VERSION" - if [ "${NO_SKIP:-0}" != "1" ] && exists "$img"; then echo "[build-image] $img exists; skipping."; exit 0; fi - guard_ref NVSNAP_UVLOOP_REF "$NVSNAP_UVLOOP_REF" - tmp="$(mktemp -d)"; trap 'rm -rf "$tmp"' EXIT - clone_fork "$tmp/src" "$NVSNAP_UVLOOP_REPO" "$NVSNAP_UVLOOP_REF" - bud -f "$ROOT/docker/uvloop/Dockerfile" -t "$img" "$tmp/src" - push_img "$img" - ;; - libzmq) - img="$REG/libzmq-builder:$NVSNAP_LIBZMQ_VERSION" - if [ "${NO_SKIP:-0}" != "1" ] && exists "$img"; then echo "[build-image] $img exists; skipping."; exit 0; fi - guard_ref NVSNAP_LIBZMQ_REF "$NVSNAP_LIBZMQ_REF" - tmp="$(mktemp -d)"; trap 'rm -rf "$tmp"' EXIT - clone_fork "$tmp/src" "$NVSNAP_LIBZMQ_REPO" "$NVSNAP_LIBZMQ_REF" - bud -f "$ROOT/docker/libzmq/Dockerfile" -t "$img" "$tmp/src" - push_img "$img" - ;; - libuv) - img="$REG/libuv-builder:$NVSNAP_LIBUV_VERSION" - if [ "${NO_SKIP:-0}" != "1" ] && exists "$img"; then echo "[build-image] $img exists; skipping."; exit 0; fi - guard_ref NVSNAP_LIBUV_REF "$NVSNAP_LIBUV_REF" - tmp="$(mktemp -d)"; trap 'rm -rf "$tmp"' EXIT - clone_fork "$tmp/src" "$NVSNAP_LIBUV_REPO" "$NVSNAP_LIBUV_REF" - bud -f "$ROOT/docker/libuv/Dockerfile" -t "$img" "$tmp/src" - push_img "$img" - ;; - pyzmq) - # pyzmq builds against the libzmq source: context holds both - # libzmq-src/ and pyzmq-src/ (see docker/pyzmq/Dockerfile). - img="$REG/pyzmq-builder:$NVSNAP_PYZMQ_VERSION" - if [ "${NO_SKIP:-0}" != "1" ] && exists "$img"; then echo "[build-image] $img exists; skipping."; exit 0; fi - guard_ref NVSNAP_PYZMQ_REF "$NVSNAP_PYZMQ_REF" - guard_ref NVSNAP_LIBZMQ_REF "$NVSNAP_LIBZMQ_REF" - tmp="$(mktemp -d)"; trap 'rm -rf "$tmp"' EXIT - clone_fork "$tmp/pyzmq-src" "$NVSNAP_PYZMQ_REPO" "$NVSNAP_PYZMQ_REF" - clone_fork "$tmp/libzmq-src" "$NVSNAP_LIBZMQ_REPO" "$NVSNAP_LIBZMQ_REF" - bud -f "$ROOT/docker/pyzmq/Dockerfile" -t "$img" "$tmp" - push_img "$img" - ;; # ── agent: app layer over base + dep images (context = repo root) ─── agent) @@ -159,24 +117,10 @@ case "$COMPONENT" in trap 'rm -f "$ROOT/cuda-checkpoint-wrapper.sh"' EXIT bud -f docker/agent/Dockerfile.app \ --build-arg BASE_IMAGE="$REG/nvsnap-agent-base:$NVSNAP_BASE_VERSION" \ - --build-arg UVLOOP_IMAGE="$REG/uvloop-builder:$NVSNAP_UVLOOP_VERSION" \ - --build-arg LIBUV_IMAGE="$REG/libuv-builder:$NVSNAP_LIBUV_VERSION" \ - --build-arg LIBZMQ_IMAGE="$REG/libzmq-builder:$NVSNAP_LIBZMQ_VERSION" \ -t "$img" . push_img "$img" ;; - # ── init: assembles dep + agent images (context = repo root) ──────── - init) - build "$REG/nvsnap-init:$NVSNAP_INIT_VERSION" \ - -f docker/init/Dockerfile \ - --build-arg UVLOOP_IMAGE="$REG/uvloop-builder:$NVSNAP_UVLOOP_VERSION" \ - --build-arg LIBUV_IMAGE="$REG/libuv-builder:$NVSNAP_LIBUV_VERSION" \ - --build-arg LIBZMQ_IMAGE="$REG/libzmq-builder:$NVSNAP_LIBZMQ_VERSION" \ - --build-arg PYZMQ_IMAGE="$REG/pyzmq-builder:$NVSNAP_PYZMQ_VERSION" \ - --build-arg AGENT_IMAGE="$REG/nvsnap-agent:$NVSNAP_APP_VERSION" - ;; - *) echo "unknown component: $COMPONENT" >&2 exit 2 diff --git a/src/compute-plane-services/nvsnap/docker/agent/Dockerfile b/src/compute-plane-services/nvsnap/docker/agent/Dockerfile deleted file mode 100644 index 152f866422..0000000000 --- a/src/compute-plane-services/nvsnap/docker/agent/Dockerfile +++ /dev/null @@ -1,215 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# NVSNAP Agent Container - Multi-Stage Build -# -# Builds everything from source to ensure consistency: -# - Stage 1: Build CRIU from our fork (github.com/balajinvda/criu) -# - Stage 2: Build Go binaries (nvsnap-agent, restore-entrypoint) -# - Stage 3: Final minimal runtime image -# -# CRIU fork is cloned from GitHub with a specific tag for reproducibility. - -# ============================================================================= -# Stage 1: Build CRIU from source (cloned from GitHub fork) -# ============================================================================= -FROM ubuntu:22.04 AS criu-builder - -# CRIU fork version - change this to update CRIU -ARG CRIU_REPO=https://github.com/balajinvda/criu.git -ARG CRIU_TAG=nvsnap-v0.3.0 -ARG CRIU_LOCAL_DIR=criu-src-local - -# Install CRIU build dependencies -RUN apt-get update && apt-get install -y --no-install-recommends \ - build-essential \ - pkg-config \ - python3 \ - python3-protobuf \ - libprotobuf-dev \ - libprotobuf-c-dev \ - protobuf-c-compiler \ - protobuf-compiler \ - libnet1-dev \ - libnl-3-dev \ - libnl-genl-3-dev \ - libcap-dev \ - libgnutls28-dev \ - libselinux1-dev \ - libnftables-dev \ - libdrm-dev \ - uuid-dev \ - liburing-dev \ - asciidoctor \ - xmlto \ - git \ - ca-certificates \ - && rm -rf /var/lib/apt/lists/* - -# Prefer local CRIU fork if provided in build context -COPY ${CRIU_LOCAL_DIR}/ /criu-local/ -RUN if [ -f /criu-local/Makefile ]; then \ - cp -a /criu-local /criu-src; \ - echo "CRIU version: local (from build context)"; \ - else \ - git clone --branch ${CRIU_TAG} --depth 1 ${CRIU_REPO} /criu-src && \ - echo "CRIU version: $(cd /criu-src && git describe --tags --always)"; \ - fi - -WORKDIR /criu-src - -# Ensure no host build artifacts leak into container build -RUN make mrproper || make clean - -# Create asm symlink for x86 architecture (required for plugin build) -RUN ln -sf ../arch/x86/include/asm criu/include/asm - -# Build CRIU -RUN make -j$(nproc) - -# Build CUDA plugin -RUN make -C plugins/cuda - -# Build ZMQ plugin -RUN make -C plugins/zmq - -# Collect CRIU and its library dependencies -RUN mkdir -p /criu-bundle/lib /criu-bundle/plugins && \ - cp criu/criu /criu-bundle/ && \ - cp plugins/cuda/cuda_plugin.so /criu-bundle/plugins/ && \ - cp plugins/zmq/zmq_plugin.so /criu-bundle/plugins/ && \ - # Collect runtime library dependencies (excluding core glibc) - for lib in $(ldd /criu-bundle/criu | grep "=>" | awk '{print $3}' | sort -u); do \ - case "$lib" in \ - */libc.so*|*/ld-linux*|*/libpthread*|*/libdl*|*/librt*|*/libm.so*) \ - echo "Skipping core lib: $lib" ;; \ - *) \ - echo "Copying: $lib" && cp -L "$lib" /criu-bundle/lib/ 2>/dev/null || true ;; \ - esac \ - done - -# Create iptables shims inside the CRIU bundle. -# Restore pods run CRIU inside arbitrary workload images, which often do not include -# iptables-restore/ip6tables-restore. CRIU's network lock/unlock may exec these binaries, -# so we provide no-op static shims to keep restore generic and self-contained. -RUN printf '#include \nint main(){return 0;}\n' > /tmp/true.c && \ - gcc -static -O2 -s -o /criu-bundle/iptables-restore /tmp/true.c && \ - cp /criu-bundle/iptables-restore /criu-bundle/ip6tables-restore && \ - chmod +x /criu-bundle/iptables-restore /criu-bundle/ip6tables-restore - -# Create wrapper script that sets library path -RUN echo '#!/bin/sh' > /criu-bundle/criu-wrapper && \ - echo 'DIR="$(dirname "$(readlink -f "$0")")"' >> /criu-bundle/criu-wrapper && \ - echo 'export LD_LIBRARY_PATH="$DIR/lib:$LD_LIBRARY_PATH"' >> /criu-bundle/criu-wrapper && \ - echo 'exec "$DIR/criu" "$@"' >> /criu-bundle/criu-wrapper && \ - chmod +x /criu-bundle/criu-wrapper - -# ============================================================================= -# Stage 2a: Build libnvsnap_intercept.so (io_uring & libuv quiescence) -# ============================================================================= -FROM ubuntu:22.04 AS intercept-builder - -RUN apt-get update && apt-get install -y --no-install-recommends \ - build-essential \ - python3-dev \ - && rm -rf /var/lib/apt/lists/* - -WORKDIR /app -COPY lib/nvsnap_intercept/ . - -ARG ENABLE_LIBUV_INTERCEPT=0 -RUN make clean && make ENABLE_LIBUV_INTERCEPT=${ENABLE_LIBUV_INTERCEPT} && \ - ls -la libnvsnap_intercept.so - -# ============================================================================= -# Stage 2b: Build Go binaries -# ============================================================================= -FROM golang:1.25-bookworm AS go-builder - -WORKDIR /app - -# Copy go mod files first for caching -COPY go.mod go.sum ./ -RUN go mod download - -# Copy source -COPY . . - -# Build agent and restore-entrypoint (use -mod=mod to ignore vendor, download fresh) -RUN CGO_ENABLED=0 GOOS=linux go build -mod=mod -ldflags="-s -w" -o /bin/nvsnap-agent ./cmd/agent && \ - CGO_ENABLED=0 GOOS=linux go build -mod=mod -ldflags="-s -w" -o /bin/restore-entrypoint ./cmd/restore-entrypoint && \ - CGO_ENABLED=0 GOOS=linux go build -mod=mod -ldflags="-s -w" -o /bin/nvsnap-mount-prep ./cmd/nvsnap-mount-prep - -# ============================================================================= -# Stage 3: Final runtime image -# ============================================================================= -FROM ubuntu:22.04 - -# Install runtime dependencies for CRIU -RUN apt-get update && apt-get install -y --no-install-recommends \ - ca-certificates \ - util-linux \ - procps \ - iproute2 \ - iptables \ - libnftables1 \ - libnet1 \ - libnl-3-200 \ - libprotobuf-c1 \ - liburing2 \ - python3 \ - python3-pip \ - gdb \ - && rm -rf /var/lib/apt/lists/* - -RUN pip3 install --no-cache-dir py-spy - -# Create directories -RUN mkdir -p /var/lib/nvsnap/checkpoints /criu-bundle - -# Indicate containerized environment -RUN touch /.dockerenv - -# Copy CRIU bundle from builder -COPY --from=criu-builder /criu-bundle/ /criu-bundle/ - -# Copy Go binaries -COPY --from=go-builder /bin/nvsnap-agent /criu-bundle/nvsnap-agent -COPY --from=go-builder /bin/restore-entrypoint /criu-bundle/restore-entrypoint -COPY --from=go-builder /bin/nvsnap-mount-prep /criu-bundle/nvsnap-mount-prep - -# Copy libnvsnap_intercept.so for io_uring/libuv quiescence -COPY --from=intercept-builder /app/libnvsnap_intercept.so /criu-bundle/lib/libnvsnap_intercept.so - -# Copy cuda-checkpoint (placed in build context root by build-agent-image.sh) -COPY cuda-checkpoint /criu-bundle/cuda-checkpoint - -# Make everything executable -RUN chmod +x /criu-bundle/* - -# Add bundled libs to library path -RUN echo "/criu-bundle/lib" > /etc/ld.so.conf.d/criu.conf && ldconfig - -# Create convenience symlinks -RUN ln -sf /criu-bundle/nvsnap-agent /usr/local/bin/nvsnap-agent && \ - ln -sf /criu-bundle/criu /usr/local/sbin/criu && \ - ln -sf /criu-bundle/nvsnap-mount-prep /nvsnap-mount-prep - -# Set PATH and LD_LIBRARY_PATH -ENV PATH="/criu-bundle:$PATH" -ENV LD_LIBRARY_PATH="/criu-bundle/lib" - -# Verify build -RUN echo "=== CRIU version ===" && /criu-bundle/criu --version && \ - echo "=== Agent ===" && /criu-bundle/nvsnap-agent --help 2>&1 | head -3 || true && \ - echo "=== Bundle contents ===" && ls -la /criu-bundle/ - -EXPOSE 8081 - -ENTRYPOINT ["/criu-bundle/nvsnap-agent"] -CMD ["--listen", ":8081", \ - "--checkpoint-dir", "/var/lib/nvsnap/checkpoints", \ - "--containerd-socket", "/run/containerd/containerd.sock", \ - "--criu-path", "/criu-bundle/criu", \ - "--cuda-checkpoint-path", "/criu-bundle/cuda-checkpoint", \ - "--log-level", "debug"] diff --git a/src/compute-plane-services/nvsnap/docker/agent/Dockerfile.local b/src/compute-plane-services/nvsnap/docker/agent/Dockerfile.local index 44f4243945..774192e098 100644 --- a/src/compute-plane-services/nvsnap/docker/agent/Dockerfile.local +++ b/src/compute-plane-services/nvsnap/docker/agent/Dockerfile.local @@ -81,20 +81,6 @@ RUN printf '#include \nint main(){return 0;}\n' > /tmp/true.c && \ # ============================================================================= -# Stage 2a: Build libnvsnap_intercept.so (io_uring & libuv quiescence) -# ============================================================================= -FROM ubuntu:22.04 AS intercept-builder - -RUN apt-get update && apt-get install -y --no-install-recommends \ - build-essential \ - && rm -rf /var/lib/apt/lists/* - -WORKDIR /app -COPY lib/nvsnap_intercept/ . - -RUN make clean && make && \ - ls -la libnvsnap_intercept.so - # ============================================================================= # Stage 2b: Build Go binaries # ============================================================================= @@ -146,8 +132,6 @@ COPY --from=criu-builder /criu-bundle/ /criu-bundle/ COPY --from=go-builder /bin/nvsnap-agent /criu-bundle/nvsnap-agent COPY --from=go-builder /bin/restore-entrypoint /criu-bundle/restore-entrypoint -# Copy libnvsnap_intercept.so for io_uring/libuv quiescence -COPY --from=intercept-builder /app/libnvsnap_intercept.so /criu-bundle/lib/libnvsnap_intercept.so # Copy cuda-checkpoint (placed in build context root by build-agent-image.sh) COPY cuda-checkpoint /criu-bundle/cuda-checkpoint.real diff --git a/src/compute-plane-services/nvsnap/docs/README.md b/src/compute-plane-services/nvsnap/docs/README.md index 3718c3a9b3..525b0f8675 100644 --- a/src/compute-plane-services/nvsnap/docs/README.md +++ b/src/compute-plane-services/nvsnap/docs/README.md @@ -35,7 +35,7 @@ benchmark summary. This directory holds the operator and developer docs. - [Third-party forks + fork-maintenance policy](THIRD-PARTY-FORKS.md) - [Benchmarks](BENCHMARK.md) · [PDF benchmark matrix](PDF-BENCH-RESULTS.md) -- [Generic Python injection](GENERIC-PYTHON-INJECTION-DESIGN.md) +- [Generic Python injection](archive/GENERIC-PYTHON-INJECTION-DESIGN.md) (archived; the injection stack was removed) - [L2 per-capture PVC (CRIU)](L2-PVC-CRIU-DESIGN.md) - [Multi-GPU rootfs fan-out](MULTI-GPU-ROOTFS-FANOUT-DESIGN.md) diff --git a/src/compute-plane-services/nvsnap/docs/GENERIC-PYTHON-INJECTION-DESIGN.md b/src/compute-plane-services/nvsnap/docs/archive/GENERIC-PYTHON-INJECTION-DESIGN.md similarity index 100% rename from src/compute-plane-services/nvsnap/docs/GENERIC-PYTHON-INJECTION-DESIGN.md rename to src/compute-plane-services/nvsnap/docs/archive/GENERIC-PYTHON-INJECTION-DESIGN.md diff --git a/src/compute-plane-services/nvsnap/internal/agent/checkpoint_v2.go b/src/compute-plane-services/nvsnap/internal/agent/checkpoint_v2.go index d95f03f45f..f3c58cc959 100644 --- a/src/compute-plane-services/nvsnap/internal/agent/checkpoint_v2.go +++ b/src/compute-plane-services/nvsnap/internal/agent/checkpoint_v2.go @@ -34,7 +34,7 @@ limitations under the License. // SCOPE: this does NOT yet restore io_uring / libuv event-loop kernel // state. vLLM's uvloop aborts post-restore with io_uring enabled, so the // vllm-small manifest still sets USE_LIBUV=0 / UV_USE_IO_URING=0 and -// preloads nvsnap_cr.so (verified: with those levers removed the restored +// preloaded the interception library (verified: with those levers removed the restored // process aborts in uvloop.run, 2026-07-13). Restoring the rings at the // CRIU layer to drop those levers is tracked separately (NVCF-9641, // io_uring ring-restore work item). diff --git a/src/compute-plane-services/nvsnap/internal/webhook/mutate.go b/src/compute-plane-services/nvsnap/internal/webhook/mutate.go index ace75f07f9..d896c5afbd 100644 --- a/src/compute-plane-services/nvsnap/internal/webhook/mutate.go +++ b/src/compute-plane-services/nvsnap/internal/webhook/mutate.go @@ -85,7 +85,7 @@ type PatchOp struct { } // mergeableArrayRe matches the per-container appendable arrays whose -// bootstrap (`add []`) the auto-inject and restore builders each +// bootstrap (`add []`) the cachedir and restore builders each // compute independently from the original pod. var mergeableArrayRe = regexp.MustCompile(`^/spec/containers/\d+/(volumeMounts|env)$`) @@ -97,7 +97,7 @@ func isMergeableArray(path string) bool { // patchElementName pulls the "name" field out of a patch value (volume, // volumeMount, initContainer, env). JSON round-trip so it works for both -// map[string]any (auto-inject) and typed corev1 values (restore builders). +// map[string]any (cachedir) and typed corev1 values (restore builders). // Returns "" when there's no name (then the element can't be deduped). func patchElementName(v any) string { b, err := json.Marshal(v) @@ -113,12 +113,12 @@ func patchElementName(v any) string { return named.Name } -// mergePatchPlan reconciles the concatenation of auto-inject + restore +// mergePatchPlan reconciles the concatenation of cachedir + restore // patches so they don't clobber each other (nvsnap#93). Both sides build // array bootstraps from the ORIGINAL pod, so a pod with empty spec.volumes // (or empty main-container volumeMounts/env) gets TWO `add [..]` // ops — and under JSON Patch the second REPLACES the first, dropping the -// auto-injected nvsnap-lib volume/mount. The same arrays can also receive a +// restore-bundle volume/mount. The same arrays can also receive a // duplicate element (two nvsnap-lib volumes). // // Normalization, preserving order: @@ -420,7 +420,7 @@ func (m *Mutator) Mutate(ctx context.Context, pod *corev1.Pod) ([]PatchOp, error // No-op when cachedir mode is off. (Restore pods get the rox mount // + env from the cachedir restore path below, not this.) injectPatches = append(injectPatches, m.cacheDirCapturePatches(pod)...) - // Return whatever auto-inject + cachedir-capture produced (may be + // Return whatever the cachedir-capture patches produced (may be // nil, which is fine — pod admitted unchanged). if len(injectPatches) > 0 { m.logger().WithField("pod", pod.Namespace+"/"+pod.Name). @@ -546,9 +546,9 @@ func (m *Mutator) Mutate(ctx context.Context, pod *corev1.Pod) ([]PatchOp, error if err != nil { return nil, err } - // Prepend any auto-inject patches so they're applied first by the + // Prepend the pre-patches so they're applied first by the // API server (in-order JSON Patch application). Restore mount - // patches reference indices that auto-inject may also touch + // patches reference indices that the pre-patches may also touch // (initContainers list, env list); doing inject-then-restore keeps // path math straightforward. patches = mergePatchPlan(append(injectPatches, patches...)) diff --git a/src/compute-plane-services/nvsnap/lib/README.md b/src/compute-plane-services/nvsnap/lib/README.md index a27e4d0250..119d3e7c71 100644 --- a/src/compute-plane-services/nvsnap/lib/README.md +++ b/src/compute-plane-services/nvsnap/lib/README.md @@ -8,19 +8,16 @@ Non-Go runtime pieces that ship inside workload pods (not the agent). ## Contents -- [`nvsnap_intercept/`](nvsnap_intercept/) — **libnvsnap_intercept.so**, the - `LD_PRELOAD` C library that reinitializes io_uring (uvloop/libuv) and libzmq - epoll after a CRIU restore. See its [README](nvsnap_intercept/README.md). -- [`nvsnap_restore_helper/`](nvsnap_restore_helper/) — small C restore helper - bundled alongside the intercept library. -- [`sitecustomize/`](sitecustomize/) — `sitecustomize.py`; Python's `site.py` - auto-imports it to prepend the patched-uvloop site-packages onto `sys.path` - at runtime (no edits to the workload image). See - [docs/GENERIC-PYTHON-INJECTION-DESIGN.md](../docs/GENERIC-PYTHON-INJECTION-DESIGN.md). +- [`nvsnap_restore_helper/`](nvsnap_restore_helper/) — small C restore helper. + +The `LD_PRELOAD` interception library, its Python `sitecustomize` shim, and the +patched uvloop/libuv/libzmq that went with them were removed once criu-v2 +replaced the approach. criu-v2 handles the workload at the OS level and injects +nothing into it. The implementation is preserved at the tag +`archive/nvsnap-injection-stack`. ## Rules -These are injected into unmodified workload containers via init containers + -env (`LD_PRELOAD`, `PYTHONPATH`) — never by modifying the application image. -The C library must link against the same glibc as the workloads it preloads -into (built on ubuntu:22.04; see [CONTRIBUTING.md](../CONTRIBUTING.md)). +What ships here runs inside workload pods, never by modifying the application +image. C pieces must link against the same glibc as the workloads they run in +(built on ubuntu:22.04; see [CONTRIBUTING.md](../CONTRIBUTING.md)). diff --git a/src/compute-plane-services/nvsnap/scripts/build-vllm-image.sh b/src/compute-plane-services/nvsnap/scripts/build-vllm-image.sh deleted file mode 100755 index ccb068dae3..0000000000 --- a/src/compute-plane-services/nvsnap/scripts/build-vllm-image.sh +++ /dev/null @@ -1,98 +0,0 @@ -#!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# Build custom vLLM image with patched uvloop + libzmq baked in -# -# This eliminates the need for: -# - get-libzmq init container (libzmq baked into image) -# - C-level uvloop hacks in libnvsnap_intercept.so (uvloop handles it natively) -# -# The resulting image is a drop-in replacement for vllm/vllm-openai:VERSION - -set -euo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -PROJECT_ROOT="$(dirname "$SCRIPT_DIR")" - -# Logging -log_info() { echo "[INFO] $*"; } -log_error() { echo "[ERROR] $*" >&2; } - -# Configuration -REGISTRY="${REGISTRY:-nvcr.io/0651155215864979/ncp-dev}" -IMAGE_NAME="vllm-openai" -VLLM_BASE="${VLLM_BASE:-vllm/vllm-openai:v0.11.2}" -UVLOOP_BUILDER="${UVLOOP_BUILDER:-${REGISTRY}/uvloop-builder:v0.22.1-gpucr2}" -LIBZMQ_BUILDER="${LIBZMQ_BUILDER:-${REGISTRY}/libzmq-builder:v0.8.0-zmq}" -VERSION="${VERSION:-v0.11.2-gpucr1}" -CONTAINER_TOOL="${CONTAINER_TOOL:-docker}" - -log_info "Building custom vLLM image" -log_info "Base: $VLLM_BASE" -log_info "uvloop builder: $UVLOOP_BUILDER" -log_info "libzmq builder: $LIBZMQ_BUILDER" -log_info "Version: $VERSION" - -# Create build context -BUILD_CONTEXT=$(mktemp -d) -trap "rm -rf $BUILD_CONTEXT" EXIT - -# Create Dockerfile -cat > "$BUILD_CONTEXT/Dockerfile" < Date: Thu, 24 Sep 2026 09:40:08 -0700 Subject: [PATCH 3/3] fix(nvsnap): make the validate-test-env build hint build the expected tag The hint printed a bare build command, so following it built the default tag while validation went on checking for EXPECTED_VERSION, and the check failed again with the same message. Pass the version through. Relates to #2094 Co-Authored-By: Balaji Ganesan --- src/compute-plane-services/nvsnap/scripts/validate-test-env.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/compute-plane-services/nvsnap/scripts/validate-test-env.sh b/src/compute-plane-services/nvsnap/scripts/validate-test-env.sh index de0b9f8740..52414988f1 100755 --- a/src/compute-plane-services/nvsnap/scripts/validate-test-env.sh +++ b/src/compute-plane-services/nvsnap/scripts/validate-test-env.sh @@ -38,7 +38,7 @@ if docker image inspect "$IMAGE" >/dev/null 2>&1; then echo -e "${GREEN} OK: Image exists locally (created: $CREATED)${NC}" else echo -e "${RED} ERROR: Image $IMAGE not found locally${NC}" - echo " Run: ./scripts/build-agent.sh app" + echo " Run: NVSNAP_APP_VERSION=\"$EXPECTED_VERSION\" ./scripts/build-agent.sh app" ERRORS=$((ERRORS + 1)) fi