Repository: Azure/AgentBaker
Pull request #6067: Add license and GRID driver validation, move nvidia systemd service files
Author: login: ganeshkumarashok
Created: 2025-03-21T16:35:32Z | Merged: 2025-03-26T21:05:49Z

**What type of PR is this?**
/kind bug
<!--
Add one of the following kinds:
/kind bug
/kind test
/kind cleanup
/kind documentation
/kind feature
/kind design

Optionally add one or more of the following kinds if applicable:
/kind api-change
/kind deprecation
/kind failing-test
/kind flake
/kind regression
-->

**What this PR does / why we need it**:
- Add e2e validation for GRID driver license server state
- Validate status of nvidia-gridd systemd unit
- Update 535 GRID driver with container image that moves nvidia systemd service files to the right path.


**Which issue(s) this PR fixes**:
<!-- 
*Automatically closes linked issue when PR is merged.
Usage: `Fixes #<issue number>`, or `Fixes (paste link of issue)`.
_If PR is about `failing-tests or flakes`, please post the related issues/tests in a comment and do not use `Fixes`_*
-->
Fixes #

**Requirements**:

- [ ] uses [conventional commit messages](https://www.conventionalcommits.org/)
  <!-- Common commit types:
        build: Build 🏭
        chore: Maintenance 🔧
        ci: Continuous Integration 💜
        docs: Documentation 📘
        feat: Features 🌈
        fix: Bug Fixes 🐞
        perf: Performance Improvements 🚀
        refactor: Code Refactoring 💎
        revert: Revert Change ◀️
        style: Code Style 🎶
        security: Security Fix 🛡️
        test: Testing 💚 -->
- [ ] includes documentation
- [ ] adds unit tests
- [ ] tested upgrade from previous version

**Special notes for your reviewer**:

**Release note**:

```
none
```



Diff:
diff --git a/e2e/scenario_test.go b/e2e/scenario_test.go
index a882949eb13..6e0b526c185 100644
--- a/e2e/scenario_test.go
+++ b/e2e/scenario_test.go
@@ -1008,7 +1008,7 @@ func Test_Ubuntu2204_GPUA100(t *testing.T) {
 }
 
 func Test_Ubuntu2204_GPUA10(t *testing.T) {
-	runScenarioUbuntu2204GPU(t, "Standard_NV6ads_A10_v5")
+	runScenarioUbuntuGRID(t, "Standard_NV6ads_A10_v5")
 }
 
 // Returns config for the 'gpu' E2E scenario
@@ -1040,6 +1040,36 @@ func runScenarioUbuntu2204GPU(t *testing.T, vmSize string) {
 	})
 }
 
+func runScenarioUbuntuGRID(t *testing.T, vmSize string) {
+	RunScenario(t, &Scenario{
+		Description: fmt.Sprintf("Tests that a GPU-enabled node with VM size %s using an Ubuntu 2204 VHD can be properly bootstrapped, and that the GRID license is valid", vmSize),
+		Tags: Tags{
+			GPU: true,
+		},
+		Config: Config{
+			Cluster: ClusterKubenet,
+			VHD:     config.VHDUbuntu2204Gen2Containerd,
+			BootstrapConfigMutator: func(nbc *datamodel.NodeBootstrappingConfiguration) {
+				nbc.AgentPoolProfile.VMSize = vmSize
+				nbc.ConfigGPUDriverIfNeeded = true
+				nbc.EnableGPUDevicePluginIfNeeded = false
+				nbc.EnableNvidia = true
+			},
+			VMConfigMutator: func(vmss *armcompute.VirtualMachineScaleSet) {
+				vmss.SKU.Name = to.Ptr(vmSize)
+			},
+			Validator: func(ctx context.Context, s *Scenario) {
+				// Ensure nvidia-modprobe install does not restart kubelet and temporarily cause node to be unschedulable
+				ValidateNvidiaModProbeInstalled(ctx, s)
+				ValidateNvidiaGRIDLicenseValid(ctx, s)
+				ValidateKubeletHasNotStopped(ctx, s)
+				ValidateServicesDoNotRestartKubelet(ctx, s)
+				ValidateNvidiaPersistencedRunning(ctx, s)
+			},
+		},
+	})
+}
+
 func Test_Ubuntu2204_GPUA10_Scriptless(t *testing.T) {
 	RunScenario(t, &Scenario{
 		Description: "Tests scriptless installer that a GPU-enabled node using the Ubuntu 2204 VHD with grid driver can be properly bootstrapped",
diff --git a/e2e/validators.go b/e2e/validators.go
index 816427ae189..1dc9fec2607 100644
--- a/e2e/validators.go
+++ b/e2e/validators.go
@@ -68,9 +68,34 @@ func ValidateNvidiaModProbeInstalled(ctx context.Context, s *Scenario) {
 		"set -ex",
 		"sudo nvidia-modprobe",
 	}
-	execScriptOnVMForScenarioValidateExitCode(ctx, s, strings.Join(command, "\n"), 0, "cound not execute nvidia-modprobe command")
+	execScriptOnVMForScenarioValidateExitCode(ctx, s, strings.Join(command, "\n"), 0, "could not execute nvidia-modprobe command")
 }
 
+func ValidateNvidiaGRIDLicenseValid(ctx context.Context, s *Scenario) {
+	command := []string{
+		"set -ex",
+		// Capture the license status output, or continue silently if not found
+		"license_status=$(sudo nvidia-smi -q | grep 'License Status' | grep 'Licensed' || true)",
+		// If the output is empty, print an error message and exit with a nonzero code
+		"if [ -z \"$license_status\" ]; then echo 'License status not valid or not found'; exit 1; fi",
+		// Check that nvidia-gridd is active by capturing its is-active output
+		"active_status=$(sudo systemctl is-active nvidia-gridd)",
+		"if [ \"$active_status\" != \"active\" ]; then echo \"nvidia-gridd is not active: $active_status\"; exit 1; fi",
+	}
+	execScriptOnVMForScenarioValidateExitCode(ctx, s, strings.Join(command, "\n"), 0, "failed to validate nvidia-smi license state or nvidia-gridd service status")
+}
+
+func ValidateNvidiaPersistencedRunning(ctx context.Context, s *Scenario) {
+	command := []string{
+		"set -ex",
+		// Check that nvidia-persistenced.service is active by capturing its is-active output
+		"active_status=$(sudo systemctl is-active nvidia-persistenced.service)",
+		"if [ \"$active_status\" != \"active\" ]; then echo \"nvidia-gridd is not active: $active_status\"; exit 1; fi",
+	}
+	execScriptOnVMForScenarioValidateExitCode(ctx, s, strings.Join(command, "\n"), 0, "failed to validate nvidia-persistenced.service status")
+}
+
+
 func ValidateNonEmptyDirectory(ctx context.Context, s *Scenario, dirName string) {
 	command := []string{
 		"set -ex",
diff --git a/parts/common/components.json b/parts/common/components.json
index d17718c9101..c91ce44a5ed 100644
--- a/parts/common/components.json
+++ b/parts/common/components.json
@@ -773,7 +773,7 @@
       "downloadURL": "mcr.microsoft.com/aks/aks-gpu-grid:*",
       "gpuVersion": {
         "renovateTag": "registry=https://mcr.microsoft.com, name=aks/aks-gpu-grid",
-        "latestVersion": "535.161.08-20250317215455"
+        "latestVersion": "535.161.08-20250325114356"
       }
     }
   ],


Reviews:
user: login: cameronmeissner
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-22T17:19:32Z
user: login: ganeshkumarashok
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-25T19:18:32Z
user: login: djsly
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-25T19:49:30Z
user: login: ganeshkumarashok
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-25T19:49:53Z
user: login: ganeshkumarashok
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-25T19:51:05Z
user: login: ganeshkumarashok
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-25T19:56:18Z
user: login: ganeshkumarashok
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-25T19:57:56Z
user: login: ganeshkumarashok
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-25T20:27:26Z
user: login: djsly
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-25T21:48:04Z
user: login: ganeshkumarashok
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-26T02:55:07Z
user: login: ganeshkumarashok
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-26T21:05:24Z

Review comments:
diff_hunk: @@ -68,7 +68,21 @@ func ValidateNvidiaModProbeInstalled(ctx context.Context, s *Scenario) {
 		"set -ex",
 		"sudo nvidia-modprobe",
 	}
-	execScriptOnVMForScenarioValidateExitCode(ctx, s, strings.Join(command, "\n"), 0, "cound not execute nvidia-modprobe command")
+	execScriptOnVMForScenarioValidateExitCode(ctx, s, strings.Join(command, "\n"), 0, "could not execute nvidia-modprobe command")
+}
+
+func ValidateNvidiaGRIDLicenseValid(ctx context.Context, s *Scenario) {
+	command := []string{
+		"set -ex",
+		// Capture the license status output, or continue silently if not found
+		"license_status=$(sudo nvidia-smi -q | grep 'License Status' | grep 'Licensed' || true)",
+		// If the output is empty, print an error message and exit with a nonzero code
+		"if [ -z \"$license_status\" ]; then echo 'License status not valid or not found'; exit 1; fi",
+		// Check that nvidia-gridd is active by capturing its is-active output
+		"active_status=$(sudo systemctl is-active nvidia-gridd)",
path: e2e/validators.go
user: login: ganeshkumarashok
body: Considered adding a similar check for nvidia-topologyd but it's not clear which GPUs and driver versions should have it. GRID license should always be valid though.
created_at: 2025-03-25T19:18:32Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1039,7 +1039,7 @@ func Test_Ubuntu2204_GPUA100(t *testing.T) {
 }
 
 func Test_Ubuntu2204_GPUA10(t *testing.T) {
-	runScenarioUbuntu2204GPU(t, "Standard_NV6ads_A10_v5")
path: e2e/scenario_test.go
user: login: djsly
body: are we still using this Scenario ?
created_at: 2025-03-25T19:46:39Z
author_association: COLLABORATOR
diff_hunk: @@ -1071,6 +1071,35 @@ func runScenarioUbuntu2204GPU(t *testing.T, vmSize string) {
 	})
 }
 
+func runScenarioUbuntuGRID(t *testing.T, vmSize string) {
path: e2e/scenario_test.go
user: login: djsly
body: should we keep 2204 in the ScenarioName?  what if we want to test with 2404 to ensure that it's working with the kernel 6.8 ?
created_at: 2025-03-25T19:47:37Z
author_association: COLLABORATOR
diff_hunk: @@ -68,7 +68,21 @@ func ValidateNvidiaModProbeInstalled(ctx context.Context, s *Scenario) {
 		"set -ex",
 		"sudo nvidia-modprobe",
 	}
-	execScriptOnVMForScenarioValidateExitCode(ctx, s, strings.Join(command, "\n"), 0, "cound not execute nvidia-modprobe command")
+	execScriptOnVMForScenarioValidateExitCode(ctx, s, strings.Join(command, "\n"), 0, "could not execute nvidia-modprobe command")
+}
+
+func ValidateNvidiaGRIDLicenseValid(ctx context.Context, s *Scenario) {
+	command := []string{
+		"set -ex",
+		// Capture the license status output, or continue silently if not found
+		"license_status=$(sudo nvidia-smi -q | grep 'License Status' | grep 'Licensed' || true)",
+		// If the output is empty, print an error message and exit with a nonzero code
+		"if [ -z \"$license_status\" ]; then echo 'License status not valid or not found'; exit 1; fi",
+		// Check that nvidia-gridd is active by capturing its is-active output
+		"active_status=$(sudo systemctl is-active nvidia-gridd)",
path: e2e/validators.go
user: login: djsly
body: what about persisted ?
created_at: 2025-03-25T19:48:52Z
author_association: COLLABORATOR
diff_hunk: @@ -68,7 +68,21 @@ func ValidateNvidiaModProbeInstalled(ctx context.Context, s *Scenario) {
 		"set -ex",
 		"sudo nvidia-modprobe",
 	}
-	execScriptOnVMForScenarioValidateExitCode(ctx, s, strings.Join(command, "\n"), 0, "cound not execute nvidia-modprobe command")
+	execScriptOnVMForScenarioValidateExitCode(ctx, s, strings.Join(command, "\n"), 0, "could not execute nvidia-modprobe command")
+}
+
+func ValidateNvidiaGRIDLicenseValid(ctx context.Context, s *Scenario) {
+	command := []string{
+		"set -ex",
+		// Capture the license status output, or continue silently if not found
+		"license_status=$(sudo nvidia-smi -q | grep 'License Status' | grep 'Licensed' || true)",
path: e2e/validators.go
user: login: djsly
body: this will always be valid for the first 20min no ? 
created_at: 2025-03-25T19:49:24Z
author_association: COLLABORATOR
diff_hunk: @@ -1039,7 +1039,7 @@ func Test_Ubuntu2204_GPUA100(t *testing.T) {
 }
 
 func Test_Ubuntu2204_GPUA10(t *testing.T) {
-	runScenarioUbuntu2204GPU(t, "Standard_NV6ads_A10_v5")
path: e2e/scenario_test.go
user: login: ganeshkumarashok
body: yes, for CUDA
created_at: 2025-03-25T19:49:53Z
author_association: CONTRIBUTOR
diff_hunk: @@ -68,7 +68,21 @@ func ValidateNvidiaModProbeInstalled(ctx context.Context, s *Scenario) {
 		"set -ex",
 		"sudo nvidia-modprobe",
 	}
-	execScriptOnVMForScenarioValidateExitCode(ctx, s, strings.Join(command, "\n"), 0, "cound not execute nvidia-modprobe command")
+	execScriptOnVMForScenarioValidateExitCode(ctx, s, strings.Join(command, "\n"), 0, "could not execute nvidia-modprobe command")
+}
+
+func ValidateNvidiaGRIDLicenseValid(ctx context.Context, s *Scenario) {
+	command := []string{
+		"set -ex",
+		// Capture the license status output, or continue silently if not found
+		"license_status=$(sudo nvidia-smi -q | grep 'License Status' | grep 'Licensed' || true)",
path: e2e/validators.go
user: login: ganeshkumarashok
body: No, it won't necessarily be valid then -- workload will work for the first 20 minutes but after that, if it remains in an unlicensed state, it will get degraded.
created_at: 2025-03-25T19:51:05Z
author_association: CONTRIBUTOR
diff_hunk: @@ -68,7 +68,21 @@ func ValidateNvidiaModProbeInstalled(ctx context.Context, s *Scenario) {
 		"set -ex",
 		"sudo nvidia-modprobe",
 	}
-	execScriptOnVMForScenarioValidateExitCode(ctx, s, strings.Join(command, "\n"), 0, "cound not execute nvidia-modprobe command")
+	execScriptOnVMForScenarioValidateExitCode(ctx, s, strings.Join(command, "\n"), 0, "could not execute nvidia-modprobe command")
+}
+
+func ValidateNvidiaGRIDLicenseValid(ctx context.Context, s *Scenario) {
+	command := []string{
+		"set -ex",
+		// Capture the license status output, or continue silently if not found
+		"license_status=$(sudo nvidia-smi -q | grep 'License Status' | grep 'Licensed' || true)",
+		// If the output is empty, print an error message and exit with a nonzero code
+		"if [ -z \"$license_status\" ]; then echo 'License status not valid or not found'; exit 1; fi",
+		// Check that nvidia-gridd is active by capturing its is-active output
+		"active_status=$(sudo systemctl is-active nvidia-gridd)",
path: e2e/validators.go
user: login: ganeshkumarashok
body: persistenced is not one of the files that got moved, but we could check it too
created_at: 2025-03-25T19:56:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1071,6 +1071,35 @@ func runScenarioUbuntu2204GPU(t *testing.T, vmSize string) {
 	})
 }
 
+func runScenarioUbuntuGRID(t *testing.T, vmSize string) {
path: e2e/scenario_test.go
user: login: ganeshkumarashok
body: We should be checking both 22.04 and 24.04 in that case. How are we switching between the two OS versions?
created_at: 2025-03-25T19:57:55Z
author_association: CONTRIBUTOR
diff_hunk: @@ -68,7 +68,21 @@ func ValidateNvidiaModProbeInstalled(ctx context.Context, s *Scenario) {
 		"set -ex",
 		"sudo nvidia-modprobe",
 	}
-	execScriptOnVMForScenarioValidateExitCode(ctx, s, strings.Join(command, "\n"), 0, "cound not execute nvidia-modprobe command")
+	execScriptOnVMForScenarioValidateExitCode(ctx, s, strings.Join(command, "\n"), 0, "could not execute nvidia-modprobe command")
+}
+
+func ValidateNvidiaGRIDLicenseValid(ctx context.Context, s *Scenario) {
+	command := []string{
+		"set -ex",
+		// Capture the license status output, or continue silently if not found
+		"license_status=$(sudo nvidia-smi -q | grep 'License Status' | grep 'Licensed' || true)",
+		// If the output is empty, print an error message and exit with a nonzero code
+		"if [ -z \"$license_status\" ]; then echo 'License status not valid or not found'; exit 1; fi",
+		// Check that nvidia-gridd is active by capturing its is-active output
+		"active_status=$(sudo systemctl is-active nvidia-gridd)",
path: e2e/validators.go
user: login: ganeshkumarashok
body: Added the check for persistenced as a separate validator.
created_at: 2025-03-25T20:27:26Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1071,6 +1071,35 @@ func runScenarioUbuntu2204GPU(t *testing.T, vmSize string) {
 	})
 }
 
+func runScenarioUbuntuGRID(t *testing.T, vmSize string) {
path: e2e/scenario_test.go
user: login: djsly
body: this line ? VHD:     config.VHDUbuntu2204Gen2Containerd,
created_at: 2025-03-25T21:48:04Z
author_association: COLLABORATOR
diff_hunk: @@ -1071,6 +1071,35 @@ func runScenarioUbuntu2204GPU(t *testing.T, vmSize string) {
 	})
 }
 
+func runScenarioUbuntuGRID(t *testing.T, vmSize string) {
path: e2e/scenario_test.go
user: login: ganeshkumarashok
body: Thanks, added another test for 24.04
created_at: 2025-03-26T02:55:07Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1071,6 +1071,35 @@ func runScenarioUbuntu2204GPU(t *testing.T, vmSize string) {
 	})
 }
 
+func runScenarioUbuntuGRID(t *testing.T, vmSize string) {
path: e2e/scenario_test.go
user: login: ganeshkumarashok
body: Removed that test since there were other failures. 
created_at: 2025-03-26T21:05:24Z
author_association: CONTRIBUTOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-03-26T20:37:00Z
body: No changes to cached containers or packages on Windows VHDs

<!-- thollander/actions-comment-pull-request "WINDOWS_VHD_CONTENT_DIFF" -->
author_association: CONTRIBUTOR