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,74 @@ 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()
+
+	// List all nodes
+	nodeList := &corev1.NodeList{}
+	if err := r.List(ctx, nodeList); err != nil {
+		return fmt.Errorf("failed to list nodes: %w", err)
+	}
+
+	// Filter nodes that have the upgrade label (any value)
+	var nodesWithUpgradeLabel []corev1.Node
+	for _, node := range nodeList.Items {
+		if _, hasLabel := node.Labels[upgradeStateLabel]; hasLabel {
+			nodesWithUpgradeLabel = append(nodesWithUpgradeLabel, node)
+		}
+	}
+
+	if len(nodesWithUpgradeLabel) == 0 {
+		return nil // No nodes with upgrade labels
+	}
+
+	// Build a set of nodes being actively managed by the upgrade process
+	// We should not remove labels from these nodes as they may be mid-upgrade
+	managedNodes := make(map[string]bool)
+	for nodeName := range state.NodeStates {
+		managedNodes[nodeName] = true
+	}
+
+	// List all driver pods (including orphaned ones)
+	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 nodesWithUpgradeLabel {
+		node := &nodesWithUpgradeLabel[i]
+		// Skip nodes being actively managed by upgrade process
+		if managedNodes[node.Name] {
+			continue
+		}
+		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 {