Skip to content

fix(gpu): keep CUDA prebakes unregistered until node installation - #9568

Draft
sulixu wants to merge 1 commit into
mainfrom
sulixu/gpu-prebake-no-dkms
Draft

sulixu wants to merge 1 commit into
mainfrom
sulixu/gpu-prebake-no-dkms

Conversation

@sulixu

@sulixu sulixu commented Sep 17, 2026

Copy link
Copy Markdown
Contributor

What this PR does / why we need it:

Draft companion to Azure/aks-gpu#191 and an alternative to #9545.

Some non-GPU nodes did not complete CUDA prebake cleanup. Later kernel security
updates could discover the remaining NVIDIA DKMS registration and try to rebuild
the driver. Shared-image safety must not depend on node-time cleanup succeeding.

The companion aks-gpu change compiles and installs the VHD's NVIDIA modules with
--no-dkms. This PR checks that contract and prepares managed CUDA nodes to use
those images:

  • Reject a CUDA prebake if dkms status reports NVIDIA state or
    /var/lib/dkms/nvidia exists, including a dangling link. Check again in the final
    VHD content test, and require the marker and compiled module.
  • For a marked Ubuntu CUDA node, require a DKMS installed record for the on-disk
    driver version, running kernel, and architecture. A loadable module or an
    added/built record is not enough.
  • Use the existing normal container install action if installation is requested
    or the required DKMS state is missing, including validation-only/PIS paths.
    The normal NVIDIA installer still uses --dkms. Check the installed state
    after installation or validation, and fail provisioning if it is incomplete.
  • Keep the existing CPU/opt-out caller gates and GRID cleanup/install/validation
    paths. Do not add a duplicate GPU-node guard inside ensureGPUDrivers.
  • Add focused ShellSpec coverage, real-DKMS fixture tests, and release notes.

Unlike #9545, this design does not park or restore DKMS trees. It does not use
dkms add or enable install-skip-build.

Trade-offs and draft blockers:

  • This is not ready to ship alone. Publish the companion aks-gpu build image
    and update the builder image reference before using these build checks. No
    image tags are changed here; an old container that registers NVIDIA will fail
    the new check when NVIDIA_CUDA_PREBAKE is enabled.
  • Roll out the compatible CSE first, then the new VHDs, and rebuild PIS caches.
    Existing registered images and deployed nodes are not repaired by this change.
  • A fresh unregistered CUDA image requires full installation, even in
    validation-only mode. This can lose the startup-time saving from validating a
    loadable prebake. No startup-time improvement or measurement is claimed.
  • The existing GRID cleanup limit remains: cleanup can report incomplete work
    but return success. If the remaining driver passes validation, a validation-only
    GRID node can finish without invoking the GRID installer. A mocked control-flow
    check demonstrated this path; it has not been verified on a real GRID node.
  • Real Ubuntu 22.04/24.04 VHD builds and GPU tests remain required: failed CPU
    cleanup plus kernel patching; CUDA install/validation/PIS; GRID and opt-out;
    kernel/driver mismatches; and GPU kernel update plus reboot.

Validation:

  • Host Bash 3.2: 55 focused ShellSpec examples passed.
  • Linux Bash 5 with DKMS 2.8.7: 100 ShellSpec examples passed, including the
    existing GPU suite and 8 real-DKMS fixture tests.
  • Linux DKMS 3.0.11: all 8 fixture tests passed.
  • ShellCheck passed for the three changed production scripts with the
    repository's standard exclusions. git diff --check passed.
  • go test ./pkg/agent -count=1 passed.

The DKMS tests use fixture source/module trees. They test status parsing and
kernel-update discovery, not a real NVIDIA build or GPU installation.

Which issue(s) this PR fixes:

Addresses the failed-cleanup / unwanted NVIDIA DKMS rebuild scenario discussed
in #9545. No separate issue is linked. Requires Azure/aks-gpu#191.

@github-actions

Copy link
Copy Markdown
Contributor

Code Coverage

Package Line Rate Health
github.com/Azure/agentbaker/aks-node-controller 72%
github.com/Azure/agentbaker/aks-node-controller/common 100%
github.com/Azure/agentbaker/aks-node-controller/helpers 71%
github.com/Azure/agentbaker/aks-node-controller/parser 89%
github.com/Azure/agentbaker/aks-node-controller/pkg/gpu 100%
github.com/Azure/agentbaker/aks-node-controller/pkg/nodeconfigutils 66%
github.com/Azure/agentbaker/aks-node-controller/utils 0%
github.com/Azure/agentbaker/apiserver 25%
github.com/Azure/agentbaker/cmd 0%
github.com/Azure/agentbaker/cmd/starter 0%
github.com/Azure/agentbaker/fuzz/api 0%
github.com/Azure/agentbaker/hotfix/render-nodecustomdata 0%
github.com/Azure/agentbaker/pkg/agent 76%
github.com/Azure/agentbaker/pkg/agent/datamodel 76%
github.com/Azure/agentbaker/pkg/agent/toggles 0%
github.com/Azure/agentbaker/pkg/vhdbuilder/datamodel 88%
Summary 74% (6311 / 8577)

@github-actions

Copy link
Copy Markdown
Contributor

Windows Unit Test Results

  3 files   14 suites   50s ⏱️
437 tests 437 ✅ 0 💤 0 ❌
440 runs  440 ✅ 0 💤 0 ❌

Results for commit 0f11354.

@github-actions

Copy link
Copy Markdown
Contributor

Code Coverage

Package Line Rate Health
staging_cse_windows Package 1 56%
debug 0%
provisioningscripts 2%
parts_windows Package 1 76%
vhdbuilder_packer_windows Package 1 20%
Summary 38% (1368 / 6070)

@github-actions

Copy link
Copy Markdown
Contributor

Code Coverage

Package Line Rate Health
shellspec spec 20%
Summary 20% (2839 / 13892)

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant