Skip to content

fix: keep build-only NVIDIA prebakes unregistered with DKMS #191

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
sulixu wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

fix: keep build-only NVIDIA prebakes unregistered with DKMS #191

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion entrypoint.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ 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 + node initialization (legacy behaviour)
# build-only -> compile/cache the kernel module only (VHD build, no GPU)
# build-only -> compile without DKMS registration (shared VHD build, no GPU)
# install-skip-build -> node initialization only, reusing the module prebuilt into the VHD
GPU_INSTALL_MODE_ENV=""
case "${1}" in
Expand Down
45 changes: 34 additions & 11 deletions install.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
PS4='+ $(date -u -I"seconds" | cut -c1-19) '

# Install mode flags (set by entrypoint.sh based on the requested action):
# AKSGPU_BUILD_ONLY=1 -> compile/cache the kernel module + userspace libs only.
# AKSGPU_BUILD_ONLY=1 -> compile the kernel module + userspace libs without DKMS registration.
# Runs on a GPU-less host (e.g. the Packer VHD builder).
# Skips every device-dependent step (modprobe, nvidia-smi,
# fabric manager, persistence) and writes a marker.
Expand All @@ -20,9 +20,11 @@ PS4='+ $(date -u -I"seconds" | cut -c1-19) '
AKSGPU_BUILD_ONLY="${AKSGPU_BUILD_ONLY:-0}"
AKSGPU_SKIP_KERNEL_BUILD="${AKSGPU_SKIP_KERNEL_BUILD:-0}"

# Host-side marker describing what was baked into the VHD at build time. AgentBaker reads
# this (plus its own image-digest record) to decide whether the boot-time fast path is safe.
# Host-side marker describing what was baked into the VHD at build time.
# AgentBaker still uses the normal install action, not install-skip-build.
DKMS_MARKER_FILE="/opt/azure/aks-gpu/dkms-marker"
NVIDIA_DKMS_DIR="/var/lib/dkms/nvidia"
GPU_CACHE_DIR="/opt/gpu"

KERNEL_NAME=$(uname -r)
LOG_FILE_NAME="/var/log/nvidia-installer-$(date +%s).log"
Expand Down Expand Up @@ -100,6 +102,12 @@ install_nvidia_container_toolkit() {
# userspace libraries. It performs NO device access, so it is safe to run at VHD build time on
# a host without a GPU.
build_kernel_module() {
# A shared image must be safe even when node-time cleanup never succeeds. Do not
# repair inherited registrations here: refuse to publish that image instead.
if [ "${AKSGPU_BUILD_ONLY}" = "1" ] && { [ -e "${NVIDIA_DKMS_DIR}" ] || [ -L "${NVIDIA_DKMS_DIR}" ]; }; then
echo "aks-gpu: refusing to prebake over an existing NVIDIA DKMS registration" >&2
return 1
fi
# blacklist nouveau driver, nvidia driver dependency
cp /opt/gpu/blacklist-nouveau.conf /etc/modprobe.d/blacklist-nouveau.conf
update-initramfs -u
Expand All @@ -117,10 +125,15 @@ build_kernel_module() {

resolve_runfile

# install nvidia drivers (DKMS build is the dominant cost we are hoisting to VHD build time)
pushd /opt/gpu
# Keep the installed module in the path used by AgentBaker's CPU/GRID cleanup.
# --no-dkms otherwise uses the runfile installer's different default path.
local dkms_options=(--dkms)
if [ "${AKSGPU_BUILD_ONLY}" = "1" ]; then
dkms_options=(--no-dkms "--kernel-install-path=/lib/modules/${KERNEL_NAME}/updates/dkms")
fi
pushd "${GPU_CACHE_DIR}"
local installer_rc=0
/opt/gpu/${RUNFILE}/nvidia-installer -s -k=$KERNEL_NAME --log-file-name=${LOG_FILE_NAME} -a --no-drm --dkms || installer_rc=$?
"${GPU_CACHE_DIR}/${RUNFILE}/nvidia-installer" -s -k="$KERNEL_NAME" --log-file-name="${LOG_FILE_NAME}" -a --no-drm "${dkms_options[@]}" || installer_rc=$?
popd
if [ "${installer_rc}" -ne 0 ]; then
echo "aks-gpu: nvidia-installer failed (rc=${installer_rc}) for kernel ${KERNEL_NAME}; installer log follows:"
Expand All @@ -145,8 +158,16 @@ build_kernel_module() {

cleanup_overlay

# validate that the kernel module was built and registered (no device access required)
dkms status
if [ "${AKSGPU_BUILD_ONLY}" = "1" ]; then
local status
status="$(dkms status -m nvidia)" || return 1
if [ -n "${status}" ] || [ -e "${NVIDIA_DKMS_DIR}" ] || [ -L "${NVIDIA_DKMS_DIR}" ]; then
echo "aks-gpu: prebake left an NVIDIA DKMS registration" >&2
return 1
fi
else
dkms status
fi
modinfo -k "$KERNEL_NAME" nvidia
}

Expand Down Expand Up @@ -211,8 +232,7 @@ EOF
}

# baked_marker_matches returns success only when the VHD's baked driver exactly matches what
# this node needs (kernel + driver_version + driver_kind). AgentBaker requests skip-build based
# only on the marker's presence and delegates the actual match check here, so a CUDA-baked VHD
# this node needs (kernel + driver_version + driver_kind). For install-skip-build, a CUDA-baked VHD
# booting a GRID node -- or a driver-version bump since bake -- fails this check and falls back
# to a full build.
baked_marker_matches() {
Expand All @@ -226,12 +246,15 @@ baked_marker_matches() {
[ "${m_kind}" = "${DRIVER_KIND}" ]
}

# build_and_mark compiles + DKMS-registers the module, then records exactly what was built so
# build_and_mark compiles the module, then records exactly what was built so
# the marker always reflects on-disk reality. Writing the marker after every build (not just at
# VHD-bake time) means a boot-time fallback build also refreshes the marker, so the *next* boot
# takes the fast path instead of rebuilding forever.
build_and_mark() {
# Do not put build_kernel_module in an `if`/`||`: that disables errexit inside it.
build_kernel_module
local build_status=$?
[ "${build_status}" -eq 0 ] || return "${build_status}"
write_dkms_marker
}

Expand Down
139 changes: 128 additions & 11 deletions test/install.bats
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,8 @@ EOF
DRIVER_KIND="cuda"
ARCH="x86_64"
DKMS_MARKER_FILE="${TEST_TMP}/dkms-marker"
NVIDIA_DKMS_DIR="${TEST_TMP}/dkms/nvidia"
GPU_CACHE_DIR="${TEST_TMP}/gpu"
}

teardown() {
Expand Down Expand Up @@ -75,6 +77,62 @@ _stub_dispatch() {
purge_gpu_cache() { :; }
}

# Stub read-only module checks; registration comes only from the fake installer.
_stub_fast_path() {
ldconfig() { :; }
modinfo() {
[ "${MODINFO_FAIL:-0}" = 0 ] || return 1
case "$*" in
*'-F version'*) echo "${MODULE_VERSION:-${DRIVER_VERSION}}" ;;
*'-F vermagic'*) echo "${MODULE_KERNEL:-${KERNEL_NAME}} SMP mod_unload" ;;
*) echo 'module information' ;;
esac
}
dkms() {
echo "$*" >> "${TEST_TMP}/dkms-calls"
case "$1" in
status)
[ "${DKMS_STATUS_FAIL:-0}" = 0 ] || return 1
if [ -s "${NVIDIA_DKMS_DIR}/${DRIVER_VERSION}/source/dkms.conf" ]; then
echo "nvidia/${DRIVER_VERSION}, ${KERNEL_NAME}, ${ARCH}: installed"
fi
;;
*) return 99 ;;
esac
}
}

_stub_build() {
_stub_fast_path
export TEST_TMP NVIDIA_DKMS_DIR DRIVER_VERSION
command mkdir -p "${GPU_CACHE_DIR}/fake"
cat > "${GPU_CACHE_DIR}/fake/nvidia-installer" <<'EOF'
#!/usr/bin/env bash
printf '%s\n' "$@" > "${TEST_TMP}/installer-args"
for arg in "$@"; do
if [ "$arg" = --dkms ]; then INSTALLER_REGISTERS=1; fi
done
if [ "${INSTALLER_REGISTERS:-0}" = 1 ]; then
mkdir -p "${NVIDIA_DKMS_DIR}/${DRIVER_VERSION}/source"
echo 'PACKAGE_NAME="nvidia"' > "${NVIDIA_DKMS_DIR}/${DRIVER_VERSION}/source/dkms.conf"
fi
exit "${INSTALLER_RC:-0}"
EOF
chmod +x "${GPU_CACHE_DIR}/fake/nvidia-installer"
resolve_runfile() { RUNFILE=fake; }
cleanup_overlay() { :; }
cp() { :; }
rm() { :; }
mount() { :; }
mkdir() { case "$*" in *"${TEST_TMP}"*) command mkdir "$@" ;; *) : ;; esac; }
update-initramfs() { :; }
GPU_DEST="${TEST_TMP}/userspace"
# Redirect the build function's fixed linker config write for these host-side tests.
build_kernel_module_for_test="$(declare -f build_kernel_module)"
build_kernel_module_for_test="${build_kernel_module_for_test//\/etc\/ld.so.conf.d\/nvidia.conf/${TEST_TMP}/nvidia.conf}"
eval "${build_kernel_module_for_test}"
}

# --- marker: write -------------------------------------------------------

@test "write_dkms_marker records kernel/version/kind/arch" {
Expand Down Expand Up @@ -131,27 +189,76 @@ _stub_dispatch() {

# --- fast-path fallback --------------------------------------------------

@test "fast_path_ok succeeds when ldconfig+dkms+modinfo all pass" {
_stub_bin ldconfig 0; _stub_bin dkms 0; _stub_bin modinfo 0
PATH="${TEST_TMP}/bin:$PATH"
@test "fast_path_ok succeeds when dkms and modinfo succeed" {
_stub_fast_path
run fast_path_ok
[ "$status" -eq 0 ]
}

@test "fast_path_ok fails (-> full build) when modinfo reports the module is unusable" {
_stub_bin ldconfig 0; _stub_bin dkms 0; _stub_bin modinfo 1
PATH="${TEST_TMP}/bin:$PATH"
_stub_fast_path
MODINFO_FAIL=1
run fast_path_ok
[ "$status" -ne 0 ]
}

@test "fast_path_ok fails (-> full build) when dkms status fails" {
_stub_bin ldconfig 0; _stub_bin dkms 1; _stub_bin modinfo 0
PATH="${TEST_TMP}/bin:$PATH"
_stub_fast_path
DKMS_STATUS_FAIL=1
run fast_path_ok
[ "$status" -ne 0 ]
}

@test "build-only uses no-dkms, preserves the cleanup path, and leaves no registration" {
_stub_build
AKSGPU_BUILD_ONLY=1
run build_and_mark
[ "$status" -eq 0 ]
[ -f "${DKMS_MARKER_FILE}" ]
[ ! -e "${NVIDIA_DKMS_DIR}" ]
run cat "${TEST_TMP}/installer-args"
[[ "$output" == *'--no-dkms'* ]]
[[ "$output" == *"--kernel-install-path=/lib/modules/${KERNEL_NAME}/updates/dkms"* ]]
[[ "$output" != *$'\n--dkms'* ]]
}

@test "normal installation retains DKMS registration" {
_stub_build
AKSGPU_BUILD_ONLY=0
run build_and_mark
[ "$status" -eq 0 ]
[ -s "${NVIDIA_DKMS_DIR}/${DRIVER_VERSION}/source/dkms.conf" ]
run cat "${TEST_TMP}/installer-args"
[[ "$output" == *'--dkms'* ]]
[[ "$output" != *'--no-dkms'* ]]
}

@test "build-only refuses inherited registrations, including dangling symlinks" {
_stub_build
AKSGPU_BUILD_ONLY=1
command mkdir -p "$(dirname "${NVIDIA_DKMS_DIR}")"
ln -s "${TEST_TMP}/missing" "${NVIDIA_DKMS_DIR}"
run build_and_mark
[ "$status" -ne 0 ]
[ ! -e "${TEST_TMP}/installer-args" ]
[ ! -e "${DKMS_MARKER_FILE}" ]
}

@test "build-only fails if the installer leaves registration or fails itself" {
_stub_build
AKSGPU_BUILD_ONLY=1
export INSTALLER_REGISTERS=1
run build_and_mark
[ "$status" -ne 0 ]
[ ! -e "${DKMS_MARKER_FILE}" ]
unset INSTALLER_REGISTERS
command rm -rf "${NVIDIA_DKMS_DIR}"
export INSTALLER_RC=42
run build_and_mark
[ "$status" -eq 42 ]
[ ! -e "${DKMS_MARKER_FILE}" ]
}

# --- target kernel selection --------------------------------------------

@test "target_build_kernel picks the newest kernel that has a build tree" {
Expand Down Expand Up @@ -187,8 +294,7 @@ _stub_dispatch() {

@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"
_stub_fast_path
KERNEL_NAME="5.15.0-1114-azure"; DRIVER_VERSION="580.0.0"; DRIVER_KIND="cuda"; ARCH="x86_64"
write_dkms_marker
AKSGPU_BUILD_ONLY=0; AKSGPU_SKIP_KERNEL_BUILD=1
Expand All @@ -205,8 +311,7 @@ _stub_dispatch() {

@test "dispatch install-skip-build with mismatched marker: falls back to a full build" {
_stub_dispatch
_stub_bin ldconfig 0; _stub_bin dkms 0; _stub_bin modinfo 0
PATH="${TEST_TMP}/bin:$PATH"
_stub_fast_path
KERNEL_NAME="5.15.0-1114-azure"; DRIVER_VERSION="580.0.0"; DRIVER_KIND="cuda"; ARCH="x86_64"
write_dkms_marker
DRIVER_VERSION="999.0.0" # node now needs a different version than the baked marker
Expand All @@ -219,6 +324,18 @@ _stub_dispatch() {
[ -f "${TEST_TMP}/configure_nvidia_container_runtime.ran" ]
}

@test "normal install still builds even when a matching prebake marker exists" {
_stub_dispatch
_stub_fast_path
write_dkms_marker
AKSGPU_BUILD_ONLY=0; AKSGPU_SKIP_KERNEL_BUILD=0
run main
[ "$status" -eq 0 ]
[ -f "${TEST_TMP}/build_and_mark.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" {
Expand Down