Skip to content

Commit 922cbea

Browse files
ganeshkumarashokCopilot
ganeshkumarashok
and
Copilot
authored
feat: build-only / install-skip-build modes for VHD-prebuilt GPU kernel module (#162)
* feat: add build-only and install-skip-build modes to support VHD-prebuilt kernel module Split the host-side driver install into two phases so the NVIDIA kernel module can be DKMS-compiled into the VHD at image build time and the boot-time install can skip straight to device init: - install.sh: refactor into build_kernel_module() (compile + stage userspace libs, no device access) and device_init() (modprobe, nvidia-smi, fabric manager, containerd config, udev). Add AKSGPU_BUILD_ONLY and AKSGPU_SKIP_KERNEL_BUILD modes, an overlay cleanup trap, and a dkms-marker (/opt/azure/aks-gpu/dkms-marker) recording kernel, driver_version, driver_kind and arch so the consumer (AgentBaker CSE) can validate an exact match before taking the skip-build fast path. - entrypoint.sh: add build-only and install-skip-build actions and pass the mode through to the host via nsenter. The default install action is unchanged. This is the aks-gpu half of the AgentBaker change that prebuilds the GPU kernel module into the VHD to reduce node provisioning time. Secure Boot module signing and GPU e2e validation are still required. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * harden skip-build: gate on baked marker + clean stale shared-VHD driver The skip-kernel-build fast path now requires a matching dkms-marker (kernel + driver_version + driver_kind); on mismatch it falls back to a full build. Before any full (re)build, remove_stale_baked_driver unloads loaded nvidia modules and removes a stale registered DKMS tree and its relocated libs/loader config, so a CUDA-baked VHD booting on a GRID node (or a version skew) cannot collide with the boot-time nvidia-installer. No-op on today's VHDs (nothing baked / nothing registered). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * refactor: scope stale-driver cleanup to prebaked VHDs; dedup entrypoint; fix EXIT trap - entrypoint.sh: collapse the per-action cp/echo duplication into a single staging block; the case now only maps action -> install-mode env var. - install.sh: the second `trap ... EXIT` (cleanup_overlay) was silently replacing the earlier `trap 'PS4="+ "' exit`; fold both into one EXIT trap. - install.sh: only run remove_stale_baked_driver when a baked marker is present, so the default `install` path (non-prebaked VHD, no marker) is byte-for-byte the legacy behaviour and never touches DKMS cleanup. The cleanup still runs for the reachable mismatch case (marker present but kind/version/kernel differs). - Correct comments that framed the mismatch path as defensive / "no-op on today's VHDs"; it is reachable because AgentBaker gates skip-build on marker presence alone. No functional change to build-only or the matching skip-build fast path. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * refactor(install): simplify prebake install.sh; address PR #162 review - Drop remove_stale_baked_driver + unload_nvidia_modules. nvidia-installer -s already uninstalls a prior runfile/DKMS driver before installing, so the bespoke stale-driver teardown is redundant. Removes the dkms-status parser bug, the swallowed unload-failure path, and the early-return cleanup gap. - Opportunistic fast path: validate via fast_path_ok and fall back to a full build instead of hard-failing under set -e (fixes "corrupt prebake bricks the node"). - Always write the marker after a build (build_and_mark) so a boot-time fallback refreshes it and later boots take the fast path (no rebuild loop). - Overlay cleanup is mount-state driven (mountpoint -q): a partial mount is always torn down and state is never cleared while a mount is live. - Write the DKMS marker atomically (temp + mv). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix(prebake): compile build-only module against the VHD's target kernel, not the builder's running kernel At VHD build time the Packer builder still runs the image's previous kernel while the build has already upgraded the image to a newer kernel and installed only that target kernel's headers (the builder kernel's headers are purged). Compiling against $(uname -r) therefore fails (no headers) and, even if it succeeded, the dkms-marker would not match the kernel the node boots. In build-only mode, select the newest kernel that has an installed headers/build tree (the VHD's target kernel); fall back to uname -r otherwise (e.g. normal node boot). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * prebake: surface nvidia-installer log on failure + log target/running kernel and available headers Improves debuggability of the build-only VHD prebake: prints the selected target kernel vs the builder's running kernel and the kernels with installed headers, and dumps the nvidia-installer log if the DKMS compile fails. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix(prebake): clear stale staged libs before a full build (GRID-on-CUDA-VHD) When a full driver build runs over a prebaked VHD whose driver kind/version differs -- e.g. a GRID node booting a CUDA-prebaked image, which the AgentBaker driver_kind guard correctly routes to a full install -- the baked driver's userspace libraries are already staged under ${GPU_DEST}/lib64/lib64. The plain 'cp -a' then MERGED both versions (e.g. libnvidia-ml.so.570.* alongside .580.*), and device_init copied the mix into /usr/lib/<arch>-linux-gnu. nvidia-smi resolved one version and worked, but NVML consumers like nvidia-device-plugin tripped on the mismatch ('Failed to initialize NVML: Driver/library version mismatch') and never recovered. Remove the staged lib dir before copying so it is authoritative for exactly the driver we just built. The skip-build fast path never reaches build_kernel_module, so a matching prebaked VHD keeps its baked libs untouched. Caught by AgentBaker GPU e2e (build 168618935): all remaining failures were GRID A10 nodes on a CUDA-prebaked VHD. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent 511547c commit 922cbea

3 files changed

Lines changed: 256 additions & 86 deletions

File tree

‎.dockerignore‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
.git
2+
__pycache__
3+
*.pyc

‎entrypoint.sh‎

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,8 +23,19 @@ if [[ "${1}" == "copy" ]]; then
2323
exit 0
2424
fi
2525

26-
if [[ "${1}" == "install" ]]; then
27-
echo "copying gpu cache files"
26+
# Map the requested action to the install mode passed to install.sh. All three install
27+
# variants stage the same gpu cache files; only the env var handed to install.sh differs.
28+
# install -> full compile + device init (legacy behaviour)
29+
# build-only -> compile/cache the kernel module only (VHD build, no GPU)
30+
# install-skip-build -> device init only, reusing the module prebuilt into the VHD
31+
GPU_INSTALL_MODE_ENV=""
32+
case "${1}" in
33+
build-only) GPU_INSTALL_MODE_ENV="AKSGPU_BUILD_ONLY=1" ;;
34+
install-skip-build) GPU_INSTALL_MODE_ENV="AKSGPU_SKIP_KERNEL_BUILD=1" ;;
35+
esac
36+
37+
if [[ "${1}" == "install" || -n "${GPU_INSTALL_MODE_ENV}" ]]; then
38+
echo "copying gpu cache files (${1})"
2839
cp -a /opt/gpu/. /mnt/gpu/
2940
echo "copied successfully!"
3041
fi
@@ -46,7 +57,11 @@ cp -R /opt/actions/. /mnt/actions
4657

4758
echo "Executing nsenter"
4859

49-
nsenter -t 1 -m bash "${ACTION_FILE}"
60+
if [[ -n "${GPU_INSTALL_MODE_ENV}" ]]; then
61+
nsenter -t 1 -m env "${GPU_INSTALL_MODE_ENV}" bash "${ACTION_FILE}"
62+
else
63+
nsenter -t 1 -m bash "${ACTION_FILE}"
64+
fi
5065
RESULT="${PIPESTATUS[0]}"
5166

5267
if [ $RESULT -eq 0 ]; then

‎install.sh‎

Lines changed: 235 additions & 83 deletions
Original file line numberDiff line numberDiff line change
@@ -4,21 +4,82 @@ set -euxo pipefail
44
source /opt/gpu/config.sh
55
source /opt/gpu/package_manager_helpers.sh
66

7-
trap 'PS4="+ "' exit
87
PS4='+ $(date -u -I"seconds" | cut -c1-19) '
98

9+
# Install mode flags (set by entrypoint.sh based on the requested action):
10+
# AKSGPU_BUILD_ONLY=1 -> compile/cache the kernel module + userspace libs only.
11+
# Runs on a GPU-less host (e.g. the Packer VHD builder).
12+
# Skips every device-dependent step (modprobe, nvidia-smi,
13+
# fabric manager, persistence) and writes a marker.
14+
# AKSGPU_SKIP_KERNEL_BUILD=1 -> the kernel module + libs were prebuilt into the VHD for
15+
# this exact kernel+driver; skip recompilation and only run
16+
# the device-dependent steps at node boot.
17+
# (neither set) -> legacy behaviour: full compile + device init in one shot.
18+
AKSGPU_BUILD_ONLY="${AKSGPU_BUILD_ONLY:-0}"
19+
AKSGPU_SKIP_KERNEL_BUILD="${AKSGPU_SKIP_KERNEL_BUILD:-0}"
20+
21+
# Host-side marker describing what was baked into the VHD at build time. AgentBaker reads
22+
# this (plus its own image-digest record) to decide whether the boot-time fast path is safe.
23+
DKMS_MARKER_FILE="/opt/azure/aks-gpu/dkms-marker"
24+
1025
KERNEL_NAME=$(uname -r)
1126
LOG_FILE_NAME="/var/log/nvidia-installer-$(date +%s).log"
1227
ARCH=$(uname -m)
1328

14-
set +euo pipefail
15-
open_devices="$(lsof /dev/nvidia* 2>/dev/null)"
16-
echo "Open devices: $open_devices"
29+
# target_build_kernel returns the kernel the VHD will actually boot. At VHD build time the Packer
30+
# builder is still running the image's PREVIOUS kernel, while the build has already upgraded the
31+
# image to a NEWER kernel and installed only that target kernel's headers (the builder kernel's
32+
# headers are purged). Compiling against `uname -r` would therefore fail (no headers) and, even if
33+
# it succeeded, the dkms-marker would not match the kernel the node boots. Select the newest kernel
34+
# that has a usable headers/build tree; fall back to the running kernel when none is found (e.g. a
35+
# normal node boot, where uname -r is already correct).
36+
target_build_kernel() {
37+
local d k
38+
# newest installed kernel that has a headers/build tree (the VHD's target kernel)
39+
k=$(for d in /lib/modules/*/build; do
40+
[ -d "$d" ] || continue
41+
d=${d%/build}
42+
echo "${d#/lib/modules/}"
43+
done | sort -V | tail -n1)
44+
if [ -n "$k" ]; then echo "$k"; else uname -r; fi
45+
}
1746

18-
open_gridd="$(lsof /usr/bin/nvidia-gridd 2>/dev/null)"
19-
echo "Open gridd: $open_gridd"
47+
# cleanup_overlay tears down the tmpfs+overlay scaffold. It is driven entirely by the live mount
48+
# state (mountpoint -q) rather than a flag, so it is idempotent and correct no matter where a
49+
# failure occurred -- including a partial setup where only the tmpfs mounted but the
50+
# overlay-on-/usr/lib mount failed. This guarantees a build-time exit never freezes a dangling
51+
# mount into the VHD, and never clears state while a mount is still live.
52+
cleanup_overlay() {
53+
set +e
54+
if mountpoint -q "/usr/lib/${ARCH}-linux-gnu"; then
55+
umount -l "/usr/lib/${ARCH}-linux-gnu"
56+
fi
57+
if mountpoint -q /tmp/overlay; then
58+
umount /tmp/overlay
59+
fi
60+
if ! mountpoint -q /tmp/overlay 2>/dev/null; then
61+
rm -rf /tmp/overlay
62+
fi
63+
set -e
64+
}
65+
# Reset PS4 on exit alongside the overlay cleanup (a single EXIT trap, since a second
66+
# `trap ... EXIT` would replace the first rather than chain).
67+
trap 'cleanup_overlay; PS4="+ "' EXIT
2068

21-
set -euo pipefail
69+
resolve_runfile() {
70+
if [[ "${DRIVER_KIND}" == "cuda" ]]; then
71+
RUNFILE="NVIDIA-Linux-${ARCH}-${DRIVER_VERSION}"
72+
elif [[ "${DRIVER_KIND}" == "grid" ]]; then
73+
if [[ "${ARCH}" != "x86_64" ]]; then
74+
echo "GRID driver is only supported on x86_64 architecture"
75+
exit 1
76+
fi
77+
RUNFILE="NVIDIA-Linux-x86_64-${DRIVER_VERSION}-grid-azure"
78+
else
79+
echo "Invalid driver kind: ${DRIVER_KIND}"
80+
exit 1
81+
fi
82+
}
2283

2384
# install cached nvidia debian packages for container runtime compatibility
2485
install_cached_nvidia_packages() {
@@ -27,87 +88,178 @@ for apt_package in $NVIDIA_PACKAGES; do
2788
done
2889
}
2990

30-
use_package_manager_with_retries wait_for_dpkg_lock install_cached_nvidia_packages 10 3
31-
32-
# blacklist nouveau driver, nvidia driver dependency
33-
cp /opt/gpu/blacklist-nouveau.conf /etc/modprobe.d/blacklist-nouveau.conf
34-
update-initramfs -u
35-
36-
# clean up lingering files from previous install
37-
set +e
38-
umount -l /usr/lib/$(uname -m)-linux-gnu || true
39-
umount -l /tmp/overlay || true
40-
rm -r /tmp/overlay || true
41-
set -e
42-
43-
# set up overlayfs to change install location of nvidia libs from /usr/lib/$ARCH-linux-gnu to /usr/local/nvidia
44-
# add an extra layer of indirection via tmpfs because it's not possible to have an overlayfs on an overlayfs (i.e., inside a container)
45-
mkdir /tmp/overlay
46-
mount -t tmpfs tmpfs /tmp/overlay
47-
mkdir /tmp/overlay/{workdir,lib64}
48-
mkdir -p ${GPU_DEST}/lib64
49-
mount -t overlay overlay -o lowerdir=/usr/lib/$(uname -m)-linux-gnu,upperdir=/tmp/overlay/lib64,workdir=/tmp/overlay/workdir /usr/lib/$(uname -m)-linux-gnu
50-
51-
if [[ "${DRIVER_KIND}" == "cuda" ]]; then
52-
RUNFILE="NVIDIA-Linux-$(uname -m)-${DRIVER_VERSION}"
53-
elif [[ "${DRIVER_KIND}" == "grid" ]]; then
54-
if [[ $(uname -m) != "x86_64" ]]; then
55-
echo "GRID driver is only supported on x86_64 architecture"
56-
exit 1
91+
install_nvidia_container_toolkit() {
92+
use_package_manager_with_retries wait_for_dpkg_lock install_cached_nvidia_packages 10 3
93+
}
94+
95+
# build_kernel_module compiles the NVIDIA kernel module (the expensive step) and stages the
96+
# userspace libraries. It performs NO device access, so it is safe to run at VHD build time on
97+
# a host without a GPU.
98+
build_kernel_module() {
99+
# blacklist nouveau driver, nvidia driver dependency
100+
cp /opt/gpu/blacklist-nouveau.conf /etc/modprobe.d/blacklist-nouveau.conf
101+
update-initramfs -u
102+
103+
# clear any lingering mounts from a previous attempt (idempotent, mount-state driven)
104+
cleanup_overlay
105+
106+
# set up overlayfs to relocate nvidia libs from /usr/lib/$ARCH-linux-gnu into GPU_DEST/lib64.
107+
# add an extra layer of indirection via tmpfs because it's not possible to have an overlayfs on an overlayfs (i.e., inside a container)
108+
mkdir -p /tmp/overlay
109+
mount -t tmpfs tmpfs /tmp/overlay
110+
mkdir -p /tmp/overlay/workdir /tmp/overlay/lib64
111+
mkdir -p ${GPU_DEST}/lib64
112+
mount -t overlay overlay -o lowerdir="/usr/lib/${ARCH}-linux-gnu",upperdir=/tmp/overlay/lib64,workdir=/tmp/overlay/workdir "/usr/lib/${ARCH}-linux-gnu"
113+
114+
resolve_runfile
115+
116+
# install nvidia drivers (DKMS build is the dominant cost we are hoisting to VHD build time)
117+
pushd /opt/gpu
118+
local installer_rc=0
119+
/opt/gpu/${RUNFILE}/nvidia-installer -s -k=$KERNEL_NAME --log-file-name=${LOG_FILE_NAME} -a --no-drm --dkms || installer_rc=$?
120+
popd
121+
if [ "${installer_rc}" -ne 0 ]; then
122+
echo "aks-gpu: nvidia-installer failed (rc=${installer_rc}) for kernel ${KERNEL_NAME}; installer log follows:"
123+
cat "${LOG_FILE_NAME}" 2>/dev/null || true
124+
return "${installer_rc}"
57125
fi
58-
RUNFILE="NVIDIA-Linux-x86_64-${DRIVER_VERSION}-grid-azure"
59-
else
60-
echo "Invalid driver kind: ${DRIVER_KIND}"
61-
exit 1
62-
fi
63126

64-
# install nvidia drivers
65-
pushd /opt/gpu
66-
/opt/gpu/${RUNFILE}/nvidia-installer -s -k=$KERNEL_NAME --log-file-name=${LOG_FILE_NAME} -a --no-drm --dkms
67-
nvidia-smi
68-
popd
69-
70-
# move nvidia libs to correct location from temporary overlayfs
71-
cp -a /tmp/overlay/lib64 ${GPU_DEST}/lib64
72-
73-
# configure system to know about nvidia lib paths
74-
echo "${GPU_DEST}/lib64" > /etc/ld.so.conf.d/nvidia.conf
75-
ldconfig
76-
77-
# unmount, cleanup
78-
set +e
79-
umount -l /usr/lib/$(uname -m)-linux-gnu
80-
umount /tmp/overlay
81-
rm -r /tmp/overlay
82-
set -e
83-
84-
# validate that nvidia driver is working
85-
dkms status
86-
nvidia-modprobe -u -c0
87-
88-
# configure persistence daemon
89-
# decreases latency for later driver loads
90-
# reduces nvidia-smi invocation time 10x from 30 to 2 sec
91-
# notable on large VM sizes with multiple GPUs
92-
# especially when nvidia-smi process is in CPU cgroup
93-
cp -r /usr/bin/lib64/lib64/* /usr/lib/$(uname -m)-linux-gnu/
94-
nvidia-smi
95-
96-
# install fabricmanager for nvlink based systems
97-
if [[ "${DRIVER_KIND}" == "cuda" ]]; then
98-
NVIDIA_FM_ARCH=$(uname -m)
99-
if [ $NVIDIA_FM_ARCH = "arm64" ]; then
100-
# NVIDIA uses the name "SBSA" for ARM64 platforms for the fabric manager. See https://en.wikipedia.org/wiki/Server_Base_System_Architecture
101-
NVIDIA_FM_ARCH="sbsa"
127+
# move nvidia libs to correct location from temporary overlayfs. Clear any pre-existing libs
128+
# first: when a full build runs over a prebaked VHD whose driver kind/version differs (e.g. a
129+
# GRID node booting a CUDA-prebaked image, routed here by the driver_kind guard), the baked
130+
# driver's userspace libs are already staged under ${GPU_DEST}/lib64. A plain copy would MERGE
131+
# both versions (e.g. libnvidia-ml.so.570.* and .580.*), which breaks NVML for consumers like the
132+
# device plugin ("Driver/library version mismatch"). Removing the directory first makes the
133+
# staged libs authoritative for exactly the driver we just built. (The skip-build fast path never
134+
# reaches here, so a matching prebaked VHD keeps its baked libs untouched.)
135+
rm -rf "${GPU_DEST}/lib64/lib64"
136+
cp -a /tmp/overlay/lib64 ${GPU_DEST}/lib64
137+
138+
# configure system to know about nvidia lib paths
139+
echo "${GPU_DEST}/lib64" > /etc/ld.so.conf.d/nvidia.conf
140+
ldconfig
141+
142+
cleanup_overlay
143+
144+
# validate that the kernel module was built and registered (no device access required)
145+
dkms status
146+
modinfo -k "$KERNEL_NAME" nvidia
147+
}
148+
149+
# device_init runs the steps that require the physical GPU and therefore must execute at node
150+
# boot, regardless of whether the kernel module was prebuilt into the VHD.
151+
device_init() {
152+
nvidia-modprobe -u -c0
153+
154+
# configure persistence daemon
155+
# decreases latency for later driver loads
156+
# reduces nvidia-smi invocation time 10x from 30 to 2 sec
157+
# notable on large VM sizes with multiple GPUs
158+
# especially when nvidia-smi process is in CPU cgroup
159+
cp -r /usr/bin/lib64/lib64/* "/usr/lib/${ARCH}-linux-gnu/"
160+
nvidia-smi
161+
162+
# install fabricmanager for nvlink based systems
163+
if [[ "${DRIVER_KIND}" == "cuda" ]]; then
164+
NVIDIA_FM_ARCH=$ARCH
165+
if [ "$NVIDIA_FM_ARCH" = "arm64" ]; then
166+
# NVIDIA uses the name "SBSA" for ARM64 platforms for the fabric manager. See https://en.wikipedia.org/wiki/Server_Base_System_Architecture
167+
NVIDIA_FM_ARCH="sbsa"
168+
fi
169+
bash /opt/gpu/fabricmanager-linux-${NVIDIA_FM_ARCH}-${DRIVER_VERSION}/sbin/fm_run_package_installer.sh
102170
fi
103-
bash /opt/gpu/fabricmanager-linux-${NVIDIA_FM_ARCH}-${DRIVER_VERSION}/sbin/fm_run_package_installer.sh
171+
172+
mkdir -p /etc/containerd/config.d
173+
cp /opt/gpu/10-nvidia-runtime.toml /etc/containerd/config.d/10-nvidia-runtime.toml
174+
175+
mkdir -p "$(dirname /lib/udev/rules.d/71-nvidia-dev-char.rules)"
176+
cp /opt/gpu/71-nvidia-char-dev.rules /lib/udev/rules.d/71-nvidia-dev-char.rules
177+
/usr/bin/nvidia-ctk system create-dev-char-symlinks --create-all
178+
}
179+
180+
write_dkms_marker() {
181+
mkdir -p "$(dirname "${DKMS_MARKER_FILE}")"
182+
local tmp_marker="${DKMS_MARKER_FILE}.tmp.$$"
183+
cat > "${tmp_marker}" <<EOF
184+
kernel=${KERNEL_NAME}
185+
driver_version=${DRIVER_VERSION}
186+
driver_kind=${DRIVER_KIND}
187+
arch=${ARCH}
188+
EOF
189+
# atomic publish: a partially-written marker must never be observed by the boot-time check
190+
mv -f "${tmp_marker}" "${DKMS_MARKER_FILE}"
191+
}
192+
193+
# baked_marker_matches returns success only when the VHD's baked driver exactly matches what
194+
# this node needs (kernel + driver_version + driver_kind). AgentBaker requests skip-build based
195+
# only on the marker's presence and delegates the actual match check here, so a CUDA-baked VHD
196+
# booting a GRID node -- or a driver-version bump since bake -- fails this check and falls back
197+
# to a full build.
198+
baked_marker_matches() {
199+
[ -f "${DKMS_MARKER_FILE}" ] || return 1
200+
local m_kernel m_version m_kind
201+
m_kernel="$(sed -n 's/^kernel=//p' "${DKMS_MARKER_FILE}" | head -n1)"
202+
m_version="$(sed -n 's/^driver_version=//p' "${DKMS_MARKER_FILE}" | head -n1)"
203+
m_kind="$(sed -n 's/^driver_kind=//p' "${DKMS_MARKER_FILE}" | head -n1)"
204+
[ "${m_kernel}" = "${KERNEL_NAME}" ] && \
205+
[ "${m_version}" = "${DRIVER_VERSION}" ] && \
206+
[ "${m_kind}" = "${DRIVER_KIND}" ]
207+
}
208+
209+
# build_and_mark compiles + DKMS-registers the module, then records exactly what was built so
210+
# the marker always reflects on-disk reality. Writing the marker after every build (not just at
211+
# VHD-bake time) means a boot-time fallback build also refreshes the marker, so the *next* boot
212+
# takes the fast path instead of rebuilding forever.
213+
build_and_mark() {
214+
build_kernel_module
215+
write_dkms_marker
216+
}
217+
218+
# fast_path_ok opportunistically validates a prebaked module without recompiling. It returns
219+
# non-zero (instead of aborting under `set -e`) so a corrupt or incomplete prebake falls back to
220+
# a full build rather than bricking the node. The authoritative device check still happens later
221+
# in device_init (nvidia-modprobe + nvidia-smi).
222+
fast_path_ok() {
223+
ldconfig || return 1
224+
dkms status >/dev/null 2>&1 || return 1
225+
modinfo -k "${KERNEL_NAME}" nvidia >/dev/null 2>&1 || return 1
226+
}
227+
228+
set +euo pipefail
229+
open_devices="$(lsof /dev/nvidia* 2>/dev/null)"
230+
echo "Open devices: $open_devices"
231+
232+
open_gridd="$(lsof /usr/bin/nvidia-gridd 2>/dev/null)"
233+
echo "Open gridd: $open_gridd"
234+
set -euo pipefail
235+
236+
if [ "${AKSGPU_BUILD_ONLY}" = "1" ]; then
237+
# VHD build time: compile + cache + marker only, no device access. Target the kernel the VHD
238+
# will boot (not the builder's running kernel) so the prebuilt module + marker match at boot.
239+
KERNEL_NAME="$(target_build_kernel)"
240+
echo "aks-gpu: build-only mode (prebuilding kernel module for kernel ${KERNEL_NAME}; builder running $(uname -r))"
241+
echo "aks-gpu: kernels with installed headers (build trees):"; ls -ld /lib/modules/*/build 2>/dev/null || echo " (none found)"
242+
build_and_mark
243+
rm -r /opt/gpu
244+
exit 0
104245
fi
105246

106-
mkdir -p /etc/containerd/config.d
107-
cp /opt/gpu/10-nvidia-runtime.toml /etc/containerd/config.d/10-nvidia-runtime.toml
247+
install_nvidia_container_toolkit
248+
249+
if [ "${AKSGPU_SKIP_KERNEL_BUILD}" = "1" ] && baked_marker_matches && fast_path_ok; then
250+
# Prebuilt module is present and valid for this kernel+driver: skip the ~100s recompile.
251+
echo "aks-gpu: using kernel module prebuilt in the VHD for kernel ${KERNEL_NAME} (recompile skipped)"
252+
else
253+
if [ "${AKSGPU_SKIP_KERNEL_BUILD}" = "1" ]; then
254+
echo "aks-gpu: prebuilt module missing/invalid for ${DRIVER_KIND} ${DRIVER_VERSION} on ${KERNEL_NAME}; building from source"
255+
fi
256+
# No bespoke stale-driver teardown is needed: `nvidia-installer -s` automatically uninstalls
257+
# any previously runfile-installed driver -- including a mismatched prebaked one -- and its
258+
# DKMS registration before installing the new one. build_and_mark then refreshes the marker
259+
# to match what we just built, so subsequent boots take the fast path.
260+
build_and_mark
261+
fi
108262

109-
mkdir -p "$(dirname /lib/udev/rules.d/71-nvidia-dev-char.rules)"
110-
cp /opt/gpu/71-nvidia-char-dev.rules /lib/udev/rules.d/71-nvidia-dev-char.rules
111-
/usr/bin/nvidia-ctk system create-dev-char-symlinks --create-all
263+
device_init
112264

113265
rm -r /opt/gpu

0 commit comments

Comments
 (0)