[GB300 test only] Enable managed GPU driver install on arm64 (GB200/GB300)#8950
Open
xuexu6666 wants to merge 4 commits into
Open
[GB300 test only] Enable managed GPU driver install on arm64 (GB200/GB300)#8950xuexu6666 wants to merge 4 commits into
xuexu6666 wants to merge 4 commits into
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
+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" | ||
| ) |
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
Comments suppressed due to low confidence (1)
parts/linux/cloud-init/artifacts/cse_config.sh:1361
- This changes arm64 behavior in ensureGPUDrivers (Ubuntu arm64 no longer returns early), but there’s no ShellSpec coverage to prevent regressions. The existing spec suite for cse_config.sh doesn’t exercise ensureGPUDrivers’ arm64/OS gating, so a future refactor could accidentally reintroduce the no-op on Ubuntu arm64.
# 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
return
fi
Comment on lines
+1987
to
+1999
| # === TEST-ONLY (DO NOT MERGE): Step 1 of managed ComputeDomains on GB300 === | ||
| # GB (arm64) MNNVL SKUs also need the node-local compute-domain kubelet-plugin | ||
| # (device class compute-domain.nvidia.com) for cross-node IMEX. The | ||
| # dra-driver-nvidia-gpu deb ships ONLY gpu-kubelet-plugin, so for validation we | ||
| # side-load the version-matched compute-domain-kubelet-plugin binary (opt-in via | ||
| # COMPUTE_DOMAIN_PLUGIN_URL; empty on normal nodes => no-op) and run it as its own | ||
| # systemd unit. The cluster-side controller comes from the upstream Helm chart for | ||
| # now (Step 2); an AKS-managed controller is Step 3. Productization also moves the | ||
| # binary into the deb and grants the plugin's RBAC. See wiki "Fully Managed Experience". | ||
| if [ "$(isARM64)" -eq 1 ] && [ -n "${COMPUTE_DOMAIN_PLUGIN_URL:-}" ]; then | ||
| if [ ! -x /usr/bin/compute-domain-kubelet-plugin ]; then | ||
| retrycmd_if_failure 10 5 60 curl -fsSL -o /usr/bin/compute-domain-kubelet-plugin "${COMPUTE_DOMAIN_PLUGIN_URL}" || exit $ERR_DRA_DRIVER_START_FAIL | ||
| chmod 0755 /usr/bin/compute-domain-kubelet-plugin |
Comment on lines
+84
to
+86
| osd.ManagedDisk = &armcompute.VirtualMachineScaleSetManagedDiskParameters{ | ||
| StorageAccountType: to.Ptr(armcompute.StorageAccountTypesStandardLRS), | ||
| } |
Comment on lines
+29
to
+31
| // Run with KEEP_VMSS=true to leave the rack up for on-node inspection: | ||
| // | ||
| // KEEP_VMSS=true go test -run '^Test_Ubuntu2404Arm64_GB300_ManagedDriver$' -count=1 -v -timeout 120m |
…GB300 DRA Run compute-domain-kubelet-plugin as its own systemd unit (compute-domain-nvidia-gpu.service) on arm64 (GB) DRA nodes, alongside dra-driver-nvidia-gpu.service. This provides the node half of DRA ComputeDomains (device class compute-domain.nvidia.com) needed for cross-node IMEX / MNNVL. The dra-driver-nvidia-gpu deb ships only gpu-kubelet-plugin, so the version-matched compute-domain-kubelet-plugin binary is fetched from COMPUTE_DOMAIN_PLUGIN_URL when set (empty => no-op, so nodes without ComputeDomains are unaffected). The unit uses the node kubeconfig and mirrors the GPU plugin's flags. Follow-ups: ship the binary in the dra-driver-nvidia-gpu deb, grant the plugin's RBAC (computedomains/computedomaincliques/deployments), and regenerate snapshots.
xuexu6666
force-pushed
the
xuxue/gb300-arm64-managed-gpu-driver
branch
from
July 18, 2026 20:05
dfae58b to
644f50a
Compare
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
Comments suppressed due to low confidence (2)
parts/linux/cloud-init/artifacts/cse_config.sh:1998
- The compute-domain-kubelet-plugin download is (1) not updated when COMPUTE_DOMAIN_PLUGIN_URL changes (because it only downloads when /usr/bin/compute-domain-kubelet-plugin is missing), and (2) pulls an executable from an arbitrary URL without any allowlisting/integrity check. That can leave nodes running a version-mismatched plugin after upgrades and increases supply-chain risk. Consider always re-downloading when COMPUTE_DOMAIN_PLUGIN_URL is set, and at minimum restrict the URL to an allowed domain (e.g., packages.aks.azure.com) before downloading.
if [ "$(isARM64)" -eq 1 ] && [ -n "${COMPUTE_DOMAIN_PLUGIN_URL:-}" ]; then
if [ ! -x /usr/bin/compute-domain-kubelet-plugin ]; then
retrycmd_if_failure 10 5 60 curl -fsSL -o /usr/bin/compute-domain-kubelet-plugin "${COMPUTE_DOMAIN_PLUGIN_URL}" || exit $ERR_DRA_DRIVER_START_FAIL
chmod 0755 /usr/bin/compute-domain-kubelet-plugin
fi
parts/linux/cloud-init/artifacts/cse_config.sh:1359
- This changes arm64 behavior in ensureGPUDrivers() (Ubuntu now proceeds into the install path). There is existing ShellSpec coverage for cse_config.sh, but no spec currently asserts ensureGPUDrivers' arm64/Ubuntu gating, so this is untested and could regress silently. Please add/extend ShellSpec tests to cover: (a) arm64+Ubuntu does not early-return, (b) arm64+non-Ubuntu still early-returns.
# 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
+38
to
+48
| Config: Config{ | ||
| Cluster: ClusterKubenet, | ||
| // The plain 2404gen2arm64containerd definition is the UNIFIED dual-kernel | ||
| // arm64 24.04 image: its release notes bake both linux-azure (6.8 LTS) and | ||
| // linux-azure-nvidia (6.14.0-1007) — the same kernel a real GB300 node runs | ||
| // (verified: gb300-test-rg/gb300 booted this definition on 6.14.0-1007-azure-nvidia). | ||
| // grub's 10_azure_nvidia selects the -azure-nvidia kernel on GB hardware. | ||
| // Unpinned → resolves in-tenant (test gallery, latest branch=main); no prod | ||
| // cross-tenant gallery, so no InvalidAuthenticationTokenTenant. | ||
| VHD: config.VHDUbuntu2404ArmContainerd, | ||
| BootstrapConfigMutator: func(_ *Cluster, nbc *datamodel.NodeBootstrappingConfiguration) { |
On arm64 Ubuntu (Grace-Blackwell), ensureGPUDrivers could dispatch to validateGPUDrivers when CONFIG_GPU_DRIVER_IF_NEEDED=false. validateGPUDrivers returns immediately on arm64 (no arm64 validation path), so no driver gets installed, yet nvidia-modprobe is then started unconditionally -> hard failure (ERR_GPU_DRIVERS_START_FAIL). arm64 only has the container-image install path, so always run configGPUDrivers on arm64. Addresses Copilot review feedback on PR #8950.
Guards the arm64 driver-install gating: arm64 non-Ubuntu skips entirely; arm64 Ubuntu always takes configGPUDrivers (even when CONFIG_GPU_DRIVER_IF_NEEDED=false, the regression the previous commit fixed); amd64 behavior (validateGPUDrivers) unchanged. 4 examples, all passing. Addresses Copilot review feedback on PR #8950.
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
Comments suppressed due to low confidence (3)
parts/linux/cloud-init/artifacts/cse_config.sh:2000
- This code downloads and executes a root-owned binary from
COMPUTE_DOMAIN_PLUGIN_URLwith no integrity/authenticity verification (checksum/signature) and no source allowlist. Given this runs in the node bootstrap path, please ensure the URL is constrained to an approved source and the artifact is verifiably pinned (e.g., checksum in components metadata + verification before install).
if [ "$(isARM64)" -eq 1 ] && [ -n "${COMPUTE_DOMAIN_PLUGIN_URL:-}" ]; then
if [ ! -x /usr/bin/compute-domain-kubelet-plugin ]; then
retrycmd_if_failure 10 5 60 curl -fsSL -o /usr/bin/compute-domain-kubelet-plugin "${COMPUTE_DOMAIN_PLUGIN_URL}" || exit $ERR_DRA_DRIVER_START_FAIL
chmod 0755 /usr/bin/compute-domain-kubelet-plugin
e2e/scenario_gb300_managed_driver_test.go:33
- This scenario provisions an 18-node GB300 rack and will run whenever the full e2e suite is executed (RunScenario always runs all scenarios unless TAGS_TO_RUN/TAGS_TO_SKIP is set). To avoid unexpected long-running/costly allocations in short-mode runs, add a
testing.Short()guard like other integration-style tests in this package.
func Test_Ubuntu2404Arm64_GB300_ManagedDriver(t *testing.T) {
RunScenario(t, &Scenario{
parts/linux/cloud-init/artifacts/cse_config.sh:1996
- This hunk adds new behavior beyond the PR’s stated goal of enabling managed GPU driver install on arm64: it conditionally downloads and runs an additional compute-domain kubelet plugin (plus a new systemd unit). Please either document this explicitly in the PR description (motivation/rollout/expected variables) or split it into a separate PR so the driver-install change can be reviewed/rolled out independently.
# GB (arm64) MNNVL SKUs also need the node-local compute-domain kubelet-plugin
# (device class compute-domain.nvidia.com) for cross-node IMEX. The
# dra-driver-nvidia-gpu deb ships only gpu-kubelet-plugin, so the version-matched
# compute-domain-kubelet-plugin binary is fetched from COMPUTE_DOMAIN_PLUGIN_URL when
# set (empty on nodes without ComputeDomains => no-op) and run as its own systemd
# unit. The cluster-side ComputeDomain controller is deployed separately. Follow-up:
# ship the binary in the deb and grant the plugin's RBAC.
Comment on lines
+1998
to
+2001
| if [ ! -x /usr/bin/compute-domain-kubelet-plugin ]; then | ||
| retrycmd_if_failure 10 5 60 curl -fsSL -o /usr/bin/compute-domain-kubelet-plugin "${COMPUTE_DOMAIN_PLUGIN_URL}" || exit $ERR_DRA_DRIVER_START_FAIL | ||
| chmod 0755 /usr/bin/compute-domain-kubelet-plugin | ||
| fi |
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 3 out of 3 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (1)
parts/linux/cloud-init/artifacts/cse_config.sh:2001
- This introduces a root-run download+exec of /usr/bin/compute-domain-kubelet-plugin from the externally-provided COMPUTE_DOMAIN_PLUGIN_URL with no allowlist and no integrity verification (checksum/signature). That bypasses the repo’s external-dependency tracking (components.json / packages.aks.azure.com) and is a supply-chain risk: a compromised/incorrect URL would yield arbitrary code execution as root during provisioning. Consider shipping this binary in the dra-driver-nvidia-gpu package or referencing it via parts/common/components.json and downloading only from an allowlisted host (e.g., packages.aks.azure.com) with a pinned digest/sha256 check. Also consider wrapping the download step in logs_to_events for better telemetry on failures.
if [ "$(isARM64)" -eq 1 ] && [ -n "${COMPUTE_DOMAIN_PLUGIN_URL:-}" ]; then
if [ ! -x /usr/bin/compute-domain-kubelet-plugin ]; then
retrycmd_if_failure 10 5 60 curl -fsSL -o /usr/bin/compute-domain-kubelet-plugin "${COMPUTE_DOMAIN_PLUGIN_URL}" || exit $ERR_DRA_DRIVER_START_FAIL
chmod 0755 /usr/bin/compute-domain-kubelet-plugin
fi
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.