fix: initialize NVIDIA driver before CDI toolkit

entrypoint.sh

@@ -25,9 +25,9 @@ fi
 
 # Map the requested action to the install mode passed to install.sh. All three install
 # variants stage the same gpu cache files; only the env var handed to install.sh differs.
-#   install            -> full compile + device init (legacy behaviour)
+#   install            -> full compile + node initialization (legacy behaviour)
 #   build-only         -> compile/cache the kernel module only (VHD build, no GPU)
-#   install-skip-build -> device init only, reusing the module prebuilt into the VHD
+#   install-skip-build -> node initialization only, reusing the module prebuilt into the VHD
 GPU_INSTALL_MODE_ENV=""
 case "${1}" in
     build-only)         GPU_INSTALL_MODE_ENV="AKSGPU_BUILD_ONLY=1" ;;

install.sh

@@ -16,7 +16,7 @@ PS4='+ $(date -u -I"seconds" | cut -c1-19) '
 #   AKSGPU_SKIP_KERNEL_BUILD=1 -> the kernel module + libs were prebuilt into the VHD for
 #                                 this exact kernel+driver; skip recompilation and only run
 #                                 the device-dependent steps at node boot.
-#   (neither set)              -> legacy behaviour: full compile + device init in one shot.
+#   (neither set)              -> legacy behaviour: full compile + node initialization in one shot.
 AKSGPU_BUILD_ONLY="${AKSGPU_BUILD_ONLY:-0}"
 AKSGPU_SKIP_KERNEL_BUILD="${AKSGPU_SKIP_KERNEL_BUILD:-0}"
 
@@ -150,9 +150,11 @@ build_kernel_module() {
     modinfo -k "$KERNEL_NAME" nvidia
 }
 
-# device_init runs the steps that require the physical GPU and therefore must execute at node
-# boot, regardless of whether the kernel module was prebuilt into the VHD.
-device_init() {
+# initialize_nvidia_driver runs the steps that require the physical GPU and therefore must
+# execute at node boot, regardless of whether the kernel module was prebuilt into the VHD.
+# Keep this ahead of the container toolkit install: the toolkit package immediately starts
+# nvidia-cdi-refresh, whose nvidia-smi readiness check requires these userspace libraries.
+initialize_nvidia_driver() {
     nvidia-modprobe -u -c0
 
     # configure persistence daemon
@@ -161,6 +163,7 @@ device_init() {
     # notable on large VM sizes with multiple GPUs
     # especially when nvidia-smi process is in CPU cgroup
     cp -r /usr/bin/lib64/lib64/* "/usr/lib/${ARCH}-linux-gnu/"
+    ldconfig
     nvidia-smi
 
     # install fabricmanager for nvlink based systems
@@ -172,13 +175,26 @@ device_init() {
         fi
         bash /opt/gpu/fabricmanager-linux-${NVIDIA_FM_ARCH}-${DRIVER_VERSION}/sbin/fm_run_package_installer.sh
     fi
+}
+
+configure_nvidia_container_runtime() {
+    local nvidia_ctk_bin="${AKSGPU_NVIDIA_CTK_BIN:-/usr/bin/nvidia-ctk}"
+
+    install_nvidia_container_toolkit
 
     mkdir -p /etc/containerd/config.d
     cp /opt/gpu/10-nvidia-runtime.toml /etc/containerd/config.d/10-nvidia-runtime.toml
 
     mkdir -p "$(dirname /lib/udev/rules.d/71-nvidia-dev-char.rules)"
     cp /opt/gpu/71-nvidia-char-dev.rules /lib/udev/rules.d/71-nvidia-dev-char.rules
-    /usr/bin/nvidia-ctk system create-dev-char-symlinks --create-all
+    "$nvidia_ctk_bin" system create-dev-char-symlinks --create-all
+
+    # The package post-install starts this oneshot too, but require a final successful run after
+    # all runtime/device setup. A successful oneshot is inactive afterward, so its exit status
+    # rather than `is-active` is the health signal. Clear the package-time start counter first:
+    # the post-install and path unit can otherwise consume the service's five-start budget.
+    systemctl reset-failed nvidia-cdi-refresh.service nvidia-cdi-refresh.path 2>/dev/null || true
+    systemctl restart nvidia-cdi-refresh.service
 }
 
 write_dkms_marker() {
@@ -222,7 +238,7 @@ build_and_mark() {
 # fast_path_ok opportunistically validates a prebaked module without recompiling. It returns
 # non-zero (instead of aborting under `set -e`) so a corrupt or incomplete prebake falls back to
 # a full build rather than bricking the node. The authoritative device check still happens later
-# in device_init (nvidia-modprobe + nvidia-smi).
+# in initialize_nvidia_driver (nvidia-modprobe + nvidia-smi).
 fast_path_ok() {
     ldconfig || return 1
     dkms status >/dev/null 2>&1 || return 1
@@ -267,8 +283,6 @@ main() {
         exit 0
     fi
 
-    install_nvidia_container_toolkit
-
     if [ "${AKSGPU_SKIP_KERNEL_BUILD}" = "1" ] && baked_marker_matches && fast_path_ok; then
         # Prebuilt module is present and valid for this kernel+driver: skip the ~100s recompile.
         echo "aks-gpu: using kernel module prebuilt in the VHD for kernel ${KERNEL_NAME} (recompile skipped)"
@@ -283,7 +297,8 @@ main() {
         build_and_mark
     fi
 
-    device_init
+    initialize_nvidia_driver
+    configure_nvidia_container_runtime
 
     purge_gpu_cache
 }

test/install.bats

@@ -61,9 +61,17 @@ _has_sort_v() { printf '1\n2\n' | sort -V >/dev/null 2>&1; }
 # Each stub records that it ran via a sentinel file.
 _stub_dispatch() {
     cleanup_overlay() { :; }
-    install_nvidia_container_toolkit() { :; }
     build_and_mark() { echo "BUILD_AND_MARK"; touch "${TEST_TMP}/build_and_mark.ran"; }
-    device_init() { echo "DEVICE_INIT"; touch "${TEST_TMP}/device_init.ran"; }
+    initialize_nvidia_driver() {
+        echo "INITIALIZE_NVIDIA_DRIVER"
+        echo "initialize_nvidia_driver" >> "${TEST_TMP}/dispatch-order"
+        touch "${TEST_TMP}/initialize_nvidia_driver.ran"
+    }
+    configure_nvidia_container_runtime() {
+        echo "CONFIGURE_NVIDIA_CONTAINER_RUNTIME"
+        echo "configure_nvidia_container_runtime" >> "${TEST_TMP}/dispatch-order"
+        touch "${TEST_TMP}/configure_nvidia_container_runtime.ran"
+    }
     purge_gpu_cache() { :; }
 }
 
@@ -164,7 +172,7 @@ _stub_dispatch() {
 
 # --- mode dispatch -------------------------------------------------------
 
-@test "dispatch build-only: builds+marks then skips device init" {
+@test "dispatch build-only: builds+marks then skips node-time driver and runtime initialization" {
     _has_sort_v || skip "build-only resolves target kernel via GNU sort -V (runs on the Linux CI)"
     _stub_dispatch
     mkdir -p "${AKSGPU_MODULES_ROOT}/6.8.0-1059-azure/build"
@@ -173,10 +181,11 @@ _stub_dispatch() {
     [ "$status" -eq 0 ]
     [[ "$output" == *"build-only mode"* ]]
     [ -f "${TEST_TMP}/build_and_mark.ran" ]
-    [ ! -f "${TEST_TMP}/device_init.ran" ]
+    [ ! -f "${TEST_TMP}/initialize_nvidia_driver.ran" ]
+    [ ! -f "${TEST_TMP}/configure_nvidia_container_runtime.ran" ]
 }
 
-@test "dispatch install-skip-build with matching marker: skips recompile, runs device init" {
+@test "dispatch install-skip-build with matching marker: initializes the driver before the runtime" {
     _stub_dispatch
     _stub_bin ldconfig 0; _stub_bin dkms 0; _stub_bin modinfo 0
     PATH="${TEST_TMP}/bin:$PATH"
@@ -187,7 +196,11 @@ _stub_dispatch() {
     [ "$status" -eq 0 ]
     [[ "$output" == *"recompile skipped"* ]]
     [ ! -f "${TEST_TMP}/build_and_mark.ran" ]
-    [ -f "${TEST_TMP}/device_init.ran" ]
+    [ -f "${TEST_TMP}/initialize_nvidia_driver.ran" ]
+    [ -f "${TEST_TMP}/configure_nvidia_container_runtime.ran" ]
+    run cat "${TEST_TMP}/dispatch-order"
+    [ "$status" -eq 0 ]
+    [ "$output" = $'initialize_nvidia_driver\nconfigure_nvidia_container_runtime' ]
 }
 
 @test "dispatch install-skip-build with mismatched marker: falls back to a full build" {
@@ -202,5 +215,63 @@ _stub_dispatch() {
     [ "$status" -eq 0 ]
     [[ "$output" == *"building from source"* ]]
     [ -f "${TEST_TMP}/build_and_mark.ran" ]
-    [ -f "${TEST_TMP}/device_init.ran" ]
+    [ -f "${TEST_TMP}/initialize_nvidia_driver.ran" ]
+    [ -f "${TEST_TMP}/configure_nvidia_container_runtime.ran" ]
+}
+
+# --- CDI lifecycle ordering ---------------------------------------------
+
+@test "initialize_nvidia_driver refreshes the linker cache before nvidia-smi" {
+    DRIVER_KIND="grid"
+    nvidia-modprobe() { echo "nvidia-modprobe" >> "${TEST_TMP}/driver-order"; }
+    cp() { echo "copy-libraries" >> "${TEST_TMP}/driver-order"; }
+    ldconfig() { echo "ldconfig" >> "${TEST_TMP}/driver-order"; }
+    nvidia-smi() { echo "nvidia-smi" >> "${TEST_TMP}/driver-order"; }
+
+    run initialize_nvidia_driver
+
+    [ "$status" -eq 0 ]
+    run cat "${TEST_TMP}/driver-order"
+    [ "$status" -eq 0 ]
+    [ "$output" = $'nvidia-modprobe\ncopy-libraries\nldconfig\nnvidia-smi' ]
+}
+
+@test "configure_nvidia_container_runtime installs the toolkit before the final CDI refresh" {
+    install_nvidia_container_toolkit() { echo "install-toolkit" >> "${TEST_TMP}/runtime-order"; }
+    mkdir() { :; }
+    cp() {
+        case "$1" in
+            */10-nvidia-runtime.toml) echo "containerd-config" >> "${TEST_TMP}/runtime-order" ;;
+            */71-nvidia-char-dev.rules) echo "udev-rule" >> "${TEST_TMP}/runtime-order" ;;
+        esac
+    }
+    dirname() { command dirname "$@"; }
+    systemctl() { echo "systemctl $*" >> "${TEST_TMP}/runtime-order"; }
+    export AKSGPU_NVIDIA_CTK_BIN="${TEST_TMP}/bin/nvidia-ctk"
+    cat > "${AKSGPU_NVIDIA_CTK_BIN}" <<EOF
+#!/usr/bin/env bash
+echo "nvidia-ctk" >> "${TEST_TMP}/runtime-order"
+EOF
+    chmod +x "${AKSGPU_NVIDIA_CTK_BIN}"
+
+    run configure_nvidia_container_runtime
+
+    [ "$status" -eq 0 ]
+    run cat "${TEST_TMP}/runtime-order"
+    [ "$status" -eq 0 ]
+    [ "$output" = $'install-toolkit\ncontainerd-config\nudev-rule\nnvidia-ctk\nsystemctl reset-failed nvidia-cdi-refresh.service nvidia-cdi-refresh.path\nsystemctl restart nvidia-cdi-refresh.service' ]
+}
+
+@test "configure_nvidia_container_runtime propagates a failed final CDI refresh" {
+    install_nvidia_container_toolkit() { :; }
+    mkdir() { :; }
+    cp() { :; }
+    dirname() { command dirname "$@"; }
+    systemctl() { return 1; }
+    export AKSGPU_NVIDIA_CTK_BIN="${TEST_TMP}/bin/nvidia-ctk"
+    _stub_bin nvidia-ctk 0
+
+    run configure_nvidia_container_runtime
+
+    [ "$status" -ne 0 ]
 }