cdi install logic change, test(install): add shellcheck + bats harness for prebake logic

.github/workflows/ci.yaml

@@ -175,3 +175,15 @@ jobs:
             rm -rf /tmp/.buildx-cache
             mv /tmp/.buildx-cache-new /tmp/.buildx-cache
 
+  lint-and-test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+      - name: Shellcheck prebake scripts
+        run: shellcheck -S warning install.sh entrypoint.sh
+      - name: Install bats
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y bats
+      - name: Run install.sh unit tests
+        run: bats test/install.bats

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

@@ -1,8 +1,10 @@
 #!/usr/bin/env bash
-set -euxo pipefail
 
-source /opt/gpu/config.sh
-source /opt/gpu/package_manager_helpers.sh
+# NOTE: `set -euxo pipefail` and the `source`s of the gpu config + package-manager helpers are
+# applied inside main() rather than at top level, so this script can be sourced by the unit tests
+# (test/install.bats) to exercise the individual functions on a GPU-less host without running the
+# install or requiring /opt/gpu to exist. main() is invoked only when the script is executed
+# directly (see the guard at the bottom).
 
 PS4='+ $(date -u -I"seconds" | cut -c1-19) '
 
@@ -14,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}"
 
@@ -35,11 +37,13 @@ ARCH=$(uname -m)
 # normal node boot, where uname -r is already correct).
 target_build_kernel() {
     local d k
-    # newest installed kernel that has a headers/build tree (the VHD's target kernel)
-    k=$(for d in /lib/modules/*/build; do
+    # newest installed kernel that has a headers/build tree (the VHD's target kernel). The modules
+    # root is overridable (AKSGPU_MODULES_ROOT) so the unit tests can point it at a fixture dir.
+    local modules_root="${AKSGPU_MODULES_ROOT:-/lib/modules}"
+    k=$(for d in "${modules_root}"/*/build; do
             [ -d "$d" ] || continue
             d=${d%/build}
-            echo "${d#/lib/modules/}"
+            echo "${d#"${modules_root}"/}"
         done | sort -V | tail -n1)
     if [ -n "$k" ]; then echo "$k"; else uname -r; fi
 }
@@ -62,9 +66,9 @@ cleanup_overlay() {
     fi
     set -e
 }
-# Reset PS4 on exit alongside the overlay cleanup (a single EXIT trap, since a second
-# `trap ... EXIT` would replace the first rather than chain).
-trap 'cleanup_overlay; PS4="+ "' EXIT
+# The EXIT trap that runs cleanup_overlay (and resets PS4) is installed at the start of main(),
+# not here, so that sourcing this script for unit tests neither registers a trap nor tears down
+# mounts when the test process exits.
 
 resolve_runfile() {
     if [[ "${DRIVER_KIND}" == "cuda" ]]; then
@@ -146,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
@@ -157,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
@@ -168,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() {
@@ -218,48 +238,71 @@ 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
     modinfo -k "${KERNEL_NAME}" nvidia >/dev/null 2>&1 || return 1
 }
 
-set +euo pipefail
-open_devices="$(lsof /dev/nvidia* 2>/dev/null)"
-echo "Open devices: $open_devices"
-
-open_gridd="$(lsof /usr/bin/nvidia-gridd 2>/dev/null)"
-echo "Open gridd: $open_gridd"
-set -euo pipefail
-
-if [ "${AKSGPU_BUILD_ONLY}" = "1" ]; then
-    # VHD build time: compile + cache + marker only, no device access. Target the kernel the VHD
-    # will boot (not the builder's running kernel) so the prebuilt module + marker match at boot.
-    KERNEL_NAME="$(target_build_kernel)"
-    echo "aks-gpu: build-only mode (prebuilding kernel module for kernel ${KERNEL_NAME}; builder running $(uname -r))"
-    echo "aks-gpu: kernels with installed headers (build trees):"; ls -ld /lib/modules/*/build 2>/dev/null || echo "  (none found)"
-    build_and_mark
+# purge_gpu_cache removes the gpu cache that entrypoint.sh staged under /opt/gpu once install.sh
+# has consumed it. Factored out of the two former inline `rm -r /opt/gpu` calls so the unit tests
+# can stub it when exercising main()'s dispatch.
+purge_gpu_cache() {
     rm -r /opt/gpu
-    exit 0
-fi
+}
 
-install_nvidia_container_toolkit
+# main is the install entrypoint. It is run only when the script is executed directly (the guard
+# at the bottom), so sourcing the script for unit tests loads the functions above without running
+# any install steps. The set-flags and the gpu config/helper sources live here (not at top level)
+# for the same reason.
+main() {
+    set -euxo pipefail
+    # shellcheck source=/dev/null
+    source "${AKSGPU_CONFIG_PATH:-/opt/gpu/config.sh}"
+    # shellcheck source=/dev/null
+    source "${AKSGPU_PMH_PATH:-/opt/gpu/package_manager_helpers.sh}"
+    trap 'cleanup_overlay; PS4="+ "' EXIT
+
+    set +euo pipefail
+    open_devices="$(lsof /dev/nvidia* 2>/dev/null)"
+    echo "Open devices: $open_devices"
+
+    open_gridd="$(lsof /usr/bin/nvidia-gridd 2>/dev/null)"
+    echo "Open gridd: $open_gridd"
+    set -euo pipefail
+
+    if [ "${AKSGPU_BUILD_ONLY}" = "1" ]; then
+        # VHD build time: compile + cache + marker only, no device access. Target the kernel the VHD
+        # will boot (not the builder's running kernel) so the prebuilt module + marker match at boot.
+        KERNEL_NAME="$(target_build_kernel)"
+        echo "aks-gpu: build-only mode (prebuilding kernel module for kernel ${KERNEL_NAME}; builder running $(uname -r))"
+        echo "aks-gpu: kernels with installed headers (build trees):"; ls -ld "${AKSGPU_MODULES_ROOT:-/lib/modules}"/*/build 2>/dev/null || echo "  (none found)"
+        build_and_mark
+        purge_gpu_cache
+        exit 0
+    fi
 
-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)"
-else
-    if [ "${AKSGPU_SKIP_KERNEL_BUILD}" = "1" ]; then
-        echo "aks-gpu: prebuilt module missing/invalid for ${DRIVER_KIND} ${DRIVER_VERSION} on ${KERNEL_NAME}; building from source"
+    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)"
+    else
+        if [ "${AKSGPU_SKIP_KERNEL_BUILD}" = "1" ]; then
+            echo "aks-gpu: prebuilt module missing/invalid for ${DRIVER_KIND} ${DRIVER_VERSION} on ${KERNEL_NAME}; building from source"
+        fi
+        # No bespoke stale-driver teardown is needed: `nvidia-installer -s` automatically uninstalls
+        # any previously runfile-installed driver -- including a mismatched prebaked one -- and its
+        # DKMS registration before installing the new one. build_and_mark then refreshes the marker
+        # to match what we just built, so subsequent boots take the fast path.
+        build_and_mark
     fi
-    # No bespoke stale-driver teardown is needed: `nvidia-installer -s` automatically uninstalls
-    # any previously runfile-installed driver -- including a mismatched prebaked one -- and its
-    # DKMS registration before installing the new one. build_and_mark then refreshes the marker
-    # to match what we just built, so subsequent boots take the fast path.
-    build_and_mark
-fi
 
-device_init
+    initialize_nvidia_driver
+    configure_nvidia_container_runtime
 
-rm -r /opt/gpu
+    purge_gpu_cache
+}
+
+if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
+    main
+fi