Merge branch 'main' into xid-analyzer-recovery-actions

Author: lalitadithya

distros/kubernetes/nvsentinel/charts/node-drainer/templates/configmap.yaml

@@ -25,6 +25,7 @@ data:
     systemNamespaces = {{ .Values.systemNamespaces | quote }}
     deleteAfterTimeoutMinutes = {{ .Values.deleteAfterTimeoutMinutes }}
     notReadyTimeoutMinutes = {{ .Values.notReadyTimeoutMinutes }}
+    drainGPUPods = {{ .Values.drainGPUPods }}
     partialDrainEnabled = {{ .Values.partialDrainEnabled }}
     
     {{- range .Values.userNamespaces }}

distros/kubernetes/nvsentinel/charts/node-drainer/values.yaml

@@ -55,6 +55,12 @@ deleteAfterTimeoutMinutes: 60
 # Default: 5 minutes if not specified (validated in config.go)
 notReadyTimeoutMinutes: 5
 
+# Flag to restrict draining to GPU workloads
+# If enabled, only pods with the metadata-collector device annotation
+# (indicating assigned GPU devices) are eligible for draining
+# Default: false if not specified
+drainGPUPods: false
+
 # User namespace configuration with eviction modes
 # Defines how pods in different namespaces should be evicted during node drain
 # Each entry specifies a namespace pattern and its corresponding eviction mode

distros/kubernetes/nvsentinel/values-full.yaml

@@ -590,6 +590,12 @@ node-drainer:
   # Default: 5 minutes
   notReadyTimeoutMinutes: 5
 
+  # Flag to restrict draining to GPU workloads
+  # If enabled, only pods with the metadata-collector device annotation
+  # (indicating assigned GPU devices) are eligible for draining
+  # Default: false if not specified
+  drainGPUPods: false
+
   # Namespace-specific eviction strategies
   # Define how pods in different namespaces should be evicted
   # Multiple rules can be defined with namespace patterns

docs/configuration/node-drainer.md

@@ -96,6 +96,29 @@ node-drainer:
 
 When a pod has been in NotReady state for longer than this timeout, it is excluded from the list of pods to evict. This prevents attempting to evict pods that are already unhealthy and unlikely to respond to eviction requests.
 
+### GPU-Only Draining
+
+If enabled, the node-drainer filters pod eviction to only target workloads that request GPU resources.
+
+```yaml
+node-drainer:
+  drainGPUPods: false
+```
+
+The node-drainer detects GPU resource requests through device annotations added to pods by the metadata-collector. Pods with device annotations are identified as GPU workloads and eligible for eviction.
+
+Device annotations are added to pods requesting GPU resources by metadata-collector with the format:
+```yaml
+annotations:
+	  dgxc.nvidia.com/devices: '{"devices":{"nvidia.com/gpu":["GPU-123"]}}'
+```
+
+#### Behavior
+
+- **When enabled (`true`)**: Only pods with GPU device annotations are evicted during drain operations
+- **When disabled (`false`)**: All eligible pods in configured namespaces are evicted (default behavior)
+- Pods without GPU requests are preserved, maintaining critical infrastructure services
+
 ## User Namespaces
 
 Defines eviction behavior for user workloads based on namespace patterns.

docs/node-drainer.md

@@ -49,6 +49,7 @@ node-drainer:
   systemNamespaces: "^(nvsentinel|kube-system|gpu-operator)$"  # Namespaces to skip
   deleteAfterTimeoutMinutes: 60      # Force delete after this timeout
   notReadyTimeoutMinutes: 5          # Timeout for stuck pods
+  drainGPUPods: false                # Only drain pods requesting GPU
 
   userNamespaces:
     - name: "*"                      # Pattern matching namespaces
@@ -81,6 +82,7 @@ The module supports three eviction modes for different workload types:
 - **System Namespaces**: Regex pattern for namespaces to skip (system pods)
 - **Delete Timeout**: Minutes to wait before force deleting pods
 - **Not Ready Timeout**: Minutes before considering a pod stuck
+- **Drain GPU Pods**: When enabled, only drains pods requesting GPU resources; CPU-only workloads remain on the node
 - **User Namespaces**: Define eviction mode per namespace pattern (supports `*` wildcard)
 - **Partial Drain**: Enable or disable partial drain functionality
 
@@ -108,4 +110,7 @@ Multiple timeout mechanisms prevent stuck drains:
 Automatically resumes drain operations after restarts - queries datastore for in-progress drains and continues from where it left off.
 
 ### Partial Drain Functionality
-For GPU faults that can be remediated with a GPU reset, the Node Drainer will only drain pods which are leveraging the unhealthy GPU. For GPU faults that require a node reboot, all pods on the given node in the configured namespaces will be drained.
+For GPU faults that can be remediated with a GPU reset, the Node Drainer will only drain pods which are leveraging the unhealthy GPU. For GPU faults that require a node reboot, all pods on the given node in the configured namespaces will be drained.
+
+### GPU-Only Draining
+When `drainGPUPods: true` is set, the Node Drainer filters pod eviction to only target workloads that request GPU resources. The feature detects GPU resources using device annotations provided by the Metadata Collector, which tracks GPU allocation across the cluster. Default is `false`.

node-drainer/pkg/config/config.go

@@ -62,6 +62,7 @@ type TomlConfig struct {
 	DeleteAfterTimeoutMinutes int      `toml:"deleteAfterTimeoutMinutes"`
 	// NotReadyTimeoutMinutes is the time after which a pod in NotReady state is considered stuck
 	NotReadyTimeoutMinutes int               `toml:"notReadyTimeoutMinutes"`
+	DrainGPUPods           bool              `toml:"drainGPUPods"`
 	UserNamespaces         []UserNamespace   `toml:"userNamespaces"`
 	CustomDrain            CustomDrainConfig `toml:"customDrain"`
 	PartialDrainEnabled    bool              `toml:"partialDrainEnabled"`

node-drainer/pkg/evaluator/evaluator_integration_test.go

@@ -103,11 +103,12 @@ func setupDirectTest(t *testing.T, userNamespaces []config.UserNamespace, dryRun
 		SystemNamespaces:          "kube-*",
 		DeleteAfterTimeoutMinutes: 5,
 		NotReadyTimeoutMinutes:    2,
+		DrainGPUPods:              false,
 		UserNamespaces:            userNamespaces,
 		PartialDrainEnabled:       partialDrainEnabled,
 	}
 
-	informersInstance, err := informers.NewInformers(client, 1*time.Minute, ptr.To(2), dryRun)
+	informersInstance, err := informers.NewInformers(client, 1*time.Minute, ptr.To(2), false, dryRun)
 	require.NoError(t, err)
 	go func() { _ = informersInstance.Run(ctx) }()
 	require.Eventually(t, informersInstance.HasSynced, 30*time.Second, 1*time.Second)

node-drainer/pkg/informers/informers.go

@@ -53,12 +53,13 @@ type Informers struct {
 	nodeInformer           cache.SharedIndexInformer
 	clientset              kubernetes.Interface
 	notReadyTimeoutMinutes *int
+	drainGPUPods           bool
 	dryRunMode             []string
 	namespace              string
 }
 
 func NewInformers(clientset kubernetes.Interface, resyncPeriod time.Duration,
-	notReadyTimeoutMinutes *int, dryRun bool) (*Informers, error) {
+	notReadyTimeoutMinutes *int, drainGPUPods bool, dryRun bool) (*Informers, error) {
 	informerFactory := informers.NewSharedInformerFactoryWithOptions(
 		clientset,
 		resyncPeriod,
@@ -104,6 +105,7 @@ func NewInformers(clientset kubernetes.Interface, resyncPeriod time.Duration,
 		eventInformer:          eventInformer,
 		nodeInformer:           nodeInformer,
 		notReadyTimeoutMinutes: notReadyTimeoutMinutes,
+		drainGPUPods:           drainGPUPods,
 		dryRunMode:             dryRunMode,
 		namespace:              metav1.NamespaceDefault,
 	}, nil
@@ -189,6 +191,10 @@ func (i *Informers) FindEvictablePodsInNamespaceAndNode(namespace, nodeName stri
 
 	pods = i.filterEvictablePods(pods)
 
+	if i.drainGPUPods && partialDrainEntity == nil {
+		pods = i.filterPodsWithGPURequests(pods)
+	}
+
 	pods, err = i.filterPodsUsingEntity(pods, partialDrainEntity, nodeName)
 	if err != nil {
 		return nil, fmt.Errorf("failed to filter pods using entity: %w", err)
@@ -353,6 +359,26 @@ func (i *Informers) filterEvictablePods(pods []*v1.Pod) []*v1.Pod {
 	return filteredPods
 }
 
+func (i *Informers) filterPodsWithGPURequests(pods []*v1.Pod) []*v1.Pod {
+	filteredPods := []*v1.Pod{}
+
+	for _, pod := range pods {
+		_, podHasDeviceAnnotation := pod.Annotations[model.PodDeviceAnnotationName]
+
+		// If the pod has been assigned GPU it must have device annotation
+		if podHasDeviceAnnotation {
+			slog.Info("Pod is eligible for draining as it is requesting GPU",
+				"pod", pod.Name,
+				"namespace", pod.Namespace,
+				"node", pod.Spec.NodeName,
+			)
+			filteredPods = append(filteredPods, pod)
+		}
+	}
+
+	return filteredPods
+}
+
 func (i *Informers) isDaemonSetPod(pod *v1.Pod) bool {
 	for _, owner := range pod.OwnerReferences {
 		if owner.Kind == "DaemonSet" {

node-drainer/pkg/initializer/init.go

@@ -98,7 +98,9 @@ func InitializeAll(ctx context.Context, params InitializationParams) (*Component
 		return nil, fmt.Errorf("failed to initialize dynamic client and mapper: %w", err)
 	}
 
-	informersInstance, err := initializeInformers(clientSet, &configs.tomlCfg.NotReadyTimeoutMinutes, params.DryRun)
+	informersInstance, err := initializeInformers(
+		clientSet, &configs.tomlCfg.NotReadyTimeoutMinutes, configs.tomlCfg.DrainGPUPods, params.DryRun,
+	)
 	if err != nil {
 		return nil, fmt.Errorf("error while initializing informers: %w", err)
 	}
@@ -286,8 +288,8 @@ func initializeKubernetesClient(kubeconfigPath string) (kubernetes.Interface, *r
 }
 
 func initializeInformers(clientset kubernetes.Interface,
-	notReadyTimeoutMinutes *int, dryRun bool) (*informers.Informers, error) {
-	return informers.NewInformers(clientset, time.Hour, notReadyTimeoutMinutes, dryRun)
+	notReadyTimeoutMinutes *int, drainGPUPods bool, dryRun bool) (*informers.Informers, error) {
+	return informers.NewInformers(clientset, time.Hour, notReadyTimeoutMinutes, drainGPUPods, dryRun)
 }
 
 func initializeStateManager(clientSet kubernetes.Interface) statemanager.StateManager {

node-drainer/pkg/reconciler/reconciler_integration_test.go

@@ -229,6 +229,7 @@ func TestReconciler_ProcessEvent(t *testing.T) {
 		expectError                     bool
 		expectedNodeLabel               *string
 		numReconciles                   int
+		enableDrainGPUPods              bool
 		validateFunc                    func(t *testing.T, client kubernetes.Interface, ctx context.Context, nodeName string, err error)
 	}{
 		{
@@ -260,6 +261,114 @@ func TestReconciler_ProcessEvent(t *testing.T) {
 				}, 30*time.Second, 1*time.Second, "pods should be evicted")
 			},
 		},
+		{
+			name:               "DrainGPUPods enabled only evicts GPU workloads",
+			nodeName:           "gpu-drain-node",
+			namespaces:         []string{"immediate-test"},
+			nodeQuarantined:    model.Quarantined,
+			enableDrainGPUPods: true,
+			existingNodeLabels: map[string]string{
+				statemanager.NVSentinelStateLabelKey: string(statemanager.QuarantinedLabelValue),
+			},
+			pods: []*v1.Pod{
+				{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "gpu-pod-1",
+						Namespace: "immediate-test",
+						Annotations: map[string]string{
+							model.PodDeviceAnnotationName: "{\"devices\":{\"nvidia.com/gpu\":[\"GPU-123\"]}}",
+						},
+					},
+					Spec: v1.PodSpec{
+						NodeName: "gpu-drain-node",
+						Containers: []v1.Container{{
+							Name:  "gpu-container",
+							Image: "nvidia/cuda:latest",
+							Resources: v1.ResourceRequirements{
+								Limits: v1.ResourceList{
+									v1.ResourceName("nvidia.com/gpu"): resource.MustParse("1"),
+								},
+							},
+						}},
+					},
+					Status: v1.PodStatus{Phase: v1.PodRunning},
+				},
+				{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "cpu-pod-1",
+						Namespace: "immediate-test",
+					},
+					Spec: v1.PodSpec{
+						NodeName: "gpu-drain-node",
+						Containers: []v1.Container{{
+							Name:  "cpu-container",
+							Image: "nginx",
+						}},
+					},
+					Status: v1.PodStatus{Phase: v1.PodRunning},
+				},
+				{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "gpu-pod-2",
+						Namespace: "immediate-test",
+						Annotations: map[string]string{
+							model.PodDeviceAnnotationName: "{\"devices\":{\"nvidia.com/gpu\":[\"GPU-456\"]}}",
+						},
+					},
+					Spec: v1.PodSpec{
+						NodeName: "gpu-drain-node",
+						Containers: []v1.Container{{
+							Name:  "gpu-container",
+							Image: "pytorch:latest",
+							Resources: v1.ResourceRequirements{
+								Limits: v1.ResourceList{
+									v1.ResourceName("nvidia.com/gpu"): resource.MustParse("2"),
+								},
+							},
+						}},
+					},
+					Status: v1.PodStatus{Phase: v1.PodRunning},
+				},
+			},
+			expectError:       true,
+			expectedNodeLabel: ptr.To(string(statemanager.DrainingLabelValue)),
+			validateFunc: func(t *testing.T, client kubernetes.Interface, ctx context.Context, nodeName string, err error) {
+				assert.Error(t, err)
+				assert.Contains(t, err.Error(), "immediate eviction completed, requeuing for status verification")
+
+				// Verify GPU pods are evicted but CPU pod is preserved
+				expectDeletedForPods := map[string]bool{
+					"gpu-pod-1": true,  // Has GPU-123, eligible to be evicted
+					"cpu-pod-1": false, // preserved
+					"gpu-pod-2": true,  // Has GPU-456, eligible to be evicted
+				}
+
+				for podName, expectDeleted := range expectDeletedForPods {
+					pod, err := client.CoreV1().Pods("immediate-test").Get(ctx, podName, metav1.GetOptions{})
+					require.NoError(t, err, "Pod %s should exist", podName)
+
+					if expectDeleted {
+						assert.NotNil(t, pod.DeletionTimestamp, "GPU pod %s should be marked for deletion", podName)
+					} else {
+						assert.Nil(t, pod.DeletionTimestamp, "CPU pod %s should NOT be marked for deletion", podName)
+					}
+				}
+
+				// Verify that eventually only CPU pod remains (GPU pods evicted)
+				require.Eventually(t, func() bool {
+					pods, _ := client.CoreV1().Pods("immediate-test").List(ctx, metav1.ListOptions{})
+					remainingPods := 0
+					for _, pod := range pods.Items {
+						if pod.DeletionTimestamp == nil {
+							remainingPods++
+							// Verify remaining pod is CPU pod
+							assert.Equal(t, "cpu-pod-1", pod.Name, "Only CPU pod should remain")
+						}
+					}
+					return remainingPods == 1
+				}, 30*time.Second, 1*time.Second, "only CPU pod should remain after GPU pods evicted")
+			},
+		},
 		{
 			name:            "DrainOverrides.Force overrides all namespace modes",
 			nodeName:        "force-node",
@@ -737,7 +846,7 @@ func TestReconciler_ProcessEvent(t *testing.T) {
 				{Name: "immediate-*", Mode: config.ModeImmediateEvict},
 				{Name: "completion-*", Mode: config.ModeAllowCompletion},
 				{Name: "timeout-*", Mode: config.ModeDeleteAfterTimeout},
-			}, false)
+			}, false, tt.enableDrainGPUPods)
 
 			nodeLabels := tt.existingNodeLabels
 			if nodeLabels == nil {
@@ -1198,7 +1307,7 @@ type testSetup struct {
 	mockDB            *mockDataStore
 }
 
-func setupDirectTest(t *testing.T, userNamespaces []config.UserNamespace, dryRun bool) *testSetup {
+func setupDirectTest(t *testing.T, userNamespaces []config.UserNamespace, dryRun bool, drainGPUPods ...bool) *testSetup {
 	t.Helper()
 	ctx := t.Context()
 
@@ -1210,11 +1319,17 @@ func setupDirectTest(t *testing.T, userNamespaces []config.UserNamespace, dryRun
 	client, err := kubernetes.NewForConfig(cfg)
 	require.NoError(t, err)
 
+	enableDrainGPUPods := false
+	if len(drainGPUPods) > 0 {
+		enableDrainGPUPods = drainGPUPods[0]
+	}
+
 	tomlConfig := config.TomlConfig{
 		EvictionTimeoutInSeconds:  config.Duration{Duration: 30 * time.Second},
 		SystemNamespaces:          "kube-*",
 		DeleteAfterTimeoutMinutes: 5,
 		NotReadyTimeoutMinutes:    2,
+		DrainGPUPods:              enableDrainGPUPods,
 		UserNamespaces:            userNamespaces,
 		PartialDrainEnabled:       true,
 	}
@@ -1237,7 +1352,7 @@ func setupDirectTest(t *testing.T, userNamespaces []config.UserNamespace, dryRun
 		StateManager: statemanager.NewStateManager(client),
 	}
 
-	informersInstance, err := informers.NewInformers(client, 1*time.Minute, ptr.To(2), dryRun)
+	informersInstance, err := informers.NewInformers(client, 1*time.Minute, ptr.To(2), enableDrainGPUPods, dryRun)
 	require.NoError(t, err)
 
 	go func() { _ = informersInstance.Run(ctx) }()
@@ -1311,6 +1426,7 @@ func setupCustomDrainTest(t *testing.T, customDrainConfig config.CustomDrainConf
 		SystemNamespaces:          "kube-*",
 		DeleteAfterTimeoutMinutes: 5,
 		NotReadyTimeoutMinutes:    2,
+		DrainGPUPods:              false,
 		CustomDrain:               customDrainConfig,
 		PartialDrainEnabled:       true,
 	}
@@ -1332,7 +1448,7 @@ func setupCustomDrainTest(t *testing.T, customDrainConfig config.CustomDrainConf
 		StateManager: statemanager.NewStateManager(client),
 	}
 
-	informersInstance, err := informers.NewInformers(client, 1*time.Minute, ptr.To(2), false)
+	informersInstance, err := informers.NewInformers(client, 1*time.Minute, ptr.To(2), false, false)
 	require.NoError(t, err)
 
 	go func() { _ = informersInstance.Run(ctx) }()
@@ -2041,6 +2157,7 @@ func TestReconciler_CustomDrainCRDNotFound(t *testing.T) {
 		SystemNamespaces:          "kube-*",
 		DeleteAfterTimeoutMinutes: 5,
 		NotReadyTimeoutMinutes:    2,
+		DrainGPUPods:              false,
 		CustomDrain:               customDrainCfg,
 		UserNamespaces: []config.UserNamespace{
 			{Name: "*", Mode: config.ModeImmediateEvict},
@@ -2065,7 +2182,7 @@ func TestReconciler_CustomDrainCRDNotFound(t *testing.T) {
 		StateManager: statemanager.NewStateManager(client),
 	}
 
-	informersInstance, err := informers.NewInformers(client, 1*time.Minute, ptr.To(2), false)
+	informersInstance, err := informers.NewInformers(client, 1*time.Minute, ptr.To(2), false, false)
 	require.NoError(t, err)
 
 	go func() { _ = informersInstance.Run(ctx) }()