From d4a3625f7b7325eb10c4dfdae8448c63155181fd Mon Sep 17 00:00:00 2001 From: Sylvain Boily <4981802+djsly@users.noreply.github.com> Date: Mon, 13 Jul 2026 15:58:09 -0400 Subject: [PATCH] fix: clean up mismatched prebaked GPU drivers Ensure Ubuntu managed GPU driver installation removes a prebaked GPU driver when its marker indicates a different driver family than the requested node driver. This prevents GRID scenarios from keeping the prebaked CUDA driver alongside a dynamically installed GRID driver. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 4110a607-c4e2-49f0-b023-5b5f73cfb153 --- .../linux/cloud-init/artifacts/cse_config.sh | 66 +++++++++++++++---- .../artifacts/ubuntu/cse_install_ubuntu.sh | 11 ++-- .../cloud-init/artifacts/cse_config_spec.sh | 58 ++++++++++++++++ 3 files changed, 117 insertions(+), 18 deletions(-) diff --git a/parts/linux/cloud-init/artifacts/cse_config.sh b/parts/linux/cloud-init/artifacts/cse_config.sh index a15bb501de6..c9d8b0d5940 100755 --- a/parts/linux/cloud-init/artifacts/cse_config.sh +++ b/parts/linux/cloud-init/artifacts/cse_config.sh @@ -1295,23 +1295,35 @@ validateGPUDrivers() { fi } -# logGPUDriverPrebakeReadiness emits a stage-1 observability signal on a managed GPU node: whether -# the aks-gpu prebake marker is present and matches this node's driver kind -- i.e. whether stage-2 -# (skip-build) would take the fast path. Lets the rollout confirm managed CUDA GPU nodes are ready -# before enabling consume. Observability only; no behavior change. +getGPUDriverKindFromType() { + local driver_type="${1:-}" + + case "${driver_type}" in + cuda*) echo "cuda" ;; + grid*) echo "grid" ;; + *) echo "${driver_type}" ;; + esac +} + +getPrebakedGPUDriverKind() { + local marker="${1:-${GPU_DKMS_MARKER_FILE:-/opt/azure/aks-gpu/dkms-marker}}" + + if [ ! -f "${marker}" ]; then + return 1 + fi + + sed -n 's/^driver_kind=//p' "${marker}" | head -n1 +} + +# logGPUDriverPrebakeReadiness emits whether the aks-gpu prebake marker is present and matches this +# managed GPU node's driver kind -- i.e. whether the aks-gpu installer can safely use skip-build. logGPUDriverPrebakeReadiness() { local marker="${GPU_DKMS_MARKER_FILE:-/opt/azure/aks-gpu/dkms-marker}" local marker_present=false driver_kind_match=false m_kind node_kind - # Map the AgentBaker driver-type to the aks-gpu marker's driver_kind (the container's DRIVER_KIND - # build arg): image variants "cuda-lts" and "grid-v20" bake markers as "cuda"/"grid" respectively. - case "${NVIDIA_GPU_DRIVER_TYPE}" in - cuda*) node_kind=cuda ;; - grid*) node_kind=grid ;; - *) node_kind="${NVIDIA_GPU_DRIVER_TYPE}" ;; - esac + node_kind="$(getGPUDriverKindFromType "${NVIDIA_GPU_DRIVER_TYPE:-}")" if [ -f "${marker}" ]; then marker_present=true - m_kind="$(sed -n 's/^driver_kind=//p' "${marker}" | head -n1)" + m_kind="$(getPrebakedGPUDriverKind "${marker}")" # require both sides non-empty so a marker missing driver_kind= (or an unset # NVIDIA_GPU_DRIVER_TYPE) does not falsely report a match (empty = empty). if [ -n "${m_kind}" ] && [ -n "${node_kind}" ] && [ "${m_kind}" = "${node_kind}" ]; then @@ -1321,11 +1333,41 @@ logGPUDriverPrebakeReadiness() { echo "AKS_GPU_PREBAKE event=managed_gpu driver_type=${NVIDIA_GPU_DRIVER_TYPE:-} marker_present=${marker_present} driver_kind_match=${driver_kind_match}" } +cleanUpMismatchedPrebakedGPUDriver() { + local marker="${GPU_DKMS_MARKER_FILE:-/opt/azure/aks-gpu/dkms-marker}" + local marker_kind node_kind reason + + if [ "$OS" != "$UBUNTU_OS_NAME" ] || [ ! -f "${marker}" ]; then + return 0 + fi + + node_kind="$(getGPUDriverKindFromType "${NVIDIA_GPU_DRIVER_TYPE:-}")" + marker_kind="$(getPrebakedGPUDriverKind "${marker}")" + + if [ -n "${marker_kind}" ] && [ -n "${node_kind}" ]; then + if [ "${marker_kind}" = "${node_kind}" ]; then + return 0 + fi + reason="driver_kind_mismatch" + elif [ -z "${marker_kind}" ] && [ "${node_kind}" = "grid" ]; then + reason="legacy_cuda_prebake_marker_on_grid_node" + else + return 0 + fi + + echo "AKS_GPU_PREBAKE event=managed_gpu_cleanup driver_type=${NVIDIA_GPU_DRIVER_TYPE:-} marker_driver_kind=${marker_kind:-unknown} node_driver_kind=${node_kind:-unknown} reason=${reason}" + cleanUpPrebakedGPUDriver +} + ensureGPUDrivers() { if [ "$(isARM64)" -eq 1 ]; then return fi + if [ "$OS" = "$UBUNTU_OS_NAME" ]; then + logs_to_events "AKS.CSE.ensureGPUDrivers.cleanUpMismatchedPrebakedGPUDriver" cleanUpMismatchedPrebakedGPUDriver || exit $ERR_GPU_DRIVERS_START_FAIL + fi + if [ "${CONFIG_GPU_DRIVER_IF_NEEDED}" = true ]; then logs_to_events "AKS.CSE.ensureGPUDrivers.configGPUDrivers" configGPUDrivers else diff --git a/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh b/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh index 49025f22dfe..a7cab1e35f6 100755 --- a/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh +++ b/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh @@ -231,12 +231,11 @@ removeNvidiaRepos() { fi } -# cleanUpPrebakedGPUDriver removes a CUDA driver pre-baked into the shared VHD on any node that does -# NOT install the AKS-managed driver -- the cleanUpGPUDrivers path (GPU_NODE != true OR -# skip_nvidia_driver_install=true): non-GPU VMs, and GPU VMs opted out via --gpu-driver None or the -# skip toggle/tag. There the driver is dead weight (wasted disk; nvidia.ko rebuilt on every kernel -# patch) and, on an opted-out GPU node, unused attack surface. The module is never loaded on these -# nodes (ensureGPUDrivers doesn't run), so deregistration is safe. No-op unless the marker exists. +# cleanUpPrebakedGPUDriver removes a driver pre-baked into the shared VHD on nodes that should not +# keep it: non-GPU VMs, GPU VMs opted out via --gpu-driver None or skip toggles, and managed GPU +# nodes whose requested driver kind does not match the prebaked driver marker. There the driver is +# either dead weight (wasted disk; nvidia.ko rebuilt on every kernel patch), unused attack surface, or +# the wrong driver family for the node SKU. No-op unless the marker exists. cleanUpPrebakedGPUDriver() { local marker="${GPU_DKMS_MARKER_FILE:-/opt/azure/aks-gpu/dkms-marker}" if [ ! -f "${marker}" ]; then diff --git a/spec/parts/linux/cloud-init/artifacts/cse_config_spec.sh b/spec/parts/linux/cloud-init/artifacts/cse_config_spec.sh index d6d4596b063..45fc2660a50 100755 --- a/spec/parts/linux/cloud-init/artifacts/cse_config_spec.sh +++ b/spec/parts/linux/cloud-init/artifacts/cse_config_spec.sh @@ -74,6 +74,64 @@ Describe 'cse_config.sh' End End + Describe 'cleanUpMismatchedPrebakedGPUDriver' + It 'is a no-op when the prebake marker is absent' + GPU_DKMS_MARKER_FILE="$(mktemp)"; rm -f "${GPU_DKMS_MARKER_FILE}" + OS="$UBUNTU_OS_NAME" + NVIDIA_GPU_DRIVER_TYPE="grid" + # shellcheck disable=SC2329 # cleanUpMismatchedPrebakedGPUDriver invokes this stub dynamically. + cleanUpPrebakedGPUDriver() { echo "unexpected cleanup"; } + When call cleanUpMismatchedPrebakedGPUDriver + The output should equal "" + End + + It 'removes a CUDA prebake before installing a GRID driver' + marker="$(mktemp)" + printf 'driver_kind=cuda\n' > "$marker" + GPU_DKMS_MARKER_FILE="$marker" + OS="$UBUNTU_OS_NAME" + NVIDIA_GPU_DRIVER_TYPE="grid" + # shellcheck disable=SC2329 # cleanUpMismatchedPrebakedGPUDriver invokes this stub dynamically. + cleanUpPrebakedGPUDriver() { echo "mock cleanup"; } + When call cleanUpMismatchedPrebakedGPUDriver + The output should include "AKS_GPU_PREBAKE event=managed_gpu_cleanup" + The output should include "marker_driver_kind=cuda" + The output should include "node_driver_kind=grid" + The output should include "reason=driver_kind_mismatch" + The output should include "mock cleanup" + rm -f "$marker" + End + + It 'treats a legacy marker without driver_kind as CUDA prebake on a GRID node' + marker="$(mktemp)" + printf 'kernel=5.15.0-1114-azure\n' > "$marker" + GPU_DKMS_MARKER_FILE="$marker" + OS="$UBUNTU_OS_NAME" + NVIDIA_GPU_DRIVER_TYPE="grid-v20" + # shellcheck disable=SC2329 # cleanUpMismatchedPrebakedGPUDriver invokes this stub dynamically. + cleanUpPrebakedGPUDriver() { echo "mock cleanup"; } + When call cleanUpMismatchedPrebakedGPUDriver + The output should include "marker_driver_kind=unknown" + The output should include "node_driver_kind=grid" + The output should include "reason=legacy_cuda_prebake_marker_on_grid_node" + The output should include "mock cleanup" + rm -f "$marker" + End + + It 'keeps a matching CUDA prebake for a CUDA node' + marker="$(mktemp)" + printf 'driver_kind=cuda\n' > "$marker" + GPU_DKMS_MARKER_FILE="$marker" + OS="$UBUNTU_OS_NAME" + NVIDIA_GPU_DRIVER_TYPE="cuda-lts" + # shellcheck disable=SC2329 # cleanUpMismatchedPrebakedGPUDriver invokes this stub dynamically. + cleanUpPrebakedGPUDriver() { echo "unexpected cleanup"; } + When call cleanUpMismatchedPrebakedGPUDriver + The output should equal "" + rm -f "$marker" + End + End + Describe 'configureAzureJson' AZURE_JSON_PATH="azure.json" AKS_CUSTOM_CLOUD_JSON_PATH="customcloud.json"