diff --git a/.pipelines/e2e-gpu.yaml b/.pipelines/e2e-gpu.yaml index 6c7e3e39ac5..375d80c5faf 100644 --- a/.pipelines/e2e-gpu.yaml +++ b/.pipelines/e2e-gpu.yaml @@ -1,6 +1,10 @@ name: $(Date:yyyyMMdd)$(Rev:.r) +parameters: + - name: scenarioFilter + type: string + default: "gpu=true" variables: - TAGS_TO_RUN: "gpu=true" + TAGS_TO_RUN: ${{ parameters.scenarioFilter }} TAGS_TO_SKIP: "os=windows,os=azurelinux,os=azurecontainerlinux" SKIP_E2E_TESTS: false # PRs skip quota/capacity-constrained scenarios; main and scheduled runs fail them. diff --git a/e2e/README.md b/e2e/README.md index 2d7c56c02ee..2c578cf6130 100644 --- a/e2e/README.md +++ b/e2e/README.md @@ -94,6 +94,14 @@ coverage matching production. A Scriptless tag or an AKSNodeConfig mutator alone does not opt into Phase 3. Existing scriptless opt-outs, image restrictions, and VHD-caching behavior remain in effect. +### Ubuntu PIS GPU cleanup + +`Ubuntu2404_GPUDriverOptOut_VHDCaching` needs a new Ubuntu 24.04 Gen2 image +containing the baked cleanup helper; use same-run E2E VHD metadata or +`--vhd-metadata-file` with `2404gen2containerd`. The helper is baked, not +transferred in CustomData. An opt-out PIS bake without it fails explicitly; +ordinary nodes and derived PIS nodes do not require it. + ### LocalDNS fallback checks The existing LocalDNS scenarios use short-lived DNS pods on the scenario node diff --git a/e2e/scenario/provision.go b/e2e/scenario/provision.go index c0044fb1a2c..699a943fca9 100644 --- a/e2e/scenario/provision.go +++ b/e2e/scenario/provision.go @@ -70,7 +70,13 @@ func runVHDCachingScenario(ctx context.Context, name string, original *Scenario) ValidateSystemdUnitIsNotRunning(ctx, scenario, "kubelet"), ) } - return validationErr + if validationErr != nil { + return validationErr + } + if original.Config.VHDCachingPostBasePrep != nil { + return original.Config.VHDCachingPostBasePrep(ctx, scenario) + } + return nil } bakeScenario.Config.VMConfigMutator = func(vmss *armcompute.VirtualMachineScaleSet) { @@ -832,6 +838,65 @@ func RunCommand(ctx context.Context, s *Scenario, command string) (armcompute.Vi return view, runCommandScriptError(view) } +const ( + customerGPUDriverImage = "mcr.microsoft.com/aks/aks-gpu-cuda-lts" +) + +func customerGPUDriverImageReference() (string, string, error) { + if datamodel.NvidiaCudaLTSDriverVersion == "" || datamodel.AKSGPUCudaLTSVersionSuffix == "" { + return "", "", fmt.Errorf("approved NVIDIA CUDA LTS driver image version is unavailable") + } + version := datamodel.NvidiaCudaLTSDriverVersion + tag := version + "-" + datamodel.AKSGPUCudaLTSVersionSuffix + return customerGPUDriverImage + ":" + tag, version, nil +} + +// installCustomerNvidiaDriver manually runs the VHD-cached, components-approved installer after +// basePrep. This models a customer RunCommand install; the scenario leaves managed GPU installation +// opted out, so nodePrep does not invoke the installer. +func installCustomerNvidiaDriver(ctx context.Context, s *Scenario) error { + imageRef, expectedVersion, err := customerGPUDriverImageReference() + if err != nil { + return err + } + command := fmt.Sprintf(`set -euo pipefail +image_ref=%s +expected_version=%s +sudo mkdir -p /opt/gpu /opt/actions +sudo ctr -n k8s.io images ls -q | grep -Fxq "${image_ref}" +sudo ctr -n k8s.io run --privileged --rm --net-host \ + --with-ns pid:/proc/1/ns/pid \ + --mount type=bind,src=/opt/gpu,dst=/mnt/gpu,options=rbind \ + --mount type=bind,src=/opt/actions,dst=/mnt/actions,options=rbind \ + "${image_ref}" gpuinstall /entrypoint.sh copy +sudo ctr -n k8s.io images rm "${image_ref}" + +installer_dir="/opt/gpu/NVIDIA-Linux-x86_64-${expected_version}" +test -x "${installer_dir}/nvidia-installer" +cd "${installer_dir}" +sudo ./nvidia-installer -s -k="$(uname -r)" --no-drm --dkms + +sudo nvidia-modprobe -u -c0 +gpu_info="$(sudo nvidia-smi --query-gpu=driver_version,name --format=csv,noheader | head -n1)" +driver_version="${gpu_info%%,*}" +driver_version="$(printf '%%s' "${driver_version}" | tr -d '[:space:]')" +gpu_name="${gpu_info#*,}" +case "${gpu_name}" in *T4*) ;; *) echo "expected a T4 GPU, got: ${gpu_name}" >&2; exit 1 ;; esac +test "${driver_version}" = "${expected_version}" +sudo nvidia-smi -L | grep -qi 'T4' +lsmod | grep -q '^nvidia[[:space:]]' +test "$(sudo modinfo -F version nvidia)" = "${expected_version}" +`, + shellSingleQuote(imageRef), + shellSingleQuote(expectedVersion), + ) + _, err = RunCommand(ctx, s, command) + if err != nil { + return fmt.Errorf("install customer NVIDIA driver from cached approved artifact: %w", err) + } + return nil +} + func RestartVMSSVM(ctx context.Context, s *Scenario) error { if s.Runtime == nil || s.Runtime.Cluster == nil || s.Runtime.Cluster.Model == nil || s.Runtime.Cluster.Model.Properties == nil || s.Runtime.Cluster.Model.Properties.NodeResourceGroup == nil || diff --git a/e2e/scenario/scenario.go b/e2e/scenario/scenario.go index 2d5121a16fd..1101b68c480 100755 --- a/e2e/scenario/scenario.go +++ b/e2e/scenario/scenario.go @@ -3106,6 +3106,37 @@ var _ = Register(&Scenario{ }, }) +var _ = Register(&Scenario{ + Name: "Ubuntu2404_GPUDriverOptOut_VHDCaching", + Description: "Installs the VHD-cached, components-approved NVIDIA driver artifact with a post-basePrep customer RunCommand while managed driver installation remains opted out, then verifies the real driver survives PIS VHD-cached provisioning", + K8sSystemPoolSKU: "Standard_D2s_v3", + Tags: Tags{ + GPU: true, + }, + Config: Config{ + Cluster: ClusterKubenet, + VHD: config.VHDUbuntu2404Gen2Containerd, + VHDCaching: true, + VHDCachingPostBasePrep: installCustomerNvidiaDriver, + PreProvisionBootstrapConfigMutator: func(_ *Cluster, nbc *datamodel.NodeBootstrappingConfiguration) { + nbc.EnableScriptlessCSECmd = !config.Config.DisableScriptless + }, + BootstrapConfigMutator: func(_ *Cluster, nbc *datamodel.NodeBootstrappingConfiguration) { + nbc.EnableScriptlessCSECmd = false + nbc.AgentPoolProfile.VMSize = "Standard_NC4as_T4_v3" + nbc.ConfigGPUDriverIfNeeded = false + nbc.EnableGPUDevicePluginIfNeeded = false + nbc.EnableNvidia = false + }, + VMConfigMutator: func(vmss *armcompute.VirtualMachineScaleSet) { + vmss.SKU.Name = to.Ptr("Standard_NC4as_T4_v3") + }, + Validator: func(ctx context.Context, s *Scenario) error { + return ValidateCustomerNvidiaDriver(ctx, s) + }, + }, +}) + var _ = Register(&Scenario{ Name: "Ubuntu2404_AKSVMExtension_FilesystemCorruption", Description: "Tests Ubuntu 24.04 NPD filesystem-corruption reporting with the AKS VM Extension and hotfix target resolution from the NBC command when ANC JSON is absent", diff --git a/e2e/scenario/types.go b/e2e/scenario/types.go index 27a2c4456dd..cefcf180b58 100644 --- a/e2e/scenario/types.go +++ b/e2e/scenario/types.go @@ -198,6 +198,10 @@ type Config struct { // The main purpose is to validate VHD Caching logic and ensure a reboot step between basePrep and nodePrep doesn't break anything. VHDCaching bool + // VHDCachingPostBasePrep runs on the bake VM after basePrep completes and before its disk is captured. + // Use it only for VHD-caching scenarios that need to seed state between the two provisioning phases. + VHDCachingPostBasePrep func(context.Context, *Scenario) error + // ExpectedError, when set, indicates that VMSS creation is expected to fail with an error containing this substring. // The assertion is performed during the scenario run. ExpectedError string diff --git a/e2e/scenario/validators.go b/e2e/scenario/validators.go index bd318d124a6..1836b486eec 100644 --- a/e2e/scenario/validators.go +++ b/e2e/scenario/validators.go @@ -627,6 +627,33 @@ func ValidateNvidiaSMIInstalled(ctx context.Context, s *Scenario) error { return err } +func ValidateCustomerNvidiaDriver(ctx context.Context, s *Scenario) error { + _, expectedVersion, err := customerGPUDriverImageReference() + if err != nil { + return err + } + installerPath := fmt.Sprintf("/opt/gpu/NVIDIA-Linux-x86_64-%s/nvidia-installer", expectedVersion) + command := customerNvidiaDriverValidationScript(expectedVersion, installerPath) + _, err = execScriptOnVMForScenarioValidateExitCode(ctx, s, command, 0, "customer-installed NVIDIA driver did not survive PIS image capture") + return err +} + +func customerNvidiaDriverValidationScript(expectedVersion, installerPath string) string { + return fmt.Sprintf(`set -euo pipefail +expected_version=%s +test -x %s +lsmod | grep -q '^nvidia[[:space:]]' +module_version="$(sudo modinfo -F version nvidia)" +test "${module_version}" = "${expected_version}" +gpu_version="$(sudo nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -n1 | tr -d '[:space:]')" +test "${gpu_version}" = "${expected_version}" +sudo nvidia-smi +`, + shellSingleQuote(expectedVersion), + shellSingleQuote(installerPath), + ) +} + func ValidateNvidiaModProbeInstalled(ctx context.Context, s *Scenario) error { command := []string{ "set -ex", diff --git a/e2e/scenario/validators_test.go b/e2e/scenario/validators_test.go index 1b23beebaad..197508dcb8a 100644 --- a/e2e/scenario/validators_test.go +++ b/e2e/scenario/validators_test.go @@ -3,6 +3,10 @@ package scenario import ( "crypto/sha256" "encoding/hex" + "fmt" + "os" + "os/exec" + "path/filepath" "testing" "github.com/Azure/azure-sdk-for-go/sdk/azcore/to" @@ -10,6 +14,39 @@ import ( resourcev1 "k8s.io/api/resource/v1" ) +func TestCustomerNvidiaDriverValidationRequiresStagedInstaller(t *testing.T) { + const expectedVersion = "580.159.04" + binDir := t.TempDir() + installerPath := filepath.Join(binDir, "nvidia-installer") + require.NoError(t, os.WriteFile(filepath.Join(binDir, "lsmod"), []byte("#!/bin/sh\nprintf 'nvidia 123 0\\n'\n"), 0o755)) + sudo := fmt.Sprintf(`#!/bin/sh +case "$1" in + modinfo) printf '%%s\n' '%s'; exit 0 ;; + nvidia-smi) + if [ "$2" = "--query-gpu=driver_version" ]; then + printf '%%s\n' '%s' + fi + exit 0 + ;; +esac +exit 1 +`, expectedVersion, expectedVersion) + require.NoError(t, os.WriteFile(filepath.Join(binDir, "sudo"), []byte(sudo), 0o755)) + t.Setenv("PATH", binDir+string(os.PathListSeparator)+os.Getenv("PATH")) + + runValidation := func() error { + output, err := exec.Command("bash", "-x", "-c", customerNvidiaDriverValidationScript(expectedVersion, installerPath)).CombinedOutput() + if err != nil { + return fmt.Errorf("%w: %s", err, output) + } + return nil + } + require.Error(t, runValidation(), "validation must reject correct driver state when the staged installer was removed") + + require.NoError(t, os.WriteFile(installerPath, []byte("#!/bin/sh\n"), 0o755)) + require.NoError(t, runValidation(), "validation must pass when the staged installer and driver state are present") +} + func TestValidateDRAMIGAllocation(t *testing.T) { tests := []struct { name string diff --git a/e2e/scenario/vmss.go b/e2e/scenario/vmss.go index bc0b1785188..be232f77f6b 100644 --- a/e2e/scenario/vmss.go +++ b/e2e/scenario/vmss.go @@ -526,7 +526,7 @@ func createVMSSModel(ctx context.Context, s *Scenario) (armcompute.VirtualMachin return armcompute.VirtualMachineScaleSet{}, fmt.Errorf("inject customData write_files entries: %w", err) } } - if !config.Config.DisableScriptless && !scriptlessProvisioningEnabled && s.VHD.SupportsScriptless() { + if s.Runtime.NBC.EnableScriptlessCSECmd && !scriptlessProvisioningEnabled && s.VHD.SupportsScriptless() { // Validate that the custom data doesn't contain any script content, // which indicates that the scriptless CSE is working as intended decodedCustomData, err := base64.StdEncoding.DecodeString(customData) diff --git a/e2e/scenario/vmss_test.go b/e2e/scenario/vmss_test.go index 6b77e77c270..2c2cc1be07d 100644 --- a/e2e/scenario/vmss_test.go +++ b/e2e/scenario/vmss_test.go @@ -44,6 +44,67 @@ func (f vmssCreationTestPolicy) Do(req *policy.Request) (*http.Response, error) return f(req.Raw()), nil } +func TestCreateVMSSModelHonorsScriptlessNBCOverride(t *testing.T) { + oldScriptless := config.Config.DisableScriptless + oldPrepareVHD := CachedPrepareVHD + config.Config.DisableScriptless = false + CachedPrepareVHD = func(context.Context, GetVHDRequest) (config.VHDResourceID, error) { + return "fixture-image", nil + } + t.Cleanup(func() { + config.Config.DisableScriptless = oldScriptless + CachedPrepareVHD = oldPrepareVHD + }) + + for _, enableScriptlessCSECmd := range []bool{false, true} { + name := "traditional override" + if enableScriptlessCSECmd { + name = "scriptless CSE" + } + t.Run(name, func(t *testing.T) { + vhd := *config.VHDUbuntu2404Gen2Containerd + cluster := &Cluster{ + Model: &armcontainerservice.ManagedCluster{ + Name: to.Ptr("cluster"), + Location: to.Ptr("westus3"), + Properties: &armcontainerservice.ManagedClusterProperties{ + CurrentKubernetesVersion: to.Ptr("1.34.0"), + NodeResourceGroup: to.Ptr("node-rg"), + NetworkProfile: &armcontainerservice.NetworkProfile{ + NetworkPlugin: to.Ptr(armcontainerservice.NetworkPluginKubenet), + }, + }, + }, + KubeletIdentity: &armcontainerservice.UserAssignedIdentity{ + ClientID: to.Ptr("client-id"), + ResourceID: to.Ptr("/subscriptions/sub/resourceGroups/rg/providers/Microsoft.ManagedIdentity/userAssignedIdentities/kubelet"), + }, + ClusterParams: &ClusterParams{ + BootstrapToken: "token", + CACert: []byte("ca-cert"), + FQDN: "cluster.example.test", + }, + } + nbc, err := getBaseNBC(t.Context(), cluster, &vhd) + require.NoError(t, err) + nbc.EnableScriptlessCSECmd = enableScriptlessCSECmd + s := &Scenario{ + Config: Config{VHD: &vhd}, + Location: "westus3", + Runtime: &ScenarioRuntime{ + NBC: nbc, + Cluster: cluster, + VMSSName: "vmss", + }, + } + + model, err := createVMSSModel(t.Context(), s) + require.NoError(t, err) + require.NotEmpty(t, *model.Properties.VirtualMachineProfile.OSProfile.CustomData) + }) + } +} + type vmssCreationTestCase struct { name string provisionCode string diff --git a/parts/linux/cloud-init/artifacts/cse_config_gpu.sh b/parts/linux/cloud-init/artifacts/cse_config_gpu.sh index 3ad27398952..a807499d2c4 100755 --- a/parts/linux/cloud-init/artifacts/cse_config_gpu.sh +++ b/parts/linux/cloud-init/artifacts/cse_config_gpu.sh @@ -231,6 +231,7 @@ ensureGPUDrivers() { } # Install AMD AMA core SW package for MA35D (Supernova GPU SKU) +#{{if or IsMariner IsCustomImage}} dnf_install_amd_ama_core_package() { ver=$1; retries=$2; wait_sleep=$3; timeout=$4; shift && shift && shift && shift # Currently version 1.5.0 is supported. Add more versions as they become available. @@ -251,6 +252,7 @@ dnf_install_amd_ama_core_package() { done echo Executed dnf install AMD AMA core package $i times; } +#{{end}} # Install AMD AMA drivers/SW for MA35D (Supernova GPU SKU) # Note that this depends on access to download.microsoft.com, so network-isolated clusters are not supported @@ -260,6 +262,9 @@ setupAmdAma() { fi if isAzureLinux "$OS"; then + : + # Template guards leave raw VHD scripts intact; unknown custom images retain the implementation. + #{{if or IsMariner IsCustomImage}} # Install MA35D packages - currently version 1.5 and above are supported # This install script will extract the driver version/build number to find # the corresponding core/FW packages @@ -325,6 +330,7 @@ setupAmdAma() { if [ "$(systemctl is-active kubelet)" = "active" ]; then systemctl restart kubelet fi + #{{end}} fi } diff --git a/parts/linux/cloud-init/artifacts/cse_main.sh b/parts/linux/cloud-init/artifacts/cse_main.sh index 0c3ffe92253..fc8fa44727f 100755 --- a/parts/linux/cloud-init/artifacts/cse_main.sh +++ b/parts/linux/cloud-init/artifacts/cse_main.sh @@ -407,6 +407,18 @@ EOF logs_to_events "AKS.CSE.ensureContainerd.ensureArtifactStreaming" ensureArtifactStreaming || exit $ERR_ARTIFACT_STREAMING_INSTALL fi + if [ "${PRE_PROVISION_ONLY}" = "true" ] && [ "${GPU_NODE}" != "true" ] && + { [ "${OS}" = "${UBUNTU_OS_NAME}" ] || { isMarinerOrAzureLinux "$OS" "$OS_VARIANT" && ! isAzureLinuxOSGuard "$OS" "$OS_VARIANT"; }; }; then + if [ "${OS}" = "${UBUNTU_OS_NAME}" ]; then + if [ ! -r /opt/azure/containers/provision_installs_ubuntu_gpu.sh ]; then + echo "image containing PIS GPU cleanup helper required" + exit $ERR_NVIDIA_DRIVER_INSTALL + fi + source /opt/azure/containers/provision_installs_ubuntu_gpu.sh || exit $ERR_NVIDIA_DRIVER_INSTALL + fi + logs_to_events "AKS.CSE.cleanUpGPUDriversForBasePrep" cleanUpGPUDriversForBasePrep || exit $ERR_NVIDIA_DRIVER_INSTALL + fi + # Enable localdns to handle node and pod DNS traffic via a local CoreDNS instance. # If hosts plugin is enabled, enableAKSLocalDNSHostsSetup() was already called at the # very start of basePrep (before disableSystemdResolved) to give the timer a head start @@ -645,7 +657,7 @@ function nodePrep { # 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 logs_to_events "AKS.CSE.cleanUpGPUDrivers" cleanUpGPUDrivers fi @@ -753,6 +765,7 @@ logs_to_events "AKS.CSE.configureNodeTimeSync" configure_node_time_sync_or_repor if [ ! -f /opt/azure/containers/base_prep.complete ]; then basePrep + BASE_PREP_RAN=true else echo "Skipping basePrep - base_prep.complete file exists" fi diff --git a/parts/linux/cloud-init/artifacts/cse_start.sh b/parts/linux/cloud-init/artifacts/cse_start.sh index 1b28fd5abfc..ae931668d9b 100644 --- a/parts/linux/cloud-init/artifacts/cse_start.sh +++ b/parts/linux/cloud-init/artifacts/cse_start.sh @@ -115,10 +115,12 @@ upload_logs() { } # Create the marker for the completed provisioning stage. if [ "${PRE_PROVISION_ONLY}" = "true" ]; then - # Stage 1: Create marker indicating Stage 2 is needed - mkdir -p /opt/azure/containers && touch /opt/azure/containers/base_prep.complete - echo "Stage 1 complete - kubelet configuration skipped, Stage 2 required" >> /var/log/azure/cluster-provision.log - echo "Created base_prep.complete marker file" >> /var/log/azure/cluster-provision.log + if [ "$EXIT_CODE" -eq 0 ]; then + # Stage 1: Create marker indicating Stage 2 is needed + mkdir -p /opt/azure/containers && touch /opt/azure/containers/base_prep.complete + echo "Stage 1 complete - kubelet configuration skipped, Stage 2 required" >> /var/log/azure/cluster-provision.log + echo "Created base_prep.complete marker file" >> /var/log/azure/cluster-provision.log + fi else # provision.complete signals that a normal provisioning attempt finished. mkdir -p /opt/azure/containers && touch /opt/azure/containers/provision.complete diff --git a/parts/linux/cloud-init/artifacts/mariner/cse_install_mariner.sh b/parts/linux/cloud-init/artifacts/mariner/cse_install_mariner.sh index 8d19ce74b2c..1b861f33107 100755 --- a/parts/linux/cloud-init/artifacts/mariner/cse_install_mariner.sh +++ b/parts/linux/cloud-init/artifacts/mariner/cse_install_mariner.sh @@ -833,6 +833,22 @@ cleanUpGPUDrivers() { done } +cleanUpGPUDriversForBasePrep() { + local packageName cache_dir + cleanUpGPUDrivers + if [ -e "${GPU_DEST}" ] || [ -e /opt/gpu ]; then + echo "GPU basePrep cleanup incomplete: AKS driver artifacts remain" >&2 + return 1 + fi + for packageName in $(managedGPUPackageList); do + cache_dir=$(getPackageCacheDir "${packageName}") + if [ -e "${cache_dir}" ]; then + echo "GPU basePrep cleanup incomplete: ${packageName} cache remains" >&2 + return 1 + fi + done +} + installMinimalBuildDeps() { echo "installMinimalBuildDeps not implemented for mariner" } diff --git a/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu_gpu.sh b/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu_gpu.sh new file mode 100755 index 00000000000..c45a6a41076 --- /dev/null +++ b/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu_gpu.sh @@ -0,0 +1,288 @@ +#!/bin/bash + +prebakedGPUDriverArtifactsRemain() { + local binary modules + [ -d /var/lib/dkms/nvidia ] || [ -e /usr/bin/lib64 ] || [ -e /etc/ld.so.conf.d/nvidia.conf ] && return 0 + for binary in nvidia-smi nvidia-debugdump nvidia-persistenced nvidia-cuda-mps-control \ + nvidia-cuda-mps-server nvidia-modprobe nvidia-bug-report.sh nvidia-powerd \ + nvidia-ngx-updater nvidia-sleep.sh; do + [ -e "/usr/bin/${binary}" ] && return 0 + done + compgen -G '/lib/modules/*/updates/dkms/nvidia*.ko*' >/dev/null && return 0 + modules=$(lsmod) || return 0 + grep -q '^nvidia' <<< "${modules}" +} + +graceBlackwellDriverBOMEntries() { + # Packer stages this BOM only on NVIDIA_GB images; package versions must still match it before removal. + local bom="${1:-${GB_MAI_BOM_FILE:-/opt/azure/containers/gb-mai-bom.json}}" + local wave="${2:-versions-wave2}" + [ -s "${bom}" ] || return 1 + + jq -er --arg wave "${wave}" ' + .[$wave] as $packages | + select(($packages | type) == "object") | + select(any($packages | keys[]; test("^nvidia-dkms-[0-9]+-open$"))) | + select(any($packages | keys[]; test("^nvidia-driver-[0-9]+-open$"))) | + select(all($packages | keys[]; test("^(libxnvctrl0|libnvidia-[A-Za-z0-9-]+|nvidia-[A-Za-z0-9-]+|xserver-xorg-video-nvidia-[0-9]+)$"))) | + select(all($packages | values[]; type == "string" and length > 0)) | + $packages | to_entries[] | "\(.key)=\(.value)" + ' "${bom}" +} + +graceBlackwellPackageManagedDriverRemains() { + local allowed="${1:-}" package_states + package_states=$(dpkg-query -W -f='${binary:Package}\t${db:Status-Status}\n') || { + echo "unable to inspect installed NVIDIA driver packages" + return 0 + } + awk -F '\t' -v allowed="${allowed}" ' + $2 == "installed" { + name = $1 + sub(/:.*/, "", name) + if (name !~ /^(nvidia-(dkms|driver|utils|compute-utils|firmware|kernel-)|nvidia-(modprobe|persistenced)$|xserver-xorg-video-nvidia-|libxnvctrl0$|libnvidia-)/ || name ~ /^libnvidia-container/) next + count = split(allowed, owned, "\n") + for (i = 1; i <= count; i++) { + sub(/=.*/, "", owned[i]) + if (owned[i] == name) next + } + print $1 + found = 1 + exit + } + END { if (!found) exit 1 } + ' <<< "${package_states}" +} + +graceBlackwellInstalledDriverPackages() { + local entries="${1}" package expected state version series="" auto_packages bom_package + local -a installed=() + + while IFS='=' read -r package expected; do + [ -n "${package}" ] || continue + case "${package}" in + nvidia-driver-[0-9]*-open) + series="${package#nvidia-driver-}" + series="${series%-open}" + ;; + esac + state=$(dpkg-query -W -f='${db:Status-Status}\t${Version}' "${package}" 2>/dev/null) || state="" + [ "${state%%$'\t'*}" = installed ] || continue + version="${state#*$'\t'}" + if [ "${version}" != "${expected}" ]; then + echo "Grace Blackwell package ${package} version ${version} does not match the GB BOM ${expected}" >&2 + return 1 + fi + installed+=("${package}=${version}") + done <<< "${entries}" + + [ -n "${series}" ] || return 1 + auto_packages=$(apt-mark showauto) || return 1 + # BOMs do not always list apt dependencies; use each dependency's exact dpkg version, never the driver's version. + for package in "nvidia-utils-${series}" "nvidia-compute-utils-${series}" nvidia-modprobe nvidia-persistenced; do + while IFS='=' read -r bom_package expected; do + [ "${bom_package}" = "${package}" ] && continue 2 + done <<< "${entries}" + state=$(dpkg-query -W -f='${db:Status-Status}\t${Version}' "${package}" 2>/dev/null) || state="" + [ "${state%%$'\t'*}" = installed ] || continue + version="${state#*$'\t'}" + if [ -z "${version}" ] || ! grep -Eq "^${package}(:arm64)?$" <<< "${auto_packages}"; then + echo "Grace Blackwell dependency ${package} is not an auto-installed package with a verifiable version" >&2 + return 1 + fi + installed+=("${package}=${version}") + done + + printf '%s\n' "${installed[@]}" +} + +graceBlackwellConfigIsUnmodified() { + local config="${1}" expected + shift + [ -e "${config}" ] || return 0 + + expected=$(printf '%s\n' "$@") + if ! cmp -s "${config}" <(printf '%s\n' "$@") && + ! cmp -s "${config}" <(printf '%s' "${expected}"); then + echo "Grace Blackwell module configuration ${config} was modified; refusing teardown" >&2 + return 1 + fi +} + +cleanUpGraceBlackwellGPUDriver() { + local entries="${1}" bom="${2}" wave="${3}" pending="${4:-false}" + local packages package version modules service module remaining_package pending_file + pending_file="${GB_DRIVER_CLEANUP_PENDING_FILE:-/opt/azure/containers/gb-driver-cleanup.pending}" + local nouveau_config="${GB_NOUVEAU_MODPROBE_CONFIG_FILE:-/etc/modprobe.d/blacklist-nouveau.conf}" + local peermem_config="${GB_NVIDIA_PEERMEM_CONFIG_FILE:-/etc/modules-load.d/nvidia-peermem.conf}" + local nvidia_config="${GB_NVIDIA_MODPROBE_CONFIG_FILE:-/etc/modprobe.d/nvidia.conf}" + packages=$(graceBlackwellInstalledDriverPackages "${entries}") || return 1 + if remaining_package=$(graceBlackwellPackageManagedDriverRemains "${packages}"); then + echo "NVIDIA package ${remaining_package} is not owned by the Grace Blackwell BOM" >&2 + return 1 + fi + if [ "${pending}" != true ] && [ ! -f "${nvidia_config}" ]; then + modules=$(lsmod) || return 1 + # BasePrep may retry after teardown succeeded but a later step failed. + if [ -z "${packages}" ] && ! grep -q '^nvidia' <<< "${modules}" && + ! prebakedGPUDriverArtifactsRemain && + [ ! -e "${nvidia_config}" ] && [ ! -L "${nvidia_config}" ] && + [ ! -e "${peermem_config}" ] && [ ! -L "${peermem_config}" ] && + [ ! -e "${nouveau_config}" ] && [ ! -L "${nouveau_config}" ]; then + return 0 + fi + echo "Grace Blackwell NVIDIA module configuration ${nvidia_config} is unavailable; refusing teardown" >&2 + return 1 + fi + [ "${pending}" = true ] || [ -f "${nvidia_config}" ] || { + echo "Grace Blackwell NVIDIA module configuration ${nvidia_config} is unavailable; refusing teardown" >&2 + return 1 + } + graceBlackwellConfigIsUnmodified "${peermem_config}" \ + '# Explicitly load nvidia_peermem at boot for GPUDirect RDMA.' \ + "# We cannot rely on a \`softdep nvidia post: nvidia_peermem\` in modprobe.d because" \ + "# \`nvidia\` is autoloaded by the kernel via PCI modalias before the OFED/RDMA" \ + '# 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 + graceBlackwellConfigIsUnmodified "${nvidia_config}" \ + 'options nvidia NVreg_RestrictProfilingToAdminUsers=0' \ + 'options nvidia NVreg_CreateImexChannel0=1' \ + 'options nvidia NVreg_CoherentGPUMemoryMode=driver' \ + 'options nvidia NVreg_RegistryDwords="RMBug5172204War=4"' || return 1 + graceBlackwellConfigIsUnmodified "${nouveau_config}" \ + 'blacklist nouveau' 'options nouveau modeset=0' || return 1 + + local -a purge_packages=() + while IFS='=' read -r package version; do + [ -n "${package}" ] && purge_packages+=("${package}") + done <<< "${packages}" + + for service in nvidia-device-plugin nvidia-dcgm-exporter nvidia-dcgm nvidia-imex nvidia-persistenced; do + systemctlDisableAndStop "${service}" || return 1 + if systemctl is-active --quiet "${service}" || systemctl is-enabled --quiet "${service}"; then + echo "Grace Blackwell service ${service} remains active or enabled" >&2 + return 1 + fi + done + + modules=$(lsmod) || return 1 + if grep -q '^nvidia' <<< "${modules}" && systemctl cat openibd &>/dev/null; then + systemctl_stop 20 5 25 openibd || { + echo "Grace Blackwell service openibd could not be stopped" >&2 + return 1 + } + if systemctl is-active --quiet openibd; then + echo "Grace Blackwell service openibd remains active" >&2 + return 1 + fi + modules=$(lsmod) || return 1 + fi + for module in nvidia_peermem nvidia_uvm nvidia_drm nvidia_modeset nvidia; do + grep -q "^${module}[[:space:]]" <<< "${modules}" || continue + rmmod "${module}" || { echo "Failed to unload Grace Blackwell module ${module}" >&2; return 1; } + modules=$(lsmod) || return 1 + grep -q "^${module}[[:space:]]" <<< "${modules}" && { + echo "Grace Blackwell module ${module} remains loaded after removal" >&2 + return 1 + } + done + grep -q '^nvidia' <<< "${modules}" && { echo "NVIDIA module remains loaded after Grace Blackwell teardown" >&2; return 1; } + + if [ "${#purge_packages[@]}" -gt 0 ]; then + apt_get_purge 10 5 300 "${purge_packages[@]}" || { + echo "Grace Blackwell driver package purge failed" >&2 + return 1 + } + fi + + if [ "${pending}" != true ]; then + printf '%s\n' "${wave}" > "${pending_file}" || { + rm -f "${pending_file}" + echo "Unable to preserve Grace Blackwell teardown provenance" >&2 + return 1 + } + fi + rm -f "${peermem_config}" "${nvidia_config}" "${nouveau_config}" || return 1 + for package in "${peermem_config}" "${nvidia_config}" "${nouveau_config}"; do + if [ -e "${package}" ] || [ -L "${package}" ]; then + echo "Grace Blackwell module configuration ${package} remains after removal" >&2 + return 1 + fi + done + update-initramfs -u || return 1 + rm -f "${pending_file}" || return 1 + [ ! -e "${pending_file}" ] || return 1 +} + +cleanUpGPUDriversForBasePrep() { + local packageName marker="${GPU_DKMS_MARKER_FILE:-/opt/azure/aks-gpu/dkms-marker}" remaining_package + local peermem_config="${GB_NVIDIA_PEERMEM_CONFIG_FILE:-/etc/modules-load.d/nvidia-peermem.conf}" + 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 + + if [ -e "${pending_file}" ]; then + IFS= read -r pending_wave < "${pending_file}" || { + echo "GPU basePrep cleanup incomplete: invalid Grace Blackwell pending marker" >&2 + return 1 + } + case "${pending_wave}" in + versions-wave2) ;; + versions-wave1) + bom="${legacy_bom}" + wave=versions-wave1 + ;; + *) + echo "GPU basePrep cleanup incomplete: invalid Grace Blackwell pending marker" >&2 + return 1 + ;; + esac + gb_layout=true + pending_teardown=true + elif [ -e "${peermem_config}" ]; then + gb_layout=true + elif [ -e "${legacy_bom}" ] && [ ! -e "${bom}" ]; then + bom="${legacy_bom}" + wave=versions-wave1 + gb_layout=true + fi + + entries=$(graceBlackwellDriverBOMEntries "${bom}" "${wave}" 2>/dev/null) || entries="" + if [ "${gb_layout}" = true ]; then + if [ -z "${entries}" ]; then + echo "GPU basePrep cleanup incomplete: Grace Blackwell BOM is unavailable" >&2 + return 1 + fi + cleanUpGraceBlackwellGPUDriver "${entries}" "${bom}" "${wave}" "${pending_teardown}" || { + echo "GPU basePrep cleanup incomplete: Grace Blackwell driver teardown failed" >&2 + return 1 + } + elif remaining_package=$(graceBlackwellPackageManagedDriverRemains); then + echo "GPU basePrep cleanup incomplete: NVIDIA package ${remaining_package} has no verifiable GB BOM ownership" >&2 + return 1 + fi + + if remaining_package=$(graceBlackwellPackageManagedDriverRemains); then + echo "GPU basePrep cleanup incomplete: NVIDIA package ${remaining_package} remains after Grace Blackwell teardown" >&2 + return 1 + fi + + cleanUpGPUDrivers + if prebakedGPUDriverArtifactsRemain; then + echo "GPU basePrep cleanup incomplete: pre-baked GPU driver artifacts remain" >&2 + return 1 + fi + if [ -e "${GPU_DEST}" ] || [ -e /opt/gpu ] || [ -e "${marker}" ]; then + echo "GPU basePrep cleanup incomplete: AKS driver artifacts remain" >&2 + return 1 + fi + for packageName in $(managedGPUPackageList); do + if [ -e "/opt/${packageName}" ]; then + echo "GPU basePrep cleanup incomplete: ${packageName} cache remains" >&2 + return 1 + fi + done +} diff --git a/pkg/agent/gpu_pis_delivery_test.go b/pkg/agent/gpu_pis_delivery_test.go new file mode 100644 index 00000000000..74f4a990a71 --- /dev/null +++ b/pkg/agent/gpu_pis_delivery_test.go @@ -0,0 +1,84 @@ +package agent + +import ( + "encoding/base64" + "os/exec" + "path/filepath" + "testing" + + "github.com/Azure/agentbaker/pkg/agent/datamodel" + "github.com/stretchr/testify/require" + "gopkg.in/yaml.v3" +) + +func TestPISGPUCustomDataFitsAndUsesBakedHelper(t *testing.T) { + for _, mode := range []struct { + preProvisionOnly bool + phase1 bool + }{ + {preProvisionOnly: true, phase1: true}, + {preProvisionOnly: true}, + {}, + } { + config := newNodeCustomDataRenderConfig(datamodel.AKSUbuntuContainerd2404Gen2) + pool := config.AgentPoolProfile + pool.Name, pool.VMSize = "agent2", "Standard_NC4as_T4_v3" + pool.StorageProfile = "ManagedDisks" + pool.AvailabilityProfile = datamodel.VirtualMachineScaleSets + pool.VnetSubnetID = "/subscriptions/test/resourceGroups/test/providers/Microsoft.Network/virtualNetworks/test/subnets/default" + properties := config.ContainerService.Properties + properties.LinuxProfile = &datamodel.LinuxProfile{AdminUsername: "azureuser"} + properties.LinuxProfile.SSH.PublicKeys = []datamodel.PublicKey{{KeyData: "fixture-key"}} + properties.ServicePrincipalProfile = &datamodel.ServicePrincipalProfile{ClientID: "fixture", Secret: "fixture"} + config.TenantID, config.SubscriptionID = "tenantID", "subID" + config.ResourceGroupName, config.UserAssignedIdentityClientID = "resourceGroupName", "userAssignedID" + config.PrimaryScaleSetName = "aks-agent2-36873793-vmss" + config.KubeletConfig = map[string]string{"--cluster-domain": "cluster.local", "--cluster-dns": "10.0.0.10", "--cloud-provider": "azure", "--max-pods": "110"} + config.PreProvisionOnly = mode.preProvisionOnly + config.EnableScriptlessCSECmd = mode.phase1 + + payload := InitializeTemplateGenerator().getLinuxNodeBootstrappingPayload(config) + compressed, err := base64.StdEncoding.DecodeString(payload) + require.NoError(t, err) + decoded, err := getGzipDecodedValue(compressed) + require.NoError(t, err) + require.LessOrEqual(t, len(compressed), 65535) + require.LessOrEqual(t, len(payload), 87380) + + var document cloudInit + require.NoError(t, yaml.Unmarshal(decoded, &document)) + var installer bool + for _, file := range document.WriteFiles { + require.NotEqual(t, "/opt/azure/containers/provision_installs_ubuntu_gpu.sh", file.Path) + if file.Path == "/opt/azure/containers/provision_installs_distro.sh" { + installer = true + script, err := getGzipDecodedValue([]byte(file.Content)) + require.NoError(t, err) + require.Contains(t, string(script), "cleanUpGPUDrivers()") + } + } + if mode.phase1 { + require.Equal(t, "/opt/azure/containers/scriptless-cse-overrides.txt", document.WriteFiles[0].Path) + } else { + require.True(t, installer) + } + } +} + +func TestGPUAMARenderingKeepsAzureLinuxImplementation(t *testing.T) { + config := newNodeCustomDataRenderConfig(datamodel.AKSAzureLinuxV3Gen2) + payload := getBase64EncodedGzippedCustomScript(kubernetesCSEConfigGPU, config) + compressed, err := base64.StdEncoding.DecodeString(payload) + require.NoError(t, err) + rendered, err := getGzipDecodedValue(compressed) + require.NoError(t, err) + + raw := exec.Command("bash", "-c", `source "$1"; declare -f setupAmdAma dnf_install_amd_ama_core_package`, + "gpu-raw-test", filepath.Join("..", "..", "parts", kubernetesCSEConfigGPU)) + rawFunctions, err := raw.CombinedOutput() + require.NoError(t, err, string(rawFunctions)) + renderedCmd := exec.Command("bash", "-c", string(rendered)+"\ndeclare -f setupAmdAma dnf_install_amd_ama_core_package") + renderedFunctions, err := renderedCmd.CombinedOutput() + require.NoError(t, err, string(renderedFunctions)) + require.Equal(t, string(rawFunctions), string(renderedFunctions)) +} diff --git a/spec/parts/linux/cloud-init/artifacts/cse_install_mariner_spec.sh b/spec/parts/linux/cloud-init/artifacts/cse_install_mariner_spec.sh index 3e3962c4dc8..4a17d7141dc 100644 --- a/spec/parts/linux/cloud-init/artifacts/cse_install_mariner_spec.sh +++ b/spec/parts/linux/cloud-init/artifacts/cse_install_mariner_spec.sh @@ -28,6 +28,32 @@ Describe 'cse_install_mariner.sh' Include "./parts/linux/cloud-init/artifacts/cse_install.sh" Include "./parts/linux/cloud-init/artifacts/mariner/cse_install_mariner.sh" + Describe 'cleanUpGPUDriversForBasePrep' + GPU_DEST="${PWD}/.shellspec-nonexistent-gpu-dest" + managedGPUPackageList() { echo "driver-cache"; } + getPackageCacheDir() { echo "$gpu_test_root/driver-cache"; } + setup_gpu_cache() { gpu_test_root="${PWD}/.shellspec-gpu-cache-$$"; mkdir -p "$gpu_test_root"; } + teardown_gpu_cache() { rm -rf "$gpu_test_root"; } + BeforeEach 'setup_gpu_cache' + AfterEach 'teardown_gpu_cache' + + It 'fails when an AKS cache survives cleanup' + mkdir -p "$gpu_test_root/driver-cache" + cleanUpGPUDrivers() { :; } + When call cleanUpGPUDriversForBasePrep + The status should be failure + The stderr should include "driver-cache cache remains" + End + + It 'succeeds when AKS caches are removed' + mkdir -p "$gpu_test_root/driver-cache" + cleanUpGPUDrivers() { rm -rf "$gpu_test_root/driver-cache"; } + When call cleanUpGPUDriversForBasePrep + The status should be success + The path "$gpu_test_root/driver-cache" should not be exist + End + End + Describe 'installDeps' It 'installs the required packages with installDeps for Mariner 2.0' OS_VERSION="2.0" @@ -772,7 +798,7 @@ EOF End Describe 'installPackageFromCache version matching' - rpm_version_cache="/tmp/shellspec-rpm-version-cache-$$" + rpm_version_cache="${PWD}/.shellspec-rpm-version-cache-$$" setup_version_cache() { RPM_PACKAGE_CACHE_BASE_DIR="$rpm_version_cache" diff --git a/spec/parts/linux/cloud-init/artifacts/cse_install_ubuntu_gpu_pis_spec.sh b/spec/parts/linux/cloud-init/artifacts/cse_install_ubuntu_gpu_pis_spec.sh new file mode 100644 index 00000000000..de9fb823273 --- /dev/null +++ b/spec/parts/linux/cloud-init/artifacts/cse_install_ubuntu_gpu_pis_spec.sh @@ -0,0 +1,75 @@ +#!/bin/bash + +Describe 'Ubuntu PIS GPU basePrep cleanup' + Include "./parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu_gpu.sh" + fixture="${PWD}/.shellspec-gpu-cleanup-$$" + GPU_DEST="${fixture}/gpu-dest" + managedGPUPackageList() { :; } + + setup_cleanup() { + mkdir -p "${fixture}" + GPU_DKMS_MARKER_FILE="${fixture}/marker" + GB_NVIDIA_PEERMEM_CONFIG_FILE="${fixture}/peermem.conf" + GB_NVIDIA_MODPROBE_CONFIG_FILE="${fixture}/nvidia.conf" + GB_NOUVEAU_MODPROBE_CONFIG_FILE="${fixture}/nouveau.conf" + GB_DRIVER_CLEANUP_PENDING_FILE="${fixture}/cleanup.pending" + } + + BeforeEach 'setup_cleanup' + AfterEach 'rm -rf "$fixture"' + + It 'fails CUDA cleanup when the prebake marker remains' + touch "${GPU_DKMS_MARKER_FILE}" + dpkg-query() { :; } + lsmod() { :; } + cleanUpGPUDrivers() { :; } + prebakedGPUDriverArtifactsRemain() { return 1; } + When call cleanUpGPUDriversForBasePrep + The status should be failure + The stderr should include "AKS driver artifacts remain" + The path "${GPU_DKMS_MARKER_FILE}" should be exist + End + + It 'refuses GB teardown when installed NVIDIA packages are not BOM-owned' + touch "${GPU_DKMS_MARKER_FILE}" + graceBlackwellInstalledDriverPackages() { printf '%s\n' 'nvidia-driver-580-open=580.159.04-1ubuntu1'; } + dpkg-query() { printf 'nvidia-driver-580-open\tinstalled\nnvidia-driver-535\tinstalled\n'; } + When call cleanUpGraceBlackwellGPUDriver 'nvidia-driver-580-open=580.159.04-1ubuntu1' "${fixture}/bom" versions-wave2 + The status should be failure + The stderr should include 'NVIDIA package nvidia-driver-535 is not owned by the Grace Blackwell BOM' + The path "${GPU_DKMS_MARKER_FILE}" should be exist + End + + It 'retains GB cleanup provenance after initramfs failure and retries' + cp parts/linux/cloud-init/artifacts/ubuntu/gb/nvidia-peermem.conf "${GB_NVIDIA_PEERMEM_CONFIG_FILE}" + cp parts/linux/cloud-init/artifacts/ubuntu/gb/modprobe-nvidia-parameters.conf "${GB_NVIDIA_MODPROBE_CONFIG_FILE}" + printf 'blacklist nouveau\noptions nouveau modeset=0\n' > "${GB_NOUVEAU_MODPROBE_CONFIG_FILE}" + touch "${GPU_DKMS_MARKER_FILE}" + graceBlackwellDriverBOMEntries() { printf '%s\n' 'nvidia-driver-580-open=580.159.04-1ubuntu1'; } + graceBlackwellInstalledDriverPackages() { :; } + graceBlackwellPackageManagedDriverRemains() { return 1; } + systemctlDisableAndStop() { :; } + systemctl() { return 1; } + lsmod() { printf 'Module Size Used by\n'; } + update_attempts=0 + update-initramfs() { + update_attempts=$((update_attempts + 1)) + [ "${update_attempts}" -gt 1 ] + } + cleanUpGPUDrivers() { rm -f "${GPU_DKMS_MARKER_FILE}"; } + prebakedGPUDriverArtifactsRemain() { return 1; } + retry_cleanup() { + local status=0 pending_wave + cleanUpGPUDriversForBasePrep > "${fixture}/first-attempt.log" 2>&1 && status=1 + grep -q 'GPU basePrep cleanup incomplete: Grace Blackwell driver teardown failed' "${fixture}/first-attempt.log" || status=1 + IFS= read -r pending_wave < "${GB_DRIVER_CLEANUP_PENDING_FILE}" || status=1 + [ "${pending_wave}" = versions-wave2 ] || status=1 + cleanUpGPUDriversForBasePrep || status=1 + [ ! -e "${GB_DRIVER_CLEANUP_PENDING_FILE}" ] || status=1 + [ ! -e "${GPU_DKMS_MARKER_FILE}" ] || status=1 + return "${status}" + } + When call retry_cleanup + The status should be success + End +End diff --git a/spec/parts/linux/cloud-init/artifacts/cse_main_gpu_pis_spec.sh b/spec/parts/linux/cloud-init/artifacts/cse_main_gpu_pis_spec.sh new file mode 100644 index 00000000000..a828b7c0391 --- /dev/null +++ b/spec/parts/linux/cloud-init/artifacts/cse_main_gpu_pis_spec.sh @@ -0,0 +1,82 @@ +#!/bin/bash + +Describe 'GPU cleanup across PIS phases' + CSE_MAIN="./parts/linux/cloud-init/artifacts/cse_main.sh" + fixture="${PWD}/.shellspec-cse-main-gpu-pis-$$" + + setup_driver() { + driver_root="${fixture}/driver" + marker="${fixture}/base_prep.complete" + mkdir -p "${driver_root}/usr/local/nvidia" + GPU_NODE=false + PRE_PROVISION_ONLY=false + skip_nvidia_driver_install=false + BASE_PREP_RAN=false + } + + run_phases() { + local dispatch cleanup + dispatch=$(awk ' + /^if \[ ! -f \/opt\/azure\/containers\/base_prep.complete \]; then/ { reading = 1 } + reading && /^echo "Custom script finished."/ { exit } + reading { print } + ' "$CSE_MAIN") + cleanup=$(sed -n '/ # Clean up GPU drivers if not a GPU node/,/ fi/p' "$CSE_MAIN") + dispatch=${dispatch//\/opt\/azure\/containers\/base_prep.complete/$marker} + basePrep() { echo "basePrep"; } + nodePrep() { eval "$cleanup"; } + logs_to_events() { shift; "$@"; } + cleanUpGPUDrivers() { + rm -rf "${driver_root}/usr/local/nvidia" + touch "${fixture}/cleanup-called" + } + eval "$dispatch" + } + + run_cse_marker() { + local marker_script + PRE_PROVISION_ONLY=true + marker_script=$(awk ' + /^if \[ "\$\{PRE_PROVISION_ONLY\}" = "true" \]; then/ { reading = 1 } + reading && /^if \[ "\$EXIT_CODE" -ne 0 \]; then/ { exit } + reading { print } + ' ./parts/linux/cloud-init/artifacts/cse_start.sh) + marker_script=${marker_script//\/opt\/azure\/containers/$fixture} + marker_script=${marker_script//\/var\/log\/azure\/cluster-provision.log/$fixture/provision.log} + eval "$marker_script" + } + + BeforeEach 'setup_driver' + AfterEach 'rm -rf "$fixture"' + + It 'preserves customer drivers on a PIS-derived opt-out node' + touch "$marker" + When call run_phases + The status should be success + The output should include "Skipping basePrep" + The path "$driver_root/usr/local/nvidia" should be exist + The path "$fixture/cleanup-called" should not be exist + End + + It 'cleans CUDA drivers only after basePrep ran in this execution' + When call run_phases + The status should be success + The output should include "basePrep" + The path "$driver_root/usr/local/nvidia" should not be exist + The path "$fixture/cleanup-called" should be exist + End + + It 'does not write the bake marker after failed basePrep' + EXIT_CODE=224 + When call run_cse_marker + The status should be success + The path "$fixture/base_prep.complete" should not be exist + End + + It 'writes the bake marker after successful basePrep' + EXIT_CODE=0 + When call run_cse_marker + The status should be success + The path "$fixture/base_prep.complete" should be exist + End +End diff --git a/spec/vhdbuilder/packer/pis_gpu_helper_spec.sh b/spec/vhdbuilder/packer/pis_gpu_helper_spec.sh new file mode 100644 index 00000000000..cb556b84757 --- /dev/null +++ b/spec/vhdbuilder/packer/pis_gpu_helper_spec.sh @@ -0,0 +1,19 @@ +#!/bin/bash + +Describe 'Ubuntu PIS GPU cleanup helper VHD installation' + Include './vhdbuilder/packer/packer_source.sh' + + isUbuntu() { [ "$1" = UBUNTU ]; } + cpAndMode() { printf '%s\n' "$*"; } + + It 'copies the helper into Ubuntu images' + OS=UBUNTU + When call copyUbuntuGPUCleanupHelper + The output should equal '/home/packer/provision_installs_ubuntu_gpu.sh /opt/azure/containers/provision_installs_ubuntu_gpu.sh 0744' + End + + It 'calls the helper copy during Packer file installation' + When call sed -n '/^copyPackerFiles()/,/^}/p' vhdbuilder/packer/packer_source.sh + The output should include ' copyUbuntuGPUCleanupHelper' + End +End diff --git a/vhdbuilder/packer/packer_source.sh b/vhdbuilder/packer/packer_source.sh index 44a32731653..07a7bb93ebf 100644 --- a/vhdbuilder/packer/packer_source.sh +++ b/vhdbuilder/packer/packer_source.sh @@ -189,6 +189,12 @@ configureServiceExecutionTelemetryForSystemdVersion() { done } +copyUbuntuGPUCleanupHelper() { + if isUbuntu "$OS"; then + cpAndMode /home/packer/provision_installs_ubuntu_gpu.sh /opt/azure/containers/provision_installs_ubuntu_gpu.sh 0744 + fi +} + copyPackerFiles() { SYSCTL_CONFIG_SRC=/home/packer/sysctl-d-60-CIS.conf SYSCTL_CONFIG_DEST=/etc/sysctl.d/60-CIS.conf @@ -487,6 +493,7 @@ copyPackerFiles() { CSE_INSTALL_DISTRO_SRC=/home/packer/provision_installs_distro.sh CSE_INSTALL_DISTRO_DEST=/opt/azure/containers/provision_installs_distro.sh cpAndMode $CSE_INSTALL_DISTRO_SRC $CSE_INSTALL_DISTRO_DEST 0744 + copyUbuntuGPUCleanupHelper CSE_HELPERS_SRC=/home/packer/provision_source.sh CSE_HELPERS_DEST=/opt/azure/containers/provision_source.sh diff --git a/vhdbuilder/packer/vhd-image-builder-arm64-gb.json b/vhdbuilder/packer/vhd-image-builder-arm64-gb.json index a74db09e3c7..231d5382d5c 100644 --- a/vhdbuilder/packer/vhd-image-builder-arm64-gb.json +++ b/vhdbuilder/packer/vhd-image-builder-arm64-gb.json @@ -169,6 +169,11 @@ "source": "parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh", "destination": "/home/packer/provision_installs_distro.sh" }, + { + "type": "file", + "source": "parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu_gpu.sh", + "destination": "/home/packer/provision_installs_ubuntu_gpu.sh" + }, { "type": "file", "source": "parts/linux/cloud-init/artifacts/cse_helpers.sh", diff --git a/vhdbuilder/packer/vhd-image-builder-arm64-gen2.json b/vhdbuilder/packer/vhd-image-builder-arm64-gen2.json index d60eba4569f..e3456826cf0 100644 --- a/vhdbuilder/packer/vhd-image-builder-arm64-gen2.json +++ b/vhdbuilder/packer/vhd-image-builder-arm64-gen2.json @@ -154,6 +154,11 @@ "source": "parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh", "destination": "/home/packer/provision_installs_distro.sh" }, + { + "type": "file", + "source": "parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu_gpu.sh", + "destination": "/home/packer/provision_installs_ubuntu_gpu.sh" + }, { "type": "file", "source": "parts/linux/cloud-init/artifacts/cse_helpers.sh", diff --git a/vhdbuilder/packer/vhd-image-builder-base.json b/vhdbuilder/packer/vhd-image-builder-base.json index c1a7cfe1c23..d2be462828f 100644 --- a/vhdbuilder/packer/vhd-image-builder-base.json +++ b/vhdbuilder/packer/vhd-image-builder-base.json @@ -157,6 +157,11 @@ "source": "parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh", "destination": "/home/packer/provision_installs_distro.sh" }, + { + "type": "file", + "source": "parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu_gpu.sh", + "destination": "/home/packer/provision_installs_ubuntu_gpu.sh" + }, { "type": "file", "source": "parts/linux/cloud-init/artifacts/cse_helpers.sh", diff --git a/vhdbuilder/packer/vhd-image-builder-cvm.json b/vhdbuilder/packer/vhd-image-builder-cvm.json index 228b2187700..d0f9387317c 100644 --- a/vhdbuilder/packer/vhd-image-builder-cvm.json +++ b/vhdbuilder/packer/vhd-image-builder-cvm.json @@ -160,6 +160,11 @@ "source": "parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh", "destination": "/home/packer/provision_installs_distro.sh" }, + { + "type": "file", + "source": "parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu_gpu.sh", + "destination": "/home/packer/provision_installs_ubuntu_gpu.sh" + }, { "type": "file", "source": "parts/linux/cloud-init/artifacts/cse_helpers.sh",