feat: support Azure Linux ARM64 images on NVIDIA GB200/300

parts/linux/cloud-init/artifacts/10_azure_nvidia

@@ -2,30 +2,63 @@
 
 set -e
 
-. /etc/grub.d/10_linux >/dev/null 2>&1
+BOOT_DIR="${BOOT_DIR:-/boot}"
+GRUB_LINUX_SCRIPT="${GRUB_LINUX_SCRIPT:-/etc/grub.d/10_linux}"
+
+boot_device_id=
+reverse_sorted_list=
+if [ -r "$GRUB_LINUX_SCRIPT" ]; then
+	# shellcheck source=/etc/grub.d/10_linux
+	# shellcheck disable=SC1091
+	. "$GRUB_LINUX_SCRIPT" >/dev/null 2>&1
+fi
 
 NVIDIA=
 OTHER=
 
-for linux in ${reverse_sorted_list}; do
-	case $linux in
-		*-azure-nvidia) : "${NVIDIA:="${linux}"}" ;;
-		*)              : "${OTHER:="${linux}"}"  ;;
-	esac
-done
+# Azure Linux kernel filenames do not identify their track, so use RPM ownership.
+if command -v rpm >/dev/null 2>&1; then
+	for linux in "${BOOT_DIR}"/vmlinuz-*; do
+		[ -f "$linux" ] || continue
+		package=$(rpm -qf --queryformat '%{NAME}' "$linux" 2>/dev/null || true)
+
+		case $package in
+			kernel-hwe) NVIDIA=$(printf '%s\n%s\n' "$NVIDIA" "$linux" | sed '/^$/d' | sort -V | tail -n1) ;;
+			kernel) OTHER=$(printf '%s\n%s\n' "$OTHER" "$linux" | sed '/^$/d' | sort -V | tail -n1) ;;
+		esac
+	done
+fi
+
+# Ubuntu identifies its NVIDIA kernel by filename and exposes a sorted list from 10_linux.
+if [ -z "$NVIDIA" ] || [ -z "$OTHER" ]; then
+	NVIDIA=
+	OTHER=
+	# shellcheck disable=SC2086 # reverse_sorted_list is a space-separated path list from 10_linux.
+	for linux in ${reverse_sorted_list}; do
+		case $linux in
+			*-azure-nvidia) : "${NVIDIA:="${linux}"}" ;;
+			*) : "${OTHER:="${linux}"}" ;;
+		esac
+	done
+fi
 
-if [ -z "${NVIDIA}" ] || [ -z "${OTHER}" ]; then
+if [ -z "$NVIDIA" ] || [ -z "$OTHER" ]; then
 	echo "Only one image type (NVIDIA or non-NVIDIA) found" >&2
 	exit 0
 fi
 
-echo "Default NVIDIA image: ${NVIDIA}" >&2
 NVIDIA=${NVIDIA##*/}
-NVIDIA=$(echo "${NVIDIA}" | sed -e "s,^[^0-9]*-,,g")
+NVIDIA=${NVIDIA#vmlinuz-}
 
-echo "Default non-NVIDIA image: ${OTHER}" >&2
 OTHER=${OTHER##*/}
-OTHER=$(echo "${OTHER}" | sed -e "s,^[^0-9]*-,,g")
+OTHER=${OTHER#vmlinuz-}
+
+if [ -z "$boot_device_id" ]; then
+	echo "Could not determine boot_device_id from $GRUB_LINUX_SCRIPT" >&2
+	exit 0
+fi
+
+echo "Default NVIDIA image: $NVIDIA, non-NVIDIA image: $OTHER, boot_device_id: $boot_device_id" >&2
 
 cat << EOF
 insmod smbios

parts/linux/cloud-init/artifacts/51-azure-nvidia.cfg

@@ -1,5 +1,5 @@
-# 10_azure_nvidia changes the default menu entry based on a flat menu.
+# 10_azure_nvidia selects the hardware-appropriate entry from a flat menu.
 GRUB_DISABLE_SUBMENU=true
 
-# 10_azure_nvidia appends additional arguments for azure-nvidia kernels.
+# Append NVIDIA Grace arguments only when the selector sets nvidia_args.
 GRUB_CMDLINE_LINUX="$GRUB_CMDLINE_LINUX \$nvidia_args"

parts/linux/cloud-init/artifacts/cse_config.sh

@@ -1553,10 +1553,6 @@ configGPUDrivers() {
 }
 
 validateGPUDrivers() {
-    if [ "$(isARM64)" -eq 1 ]; then
-        return
-    fi
-
     retrycmd_if_failure 24 5 25 nvidia-modprobe -u -c0 && echo "gpu driver loaded" || configGPUDrivers || exit $ERR_GPU_DRIVERS_START_FAIL
 
     if which nvidia-smi; then
@@ -1635,7 +1631,9 @@ cleanUpGridNodeCudaPrebake() {
 
 ensureGPUDrivers() {
     if [ "$(isARM64)" -eq 1 ]; then
-        return
+        if [ "$OS" != "$AZURELINUX_OS_NAME" ] || [ "$OS_VERSION" != "3.0" ] || [ -n "$OS_VARIANT" ] || [ "${ENABLE_FIPS,,}" = "true" ]; then
+            return
+        fi
     fi
 
     # Tear down a mismatched cuda-lts VHD prebake before a GRID node installs its own driver, or the

parts/linux/cloud-init/artifacts/mariner/cse_install_mariner.sh

@@ -116,6 +116,7 @@ getAzureLinuxNvidiaDriverVersionFromPackage() {
     local version_with_epoch
 
     version_with_epoch=${package#"${package_prefix}"}
+    version_with_epoch=${version_with_epoch#hwe-}
     version_with_epoch=${version_with_epoch%%-*}
     echo "${version_with_epoch#*:}"
 }
@@ -127,7 +128,7 @@ getAzureLinuxNvidiaDriverReleaseNotes() {
     local cuda_open_package
     local cuda_package
     local grid_package
-    cuda_open_package=$(getLatestAzureLinuxNvidiaDriverPackageForKernel "cuda-open*" "^cuda-open-[0-9]" "${kernel_version}")
+    cuda_open_package=$(getLatestAzureLinuxNvidiaDriverPackageForKernel "cuda-open*" "^cuda-open(-hwe)?-[0-9]" "${kernel_version}")
     cuda_package=$(getLatestAzureLinuxNvidiaDriverPackageForKernel "cuda" "^cuda-[0-9]" "${kernel_version}")
     grid_package=$(getLatestAzureLinuxNvidiaDriverPackageForKernel "nvidia-vgpu-guest-driver*" "^nvidia-vgpu-guest-driver-[0-9]" "${kernel_version}")
 
@@ -155,7 +156,7 @@ downloadGPUDrivers() {
     # cuda-%{nvidia gpu driver version}_%{kernel source version}.%{kernel release version}.{mariner rpm postfix}
     #
     # 2. NVIDIA OpenRM driver:
-    # cuda-open-%{nvidia gpu driver version}_%{kernel source version}.%{kernel release version}.{mariner rpm postfix}
+    # cuda-open[-hwe]-%{nvidia gpu driver version}_%{kernel source version}.%{kernel release version}.{mariner rpm postfix}
     #
     # 3. NVIDIA GRID (vGPU guest) driver for converged GPU sizes:
     # nvidia-vgpu-guest-driver-%{version}_%{kernel version}.{mariner rpm postfix}
@@ -193,10 +194,10 @@ downloadGPUDrivers() {
         exit $ERR_MISSING_CUDA_PACKAGE
     elif [ "$driver_ret" -eq 0 ]; then
         echo "VM SKU ${VM_SKU} uses NVIDIA OpenRM driver (cuda-open)"
-        CUDA_PACKAGE=$(dnf repoquery -y --available "cuda-open*" | grep -E "^cuda-open-[0-9]+.*_${KERNEL_VERSION}" | sort -V | tail -n 1)
+      CUDA_PACKAGE=$(getLatestAzureLinuxNvidiaDriverPackageForKernel "cuda-open*" "^cuda-open(-hwe)?-[0-9]" "$KERNEL_VERSION")
     else
         echo "VM SKU ${VM_SKU} uses NVIDIA proprietary driver (cuda)"
-        CUDA_PACKAGE=$(dnf repoquery -y --available "cuda-[0-9]*" | grep -E "^cuda-[0-9]+.*_${KERNEL_VERSION}" | sort -V | tail -n 1)
+      CUDA_PACKAGE=$(getLatestAzureLinuxNvidiaDriverPackageForKernel "cuda" "^cuda-[0-9]" "$KERNEL_VERSION")
     fi
 
     if [ -z "$CUDA_PACKAGE" ]; then

pkg/agent/baker_test.go

@@ -1081,6 +1081,10 @@ var _ = Describe("GetGPUDriverType", func() {
 	It("should use cuda-lts with nc v3", func() {
 		Expect(GetGPUDriverType("standard_nc6_v3")).To(Equal("cuda-lts"))
 	})
+	It("should use cuda-lts with GB", func() {
+		Expect(GetGPUDriverType("Standard_ND128isr_NDR_GB200_v6")).To(Equal("cuda-lts"))
+		Expect(GetGPUDriverType("Standard_ND128isr_GB300_v6")).To(Equal("cuda-lts"))
+	})
 	It("should keep cuda (legacy R470) with nc v1 (K80)", func() {
 		Expect(GetGPUDriverType("standard_nc6")).To(Equal("cuda"))
 	})
@@ -1119,6 +1123,13 @@ var _ = Describe("GetAKSGPUImageSHA", func() {
 	})
 })
 
+var _ = Describe("GPUNeedsFabricManager", func() {
+	It("should not use Fabric Manager with GB", func() {
+		Expect(GPUNeedsFabricManager("Standard_ND128isr_NDR_GB200_v6")).To(BeFalse())
+		Expect(GPUNeedsFabricManager("Standard_ND128isr_GB300_v6")).To(BeFalse())
+	})
+})
+
 var _ = Describe("getLinuxNodeCSECommand", func() {
 	var (
 		templateGenerator *TemplateGenerator

spec/parts/linux/cloud-init/artifacts/azure_nvidia_spec.sh

@@ -0,0 +1,73 @@
+#!/bin/bash
+
+Describe '10_azure_nvidia'
+    SCRIPT='./parts/linux/cloud-init/artifacts/10_azure_nvidia'
+
+    Mock rpm
+        for arg do
+            kernel_path=$arg
+        done
+        case "${kernel_path##*/}" in
+            vmlinuz-1.0.*) printf 'kernel' ;;
+            vmlinuz-2.0.*) printf 'kernel-hwe' ;;
+            *) exit 1 ;;
+        esac
+    End
+
+    setup() {
+        TEST_DIR=$(mktemp -d)
+        BOOT_DIR="${TEST_DIR}/boot"
+        GRUB_LINUX_SCRIPT="${TEST_DIR}/10_linux"
+        mkdir -p "$BOOT_DIR"
+        for kernel in 1.0.9-test 1.0.10-test 2.0.9-test 2.0.10-test 9.0.0-unowned; do
+            touch "${BOOT_DIR}/vmlinuz-${kernel}"
+        done
+        printf '%s\n' 'boot_device_id=test-device' > "$GRUB_LINUX_SCRIPT"
+        export BOOT_DIR GRUB_LINUX_SCRIPT
+    }
+
+    cleanup() {
+        rm -rf "$TEST_DIR"
+    }
+
+    BeforeEach 'setup'
+    AfterEach 'cleanup'
+
+    It 'selects the newest RPM-owned kernel from each Azure Linux track'
+        When run script "$SCRIPT"
+        The status should be success
+        The output should include 'set default="gnulinux-2.0.10-test-advanced-test-device"'
+        The output should include 'set default="gnulinux-1.0.10-test-advanced-test-device"'
+        The output should not include '9.0.0-unowned'
+        The stderr should include 'Default NVIDIA image: 2.0.10-test, non-NVIDIA image: 1.0.10-test'
+    End
+
+    It 'preserves Ubuntu kernel selection by filename'
+        rm -f "${BOOT_DIR}"/vmlinuz-*
+        cat > "$GRUB_LINUX_SCRIPT" <<'EOF'
+boot_device_id=test-device
+reverse_sorted_list='/boot/vmlinuz-3.0.10-azure-nvidia /boot/vmlinuz-3.0.10-azure'
+EOF
+        When run script "$SCRIPT"
+        The status should be success
+        The output should include 'set default="gnulinux-3.0.10-azure-nvidia-advanced-test-device"'
+        The output should include 'set default="gnulinux-3.0.10-azure-advanced-test-device"'
+        The stderr should include 'Default NVIDIA image: 3.0.10-azure-nvidia, non-NVIDIA image: 3.0.10-azure'
+    End
+
+    It 'emits no override when either kernel track is missing'
+        rm -f "${BOOT_DIR}"/vmlinuz-2.0.*
+        When run script "$SCRIPT"
+        The status should be success
+        The output should eq ''
+        The stderr should include 'Only one image type (NVIDIA or non-NVIDIA) found'
+    End
+
+    It 'emits no override when 10_linux cannot derive a boot device ID'
+        printf '%s\n' 'boot_device_id=' > "$GRUB_LINUX_SCRIPT"
+        When run script "$SCRIPT"
+        The status should be success
+        The output should eq ''
+        The stderr should include "Could not determine boot_device_id from $GRUB_LINUX_SCRIPT"
+    End
+End

spec/parts/linux/cloud-init/artifacts/cse_config_spec.sh

@@ -2909,6 +2909,55 @@ OVERRIDE_EOF
         End
     End
 
+    Describe 'ARM64 GPU driver dispatch'
+        isARM64() { echo 1; }
+        logs_to_events() {
+            shift
+            "$@"
+        }
+        configGPUDrivers() { echo "configGPUDrivers called"; }
+        validateGPUDrivers() { echo "validateGPUDrivers called"; }
+
+        Parameters
+            "$AZURELINUX_OS_NAME" "3.0" "" false true  "configGPUDrivers called"
+            "$AZURELINUX_OS_NAME" "3.0" "" false false "validateGPUDrivers called"
+            "$UBUNTU_OS_NAME"     "24.04" "" false true ""
+            "$AZURELINUX_OS_NAME" "2.0" "" false true ""
+            "$AZURELINUX_OS_NAME" "3.0" "$AZURELINUX_OSGUARD_OS_VARIANT" false true ""
+            "$AZURELINUX_OS_NAME" "3.0" "" true true ""
+        End
+
+        It "dispatches ARM64 driver setup for OS=$1 version=$2 variant=$3 fips=$4 install=$5"
+            OS=$1
+            OS_VERSION=$2
+            OS_VARIANT=$3
+            ENABLE_FIPS=$4
+            CONFIG_GPU_DRIVER_IF_NEEDED=$5
+
+            When call ensureGPUDrivers
+
+            The output should equal "$6"
+        End
+    End
+
+    Describe 'ARM64 GPU driver validation'
+        isARM64() { echo 1; }
+        retrycmd_if_failure() {
+            echo "retrycmd_if_failure $*" >&2
+            return 0
+        }
+        which() { return 0; }
+
+        It 'checks modprobe and nvidia-smi'
+            When call validateGPUDrivers
+
+            The status should be success
+            The output should include "gpu driver loaded"
+            The stderr should include "retrycmd_if_failure 24 5 25 nvidia-modprobe -u -c0"
+            The stderr should include "retrycmd_if_failure 24 5 30 nvidia-smi"
+        End
+    End
+
     Describe 'configGPUDrivers'
         # Assert the per-step CSE timing event names emitted via logs_to_events,
         # without running the real (hardware/daemon) driver steps. logs_to_events

spec/parts/linux/cloud-init/artifacts/cse_install_mariner_spec.sh

@@ -27,6 +27,7 @@ Describe 'cse_install_mariner.sh'
     BeforeAll 'setup'
     Include "./parts/linux/cloud-init/artifacts/cse_install.sh"
     Include "./parts/linux/cloud-init/artifacts/mariner/cse_install_mariner.sh"
+
     Describe 'installDeps'
         It 'installs the required packages with installDeps for Mariner 2.0'
             OS_VERSION="2.0"
@@ -340,6 +341,23 @@ Describe 'cse_install_mariner.sh'
             The variable GRID_CALLED should not equal "true"
         End
 
+        It 'selects the newest HWE OpenRM package for GB200'
+            NVIDIA_GPU_DRIVER_TYPE="cuda-lts"
+            MOCK_VM_SKU="Standard_ND128isr_NDR_GB200_v6"
+            MOCK_OPEN_RET=0
+            uname() { echo "9.9.2-1.azl3"; }
+            dnf() {
+                echo "cuda-open-hwe-999.1.2-1_9.9.2.1.azl3.aarch64"
+                echo "cuda-open-hwe-999.1.2-2_9.9.2.1.azl3.aarch64"
+                echo "cuda-open-hwe-999.1.2-2_9.8.1.1.azl3.aarch64"
+            }
+
+            When call downloadGPUDrivers
+
+            The status should be success
+            The output should include "dnf install 30 1 600 cuda-open-hwe-999.1.2-2_9.9.2.1.azl3.aarch64"
+        End
+
         It 'selects proprietary cuda path for T4 when NVIDIA_GPU_DRIVER_TYPE is cuda'
             NVIDIA_GPU_DRIVER_TYPE="cuda"
             MOCK_VM_SKU="Standard_NC4as_T4_v3"
@@ -389,7 +407,7 @@ Describe 'cse_install_mariner.sh'
             dnf() {
                 case "$4" in
                     "cuda-open*")
-                        echo "cuda-open-580.126.09-2_6.6.121.1.1.azl3.x86_64"
+                        echo "cuda-open-hwe-999.1.2-2_6.6.121.1.1.azl3.x86_64"
                         ;;
                     "cuda")
                         echo "cuda-570.195.03-1_6.6.121.1.1.azl3.x86_64"
@@ -404,7 +422,7 @@ Describe 'cse_install_mariner.sh'
 
             The status should be success
             The output should include "NVIDIA GPU driver versions available at VHD build time for supported Azure Linux GPU VM sizes:"
-            The output should include "  - nvidia-cuda-open-driver version 580.126.09"
+            The output should include "  - nvidia-cuda-open-driver version 999.1.2"
             The output should include "  - nvidia-cuda-driver version 570.195.03"
             The output should include "  - nvidia-grid-driver version 570.211.01"
             The output should include "build-time snapshot only"

spec/vhdbuilder/packer/test/linux_vhd_content_test_helpers_spec.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-# shellcheck disable=SC2329
+# shellcheck disable=SC2016,SC2329
 
 # ShellSpec tests for parseAutologinSessions helper function
 
@@ -233,3 +233,19 @@ Describe 'Inspektor Gadget version helper functions'
     End
   End
 End
+
+Describe 'AgentBaker source checkout'
+  SCRIPT='./vhdbuilder/packer/test/linux-vhd-content-test.sh'
+
+  It 'downloads no-git archives from the configured repository'
+    When run grep -Fq 'AGENTBAKER_ARCHIVE_URL="${AGENTBAKER_REPOSITORY_URL%.git}/archive/${GIT_COMMIT_HASH}.tar.gz"' "$SCRIPT"
+    The status should be success
+  End
+
+  It 'fetches the source ref and checks out the captured commit'
+    When run sed -n '/git fetch --quiet origin/,/git checkout --quiet/p' "$SCRIPT"
+    The output should include 'git fetch --quiet origin "$GIT_BRANCH"'
+    The output should include 'git checkout --quiet "$GIT_COMMIT_HASH"'
+    The status should be success
+  End
+End

vhdbuilder/packer/packer_source.sh

@@ -611,6 +611,16 @@ copyPackerFiles() {
     # Mariner/AzureLinux uses system-auth and system-password instead of common-auth and common-password.
     cpAndMode $PAM_D_SYSTEM_AUTH_SRC $PAM_D_SYSTEM_AUTH_DEST 644
     cpAndMode $PAM_D_SYSTEM_PASSWORD_SRC $PAM_D_SYSTEM_PASSWORD_DEST 644
+
+    if [ "$OS" = "$AZURELINUX_OS_NAME" ] && [ "$OS_VERSION" = "3.0" ] && [ "$CPU_ARCH" = "arm64" ] && [ -z "$OS_VARIANT" ] && [ "${ENABLE_FIPS,,}" != "true" ]; then
+      GRUB_AZ_NV_SCRIPT_SRC=/home/packer/10_azure_nvidia
+      GRUB_AZ_NV_SCRIPT_DEST=/etc/grub.d/10_azure_nvidia
+      cpAndMode $GRUB_AZ_NV_SCRIPT_SRC $GRUB_AZ_NV_SCRIPT_DEST 755
+
+      GRUB_AZ_NV_ENV_SRC=/home/packer/51-azure-nvidia.cfg
+      GRUB_AZ_NV_ENV_DEST=/etc/default/grub.d/51-azure-nvidia.cfg
+      cpAndMode $GRUB_AZ_NV_ENV_SRC $GRUB_AZ_NV_ENV_DEST 644
+    fi
   elif isACL "$OS" "$OS_VARIANT"; then
     # ACL cannot share the isMarinerOrAzureLinux block because:
     # - containerd.service: ACL provides containerd via sysext.