remove driver upgrade label from nodes

this commit removes driver upgrade label from nodes which don't have any driver pod running on them

Signed-off-by: Rahul Sharma <rahulsharm@nvidia.com>

Author: rahulait

controllers/upgrade_controller.go

@@ -152,6 +152,13 @@ func (r *UpgradeReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ct
 		return ctrl.Result{}, err
 	}
 
+	// Clear stale upgrade labels from nodes that no longer have driver pods
+	// Use the built state so we can avoid removing labels from nodes actively being upgraded
+	if err := r.clearUpgradeLabelsWhereDriverNotRunning(ctx, state, driverLabel, clusterPolicyCtrl.operatorNamespace); err != nil {
+		// Log the error but continue with the upgrade process, as this is a best-effort cleanup and should not block upgrades
+		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)
 
@@ -198,6 +205,78 @@ func (r *UpgradeReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ct
 	return ctrl.Result{Requeue: true, RequeueAfter: plannedRequeueInterval}, nil
 }
 
+// clearUpgradeLabelsWhereDriverNotRunning removes upgrade labels from nodes where driver pods are no longer scheduled.
+// This handles the case where a nodeSelector change causes pods to be terminated from certain nodes,
+// but the upgrade labels remain. It skips nodes that are actively being managed by the upgrade process.
+func (r *UpgradeReconciler) clearUpgradeLabelsWhereDriverNotRunning(ctx context.Context, state *upgrade.ClusterUpgradeState, driverLabel map[string]string, namespace string) error {
+	upgradeStateLabel := upgrade.GetUpgradeStateLabelKey()
+
+	// Build a set of nodes being actively managed by the upgrade process first
+	// This allows us to skip API calls if all labeled nodes are already managed
+	managedNodes := make(map[string]bool)
+	for _, nodeStates := range state.NodeStates {
+		for _, nodeState := range nodeStates {
+			if nodeState.Node != nil {
+				managedNodes[nodeState.Node.Name] = true
+			}
+		}
+	}
+
+	// List only nodes that have the upgrade label
+	nodeList := &corev1.NodeList{}
+	if err := r.List(ctx, nodeList, client.HasLabels{upgradeStateLabel}); err != nil {
+		return fmt.Errorf("failed to list nodes with upgrade labels: %w", err)
+	}
+
+	if len(nodeList.Items) == 0 {
+		return nil
+	}
+
+	// Filter out nodes being actively managed by upgrade process
+	var nodesToCheck []corev1.Node
+	for _, node := range nodeList.Items {
+		if !managedNodes[node.Name] {
+			nodesToCheck = append(nodesToCheck, node)
+		}
+	}
+
+	if len(nodesToCheck) == 0 {
+		return nil
+	}
+
+	// List driver pods only if we have nodes to check (optimization)
+	podList := &corev1.PodList{}
+	if err := r.List(ctx, podList, client.InNamespace(namespace), client.MatchingLabels(driverLabel)); err != nil {
+		return fmt.Errorf("failed to list driver pods: %w", err)
+	}
+
+	// Create a set of nodes that have driver pods (any driver pods)
+	nodesWithPods := make(map[string]bool)
+	for _, pod := range podList.Items {
+		if pod.Spec.NodeName != "" {
+			nodesWithPods[pod.Spec.NodeName] = true
+		}
+	}
+
+	// Clear upgrade label from nodes that don't have driver pods
+	for i := range nodesToCheck {
+		node := &nodesToCheck[i]
+		if _, hasDriverPod := nodesWithPods[node.Name]; !hasDriverPod {
+			r.Log.Info("Clearing stale upgrade label from node", "node", node.Name)
+
+			nodeCopy := node.DeepCopy()
+			delete(node.Labels, upgradeStateLabel)
+			if err := r.Patch(ctx, node, client.MergeFrom(nodeCopy)); err != nil {
+				r.Log.Error(err, "Failed to clear upgrade label from node", "node", node.Name)
+				// Continue with other nodes even if one fails
+				continue
+			}
+		}
+	}
+
+	return nil
+}
+
 // 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 {