Skip to content

fix: clean up mismatched prebaked GPU drivers #8915

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鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
djsly wants to merge 1 commit into from
Closed

fix: clean up mismatched prebaked GPU drivers #8915

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
66 changes: 54 additions & 12 deletions parts/linux/cloud-init/artifacts/cse_config.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down
11 changes: 5 additions & 6 deletions parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Comment on lines +234 to +238
cleanUpPrebakedGPUDriver() {
local marker="${GPU_DKMS_MARKER_FILE:-/opt/azure/aks-gpu/dkms-marker}"
if [ ! -f "${marker}" ]; then
Expand Down
58 changes: 58 additions & 0 deletions spec/parts/linux/cloud-init/artifacts/cse_config_spec.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
Loading