diff --git a/entrypoint.sh b/entrypoint.sh index b402c56..40f8c82 100755 --- a/entrypoint.sh +++ b/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" ;; diff --git a/install.sh b/install.sh index 83125a0..adc8fcb 100644 --- a/install.sh +++ b/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 } diff --git a/test/install.bats b/test/install.bats index 26b84cc..870f3a7 100644 --- a/test/install.bats +++ b/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}" <> "${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 ] }