[improvement] : remove driver upgrade label from nodes

controllers/upgrade_controller.go

@@ -153,6 +153,10 @@ func (r *UpgradeReconciler) reconcileClusterPolicyDriverUpgrades(ctx context.Con
 		return ctrl.Result{}, err
 	}
 
+	if err := r.clearStaleUpgradeLabels(ctx, state, driverLabel, clusterPolicyCtrl.operatorNamespace); err != nil {
+		r.Log.Error(err, "Failed to clear stale upgrade labels")
+	}
+
 	reqLogger.Info("Propagate state to state manager")
 	reqLogger.V(consts.LogLevelDebug).Info("Current cluster upgrade state", "state", state)
 
@@ -240,6 +244,10 @@ func (r *UpgradeReconciler) reconcileNVIDIADriverUpgrades(ctx context.Context, r
 		return ctrl.Result{}, err
 	}
 
+	if err := r.clearStaleUpgradeLabels(ctx, clusterState, map[string]string{AppComponentLabelKey: DriverAppComponentLabelValue}, clusterPolicyCtrl.operatorNamespace); err != nil {
+		r.Log.Error(err, "Failed to clear stale upgrade labels")
+	}
+
 	// Partition the cluster upgrade state into per-NVIDIADriver buckets by reading the
 	// nvidia.com/gpu-operator.driver.owner label from each node.
 	statesByNVD := make(map[string]*upgrade.ClusterUpgradeState)
@@ -324,6 +332,54 @@ func (r *UpgradeReconciler) reconcileNVIDIADriverUpgrades(ctx context.Context, r
 	return ctrl.Result{Requeue: true, RequeueAfter: plannedRequeueInterval}, nil
 }
 
+// clearStaleUpgradeLabels removes upgrade-state labels from nodes no longer selected by a driver DaemonSet.
+func (r *UpgradeReconciler) clearStaleUpgradeLabels(ctx context.Context, state *upgrade.ClusterUpgradeState, driverLabel map[string]string, namespace string) error {
+	upgradeStateLabel := upgrade.GetUpgradeStateLabelKey()
+	managedNodes := make(map[string]bool)
+	for _, nodeStates := range state.NodeStates {
+		for _, nodeState := range nodeStates {
+			if nodeState.Node != nil {
+				managedNodes[nodeState.Node.Name] = true
+			}
+		}
+	}
+
+	nodeList := &corev1.NodeList{}
+	if err := r.List(ctx, nodeList, client.HasLabels{upgradeStateLabel}); err != nil {
+		return fmt.Errorf("list nodes with upgrade labels: %w", err)
+	}
+
+	driverDaemonSets := &appsv1.DaemonSetList{}
+	if err := r.List(ctx, driverDaemonSets, client.InNamespace(namespace), client.MatchingLabels(driverLabel)); err != nil {
+		return fmt.Errorf("list driver DaemonSets: %w", err)
+	}
+
+	for index := range nodeList.Items {
+		node := &nodeList.Items[index]
+		if managedNodes[node.Name] || selectedByDriverDaemonSet(node, driverDaemonSets.Items) {
+			continue
+		}
+
+		patch := client.MergeFrom(node.DeepCopy())
+		delete(node.Labels, upgradeStateLabel)
+		if err := r.Patch(ctx, node, patch); err != nil {
+			r.Log.Error(err, "Failed to clear stale upgrade label from node", "node", node.Name)
+		}
+	}
+
+	return nil
+}
+
+func selectedByDriverDaemonSet(node *corev1.Node, driverDaemonSets []appsv1.DaemonSet) bool {
+	for _, daemonSet := range driverDaemonSets {
+		selector := labels.SelectorFromSet(daemonSet.Spec.Template.Spec.NodeSelector)
+		if selector.Matches(labels.Set(node.Labels)) {
+			return true
+		}
+	}
+	return false
+}
+
 // removeNodeUpgradeStateLabels loops over nodes in the cluster and removes "nvidia.com/gpu-driver-upgrade-state"
 // It is used for cleanup when autoUpgrade feature gets disabled
 func (r *UpgradeReconciler) removeNodeUpgradeStateLabels(ctx context.Context) error {

controllers/upgrade_controller_test.go

@@ -17,11 +17,20 @@
 package controllers
 
 import (
+	"context"
 	"fmt"
 	"testing"
 
 	upgrade_v1alpha1 "github.com/NVIDIA/k8s-operator-libs/api/upgrade/v1alpha1"
+	"github.com/NVIDIA/k8s-operator-libs/pkg/upgrade"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	appsv1 "k8s.io/api/apps/v1"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/client/fake"
 )
 
 func TestSetDrainSpecPodSelector(t *testing.T) {
@@ -69,3 +78,52 @@ func TestSetDrainSpecPodSelector(t *testing.T) {
 		})
 	}
 }
+
+func TestClearStaleUpgradeLabels(t *testing.T) {
+	upgradeStateLabel := upgrade.GetUpgradeStateLabelKey()
+	tests := []struct {
+		name          string
+		nodeLabels    map[string]string
+		daemonSets    []appsv1.DaemonSet
+		expectRemoved bool
+	}{
+		{
+			name:          "removes label from node excluded by driver node selector",
+			nodeLabels:    map[string]string{upgradeStateLabel: "upgrade-required", "gpu": "true"},
+			daemonSets:    []appsv1.DaemonSet{{ObjectMeta: metav1.ObjectMeta{Name: "driver", Namespace: "gpu-operator", Labels: map[string]string{DriverLabelKey: DriverLabelValue}}, Spec: appsv1.DaemonSetSpec{Template: corev1.PodTemplateSpec{Spec: corev1.PodSpec{NodeSelector: map[string]string{"gpu": "true", "type": "A100"}}}}}},
+			expectRemoved: true,
+		},
+		{
+			name:          "keeps label while a driver DaemonSet still targets node",
+			nodeLabels:    map[string]string{upgradeStateLabel: "upgrade-required", "gpu": "true", "type": "A100"},
+			daemonSets:    []appsv1.DaemonSet{{ObjectMeta: metav1.ObjectMeta{Name: "driver", Namespace: "gpu-operator", Labels: map[string]string{DriverLabelKey: DriverLabelValue}}, Spec: appsv1.DaemonSetSpec{Template: corev1.PodTemplateSpec{Spec: corev1.PodSpec{NodeSelector: map[string]string{"gpu": "true", "type": "A100"}}}}}},
+			expectRemoved: false,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			scheme := runtime.NewScheme()
+			require.NoError(t, corev1.AddToScheme(scheme))
+			require.NoError(t, appsv1.AddToScheme(scheme))
+
+			node := &corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: "node-1", Labels: test.nodeLabels}}
+			objects := []runtime.Object{node}
+			for index := range test.daemonSets {
+				objects = append(objects, &test.daemonSets[index])
+			}
+			reconciler := &UpgradeReconciler{Client: fake.NewClientBuilder().WithScheme(scheme).WithRuntimeObjects(objects...).Build()}
+
+			err := reconciler.clearStaleUpgradeLabels(context.Background(), &upgrade.ClusterUpgradeState{}, map[string]string{DriverLabelKey: DriverLabelValue}, "gpu-operator")
+			require.NoError(t, err)
+
+			updatedNode := &corev1.Node{}
+			require.NoError(t, reconciler.Get(context.Background(), client.ObjectKey{Name: node.Name}, updatedNode))
+			if test.expectRemoved {
+				assert.NotContains(t, updatedNode.Labels, upgradeStateLabel)
+			} else {
+				assert.Contains(t, updatedNode.Labels, upgradeStateLabel)
+			}
+		})
+	}
+}