From c5d3eafe7efb6cec302a6e428e91ea37a125405d Mon Sep 17 00:00:00 2001 From: Rahul Sharma Date: Wed, 22 Jul 2026 14:12:27 -0700 Subject: [PATCH] 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 --- controllers/upgrade_controller.go | 56 +++++++++++++++++++++++++ controllers/upgrade_controller_test.go | 58 ++++++++++++++++++++++++++ 2 files changed, 114 insertions(+) diff --git a/controllers/upgrade_controller.go b/controllers/upgrade_controller.go index ca3627ddf9..a198460f41 100644 --- a/controllers/upgrade_controller.go +++ b/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 { diff --git a/controllers/upgrade_controller_test.go b/controllers/upgrade_controller_test.go index 3b72da082f..faf765093e 100644 --- a/controllers/upgrade_controller_test.go +++ b/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) + } + }) + } +}