Enable managed GPU driver install on arm64 (GB200/GB300)#8950
Open
xuexu6666 wants to merge 1 commit into
Open
Conversation
ensureGPUDrivers() returns immediately on arm64, so the managed driver install (--gpu-driver Install / ConfigGPUDriverIfNeeded) is a no-op on the Grace-Blackwell arm64 SKUs. On a GB300 node CSE logs 'Start configuring GPU drivers' and 'End configuring GPU drivers' in the same second and nothing is installed (no /usr/local/nvidia, no DKMS module, nvidia-smi absent), even though aks-gpu-cuda-lts now publishes an arm64 image and CSE resolves it. Scope the early-return so arm64 Ubuntu (the GB SKUs) proceeds through the container-based install path in configGPUDrivers, which pulls the arm64 aks-gpu-cuda-lts image and runs nvidia-modprobe/nvidia-smi. Non-Ubuntu arm64 still skips (no arm64 GPU driver path there). Adds e2e Test_Ubuntu2404Arm64_GB300_ManagedDriver that provisions an 18-node GB300 arm64 rack and validates the R580 open kernel module on-node.
xuexu6666
requested review from
AbelHu,
Devinwong,
SriHarsha001,
awesomenix,
calvin197,
cameronmeissner,
djsly,
ganeshkumarashok,
karenychen,
lilypan26,
mxj220,
pdamianov-dev,
phealy,
r2k1,
runzhen,
sulixu,
surajssd,
timmy-wright,
titilambert and
zachary-bailey
as code owners
July 15, 2026 03:57
Contributor
PR Title Lint Failed ❌Current Title: Your PR title doesn't follow the expected format. Please update your PR title to follow one of these patterns: Conventional Commits Format:
Guidelines:
Examples:
Please update your PR title and the lint check will run again automatically. |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR updates the Linux CSE GPU driver flow to allow arm64 Ubuntu nodes (GB200/GB300) to proceed with the existing container-based managed NVIDIA driver install, instead of short-circuiting on all arm64. It also adds a new GB300-focused E2E scenario to validate the managed install produces a working R580 (580.x) open kernel module on Ubuntu 24.04 arm64.
Changes:
- Adjust
ensureGPUDrivers()to skip arm64 only for non-Ubuntu OSes, enabling managed driver install on Ubuntu arm64. - Add an E2E scenario that provisions an 18-node GB300 rack and validates
nvidia-smiplus open-module driver flavor/version.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| parts/linux/cloud-init/artifacts/cse_config.sh | Refines arm64 early-return gating so Ubuntu arm64 can run the managed/container GPU driver install path. |
| e2e/scenario_gb300_managed_driver_test.go | Adds a GB300 (Ubuntu 24.04 arm64) managed-driver E2E scenario and validator for R580 open kernel module. |
Comment on lines
+1359
to
1361
| if [ "$(isARM64)" -eq 1 ] && [ "$OS" != "$UBUNTU_OS_NAME" ]; then | ||
| return | ||
| fi |
Comment on lines
+1355
to
+1359
| # arm64 GPU nodes are the Grace-Blackwell SKUs (GB200/GB300). They install the | ||
| # driver from the container image via the Ubuntu path in configGPUDrivers - | ||
| # aks-gpu-cuda-lts now publishes an arm64 variant. Other OSes have no arm64 GPU | ||
| # driver install path, so continue skipping arm64 there. | ||
| if [ "$(isARM64)" -eq 1 ] && [ "$OS" != "$UBUNTU_OS_NAME" ]; then |
Comment on lines
+32
to
+33
| func Test_Ubuntu2404Arm64_GB300_ManagedDriver(t *testing.T) { | ||
| RunScenario(t, &Scenario{ |
Comment on lines
+3
to
+12
| import ( | ||
| "context" | ||
| "strings" | ||
| "testing" | ||
|
|
||
| "github.com/Azure/agentbaker/e2e/config" | ||
| "github.com/Azure/agentbaker/pkg/agent/datamodel" | ||
| "github.com/Azure/azure-sdk-for-go/sdk/azcore/to" | ||
| "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v7" | ||
| ) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Enable the managed GPU driver install on arm64 (Grace-Blackwell: GB200 / GB300).
ensureGPUDrivers()returns immediately on arm64, so--gpu-driver Install(
ConfigGPUDriverIfNeeded+EnableNvidia) is a no-op on GB300. This PR scopesthe early-return so arm64 Ubuntu proceeds through the existing container-based
install path in
configGPUDrivers(which pulls the now-publishedaks-gpu-cuda-ltsarm64 image and runs
nvidia-modprobe/nvidia-smi). Non-Ubuntu arm64 keeps skipping(no arm64 GPU driver path there).
ensureGPUDrivers() { - if [ "$(isARM64)" -eq 1 ]; then + if [ "$(isARM64)" -eq 1 ] && [ "$OS" != "$UBUNTU_OS_NAME" ]; then return fiWhy / evidence
Traced on a live GB300 node (
Standard_ND128isr_GB300_v6, kernel6.14.0-1007-azure-nvidia) provisioned via the new e2e:NVIDIA_DRIVER_IMAGE=mcr.microsoft.com/aks/aks-gpu-cuda-lts)and the gate passes (
GPU_NODE=true,SkipGpuDriverInstallunset).ensureGPUDriverslogsStart configuring GPU driversandEnd configuring GPU driversin the same second — it hitsisARM64 → return./usr/local/nvidia, no DKMS module,lsmodempty,nvidia-smiabsent.The
aks-gpu-cuda-ltsarm64 image exists and pulls on the node, so the only blockeris this CSE short-circuit.
Test
Adds
e2e/scenario_gb300_managed_driver_test.go→Test_Ubuntu2404Arm64_GB300_ManagedDriver: provisions an 18-node GB300 arm64 rack(dual-kernel
2404gen2arm64containerdimage → bootslinux-azure-nvidia), drives themanaged install, and validates
nvidia-smi+580.x+ the NVIDIA open kernel module(
/proc/driver/nvidia/version).GB300-specific VMSS wiring discovered while getting the raw VMSS to allocate (all matching
a real AKS GB300 nodepool VMSS):
SinglePlacementGroup=false(Vertical Connect / NVLinkrack), ephemeral OS disk on CacheDisk (1024 GB,
Standard_LRS), and the image'sDiskControllerTypes: SCSI,NVMefeature.Notes / follow-ups (not in this PR)
--gpu-driver Installfor GB300(
VMSizeDoesNotSupportGPU, because Compute SKU metadata lacks theGPUscapability), sothis CSE change stays dormant in production until that gate is lifted. The e2e sets
EnableNvidiadirectly to bypass it.validateGPUDrivers()(the prebaked-driver path, used whenConfigGPUDriverIfNeeded=false) still early-returns on arm64 — left as-is since the GBimages aren't prebaked; worth revisiting if that changes.
localdns.servicefails to start on GB300 and CSE retries it ~100× (separate issue); thee2e disables localdns to keep the run in budget.
make generate: no testdata changes. Itsvalidate-shellstep fails on a pre-existing,unrelated SC3014 warning in
disk_queue.sh(not touched here).Needs an image/VHD build + a real GB300 run to validate end-to-end.