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/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 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.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/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/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/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/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..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: @@ -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 == "" { @@ -426,12 +420,12 @@ 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). WithField("patches", len(injectPatches)). - Info("auto-inject only") + Info("capture pod: cache-dir patches only") } return injectPatches, nil } @@ -552,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...)) @@ -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/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/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/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" < - -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/retag-and-push-to-ncp-dev.sh b/src/compute-plane-services/nvsnap/scripts/retag-and-push-to-ncp-dev.sh index 4a52bc2244..e7d305e00a 100755 --- a/src/compute-plane-services/nvsnap/scripts/retag-and-push-to-ncp-dev.sh +++ b/src/compute-plane-services/nvsnap/scripts/retag-and-push-to-ncp-dev.sh @@ -26,11 +26,6 @@ declare -a IMAGES=( "nvsnap-agent:v0.24.16-ensure-capture-endpoint" "nvsnap-server:v0.9.0-cross-node-restore" "nvsnap-blobstore:v0.2.0-stats-captures" - "nvsnap-init:v0.24.16-ensure-capture-endpoint" - "uvloop-builder:v0.22.1-multipy1" - "libuv-builder:v1.48.0-criu-v3" - "libzmq-builder:v4.3.6-criu-epoll-v12" - "pyzmq-builder:v27.2.0-gpucr3" ) for entry in "${IMAGES[@]}"; do 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-agent-driven-e2e.sh b/src/compute-plane-services/nvsnap/scripts/test-agent-driven-e2e.sh index ee748a7c6f..ab62324a74 100755 --- a/src/compute-plane-services/nvsnap/scripts/test-agent-driven-e2e.sh +++ b/src/compute-plane-services/nvsnap/scripts/test-agent-driven-e2e.sh @@ -234,7 +234,7 @@ RESTORED_PID=$(echo "$RESTORE_RESP" | sed -nE 's/.*"restoredPid":([0-9]+).*/\1/p step_end "Restore" OK log " restoredPid=$RESTORED_PID" -# Give the workload a moment to settle (libnvsnap_intercept reinit + wakeRestoredThreads). +# Give the workload a moment to settle (wakeRestoredThreads). sleep 5 # Step 6: post-restore inference via nsenter into the restored process's netns. 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..52414988f1 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: NVSNAP_APP_VERSION=\"$EXPECTED_VERSION\" ./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}"