diff --git a/aks-node-controller/parser/parser.go b/aks-node-controller/parser/parser.go index 0eebd3e542b..cf317936fa5 100644 --- a/aks-node-controller/parser/parser.go +++ b/aks-node-controller/parser/parser.go @@ -197,6 +197,7 @@ func getCSEEnv(config *aksnodeconfigv1.Configuration) map[string]string { "SKIP_WAAGENT_HOLD": "true", "NETWORK_ISOLATED_CLUSTER_TEST_MODE": "false", // temp: needs to be added to config "STANDARD_SECONDARY_NIC_COUNT": fmt.Sprintf("%d", config.GetNetworkConfig().GetStandardSecondaryNicCount()), + "ENABLE_MANAGED_GPU_DRA": "false", // TODO: add protobuf field } for i, cert := range config.CustomCaCerts { diff --git a/e2e/scenario_gpu_managed_experience_test.go b/e2e/scenario_gpu_managed_experience_test.go index 07d881bf019..d8b57067e2c 100644 --- a/e2e/scenario_gpu_managed_experience_test.go +++ b/e2e/scenario_gpu_managed_experience_test.go @@ -783,3 +783,41 @@ func Test_Ubuntu2404_NvidiaDevicePluginRunning_MIG_Mixed(t *testing.T) { }, }) } + +func Test_Ubuntu2404_DraDriverNvidiaGpuRunning(t *testing.T) { + RunScenario(t, &Scenario{ + Description: "Tests DRA driver works on Ubuntu 24.04 VHD with containerd v2", + Tags: Tags{ + GPU: true, + }, + + Config: Config{ + Cluster: ClusterKubenet, + SkipScriptlessNBC: true, + VHD: config.VHDUbuntu2404Gen2Containerd, + BootstrapConfigMutator: func(_ *Cluster, nbc *datamodel.NodeBootstrappingConfiguration) { + nbc.AgentPoolProfile.VMSize = "Standard_NV6ads_A10_v5" + nbc.ConfigGPUDriverIfNeeded = true + nbc.EnableNvidia = true + nbc.EnableManagedGPUDRA = true + }, + VMConfigMutator: func(vmss *armcompute.VirtualMachineScaleSet) { + vmss.SKU.Name = to.Ptr("Standard_NV6ads_A10_v5") + + // Enable the AKS VM extension for GPU nodes + extension, err := createVMExtensionLinuxAKSNode(t.Context(), vmss.Location) + require.NoError(t, err, "creating AKS VM extension") + vmss.Properties = addVMExtensionToVMSS(vmss.Properties, extension) + }, + Validator: func(ctx context.Context, s *Scenario) { + containerdVersions := components.GetExpectedPackageVersions("containerd", "ubuntu", "r2404") + runcVersions := components.GetExpectedPackageVersions("runc", "ubuntu", "r2404") + ValidateContainerd2Properties(ctx, s, containerdVersions) + ValidateRuncVersion(ctx, s, runcVersions) + ValidateContainerRuntimePlugins(ctx, s) + ValidateDraDriverNvidiaGpuServiceRunning(ctx, s) + ValidateDRAWorkloadSchedulable(ctx, s) + }, + }, + }) +} diff --git a/e2e/validators.go b/e2e/validators.go index 83683053887..9e66c0f6ee2 100644 --- a/e2e/validators.go +++ b/e2e/validators.go @@ -28,6 +28,8 @@ import ( "github.com/stretchr/testify/require" certv1 "k8s.io/api/certificates/v1" corev1 "k8s.io/api/core/v1" + resourcev1 "k8s.io/api/resource/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" "k8s.io/apimachinery/pkg/api/resource" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/util/wait" @@ -3111,3 +3113,97 @@ func ValidateSecondaryNICDualStack(ctx context.Context, s *Scenario, ifaceName s require.Contains(s.T, result.stdout, "scope global", "expected interface %s to have a global IPv6 address (not just link-local), got:\n%s", ifaceName, result.stdout) } + +func ValidateDraDriverNvidiaGpuServiceRunning(ctx context.Context, s *Scenario) { + s.T.Helper() + s.T.Logf("validating DRA driver NVIDIA GPU systemd service is running") + + command := []string{ + "set -ex", + "systemctl is-active dra-driver-nvidia-gpu.service", + "systemctl is-enabled dra-driver-nvidia-gpu.service", + } + execScriptOnVMForScenarioValidateExitCode(ctx, s, strings.Join(command, "\n"), 0, "DRA driver NVIDIA GPU systemd service should be active and enabled") +} + +func ValidateDRAWorkloadSchedulable(ctx context.Context, s *Scenario) { + s.T.Helper() + s.T.Logf("validating that DRA workloads can be scheduled") + + time.Sleep(20 * time.Second) // Same delay as existing GPU tests + + baseName := strings.ToLower(s.Runtime.VM.KubeName) + if len(baseName) > 40 { + baseName = baseName[:40] + } + baseName = strings.TrimRight(baseName, "-") + deviceClassName := fmt.Sprintf("gpu-nvidia-%s", baseName) + claimName := fmt.Sprintf("single-gpu-%s", baseName) + podClaimRefName := "gpu-claim" + + _, err := s.Runtime.Kube.Typed.ResourceV1().DeviceClasses().Create(ctx, &resourcev1.DeviceClass{ + ObjectMeta: metav1.ObjectMeta{ + Name: deviceClassName, + }, + Spec: resourcev1.DeviceClassSpec{}, + }, metav1.CreateOptions{}) + require.Truef(s.T, err == nil || apierrors.IsAlreadyExists(err), "failed to create DeviceClass %q: %v", deviceClassName, err) + + _, err = s.Runtime.Kube.Typed.ResourceV1().ResourceClaims("default").Create(ctx, &resourcev1.ResourceClaim{ + ObjectMeta: metav1.ObjectMeta{ + Name: claimName, + Namespace: "default", + }, + Spec: resourcev1.ResourceClaimSpec{ + Devices: resourcev1.DeviceClaim{ + Requests: []resourcev1.DeviceRequest{ + { + Name: "gpu", + Exactly: &resourcev1.ExactDeviceRequest{ + DeviceClassName: deviceClassName, + }, + }, + }, + }, + }, + }, metav1.CreateOptions{}) + require.Truef(s.T, err == nil || apierrors.IsAlreadyExists(err), "failed to create ResourceClaim %q: %v", claimName, err) + + // Create a DRA test pod that consumes the ResourceClaim. + pod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: fmt.Sprintf("%s-dra-test", s.Runtime.VM.KubeName), + Namespace: "default", + }, + Spec: corev1.PodSpec{ + ResourceClaims: []corev1.PodResourceClaim{ + { + Name: podClaimRefName, + ResourceClaimName: &claimName, + }, + }, + Containers: []corev1.Container{ + { + Name: "dra-test-container", + Image: "mcr.microsoft.com/azuredocs/samples-tf-mnist-demo:gpu", + Args: []string{ + "--max-steps", "1", + }, + Resources: corev1.ResourceRequirements{ + Claims: []corev1.ResourceClaim{ + { + Name: podClaimRefName, + }, + }, + }, + }, + }, + NodeSelector: map[string]string{ + "kubernetes.io/hostname": s.Runtime.VM.KubeName, + }, + }, + } + ValidatePodRunning(ctx, s, pod) + + s.T.Logf("GPU workload is schedulable and runs successfully") +} diff --git a/parts/linux/cloud-init/artifacts/cse_cmd.sh b/parts/linux/cloud-init/artifacts/cse_cmd.sh index 52ffb72de76..12ef76a049d 100644 --- a/parts/linux/cloud-init/artifacts/cse_cmd.sh +++ b/parts/linux/cloud-init/artifacts/cse_cmd.sh @@ -80,6 +80,7 @@ CONFIG_GPU_DRIVER_IF_NEEDED={{GetVariable "configGPUDriverIfNeeded"}} ENABLE_GPU_DEVICE_PLUGIN_IF_NEEDED={{GetVariable "enableGPUDevicePluginIfNeeded"}} MANAGED_GPU_EXPERIENCE_AFEC_ENABLED="{{IsManagedGPUExperienceAFECEnabled}}" ENABLE_MANAGED_GPU="{{IsEnableManagedGPU}}" +ENABLE_MANAGED_GPU_DRA="{{IsEnableManagedGPUDRA}}" NVIDIA_MIG_STRATEGY="{{GetMigStrategy}}" CREDENTIAL_PROVIDER_DOWNLOAD_URL={{GetParameter "linuxCredentialProviderURL"}} CONTAINERD_VERSION={{GetParameter "containerdVersion"}} diff --git a/parts/linux/cloud-init/artifacts/cse_config.sh b/parts/linux/cloud-init/artifacts/cse_config.sh index ff146905b4b..28288dc4062 100755 --- a/parts/linux/cloud-init/artifacts/cse_config.sh +++ b/parts/linux/cloud-init/artifacts/cse_config.sh @@ -1825,6 +1825,10 @@ configureManagedGPUExperience() { if [ "${GPU_NODE}" != "true" ] || [ "${skip_nvidia_driver_install}" = "true" ]; then return fi + if [ "${ENABLE_MANAGED_GPU_EXPERIENCE}" = "true" ] && [ "${ENABLE_MANAGED_GPU_EXPERIENCE_DRA}" = "true" ]; then + echo "Error: ENABLE_MANAGED_GPU_EXPERIENCE and ENABLE_MANAGED_GPU_EXPERIENCE_DRA cannot both be true" + exit $ERR_ENABLE_MANAGED_GPU_EXPERIENCE + fi local managed_gpu_marker="/opt/azure/containers/managed-gpu-experience.enabled" if [ "${ENABLE_MANAGED_GPU_EXPERIENCE}" = "true" ]; then logs_to_events "AKS.CSE.installNvidiaManagedExpPkgFromCache" "installNvidiaManagedExpPkgFromCache" || exit $ERR_NVIDIA_DCGM_INSTALL @@ -1832,10 +1836,19 @@ configureManagedGPUExperience() { addKubeletNodeLabel "kubernetes.azure.com/dcgm-exporter=enabled" mkdir -p "$(dirname "${managed_gpu_marker}")" touch "${managed_gpu_marker}" + elif [ "${ENABLE_MANAGED_GPU_EXPERIENCE_DRA}" = "true" ]; then + logs_to_events "AKS.CSE.installNvidiaManagedExpPkgFromCache" "installNvidiaManagedExpPkgFromCache" || exit $ERR_NVIDIA_DCGM_INSTALL + # defer startNvidiaManagedExpServices() after kubelet starts + addKubeletNodeLabel "kubernetes.azure.com/dcgm-exporter=enabled" + mkdir -p "$(dirname "${managed_gpu_marker}")" + touch "${managed_gpu_marker}" else # EnableManagedGPUExperience is mutable, so services may have been # installed on a previous CSE run. Stop them if they exist. + # systemctlDisableAndStop check if the service exists before attempting to stop it, + # so this is safe to call even if the services were never installed. logs_to_events "AKS.CSE.stop.nvidia-device-plugin" "systemctlDisableAndStop nvidia-device-plugin" + logs_to_events "AKS.CSE.stop.dra-driver-nvidia-gpu" "systemctlDisableAndStop dra-driver-nvidia-gpu" logs_to_events "AKS.CSE.stop.nvidia-dcgm" "systemctlDisableAndStop nvidia-dcgm" logs_to_events "AKS.CSE.stop.nvidia-dcgm-exporter" "systemctlDisableAndStop nvidia-dcgm-exporter" rm -f "${managed_gpu_marker}" @@ -1843,46 +1856,65 @@ configureManagedGPUExperience() { } startNvidiaManagedExpServices() { - # 1. Start the nvidia-device-plugin service. - # Create systemd override directory to configure device plugin - NVIDIA_DEVICE_PLUGIN_OVERRIDE_DIR="/etc/systemd/system/nvidia-device-plugin.service.d" - mkdir -p "${NVIDIA_DEVICE_PLUGIN_OVERRIDE_DIR}" - - if [ "${MIG_NODE}" = "true" ]; then - # Configure with MIG strategy for MIG nodes. - # MIG strategy controls how nvidia-device-plugin exposes MIG instances to Kubernetes: - # - "single": All MIG devices exposed as generic nvidia.com/gpu resources - # - "mixed": MIG devices exposed with specific types like nvidia.com/mig-1g.5gb - # - # We only use "mixed" when explicitly specified via NVIDIA_MIG_STRATEGY. - # Otherwise, we default to "single" which is the safer/simpler option. - # Note: NVIDIA_MIG_STRATEGY values from RP are "None", "Single", "Mixed". - # "None" and "Single" both result in using the "single" strategy. - if [ "${NVIDIA_MIG_STRATEGY}" = "Mixed" ]; then - MIG_STRATEGY_FLAG="--mig-strategy mixed" - else - # Default to "single" for "Single", "None", empty, or any other value - MIG_STRATEGY_FLAG="--mig-strategy single" - fi + # 1. Start the nvidia-device-plugin service or dra-driver-nvidia-gpu service. + if [ "${ENABLE_MANAGED_GPU_EXPERIENCE}" = "true" ]; then + # Create systemd override directory to configure device plugin + NVIDIA_DEVICE_PLUGIN_OVERRIDE_DIR="/etc/systemd/system/nvidia-device-plugin.service.d" + mkdir -p "${NVIDIA_DEVICE_PLUGIN_OVERRIDE_DIR}" + + if [ "${MIG_NODE}" = "true" ]; then + # Configure with MIG strategy for MIG nodes. + # MIG strategy controls how nvidia-device-plugin exposes MIG instances to Kubernetes: + # - "single": All MIG devices exposed as generic nvidia.com/gpu resources + # - "mixed": MIG devices exposed with specific types like nvidia.com/mig-1g.5gb + # + # We only use "mixed" when explicitly specified via NVIDIA_MIG_STRATEGY. + # Otherwise, we default to "single" which is the safer/simpler option. + # Note: NVIDIA_MIG_STRATEGY values from RP are "None", "Single", "Mixed". + # "None" and "Single" both result in using the "single" strategy. + if [ "${NVIDIA_MIG_STRATEGY}" = "Mixed" ]; then + MIG_STRATEGY_FLAG="--mig-strategy mixed" + else + # Default to "single" for "Single", "None", empty, or any other value + MIG_STRATEGY_FLAG="--mig-strategy single" + fi - tee "${NVIDIA_DEVICE_PLUGIN_OVERRIDE_DIR}/10-device-plugin-config.conf" > /dev/null < /dev/null < /dev/null <<'EOF' + else + # Configure with pass-device-specs for non-MIG nodes + tee "${NVIDIA_DEVICE_PLUGIN_OVERRIDE_DIR}/10-device-plugin-config.conf" > /dev/null <<'EOF' [Service] ExecStart= ExecStart=/usr/bin/nvidia-device-plugin --pass-device-specs EOF - fi + fi - # Reload systemd to pick up the override - systemctl daemon-reload + # Reload systemd to pick up the override + systemctl daemon-reload + + logs_to_events "AKS.CSE.start.nvidia-device-plugin" "systemctlEnableAndStart nvidia-device-plugin 30" || exit $ERR_GPU_DEVICE_PLUGIN_START_FAIL + elif [ "${ENABLE_MANAGED_GPU_EXPERIENCE_DRA}" = "true" ]; then + DRA_DRIVER_NVIDIA_GPU_OVERRIDE_DIR="/etc/systemd/system/dra-driver-nvidia-gpu.service.d" + mkdir -p "${DRA_DRIVER_NVIDIA_GPU_OVERRIDE_DIR}" + + tee "${DRA_DRIVER_NVIDIA_GPU_OVERRIDE_DIR}/10-dra-driver-nvidia-gpu.conf" > /dev/null </dev/null; then systemctlEnableAndStartNoBlock aks-log-collector.timer 30 || echo "Warning: Could not start aks-log-collector.timer" else 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 7e5c8a0e7f0..554f7511ebf 100755 --- a/parts/linux/cloud-init/artifacts/mariner/cse_install_mariner.sh +++ b/parts/linux/cloud-init/artifacts/mariner/cse_install_mariner.sh @@ -418,12 +418,18 @@ isPackageInstalled() { } managedGPUPackageList() { - packages=( - nvidia-device-plugin + local packages=( datacenter-gpu-manager-4-core datacenter-gpu-manager-4-proprietary dcgm-exporter ) + + if [ "${ENABLE_MANAGED_GPU_EXPERIENCE:-false}" = "true" ]; then + packages+=(nvidia-device-plugin) + elif [ "${ENABLE_MANAGED_GPU_EXPERIENCE_DRA:-false}" = "true" ]; then + packages+=(dra-driver-nvidia-gpu) + fi + echo "${packages[@]}" } @@ -435,6 +441,8 @@ installNvidiaManagedExpPkgFromCache() { # Ensure kubelet device-plugins directory exists BEFORE package installation mkdir -p /var/lib/kubelet/device-plugins + mkdir -p /var/lib/kubelet/plugins_registry + mkdir -p /var/lib/kubelet/plugins for packageName in $(managedGPUPackageList); do downloadDir="$(getPackageDownloadDir "${packageName}")" diff --git a/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh b/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh index dfe68bfa8e5..11e964fcd39 100755 --- a/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh +++ b/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh @@ -179,18 +179,26 @@ isPackageInstalled() { } managedGPUPackageList() { - packages=( - nvidia-device-plugin + local packages=( datacenter-gpu-manager-4-core datacenter-gpu-manager-4-proprietary dcgm-exporter ) + + if [ "${ENABLE_MANAGED_GPU_EXPERIENCE:-false}" = "true" ]; then + packages+=(nvidia-device-plugin) + elif [ "${ENABLE_MANAGED_GPU_EXPERIENCE_DRA:-false}" = "true" ]; then + packages+=(dra-driver-nvidia-gpu) + fi + echo "${packages[@]}" } installNvidiaManagedExpPkgFromCache() { # Ensure kubelet device-plugins directory exists BEFORE package installation mkdir -p /var/lib/kubelet/device-plugins + mkdir -p /var/lib/kubelet/plugins_registry + mkdir -p /var/lib/kubelet/plugins for packageName in $(managedGPUPackageList); do downloadDir="/opt/${packageName}/downloads" diff --git a/pkg/agent/baker.go b/pkg/agent/baker.go index bdc69f59aed..571af24d3d8 100644 --- a/pkg/agent/baker.go +++ b/pkg/agent/baker.go @@ -1369,6 +1369,9 @@ func getContainerServiceFuncMap(config *datamodel.NodeBootstrappingConfiguration "IsEnableManagedGPU": func() bool { return config.EnableManagedGPU }, + "IsEnableManagedGPUDRA": func() bool { + return config.EnableManagedGPUDRA + }, "EnableIMDSRestriction": func() bool { return config.EnableIMDSRestriction }, diff --git a/pkg/agent/datamodel/types.go b/pkg/agent/datamodel/types.go index 2f3acde27e4..dd1597ee55b 100644 --- a/pkg/agent/datamodel/types.go +++ b/pkg/agent/datamodel/types.go @@ -1754,6 +1754,7 @@ type NodeBootstrappingConfiguration struct { EnableAMDGPU bool ManagedGPUExperienceAFECEnabled bool EnableManagedGPU bool + EnableManagedGPUDRA bool MigStrategy string EnableArtifactStreaming bool ContainerdVersion string diff --git a/spec/parts/linux/cloud-init/artifacts/cse_config_spec.sh b/spec/parts/linux/cloud-init/artifacts/cse_config_spec.sh index bc98214c428..af1917d7815 100755 --- a/spec/parts/linux/cloud-init/artifacts/cse_config_spec.sh +++ b/spec/parts/linux/cloud-init/artifacts/cse_config_spec.sh @@ -2024,7 +2024,7 @@ SETUP_EOF echo "systemctl $@" } - BeforeEach 'MIG_NODE="false"' + BeforeEach 'MIG_NODE="false"; ENABLE_MANAGED_GPU_EXPERIENCE="true"; ENABLE_MANAGED_GPU_EXPERIENCE_DRA="false"' It 'starts the device-plugin blocking but dcgm and dcgm-exporter off the critical path' When call startNvidiaManagedExpServices @@ -2207,4 +2207,49 @@ EOF The variable AMD_AMA_DRIVER_VERSION should equal "1.5.0" End End + + Describe 'managedGPUPackageList on Ubuntu' + Include "./parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh" + + BeforeEach 'setup' + setup() { + ENABLE_MANAGED_GPU_EXPERIENCE="" + ENABLE_MANAGED_GPU_EXPERIENCE_DRA="" + } + + It 'returns base managed GPU packages by default' + When call managedGPUPackageList + + The status should be success + The output should equal 'datacenter-gpu-manager-4-core datacenter-gpu-manager-4-proprietary dcgm-exporter' + The output should not include 'nvidia-device-plugin' + The output should not include 'dra-driver-nvidia-gpu' + End + + It 'includes nvidia-device-plugin when managed GPU experience is enabled' + ENABLE_MANAGED_GPU_EXPERIENCE="true" + + When call managedGPUPackageList + + The status should be success + The output should include 'datacenter-gpu-manager-4-core' + The output should include 'datacenter-gpu-manager-4-proprietary' + The output should include 'dcgm-exporter' + The output should include 'nvidia-device-plugin' + The output should not include 'dra-driver-nvidia-gpu' + End + + It 'includes dra-driver-nvidia-gpu when DRA mode is enabled' + ENABLE_MANAGED_GPU_EXPERIENCE_DRA="true" + + When call managedGPUPackageList + + The status should be success + The output should include 'datacenter-gpu-manager-4-core' + The output should include 'datacenter-gpu-manager-4-proprietary' + The output should include 'dcgm-exporter' + The output should include 'dra-driver-nvidia-gpu' + The output should not include 'nvidia-device-plugin' + End + End End 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 3a6f840ca7a..ea51e77c6f6 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 @@ -413,4 +413,47 @@ Describe 'cse_install_mariner.sh' The status should equal 242 End End + + Describe 'managedGPUPackageList on Mariner' + BeforeEach 'setup' + setup() { + ENABLE_MANAGED_GPU_EXPERIENCE="" + ENABLE_MANAGED_GPU_EXPERIENCE_DRA="" + } + + It 'returns base managed GPU packages by default' + When call managedGPUPackageList + + The status should be success + The output should equal 'datacenter-gpu-manager-4-core datacenter-gpu-manager-4-proprietary dcgm-exporter' + The output should not include 'nvidia-device-plugin' + The output should not include 'dra-driver-nvidia-gpu' + End + + It 'includes nvidia-device-plugin when managed GPU experience is enabled' + ENABLE_MANAGED_GPU_EXPERIENCE="true" + + When call managedGPUPackageList + + The status should be success + The output should include 'datacenter-gpu-manager-4-core' + The output should include 'datacenter-gpu-manager-4-proprietary' + The output should include 'dcgm-exporter' + The output should include 'nvidia-device-plugin' + The output should not include 'dra-driver-nvidia-gpu' + End + + It 'includes dra-driver-nvidia-gpu when DRA mode is enabled' + ENABLE_MANAGED_GPU_EXPERIENCE_DRA="true" + + When call managedGPUPackageList + + The status should be success + The output should include 'datacenter-gpu-manager-4-core' + The output should include 'datacenter-gpu-manager-4-proprietary' + The output should include 'dcgm-exporter' + The output should include 'dra-driver-nvidia-gpu' + The output should not include 'nvidia-device-plugin' + End + End End