Skip to content

Commit ca8a90e

Browse files
ganeshkumarashokganeshkumarashok
authored
test: add e2e test for NVIDIA device plugin as DaemonSet (#7964)
1 parent 510c2a3 commit ca8a90e

1 file changed

Lines changed: 250 additions & 0 deletions

File tree

‎e2e/scenario_gpu_daemonset_test.go‎

Lines changed: 250 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,250 @@
1+
package e2e
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"strings"
7+
"testing"
8+
"time"
9+
10+
"github.com/Azure/agentbaker/e2e/config"
11+
"github.com/Azure/agentbaker/pkg/agent/datamodel"
12+
"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
13+
"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v7"
14+
"github.com/stretchr/testify/require"
15+
appsv1 "k8s.io/api/apps/v1"
16+
corev1 "k8s.io/api/core/v1"
17+
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
18+
)
19+
20+
const (
21+
// nvidiaDevicePluginImage is the upstream NVIDIA device plugin image from MCR.
22+
// This is intentionally different from components.json which tracks the systemd-packaged version.
23+
// This test validates the upstream container-based deployment model.
24+
// Update this when a new version is available in MCR.
25+
nvidiaDevicePluginImage = "mcr.microsoft.com/oss/v2/nvidia/k8s-device-plugin:v0.18.2"
26+
)
27+
28+
// Test_Ubuntu2204_NvidiaDevicePlugin_Daemonset tests that a GPU node can function correctly
29+
// with the NVIDIA device plugin deployed as a Kubernetes DaemonSet instead of a systemd service.
30+
// This is the "upstream" deployment model commonly used by customers who manage their own
31+
// NVIDIA device plugin deployment.
32+
func Test_Ubuntu2204_NvidiaDevicePlugin_Daemonset(t *testing.T) {
33+
RunScenario(t, &Scenario{
34+
Description: "Tests that NVIDIA device plugin works when deployed as a DaemonSet (not systemd service)",
35+
Tags: Tags{
36+
GPU: true,
37+
},
38+
Config: Config{
39+
Cluster: ClusterKubenet,
40+
VHD: config.VHDUbuntu2204Gen2Containerd,
41+
BootstrapConfigMutator: func(nbc *datamodel.NodeBootstrappingConfiguration) {
42+
nbc.AgentPoolProfile.VMSize = "Standard_NV6ads_A10_v5"
43+
nbc.ConfigGPUDriverIfNeeded = true
44+
// Don't enable the managed GPU experience - we'll deploy the device plugin as a DaemonSet instead.
45+
// By not setting EnableManagedGPU=true or the VMSS tag, the systemd-based device plugin won't start.
46+
nbc.EnableGPUDevicePluginIfNeeded = false
47+
nbc.EnableNvidia = true
48+
},
49+
VMConfigMutator: func(vmss *armcompute.VirtualMachineScaleSet) {
50+
vmss.SKU.Name = to.Ptr("Standard_NV6ads_A10_v5")
51+
},
52+
Validator: func(ctx context.Context, s *Scenario) {
53+
// First, validate that GPU drivers are installed
54+
ValidateNvidiaModProbeInstalled(ctx, s)
55+
56+
// Verify that the systemd-based device plugin is NOT running
57+
// (managed GPU experience is not enabled, so the service should not be active)
58+
validateNvidiaDevicePluginServiceNotRunning(ctx, s)
59+
60+
// Deploy the NVIDIA device plugin as a DaemonSet
61+
deployNvidiaDevicePluginDaemonset(ctx, s)
62+
63+
// Wait for the DaemonSet pod to be running on our node
64+
waitForNvidiaDevicePluginDaemonsetReady(ctx, s)
65+
66+
// Validate that GPU resources are advertised by the device plugin
67+
ValidateNodeAdvertisesGPUResources(ctx, s, 1, "nvidia.com/gpu")
68+
69+
// Validate that GPU workloads can be scheduled
70+
ValidateGPUWorkloadSchedulable(ctx, s, 1)
71+
72+
s.T.Logf("NVIDIA device plugin DaemonSet is functioning correctly")
73+
},
74+
},
75+
})
76+
}
77+
78+
// validateNvidiaDevicePluginServiceNotRunning verifies that the systemd-based
79+
// NVIDIA device plugin service is not running (since we're testing the DaemonSet model).
80+
func validateNvidiaDevicePluginServiceNotRunning(ctx context.Context, s *Scenario) {
81+
s.T.Helper()
82+
s.T.Logf("Verifying that nvidia-device-plugin.service is not running...")
83+
84+
// Check if the service exists and is inactive
85+
// Using "is-active" which returns non-zero if not active
86+
result := execScriptOnVMForScenario(ctx, s, "systemctl is-active nvidia-device-plugin.service 2>/dev/null || echo 'not-running'")
87+
output := strings.TrimSpace(result.stdout)
88+
89+
// The service should either not exist or be inactive
90+
if output == "active" {
91+
s.T.Fatalf("nvidia-device-plugin.service is unexpectedly running - this test requires the systemd service to be disabled")
92+
}
93+
s.T.Logf("Confirmed nvidia-device-plugin.service is not active (status: %s)", output)
94+
}
95+
96+
// nvidiaDevicePluginDaemonsetName returns a unique DaemonSet name for the given node.
97+
// The name is truncated to fit within Kubernetes' 63-character limit for resource names.
98+
func nvidiaDevicePluginDaemonsetName(nodeName string) string {
99+
prefix := "nvdp-" // Short prefix to leave room for node name
100+
maxLen := 63
101+
name := prefix + nodeName
102+
if len(name) > maxLen {
103+
name = name[:maxLen]
104+
}
105+
return name
106+
}
107+
108+
// nvidiaDevicePluginDaemonset returns the NVIDIA device plugin DaemonSet spec
109+
// based on the official upstream deployment from:
110+
// https://github.com/NVIDIA/k8s-device-plugin/blob/main/deployments/static/nvidia-device-plugin.yml
111+
//
112+
// The DaemonSet name includes the node name to avoid collisions when multiple
113+
// GPU tests run against the same shared cluster.
114+
func nvidiaDevicePluginDaemonset(nodeName string) *appsv1.DaemonSet {
115+
dsName := nvidiaDevicePluginDaemonsetName(nodeName)
116+
117+
return &appsv1.DaemonSet{
118+
TypeMeta: metav1.TypeMeta{
119+
Kind: "DaemonSet",
120+
APIVersion: "apps/v1",
121+
},
122+
ObjectMeta: metav1.ObjectMeta{
123+
Name: dsName,
124+
Namespace: "kube-system",
125+
},
126+
Spec: appsv1.DaemonSetSpec{
127+
Selector: &metav1.LabelSelector{
128+
MatchLabels: map[string]string{
129+
"name": dsName,
130+
},
131+
},
132+
UpdateStrategy: appsv1.DaemonSetUpdateStrategy{
133+
Type: appsv1.RollingUpdateDaemonSetStrategyType,
134+
},
135+
Template: corev1.PodTemplateSpec{
136+
ObjectMeta: metav1.ObjectMeta{
137+
Labels: map[string]string{
138+
"name": dsName,
139+
},
140+
},
141+
Spec: corev1.PodSpec{
142+
// Target only our specific test node
143+
NodeSelector: map[string]string{
144+
"kubernetes.io/hostname": nodeName,
145+
},
146+
Tolerations: []corev1.Toleration{
147+
{
148+
Key: "nvidia.com/gpu",
149+
Operator: corev1.TolerationOpExists,
150+
Effect: corev1.TaintEffectNoSchedule,
151+
},
152+
},
153+
PriorityClassName: "system-node-critical",
154+
Containers: []corev1.Container{
155+
{
156+
Name: "nvidia-device-plugin-ctr",
157+
Image: nvidiaDevicePluginImage,
158+
Env: []corev1.EnvVar{
159+
{
160+
Name: "FAIL_ON_INIT_ERROR",
161+
Value: "false",
162+
},
163+
},
164+
SecurityContext: &corev1.SecurityContext{
165+
// Privileged mode is required for the device plugin to access
166+
// GPU devices and register with kubelet's device plugin framework.
167+
// This matches the upstream NVIDIA device plugin deployment spec.
168+
Privileged: to.Ptr(true),
169+
},
170+
VolumeMounts: []corev1.VolumeMount{
171+
{
172+
Name: "device-plugin",
173+
MountPath: "/var/lib/kubelet/device-plugins",
174+
},
175+
},
176+
},
177+
},
178+
Volumes: []corev1.Volume{
179+
{
180+
Name: "device-plugin",
181+
VolumeSource: corev1.VolumeSource{
182+
HostPath: &corev1.HostPathVolumeSource{
183+
Path: "/var/lib/kubelet/device-plugins",
184+
},
185+
},
186+
},
187+
},
188+
},
189+
},
190+
},
191+
}
192+
}
193+
194+
// deployNvidiaDevicePluginDaemonset creates the NVIDIA device plugin DaemonSet in the cluster
195+
// and registers cleanup to delete it when the test finishes.
196+
func deployNvidiaDevicePluginDaemonset(ctx context.Context, s *Scenario) {
197+
s.T.Helper()
198+
s.T.Logf("Deploying NVIDIA device plugin as DaemonSet...")
199+
200+
ds := nvidiaDevicePluginDaemonset(s.Runtime.VM.KubeName)
201+
202+
// Delete any existing DaemonSet from a previous failed run
203+
deleteCtx, deleteCancel := context.WithTimeout(ctx, 30*time.Second)
204+
defer deleteCancel()
205+
_ = s.Runtime.Cluster.Kube.Typed.AppsV1().DaemonSets(ds.Namespace).Delete(
206+
deleteCtx,
207+
ds.Name,
208+
metav1.DeleteOptions{},
209+
)
210+
211+
// Create the DaemonSet
212+
err := s.Runtime.Cluster.Kube.CreateDaemonset(ctx, ds)
213+
require.NoError(s.T, err, "failed to create NVIDIA device plugin DaemonSet")
214+
215+
s.T.Logf("NVIDIA device plugin DaemonSet %s/%s created successfully", ds.Namespace, ds.Name)
216+
217+
// Register cleanup to delete the DaemonSet when the test finishes
218+
s.T.Cleanup(func() {
219+
s.T.Logf("Cleaning up NVIDIA device plugin DaemonSet %s/%s...", ds.Namespace, ds.Name)
220+
cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 30*time.Second)
221+
defer cleanupCancel()
222+
deleteErr := s.Runtime.Cluster.Kube.Typed.AppsV1().DaemonSets(ds.Namespace).Delete(
223+
cleanupCtx,
224+
ds.Name,
225+
metav1.DeleteOptions{},
226+
)
227+
if deleteErr != nil {
228+
s.T.Logf("Failed to delete NVIDIA device plugin DaemonSet %s/%s: %v", ds.Namespace, ds.Name, deleteErr)
229+
}
230+
})
231+
}
232+
233+
// waitForNvidiaDevicePluginDaemonsetReady waits for the NVIDIA device plugin pod to be running on the test node.
234+
// Uses the existing WaitUntilPodRunning helper which handles CrashLoopBackOff and other failure states.
235+
func waitForNvidiaDevicePluginDaemonsetReady(ctx context.Context, s *Scenario) {
236+
s.T.Helper()
237+
238+
dsName := nvidiaDevicePluginDaemonsetName(s.Runtime.VM.KubeName)
239+
s.T.Logf("Waiting for NVIDIA device plugin DaemonSet pod to be ready on node %s...", s.Runtime.VM.KubeName)
240+
241+
_, err := s.Runtime.Cluster.Kube.WaitUntilPodRunning(
242+
ctx,
243+
"kube-system",
244+
fmt.Sprintf("name=%s", dsName),
245+
fmt.Sprintf("spec.nodeName=%s", s.Runtime.VM.KubeName),
246+
)
247+
require.NoError(s.T, err, "timed out waiting for NVIDIA device plugin DaemonSet pod to be ready")
248+
249+
s.T.Logf("NVIDIA device plugin DaemonSet pod is ready")
250+
}

0 commit comments

Comments
 (0)