Repository navigation
fix(gpu): preserve customer-installed PIS drivers - #9702
Ben Brady (benjamin-brady) wants to merge 26 commits into
Conversation
Windows Unit Test Results 3 files 18 suites 1m 5s ⏱️ Results for commit 3022615. ♻️ This comment has been updated with latest results. |
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
The E2E sentinel can survive destructive cleanup and incorrectly report that customer drivers were preserved.
Review effort: Balanced
Findings: 1
What changed in this PR
Preserves customer-installed GPU drivers on PIS-derived nodes by moving managed-driver cleanup to image preparation.
Changes:
- Adds verified distro-specific bake cleanup and success-only completion markers.
- Restricts destructive node cleanup to runs that execute basePrep.
- Adds ShellSpec checks and a GPU VHD-caching scenario.
| File | Description |
|---|---|
| spec/parts/linux/cloud-init/artifacts/cse_main_gpu_pis_spec.sh | Tests cleanup across PIS phases. |
| spec/parts/linux/cloud-init/artifacts/cse_install_ubuntu_spec.sh | Tests Ubuntu cleanup verification. |
| spec/parts/linux/cloud-init/artifacts/cse_install_mariner_spec.sh | Tests cache cleanup verification. |
| parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh | Verifies driver teardown before capture. |
| parts/linux/cloud-init/artifacts/mariner/cse_install_mariner.sh | Verifies bake-time cache removal. |
| parts/linux/cloud-init/artifacts/cse_start.sh | Marks only successful image preparation. |
| parts/linux/cloud-init/artifacts/cse_main.sh | Gates cleanup by provisioning phase. |
| e2e/scenario/types.go | Adds a post-basePrep callback. |
| e2e/scenario/scenario.go | Registers GPU driver-preservation scenario. |
| e2e/scenario/provision.go | Runs callback and installs sentinel. |
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Clean AKS-owned artifacts before PIS customization and verify teardown independently of the prebake marker. Preserve customer drivers on derived nodes while retaining non-PIS cleanup. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
2cbf2cf to
d9ddd90
Compare
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Install a cached NVIDIA driver through customer RunCommand, then check version, binary/module hashes and GPU functionality on the derived node. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Use staged GB BOM versions to identify package ownership before PIS customization. Support both GB200 and current layouts; retain failure state when module teardown or initramfs rebuild cannot complete. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
AgentBaker Linux gate detectiveRun: 184124893 Detective summary: E2E finished with VHDCaching scenarios failing only in cleanup while deleting gallery image versions; ARM returned ScopeLocked for resource group Likely cause / signature: Confidence: High for E2E cleanup/SIG resource-group lock infra flake. Recommended owner/action: E2E/VHD infra owners should remove or avoid the Strongest alternative: PR 9702 changes E2E cleanup/provisioning code, so a cleanup regression is plausible; however, the same locked resource-group signature also appears on unrelated PR 9699. Evidence: build run, failed timeline records, focused E2E task log, build metadata, PR changed files, and wiki tracker. |
Avoid Fabric Manager, which requires NVSwitch hardware. Log driver receipts on bake and derived nodes to retain preservation evidence. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Use the CI-supported default region instead of the westus2 gallery, where the pipeline identity cannot capture the PIS image. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Use baked scripts for the bake unless explicitly disabled, then traditional delivery on the derived node. Document the required new source image for affected Ubuntu PIS pools. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Keep representative preservation, teardown failure, GB ownership and retry, payload, and E2E checks. Remove the expanded test matrices and repeated fixture setup without changing production or Packer behavior. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
PIS derived nodes deliberately override the global scriptless default. Validate CustomData against the effective NBC flag instead. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Keep the default GPU suite; allow a single fresh-image PIS scenario without rerunning every GPU scenario. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
|
|
|
|
Tim Wright (@timmy-wright) Nishchay (@awesomenix) could you review this for today’s release cut? I’ve removed the general CustomData compression changes and slimmed the diff from 1,423 to 893 changed lines, keeping both CUDA and Grace Blackwell support. The fix cleans AKS-owned drivers during the PIS bake, before customer scripts, and preserves customer drivers on derived nodes. Important compatibility point: opted-out Ubuntu PIS bakes, including CPU-only pools, require a new source image containing the cleanup helper. Local regression/scenario tests pass. Fresh-image GPU PIS E2E is still pending, so this is a request to review in parallel—not a claim that it is ready to merge. The PR description has the exact validation and remaining gate. |
Review: Preserve customer-installed PIS GPU driversWhat is goodThis PR fixes the ownership boundary, not just the deletion symptom. Cleaning AKS-owned artifacts before customer customization, then preserving the resulting driver on derived nodes, is the right lifecycle model.
Review outcomeReviewed #9702 at No confirmed current reliability or performance defect was found. There is one blocking validation gate below. The design recommendations are non-blocking; the final readability comment is a nitpick. The documented source-image requirement is an intentional compatibility change, not an additional finding: opted-out Ubuntu PIS bakes, including CPU-only pools, require a source VHD containing the new helper. Ordinary and derived nodes do not require it. Release sequencing must respect that requirement. Comments, ordered by importance1. Blocking: Verify the complete lifecycle with a fresh branch-built VHDWhere: PR validation evidence; The PR description still explicitly marks fresh-image E2E validation as pending. The local tests pass, but most destructive cleanup operations are stubbed. A green GPU pipeline check alone does not establish that the run exercised the new baked helper. Before merge, attach a successful run using a branch-built Why blocking: This is the main changed lifecycle, and it combines destructive cleanup with two VM runs and an image capture. This is a verification gate, not a claim that the implementation is broken. Clear it when the evidence is available; it does not inherently require a code change. 2. Non-blocking design recommendation: Keep configuration ownership knowledge synchronizedWhere: The expected configuration, including comments, is copied from the authoritative files under Risk: A future template edit can leave the cleanup verifier stale. Even a comment-only edit to Recommendation: Keep exact ownership verification, but tie its expected content to the authoritative templates. A focused consistency test is the smallest interim safeguard. Build-time generation is an option if this duplication continues to grow. Trade-off: A consistency test retains duplication but catches drift cheaply. Generation removes manual synchronization at the cost of additional build plumbing. Evidence grade: Weak for the general DRY prescription, which is primarily practitioner consensus; the duplicated content is directly observable. This is not a demonstrated failure in the current revision. 3. Non-blocking design recommendation: Make the cleanup state and dependency contract explicitWhere: Fresh cleanup, interrupted cleanup, and already-completed cleanup are represented through combinations of files, booleans, and missing configuration. Understanding the allowed transitions requires tracing both the orchestrator and the teardown function. The helper also relies on previously sourced distro cleanup functions and globals. Risk: A later change can preserve one path while accidentally breaking retries or legacy-layout handling. Recommendation: Add a concise contract describing the entry states, the pending marker's lifetime, required sourced functions/globals, and successful postconditions. Remove the unused Trade-off: This makes the existing design easier to audit without multiplying abstractions. It does not remove the safety complexity, which is justified. Evidence grade: Moderate for control-flow and naming guidance; its applicability to shell is indirect, not proof that a particular refactor would reduce defects. 4. Non-blocking design recommendation: Separate inspection failure from a positive resultWhere: The predicates return success both when artifacts remain and when inspection fails. The package predicate also accepts an allowlist, despite its name suggesting a general presence check, and writes inspection-error text to the same output channel as package names. Risk: A future caller could interpret "could not inspect" as "found a package," or misunderstand what the allowlist excludes. The current callers remain conservative; this is an interface-design concern. Recommendation: At minimum, document the return and output contracts. If changing the interface, distinguish found, absent, and inspection-failed outcomes, send diagnostics to stderr, and explicitly retain fail-closed behavior at every caller. Trade-off: Distinct outcomes add some branching but make the safety policy explicit rather than hiding it inside a presence query. Evidence grade: Moderate for aligning names and contracts with behavior. This does not justify weakening the existing failure handling. 5. Nitpick: Make the cleanup eligibility rules read like the lifecycle policyWhere: The compound shell conditions are valid, but nested command groups and mixed Recommendation: Prefer explicit Trade-off: A few more lines in exchange for more direct control flow. No behavioral change or new helper abstraction is needed. Evidence grade: Moderate for control-flow clarity. This is a readability preference, not a demonstrated bug, and should not delay merge. Validation performed during review
The selected ShellSpec suites contain 4 phase-dispatch examples, 3 Ubuntu-cleanup examples, 59 Mariner-installer examples, and 2 Packer-helper examples. These results establish local regression coverage, not a live fresh-image preservation pass or measured provisioning performance. The inspected Linux E2E build, 187241126, reported an Azure Linux iptables validation failure rather than a GPU cleanup failure. That observation does not clear the fresh-image validation gate above. |
Tim Wright (timmy-wright)
left a comment
Tim Wright (timmy-wright)
left a comment
There was a problem hiding this comment.
The lifecycle ownership boundary and conservative cleanup checks are good. Full review report: #9702 (comment)
The inline comments below are non-blocking design recommendations and one nitpick. The report separately identifies fresh branch-built VHD E2E evidence as the blocking validation gate; this review does not assert a current correctness defect.
| '# stack is up, so the post-dep fails silently. systemd-modules-load.service' \ | ||
| '# runs later in boot, by which point modprobe can resolve the full dep chain' \ | ||
| '# (ib_core, mlx5_ib, nvidia, ...).' \ | ||
| 'nvidia_peermem' || return 1 |
There was a problem hiding this comment.
Non-blocking design recommendation - keep ownership knowledge synchronized. The expected configuration here duplicates the authoritative GB configuration files, including their comments. A future comment-only edit to nvidia-peermem.conf could therefore make legitimate AKS configuration fail the exact ownership check and block a PIS bake.
Please retain exact verification, but consider a focused consistency test tying this expected content to the checked-in templates. That is a smaller change than adding build-time generation now: it retains duplication but catches drift cheaply. This is a future-maintenance concern, not a failure in the current revision.
| local bom="${GB_MAI_BOM_FILE:-/opt/azure/containers/gb-mai-bom.json}" | ||
| local legacy_bom="${GB200_MAI_BOM_FILE:-/opt/azure/containers/gb200-mai-bom.json}" | ||
| local pending_file="${GB_DRIVER_CLEANUP_PENDING_FILE:-/opt/azure/containers/gb-driver-cleanup.pending}" | ||
| local wave=versions-wave2 entries gb_layout=false pending_teardown=false pending_wave |
There was a problem hiding this comment.
Non-blocking design recommendation - make the cleanup state contract explicit. Fresh cleanup, interrupted cleanup, and already-completed cleanup are encoded through the pending marker, layout files, booleans, and missing configuration. The recovery behavior is valuable, but its invariants currently require tracing both this function and cleanUpGraceBlackwellGPUDriver.
A concise contract describing the entry states, pending-marker lifetime, required sourced functions/globals, and successful postconditions would make this safer to maintain without introducing a state-machine framework. The unused bom parameter on cleanUpGraceBlackwellGPUDriver can also be removed to avoid implying a dependency that the function does not use.
|
|
||
| graceBlackwellPackageManagedDriverRemains() { | ||
| local allowed="${1:-}" package_states | ||
| package_states=$(dpkg-query -W -f='${binary:Package}\t${db:Status-Status}\n') || { |
There was a problem hiding this comment.
Non-blocking design recommendation - clarify inspection results versus inspection failure. This predicate returns success both when an unowned package remains and when dpkg-query fails. Its stdout can therefore contain either a package name or an inspection-error message. The current caller remains conservative, but the interface is easy for a future caller to misinterpret.
At minimum, document the return/output contract and optional allowlist. If revising the interface, distinguish found, absent, and inspection-failed outcomes, send diagnostics to stderr, and explicitly preserve fail-closed handling at every caller. The extra branching would buy a clearer safety contract; this is not a request to weaken the check.
|
|
||
| # Clean up GPU drivers if not a GPU node or if skipping driver install | ||
| if [ "${GPU_NODE}" != "true" ] || [ "${skip_nvidia_driver_install}" = "true" ]; then | ||
| if [ "${BASE_PREP_RAN}" = "true" ] && { [ "${GPU_NODE}" != "true" ] || [ "${skip_nvidia_driver_install}" = "true" ]; }; then |
There was a problem hiding this comment.
Nitpick - make the lifecycle policy explicit beside this guard. The comment above still describes only the old GPU conditions. The important new condition is that cleanup runs only when basePrep ran in this execution, so PIS-derived nodes retain customer-installed drivers.
Please consider updating that comment and initializing BASE_PREP_RAN=false beside the phase dispatcher. Explicit nested if blocks would also make the eligibility rules easier to scan than the compound shell grouping. This is readability-only and should not delay merge.
What this PR does / why we need it:
PIS-derived nodes skip
basePrep, but theirnodePreppreviously deleted customer-installed GPU drivers. Clean AKS-owned artifacts before customer customization, then preserve the customer driver on derived nodes.flowchart LR A[New AKS source VHD] --> B[basePrep: verify AKS GPU cleanup] B --> C[Customer installs driver] C --> D[Capture PIS image] D --> E[Derived node: preserve driver]base_prep.completeonly after success. Azure Linux retains distro-specific cleanup.basePrepran in the same execution. Ordinary-node ordering and legacy skip-tag behavior are unchanged.nvidia-smivalidation. The CustomData guard uses the effective NBC flag rather than the global scriptless default.scenarioFilterfor a targeted fresh-image run; its default remainsgpu=true. Missing source images still fail.Compatibility: all opted-out Ubuntu PIS bakes, including CPU-only pools, require a new source VHD containing the helper. This broader requirement is intentional and accepted; a missing helper fails explicitly. Ordinary and derived nodes do not require the helper. Old CSEs can still delete customer drivers; existing captures do not gain bake-time cleanup retroactively. GRID mismatch/validation paths are unchanged. No dependency on #9609 or #9753; general compression changes and the unrelated lint pin were removed.
Payload: the synthetic traditional Ubuntu 22.04/24.04 T4 fixture on base
f3163c3013e5(2026-10-22) measures 65,241 gzip bytes / 86,988 Base64 characters, leaving 294 bytes / 392 characters. Production/Packer code is unchanged by subsequent test/CI fixes. This narrow margin is not a guarantee for all customer/custom-cloud configurations.Review scope: 884 additions / 9 deletions (893 changed lines, down from 1,423). Both GPU implementations remain. Focused preservation, fail-closed cleanup, GB ownership/retry, payload, AMD equivalence and Packer tests remain; exhaustive matrices were removed.
cd e2e && go test ./scenario -count=1cd e2e && go test ./scenario -run '^TestCreateVMSSModelHonorsScriptlessNBCOverride$' -count=1docker run --rm -e TMPDIR=/workspace/.copilot-test-tmp -v "$PWD:/workspace" -w /workspace shellspec-docker --shell bash --format d spec/parts/linux/cloud-init/artifacts/cse_main_gpu_pis_spec.shdocker run --rm -e TMPDIR=/workspace/.copilot-test-tmp -v "$PWD:/workspace" -w /workspace shellspec-docker --shell bash --format d spec/parts/linux/cloud-init/artifacts/cse_install_ubuntu_gpu_pis_spec.shdocker run --rm -e TMPDIR=/workspace/.copilot-test-tmp -v "$PWD:/workspace" -w /workspace shellspec-docker --shell bash --format d spec/parts/linux/cloud-init/artifacts/cse_install_mariner_spec.shdocker run --rm -e TMPDIR=/workspace/.copilot-test-tmp -v "$PWD:/workspace" -w /workspace shellspec-docker --shell bash --format d spec/vhdbuilder/packer/pis_gpu_helper_spec.shgofmt -d pkg/agent/gpu_pis_delivery_test.go;git diff --checkGPU pipeline YAML parsed successfully; parameter default and missing-image failure behavior checked locally. The larger pre-slimming suites also passed in an isolated worktree at
bb5b398e95(2026-10-23 NZDT), but are not claimed as retained coverage. Independent review found no significant issues in the GPU implementation or test slimming. Repository-wide local ShellCheck previously panicked in unchangedpublish-generalized-cvm-if-required.sh; scoped checks and pre-slimming CI ShellCheck passed.Live E2E pending. GPU build 187228397 failed at the now-fixed CustomData guard and used a main image; it is not fresh-helper proof. The targeted run must use a successful branch-built
2404gen2containerdimage available in westus3. Metadata is inspected externally; this GPU pipeline selects the image byVHD_BUILD_ID/ gallerybuildIdtag, not by consuming metadata. Verify the selected resource ID equals the metadata ID, then verify the run logs before counting the result. Missing-image skips are not passes. No fresh-image preservation pass is claimed yet.Which issue(s) this PR fixes:
Fixes AB#39923593 (Don't delete user installed GPU drivers); depends on AB#39198532.