diff --git a/controllers/nodelabeling_controller.go b/controllers/nodelabeling_controller.go index 63bd1d624a..1594679747 100644 --- a/controllers/nodelabeling_controller.go +++ b/controllers/nodelabeling_controller.go @@ -57,6 +57,8 @@ type NodeLabelingReconciler struct { Scheme *runtime.Scheme Namespace string Log logr.Logger + + nvidiaDriverNodeSelectorCache nvidiaDriverNodeSelectorCache } // nodeLabelingController holds per-reconcile state so that helper methods don't need to @@ -139,6 +141,9 @@ func getNodeLabelUpdateReasons(oldLabels, newLabels map[string]string) nodeLabel // Reconcile applies GPU-Operator related labels and annotations to all cluster nodes. func (r *NodeLabelingReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) { r.Log.Info("Reconciling node labels") + if err := r.nvidiaDriverNodeSelectorCache.refresh(ctx, r.Client); err != nil { + return reconcile.Result{}, fmt.Errorf("failed to refresh NVIDIADriver node-selector cache: %w", err) + } // The ClusterPolicy (device-plugin stack) and GPUCluster (DRA stack) CRs may // coexist; neither existing means there is nothing to label. @@ -639,7 +644,7 @@ func (nlc *nodeLabelingController) labelNodesWithOrphanedDriverPods(ctx context. return nil } -// nodeOwnedByNVIDIADriver returns true when the node has an owner label matching a live NVIDIADriver. +// nodeOwnedByNVIDIADriver returns true when the node has an owner label matching a NVIDIADriver that is not being deleted. func nodeOwnedByNVIDIADriver(node *corev1.Node, nvidiaDrivers []nvidiav1alpha1.NVIDIADriver) bool { if node.Labels == nil || node.Labels[consts.NVIDIADriverOwnerLabel] == "" { return false @@ -713,15 +718,26 @@ func (r *NodeLabelingReconciler) SetupWithManager(ctx context.Context, mgr ctrl. return fmt.Errorf("error watching GPUCluster: %w", err) } - // Watch NVIDIADriver including delete events so owner labels are cleaned up promptly. - nvidiaDriverMapFn := func(ctx context.Context, nd *nvidiav1alpha1.NVIDIADriver) []reconcile.Request { - return mapToSingleton(ctx, nd) + nvidiaDriverMapFn := func(ctx context.Context, driver *nvidiav1alpha1.NVIDIADriver) []reconcile.Request { + return mapToSingleton(ctx, driver) + } + nvidiaDriverPredicate := predicate.TypedFuncs[*nvidiav1alpha1.NVIDIADriver]{ + CreateFunc: func(e event.TypedCreateEvent[*nvidiav1alpha1.NVIDIADriver]) bool { + return true + }, + UpdateFunc: func(e event.TypedUpdateEvent[*nvidiav1alpha1.NVIDIADriver]) bool { + return e.ObjectOld.GetGeneration() != e.ObjectNew.GetGeneration() || + (e.ObjectOld.GetDeletionTimestamp() == nil) != (e.ObjectNew.GetDeletionTimestamp() == nil) + }, + DeleteFunc: func(e event.TypedDeleteEvent[*nvidiav1alpha1.NVIDIADriver]) bool { + return true + }, } if err := c.Watch(source.Kind( mgr.GetCache(), &nvidiav1alpha1.NVIDIADriver{}, handler.TypedEnqueueRequestsFromMapFunc(nvidiaDriverMapFn), - predicate.TypedGenerationChangedPredicate[*nvidiav1alpha1.NVIDIADriver]{}, + nvidiaDriverPredicate, )); err != nil { return fmt.Errorf("error watching NVIDIADriver: %w", err) } @@ -739,24 +755,15 @@ func (r *NodeLabelingReconciler) SetupWithManager(ctx context.Context, mgr ctrl. reasons := getNodeLabelUpdateReasons(oldLabels, newLabels) needsUpdate := reasons.needsUpdate() - // When an NVIDIADriver daemonset pod is running on the node, check if any - // label which is configured in the NVIDIADriver's node selector has changed. nvidiaDriverNodeSelectorLabelChanged := false - if !needsUpdate && newLabels[consts.NVIDIADriverOwnerLabel] != "" { - name := newLabels[consts.NVIDIADriverOwnerLabel] - nvidiaDriver := &nvidiav1alpha1.NVIDIADriver{} - err := r.Get(ctx, types.NamespacedName{Name: name}, nvidiaDriver) + if !needsUpdate && (hasGPULabels(oldLabels) || hasGPULabels(newLabels)) { + changed, err := r.nvidiaDriverNodeSelectorCache.selectorLabelsChanged(ctx, r.Client, oldLabels, newLabels) if err != nil { - r.Log.Error(err, "failed to get NVIDIADriver object that owns this node", "name", name, "node", nodeName) - return false - } - for key := range nvidiaDriver.Spec.NodeSelector { - if oldLabels[key] != newLabels[key] { - nvidiaDriverNodeSelectorLabelChanged = true - needsUpdate = true - break - } + r.Log.Error(err, "failed to initialize NVIDIADriver selector-key cache") + return true } + nvidiaDriverNodeSelectorLabelChanged = changed + needsUpdate = nvidiaDriverNodeSelectorLabelChanged } if needsUpdate { diff --git a/controllers/nodelabeling_controller_test.go b/controllers/nodelabeling_controller_test.go index 8ec98d6d21..23e31bf6a6 100644 --- a/controllers/nodelabeling_controller_test.go +++ b/controllers/nodelabeling_controller_test.go @@ -1031,8 +1031,8 @@ func TestLabelNodesWithOrphanedDriverPods(t *testing.T) { upgradeStateLabel := upgrade.GetUpgradeStateLabelKey() - // liveDriver returns a NVIDIADriver with no deletion timestamp. - liveDriver := func() *nvidiav1alpha1.NVIDIADriver { + // nonDeletingDriver returns a NVIDIADriver with no deletion timestamp. + nonDeletingDriver := func() *nvidiav1alpha1.NVIDIADriver { return &nvidiav1alpha1.NVIDIADriver{ ObjectMeta: metav1.ObjectMeta{Name: driverName}, } @@ -1080,35 +1080,35 @@ func TestLabelNodesWithOrphanedDriverPods(t *testing.T) { }, { name: "orphaned pod on owned node, no upgrade state → labeled upgrade-required", - nvidiaDrivers: []*nvidiav1alpha1.NVIDIADriver{liveDriver()}, + nvidiaDrivers: []*nvidiav1alpha1.NVIDIADriver{nonDeletingDriver()}, nodes: []*corev1.Node{ownedNode("node-1", "")}, pods: []*corev1.Pod{orphanedPod("pod-1", "node-1")}, expectedUpgradeState: map[string]string{"node-1": upgrade.UpgradeStateUpgradeRequired}, }, { name: "orphaned pod on owned node, upgrade-done state → labeled upgrade-required", - nvidiaDrivers: []*nvidiav1alpha1.NVIDIADriver{liveDriver()}, + nvidiaDrivers: []*nvidiav1alpha1.NVIDIADriver{nonDeletingDriver()}, nodes: []*corev1.Node{ownedNode("node-1", upgrade.UpgradeStateDone)}, pods: []*corev1.Pod{orphanedPod("pod-1", "node-1")}, expectedUpgradeState: map[string]string{"node-1": upgrade.UpgradeStateUpgradeRequired}, }, { name: "orphaned pod on owned node, active upgrade state → not relabeled", - nvidiaDrivers: []*nvidiav1alpha1.NVIDIADriver{liveDriver()}, + nvidiaDrivers: []*nvidiav1alpha1.NVIDIADriver{nonDeletingDriver()}, nodes: []*corev1.Node{ownedNode("node-1", upgrade.UpgradeStatePodRestartRequired)}, pods: []*corev1.Pod{orphanedPod("pod-1", "node-1")}, expectedUpgradeState: map[string]string{"node-1": upgrade.UpgradeStatePodRestartRequired}, }, { name: "orphaned pod on owned node, failed upgrade state → not relabeled", - nvidiaDrivers: []*nvidiav1alpha1.NVIDIADriver{liveDriver()}, + nvidiaDrivers: []*nvidiav1alpha1.NVIDIADriver{nonDeletingDriver()}, nodes: []*corev1.Node{ownedNode("node-1", upgrade.UpgradeStateFailed)}, pods: []*corev1.Pod{orphanedPod("pod-1", "node-1")}, expectedUpgradeState: map[string]string{"node-1": upgrade.UpgradeStateFailed}, }, { name: "pod has owner references → skipped", - nvidiaDrivers: []*nvidiav1alpha1.NVIDIADriver{liveDriver()}, + nvidiaDrivers: []*nvidiav1alpha1.NVIDIADriver{nonDeletingDriver()}, nodes: []*corev1.Node{ownedNode("node-1", "")}, pods: []*corev1.Pod{func() *corev1.Pod { p := orphanedPod("pod-1", "node-1") @@ -1119,7 +1119,7 @@ func TestLabelNodesWithOrphanedDriverPods(t *testing.T) { }, { name: "pod not in Running phase → skipped", - nvidiaDrivers: []*nvidiav1alpha1.NVIDIADriver{liveDriver()}, + nvidiaDrivers: []*nvidiav1alpha1.NVIDIADriver{nonDeletingDriver()}, nodes: []*corev1.Node{ownedNode("node-1", "")}, pods: []*corev1.Pod{func() *corev1.Pod { p := orphanedPod("pod-1", "node-1") @@ -1130,7 +1130,7 @@ func TestLabelNodesWithOrphanedDriverPods(t *testing.T) { }, { name: "pod has no NodeName → skipped", - nvidiaDrivers: []*nvidiav1alpha1.NVIDIADriver{liveDriver()}, + nvidiaDrivers: []*nvidiav1alpha1.NVIDIADriver{nonDeletingDriver()}, nodes: []*corev1.Node{ownedNode("node-1", "")}, pods: []*corev1.Pod{func() *corev1.Pod { p := orphanedPod("pod-1", "node-1") @@ -1141,7 +1141,7 @@ func TestLabelNodesWithOrphanedDriverPods(t *testing.T) { }, { name: "node not owned by any NVIDIADriver → not labeled", - nvidiaDrivers: []*nvidiav1alpha1.NVIDIADriver{liveDriver()}, + nvidiaDrivers: []*nvidiav1alpha1.NVIDIADriver{nonDeletingDriver()}, nodes: []*corev1.Node{{ ObjectMeta: metav1.ObjectMeta{Name: "node-1"}, // no NVIDIADriverOwnerLabel }}, diff --git a/controllers/nvidiadriver_controller.go b/controllers/nvidiadriver_controller.go index 14ca610111..025562db68 100644 --- a/controllers/nvidiadriver_controller.go +++ b/controllers/nvidiadriver_controller.go @@ -20,7 +20,6 @@ import ( "context" "errors" "fmt" - "maps" "time" appsv1 "k8s.io/api/apps/v1" @@ -59,6 +58,7 @@ type NVIDIADriverReconciler struct { stateManager state.Manager nodeSelectorValidator validator.Validator conditionUpdater conditions.Updater + nodeSelectorCache nvidiaDriverNodeSelectorCache } //+kubebuilder:rbac:groups=nvidia.com,resources=nvidiadrivers,verbs=get;list;watch;create;update;patch;delete @@ -82,6 +82,7 @@ func (r *NVIDIADriverReconciler) Reconcile(ctx context.Context, req ctrl.Request instance := &nvidiav1alpha1.NVIDIADriver{} if err := r.Get(ctx, req.NamespacedName, instance); err != nil { if apierrors.IsNotFound(err) { + r.nodeSelectorCache.removeByKey(req.NamespacedName) // Request object not found, could have been deleted after reconcile request. return reconcile.Result{}, nil } @@ -94,6 +95,7 @@ func (r *NVIDIADriverReconciler) Reconcile(ctx context.Context, req ctrl.Request // Error reading the object - requeue the request. return reconcile.Result{}, wrappedErr } + r.nodeSelectorCache.replace(instance) if instance.HasDeletionTimestamp() { return reconcile.Result{}, nil } @@ -334,6 +336,47 @@ func dedupeReconcileRequests(requests []reconcile.Request) []reconcile.Request { return deduped } +func nodeDriverRenderingLabelsChanged(oldLabels, newLabels map[string]string) bool { + relevantLabels := []string{ + consts.NVIDIADriverOwnerLabel, + consts.GPUPresentLabel, + driverDeployLabelKey, + vgpuManagerDeployLabelKey, + nfdOSReleaseIDLabelKey, + nfdOSVersionIDLabelKey, + nfdKernelLabelKey, + nfdOSTreeVersionLabelKey, + } + + for _, label := range relevantLabels { + if oldLabels[label] != newLabels[label] { + return true + } + } + return hasGPULabels(oldLabels) != hasGPULabels(newLabels) +} + +func (r *NVIDIADriverReconciler) enqueueNodeDriverReconcileRequests(ctx context.Context, queue workqueue.TypedRateLimitingInterface[reconcile.Request], oldLabels, newLabels map[string]string) { + if !r.nodeLabelsAffectNVIDIADrivers(ctx, oldLabels, newLabels) { + return + } + for _, request := range r.enqueueAllNVIDIADrivers(ctx) { + queue.Add(request) + } +} + +func (r *NVIDIADriverReconciler) nodeLabelsAffectNVIDIADrivers(ctx context.Context, oldLabels, newLabels map[string]string) bool { + if nodeDriverRenderingLabelsChanged(oldLabels, newLabels) { + return true + } + changed, err := r.nodeSelectorCache.selectorLabelsChanged(ctx, r.Client, oldLabels, newLabels) + if err != nil { + log.FromContext(ctx).Error(err, "failed to initialize NVIDIADriver node-selector cache") + return true + } + return changed +} + // SetupWithManager sets up the controller with the Manager. func (r *NVIDIADriverReconciler) SetupWithManager(ctx context.Context, mgr ctrl.Manager) error { // Create state manager @@ -363,18 +406,27 @@ func (r *NVIDIADriverReconciler) SetupWithManager(ctx context.Context, mgr ctrl. return err } - // Watch for changes to NVIDIADriver CRs. Whenever an event is generated for a NVIDIADriver CR, - // enqueue a reconcile request for all NVIDIADriver instances. nvidiaDriverMapFn := func(ctx context.Context, driver *nvidiav1alpha1.NVIDIADriver) []reconcile.Request { return r.enqueueNVIDIADriverReconcilers(ctx, driver) } + nvidiaDriverPredicate := predicate.TypedFuncs[*nvidiav1alpha1.NVIDIADriver]{ + CreateFunc: func(e event.TypedCreateEvent[*nvidiav1alpha1.NVIDIADriver]) bool { + return true + }, + UpdateFunc: func(e event.TypedUpdateEvent[*nvidiav1alpha1.NVIDIADriver]) bool { + return e.ObjectOld.GetGeneration() != e.ObjectNew.GetGeneration() || + (e.ObjectOld.GetDeletionTimestamp() == nil) != (e.ObjectNew.GetDeletionTimestamp() == nil) + }, + DeleteFunc: func(e event.TypedDeleteEvent[*nvidiav1alpha1.NVIDIADriver]) bool { + return true + }, + } - // Watch for changes to the primary resource NVIDIADriver err = c.Watch(source.Kind( mgr.GetCache(), &nvidiav1alpha1.NVIDIADriver{}, handler.TypedEnqueueRequestsFromMapFunc(nvidiaDriverMapFn), - predicate.TypedGenerationChangedPredicate[*nvidiav1alpha1.NVIDIADriver]{}, + nvidiaDriverPredicate, ), ) if err != nil { @@ -387,12 +439,6 @@ func (r *NVIDIADriverReconciler) SetupWithManager(ctx context.Context, mgr ctrl. return r.enqueueAllNVIDIADrivers(ctx) } - // Watch for changes to the Nodes. Whenever an event is generated for a Node, enqueue - // a reconcile request for all NVIDIADriver instances. - nodeMapFn := func(ctx context.Context, _ *corev1.Node) []reconcile.Request { - return r.enqueueAllNVIDIADrivers(ctx) - } - err = c.Watch( source.Kind( mgr.GetCache(), @@ -421,38 +467,30 @@ func (r *NVIDIADriverReconciler) SetupWithManager(ctx context.Context, mgr ctrl. return err } - nodePredicate := predicate.TypedFuncs[*corev1.Node]{ - CreateFunc: func(e event.TypedCreateEvent[*corev1.Node]) bool { - labels := e.Object.GetLabels() - return hasGPULabels(labels) + nodeHandler := handler.TypedFuncs[*corev1.Node, reconcile.Request]{ + CreateFunc: func(ctx context.Context, e event.TypedCreateEvent[*corev1.Node], queue workqueue.TypedRateLimitingInterface[reconcile.Request]) { + if hasGPULabels(e.Object.GetLabels()) { + r.enqueueNodeDriverReconcileRequests(ctx, queue, nil, e.Object.GetLabels()) + } }, - UpdateFunc: func(e event.TypedUpdateEvent[*corev1.Node]) bool { - logger := log.FromContext(ctx) - newLabels := e.ObjectNew.GetLabels() + UpdateFunc: func(ctx context.Context, e event.TypedUpdateEvent[*corev1.Node], queue workqueue.TypedRateLimitingInterface[reconcile.Request]) { oldLabels := e.ObjectOld.GetLabels() - nodeName := e.ObjectNew.GetName() - - needsUpdate := hasGPULabels(newLabels) && !maps.Equal(newLabels, oldLabels) - - if needsUpdate { - logger.Info("Node labels have been changed", - "name", nodeName, - ) + newLabels := e.ObjectNew.GetLabels() + if hasGPULabels(oldLabels) || hasGPULabels(newLabels) { + r.enqueueNodeDriverReconcileRequests(ctx, queue, oldLabels, newLabels) } - return needsUpdate }, - DeleteFunc: func(e event.TypedDeleteEvent[*corev1.Node]) bool { - labels := e.Object.GetLabels() - return hasGPULabels(labels) + DeleteFunc: func(ctx context.Context, e event.TypedDeleteEvent[*corev1.Node], queue workqueue.TypedRateLimitingInterface[reconcile.Request]) { + if hasGPULabels(e.Object.GetLabels()) { + r.enqueueNodeDriverReconcileRequests(ctx, queue, e.Object.GetLabels(), nil) + } }, } - // Watch for changes to node labels err = c.Watch( source.Kind(mgr.GetCache(), &corev1.Node{}, - handler.TypedEnqueueRequestsFromMapFunc(nodeMapFn), - nodePredicate, + nodeHandler, ), ) if err != nil { diff --git a/controllers/nvidiadriver_controller_test.go b/controllers/nvidiadriver_controller_test.go index 3934b376d0..a8e8dae86d 100644 --- a/controllers/nvidiadriver_controller_test.go +++ b/controllers/nvidiadriver_controller_test.go @@ -40,6 +40,7 @@ import ( gpuv1 "github.com/NVIDIA/gpu-operator/api/nvidia/v1" nvidiav1alpha1 "github.com/NVIDIA/gpu-operator/api/nvidia/v1alpha1" "github.com/NVIDIA/gpu-operator/internal/conditions" + "github.com/NVIDIA/gpu-operator/internal/consts" "github.com/NVIDIA/gpu-operator/internal/state" "github.com/NVIDIA/gpu-operator/internal/validator" ) @@ -416,3 +417,93 @@ func TestEnqueueNVIDIADriverReconcilersDedupesEventDriver(t *testing.T) { require.Len(t, requests, 1) require.Equal(t, "default/driver-a", requests[0].String()) } + +func TestNodeDriverRenderingLabelsChanged(t *testing.T) { + gpuLabels := map[string]string{"feature.node.kubernetes.io/pci-10de.present": "true"} + + tests := []struct { + name string + oldLabels map[string]string + newLabels map[string]string + want bool + }{ + { + name: "unrelated label change is ignored", + oldLabels: map[string]string{"feature.node.kubernetes.io/pci-10de.present": "true", consts.NVIDIADriverOwnerLabel: "driver-a", "example.com/label": "old"}, + newLabels: map[string]string{"feature.node.kubernetes.io/pci-10de.present": "true", consts.NVIDIADriverOwnerLabel: "driver-a", "example.com/label": "new"}, + }, + { + name: "NFD kernel change is relevant", + oldLabels: map[string]string{"feature.node.kubernetes.io/pci-10de.present": "true", consts.NVIDIADriverOwnerLabel: "driver-a", nfdKernelLabelKey: "old"}, + newLabels: map[string]string{"feature.node.kubernetes.io/pci-10de.present": "true", consts.NVIDIADriverOwnerLabel: "driver-a", nfdKernelLabelKey: "new"}, + want: true, + }, + { + name: "driver deployment label change is relevant", + oldLabels: map[string]string{"feature.node.kubernetes.io/pci-10de.present": "true", consts.NVIDIADriverOwnerLabel: "driver-a", driverDeployLabelKey: "true"}, + newLabels: map[string]string{"feature.node.kubernetes.io/pci-10de.present": "true", consts.NVIDIADriverOwnerLabel: "driver-a", driverDeployLabelKey: "false"}, + want: true, + }, + { + name: "vGPU manager deployment label change is relevant", + oldLabels: map[string]string{"feature.node.kubernetes.io/pci-10de.present": "true", consts.NVIDIADriverOwnerLabel: "driver-a", vgpuManagerDeployLabelKey: "true"}, + newLabels: map[string]string{"feature.node.kubernetes.io/pci-10de.present": "true", consts.NVIDIADriverOwnerLabel: "driver-a", vgpuManagerDeployLabelKey: "false"}, + want: true, + }, + { + name: "GPU present label change is relevant", + oldLabels: map[string]string{"feature.node.kubernetes.io/pci-10de.present": "true", consts.NVIDIADriverOwnerLabel: "driver-a", consts.GPUPresentLabel: "true"}, + newLabels: map[string]string{"feature.node.kubernetes.io/pci-10de.present": "true", consts.NVIDIADriverOwnerLabel: "driver-a", consts.GPUPresentLabel: "false"}, + want: true, + }, + { + name: "ownership reassignment affects all drivers", + oldLabels: map[string]string{"feature.node.kubernetes.io/pci-10de.present": "true", consts.NVIDIADriverOwnerLabel: "driver-a"}, + newLabels: map[string]string{"feature.node.kubernetes.io/pci-10de.present": "true", consts.NVIDIADriverOwnerLabel: "driver-b"}, + want: true, + }, + { + name: "GPU label removal is relevant", + oldLabels: map[string]string{"feature.node.kubernetes.io/pci-10de.present": "true", consts.NVIDIADriverOwnerLabel: "driver-a"}, + newLabels: map[string]string{consts.NVIDIADriverOwnerLabel: "driver-a"}, + want: true, + }, + { + name: "NFD kernel label addition is relevant without an owner", + oldLabels: gpuLabels, + newLabels: map[string]string{"feature.node.kubernetes.io/pci-10de.present": "true", nfdKernelLabelKey: "new"}, + want: true, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + require.Equal(t, tc.want, nodeDriverRenderingLabelsChanged(tc.oldLabels, tc.newLabels)) + }) + } +} + +func TestNodeLabelsAffectNVIDIADrivers(t *testing.T) { + scheme := runtime.NewScheme() + require.NoError(t, nvidiav1alpha1.AddToScheme(scheme)) + + driverA := &nvidiav1alpha1.NVIDIADriver{ + ObjectMeta: metav1.ObjectMeta{Name: "driver-a"}, + } + driverB := &nvidiav1alpha1.NVIDIADriver{ + ObjectMeta: metav1.ObjectMeta{Name: "driver-b"}, + Spec: nvidiav1alpha1.NVIDIADriverSpec{NodeSelector: map[string]string{"region": "us-east"}}, + } + reconciler := &NVIDIADriverReconciler{ + Client: fake.NewClientBuilder().WithScheme(scheme).WithObjects(driverA, driverB).Build(), + } + + require.True(t, reconciler.nodeLabelsAffectNVIDIADrivers(context.Background(), + map[string]string{"feature.node.kubernetes.io/pci-10de.present": "true", consts.NVIDIADriverOwnerLabel: "driver-a"}, + map[string]string{"feature.node.kubernetes.io/pci-10de.present": "true", consts.NVIDIADriverOwnerLabel: "driver-a", "region": "us-east"}, + )) + require.False(t, reconciler.nodeLabelsAffectNVIDIADrivers(context.Background(), + map[string]string{"feature.node.kubernetes.io/pci-10de.present": "true", consts.NVIDIADriverOwnerLabel: "driver-a", "example.com/label": "old"}, + map[string]string{"feature.node.kubernetes.io/pci-10de.present": "true", consts.NVIDIADriverOwnerLabel: "driver-a", "example.com/label": "new"}, + )) +} diff --git a/controllers/nvidiadriver_nodeselector_cache.go b/controllers/nvidiadriver_nodeselector_cache.go new file mode 100644 index 0000000000..fc09a5a5a6 --- /dev/null +++ b/controllers/nvidiadriver_nodeselector_cache.go @@ -0,0 +1,149 @@ +/** +# Copyright (c) NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +**/ + +package controllers + +import ( + "context" + "sync" + + "k8s.io/apimachinery/pkg/types" + "sigs.k8s.io/controller-runtime/pkg/client" + + nvidiav1alpha1 "github.com/NVIDIA/gpu-operator/api/nvidia/v1alpha1" +) + +// nvidiaDriverNodeSelectorCache avoids listing drivers for unrelated Node updates. +type nvidiaDriverNodeSelectorCache struct { + mu sync.RWMutex + keysByDriver map[types.NamespacedName]map[string]struct{} + keyRefCounts map[string]int + initialized bool +} + +// replace updates the selector keys for one NVIDIADriver. +func (cache *nvidiaDriverNodeSelectorCache) replace(driver *nvidiav1alpha1.NVIDIADriver) { + cache.mu.Lock() + defer cache.mu.Unlock() + + key := client.ObjectKeyFromObject(driver) + if driver.HasDeletionTimestamp() { + cache.removeLocked(key) + return + } + + if cache.keysByDriver == nil { + cache.keysByDriver = make(map[types.NamespacedName]map[string]struct{}) + cache.keyRefCounts = make(map[string]int) + } + cache.removeLocked(key) + + selectorKeys := make(map[string]struct{}, len(driver.Spec.NodeSelector)) + for selectorKey := range driver.Spec.NodeSelector { + selectorKeys[selectorKey] = struct{}{} + cache.keyRefCounts[selectorKey]++ + } + cache.keysByDriver[key] = selectorKeys +} + +// removeByKey removes a NVIDIADriver's selector keys after deletion. +func (cache *nvidiaDriverNodeSelectorCache) removeByKey(key types.NamespacedName) { + cache.mu.Lock() + defer cache.mu.Unlock() + cache.removeLocked(key) +} + +// removeLocked removes a NVIDIADriver's selector keys while the cache is locked. +func (cache *nvidiaDriverNodeSelectorCache) removeLocked(key types.NamespacedName) { + for selectorKey := range cache.keysByDriver[key] { + cache.keyRefCounts[selectorKey]-- + if cache.keyRefCounts[selectorKey] == 0 { + delete(cache.keyRefCounts, selectorKey) + } + } + delete(cache.keysByDriver, key) +} + +// initialize populates the cache once when a Node event arrives before reconciliation. +func (cache *nvidiaDriverNodeSelectorCache) initialize(ctx context.Context, c client.Client) error { + cache.mu.Lock() + defer cache.mu.Unlock() + if cache.initialized { + return nil + } + return cache.refreshLocked(ctx, c) +} + +// refresh rebuilds the cache from the current NVIDIADriver list. +func (cache *nvidiaDriverNodeSelectorCache) refresh(ctx context.Context, c client.Client) error { + cache.mu.Lock() + defer cache.mu.Unlock() + return cache.refreshLocked(ctx, c) +} + +// refreshLocked rebuilds the cache while the cache is locked. +func (cache *nvidiaDriverNodeSelectorCache) refreshLocked(ctx context.Context, c client.Client) error { + drivers := &nvidiav1alpha1.NVIDIADriverList{} + if err := c.List(ctx, drivers); err != nil { + return err + } + + cache.keysByDriver = make(map[types.NamespacedName]map[string]struct{}) + cache.keyRefCounts = make(map[string]int) + for index := range drivers.Items { + driver := &drivers.Items[index] + if driver.HasDeletionTimestamp() { + continue + } + + key := client.ObjectKeyFromObject(driver) + selectorKeys := make(map[string]struct{}, len(driver.Spec.NodeSelector)) + for selectorKey := range driver.Spec.NodeSelector { + selectorKeys[selectorKey] = struct{}{} + cache.keyRefCounts[selectorKey]++ + } + cache.keysByDriver[key] = selectorKeys + } + cache.initialized = true + return nil +} + +// selectorLabelsChanged reports whether a configured node-selector key changed. +func (cache *nvidiaDriverNodeSelectorCache) selectorLabelsChanged(ctx context.Context, c client.Client, oldLabels, newLabels map[string]string) (bool, error) { + cache.mu.RLock() + initialized := cache.initialized + cache.mu.RUnlock() + if !initialized { + if err := cache.initialize(ctx, c); err != nil { + return false, err + } + } + + cache.mu.RLock() + defer cache.mu.RUnlock() + for key, oldValue := range oldLabels { + newValue, exists := newLabels[key] + if (!exists || oldValue != newValue) && cache.keyRefCounts[key] > 0 { + return true, nil + } + } + for key := range newLabels { + if _, exists := oldLabels[key]; !exists && cache.keyRefCounts[key] > 0 { + return true, nil + } + } + return false, nil +} diff --git a/controllers/nvidiadriver_nodeselector_cache_test.go b/controllers/nvidiadriver_nodeselector_cache_test.go new file mode 100644 index 0000000000..d821ec1af0 --- /dev/null +++ b/controllers/nvidiadriver_nodeselector_cache_test.go @@ -0,0 +1,137 @@ +/** +# Copyright (c) NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +**/ + +package controllers + +import ( + "context" + "sync" + "testing" + + "github.com/stretchr/testify/require" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/types" + "k8s.io/utils/ptr" + "sigs.k8s.io/controller-runtime/pkg/client/fake" + + nvidiav1alpha1 "github.com/NVIDIA/gpu-operator/api/nvidia/v1alpha1" +) + +func TestNVIDIADriverNodeSelectorCache(t *testing.T) { + ctx := context.Background() + scheme := runtime.NewScheme() + require.NoError(t, nvidiav1alpha1.AddToScheme(scheme)) + + drivers := []nvidiav1alpha1.NVIDIADriver{ + { + ObjectMeta: metav1.ObjectMeta{Name: "default"}, + Spec: nvidiav1alpha1.NVIDIADriverSpec{Default: true}, + }, + { + ObjectMeta: metav1.ObjectMeta{Name: "gold"}, + Spec: nvidiav1alpha1.NVIDIADriverSpec{ + NodeSelector: map[string]string{"region": "us-east-1"}, + }, + }, + } + client := fake.NewClientBuilder().WithScheme(scheme).WithObjects(&drivers[0], &drivers[1]).Build() + cache := &nvidiaDriverNodeSelectorCache{} + + tests := []struct { + name string + oldLabels map[string]string + newLabels map[string]string + want bool + }{ + {name: "selector label added", oldLabels: map[string]string{}, newLabels: map[string]string{"region": "us-east-1"}, want: true}, + {name: "selector label value changed", oldLabels: map[string]string{"region": "us-east-1"}, newLabels: map[string]string{"region": "us-east-2"}, want: true}, + {name: "selector label removed", oldLabels: map[string]string{"region": "us-east-1"}, newLabels: map[string]string{}, want: true}, + {name: "empty selector label removed", oldLabels: map[string]string{"region": ""}, newLabels: map[string]string{}, want: true}, + {name: "unrelated label changed", oldLabels: map[string]string{"example.com/probe": "old"}, newLabels: map[string]string{"example.com/probe": "new"}, want: false}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got, err := cache.selectorLabelsChanged(ctx, client, tc.oldLabels, tc.newLabels) + require.NoError(t, err) + require.Equal(t, tc.want, got) + }) + } + + updatedDriver := drivers[1].DeepCopy() + updatedDriver.Spec.NodeSelector = map[string]string{"zone": "us-east-1a"} + cache.replace(updatedDriver) + changed, err := cache.selectorLabelsChanged(ctx, client, + map[string]string{"region": "us-east-1"}, map[string]string{"region": "us-east-2"}) + require.NoError(t, err) + require.False(t, changed) + + changed, err = cache.selectorLabelsChanged(ctx, client, map[string]string{}, map[string]string{"zone": "us-east-1a"}) + require.NoError(t, err) + require.True(t, changed) + + otherDriver := &nvidiav1alpha1.NVIDIADriver{ + ObjectMeta: metav1.ObjectMeta{Name: "silver"}, + Spec: nvidiav1alpha1.NVIDIADriverSpec{ + NodeSelector: map[string]string{"zone": "us-east-1b"}, + }, + } + cache.replace(otherDriver) + + updatedDriver.DeletionTimestamp = ptr.To(metav1.Now()) + cache.replace(updatedDriver) + changed, err = cache.selectorLabelsChanged(ctx, client, + map[string]string{"region": "us-east-1"}, map[string]string{"region": "us-east-2"}) + require.NoError(t, err) + require.False(t, changed) + + changed, err = cache.selectorLabelsChanged(ctx, client, map[string]string{}, map[string]string{"zone": "us-east-1b"}) + require.NoError(t, err) + require.True(t, changed) + + cache.removeByKey(types.NamespacedName{Name: otherDriver.Name, Namespace: otherDriver.Namespace}) + changed, err = cache.selectorLabelsChanged(ctx, client, map[string]string{}, map[string]string{"zone": "us-east-1b"}) + require.NoError(t, err) + require.False(t, changed) +} + +func TestNVIDIADriverNodeSelectorCacheConcurrentAccess(t *testing.T) { + cache := &nvidiaDriverNodeSelectorCache{initialized: true} + driver := &nvidiav1alpha1.NVIDIADriver{ + ObjectMeta: metav1.ObjectMeta{Name: "gold"}, + Spec: nvidiav1alpha1.NVIDIADriverSpec{ + NodeSelector: map[string]string{"region": "us-east-1"}, + }, + } + + var wg sync.WaitGroup + for range 8 { + wg.Go(func() { + for range 100 { + cache.replace(driver) + cache.removeByKey(types.NamespacedName{Name: driver.Name, Namespace: driver.Namespace}) + } + }) + wg.Go(func() { + for range 100 { + _, _ = cache.selectorLabelsChanged(context.Background(), nil, + map[string]string{}, map[string]string{"region": "us-east-1"}) + } + }) + } + wg.Wait() +} diff --git a/tests/scripts/checks.sh b/tests/scripts/checks.sh index 5056d48b2f..0a6336b69c 100755 --- a/tests/scripts/checks.sh +++ b/tests/scripts/checks.sh @@ -241,13 +241,17 @@ print_driver_upgrade_debug() { } wait_for_driver_upgrade_done() { - gpu_node_count=$(kubectl get node -l nvidia.com/gpu.present --no-headers | wc -l) + local gpu_node_count local current_time=0 + local node_resource + local upgrade_state + + gpu_node_count=$(kubectl get node -l nvidia.com/gpu.present --no-headers | wc -l) echo "waiting for the gpu driver upgrade to complete" while :; do local upgraded_count=0 - for node in $(kubectl get nodes -o NAME); do - upgrade_state=$(kubectl get $node -ojsonpath='{.metadata.labels.nvidia\.com/gpu-driver-upgrade-state}') + for node_resource in $(kubectl get nodes -o NAME); do + upgrade_state=$(kubectl get "$node_resource" -ojsonpath='{.metadata.labels.nvidia\.com/gpu-driver-upgrade-state}') if [ "${upgrade_state}" = "upgrade-done" ]; then upgraded_count=$((${upgraded_count} + 1)) fi diff --git a/tests/scripts/update-nvidiadriver.sh b/tests/scripts/update-nvidiadriver.sh index d104673654..66cd5a8e5e 100755 --- a/tests/scripts/update-nvidiadriver.sh +++ b/tests/scripts/update-nvidiadriver.sh @@ -14,6 +14,9 @@ source ${SCRIPT_DIR}/checks.sh NVIDIA_DRIVER_NAME="${NVIDIA_DRIVER_NAME:-e2e-driver}" DEFAULT_NVIDIA_DRIVER_NAME="${DEFAULT_NVIDIA_DRIVER_NAME:-e2e-default-driver}" DUPLICATE_DEFAULT_NVIDIA_DRIVER_NAME="${DUPLICATE_DEFAULT_NVIDIA_DRIVER_NAME:-e2e-duplicate-default-driver}" +SELECTOR_CONFLICT_NVIDIA_DRIVER_NAME="${SELECTOR_CONFLICT_NVIDIA_DRIVER_NAME:-e2e-selector-conflict-driver}" +SELECTOR_CONFLICT_OVERLAP_NVIDIA_DRIVER_NAME="${SELECTOR_CONFLICT_OVERLAP_NVIDIA_DRIVER_NAME:-e2e-selector-conflict-overlap-driver}" +SELECTOR_CONFLICT_LABEL="${SELECTOR_CONFLICT_LABEL:-e2e.nvidia.com/nvidiadriver-selector-conflict}" get_default_nvidiadriver_name() { kubectl get nvidiadriver -o json | @@ -42,6 +45,20 @@ create_nvidiadriver_from() { ' | kubectl apply -f - } +create_selector_conflict_nvidiadriver() { + local driver_name=$1 + + kubectl get nvidiadriver/"${DEFAULT_NVIDIA_DRIVER_NAME}" -o json | + jq --arg name "${driver_name}" --arg selector_key "${SELECTOR_CONFLICT_LABEL}" ' + { + apiVersion: .apiVersion, + kind: .kind, + metadata: {name: $name}, + spec: (.spec + {default: false, nodeSelector: {($selector_key): "enabled"}}) + } + ' | kubectl apply -f - +} + unset_default_driver() { local driver_name=$1 @@ -347,7 +364,41 @@ test_multiple_default_drivers_are_not_ready() { wait_for_nvidiadriver_owner "${NVIDIA_DRIVER_NAME}" } +test_selector_label_conflict_status_on_single_gpu_node() { + local gpu_nodes node gpu_node_count + + mapfile -t gpu_nodes < <(kubectl get nodes -l nvidia.com/gpu.present=true -o jsonpath='{range .items[*]}{.metadata.name}{"\n"}{end}') + gpu_node_count=${#gpu_nodes[@]} + if [[ "${gpu_node_count}" -ne 1 ]]; then + echo "Skipping selector-label conflict status test: expected exactly one GPU node, found ${gpu_node_count}" + kubectl get nodes -l nvidia.com/gpu.present=true -o wide + return 0 + fi + node="${gpu_nodes[0]}" + + echo "Testing selector-label conflict status on GPU node ${node}" + kubectl label node "${node}" "${SELECTOR_CONFLICT_LABEL}-" >/dev/null 2>&1 || true + kubectl delete nvidiadriver/"${SELECTOR_CONFLICT_NVIDIA_DRIVER_NAME}" --ignore-not-found >/dev/null 2>&1 || true + kubectl delete nvidiadriver/"${SELECTOR_CONFLICT_OVERLAP_NVIDIA_DRIVER_NAME}" --ignore-not-found >/dev/null 2>&1 || true + + wait_for_nvidiadriver_owner "${DEFAULT_NVIDIA_DRIVER_NAME}" + create_selector_conflict_nvidiadriver "${SELECTOR_CONFLICT_NVIDIA_DRIVER_NAME}" + create_selector_conflict_nvidiadriver "${SELECTOR_CONFLICT_OVERLAP_NVIDIA_DRIVER_NAME}" + wait_for_nvidiadriver_ready "${SELECTOR_CONFLICT_NVIDIA_DRIVER_NAME}" + wait_for_nvidiadriver_ready "${SELECTOR_CONFLICT_OVERLAP_NVIDIA_DRIVER_NAME}" + kubectl label node "${node}" "${SELECTOR_CONFLICT_LABEL}=enabled" --overwrite + wait_for_nvidiadriver_condition_message "${SELECTOR_CONFLICT_NVIDIA_DRIVER_NAME}" "multiple NVIDIADrivers match the same node" + wait_for_nvidiadriver_condition_message "${SELECTOR_CONFLICT_OVERLAP_NVIDIA_DRIVER_NAME}" "multiple NVIDIADrivers match the same node" + assert_nvidiadriver_owner_count "${DEFAULT_NVIDIA_DRIVER_NAME}" + kubectl label node "${node}" "${SELECTOR_CONFLICT_LABEL}-" + wait_for_nvidiadriver_ready "${SELECTOR_CONFLICT_NVIDIA_DRIVER_NAME}" + wait_for_nvidiadriver_ready "${SELECTOR_CONFLICT_OVERLAP_NVIDIA_DRIVER_NAME}" + kubectl delete nvidiadriver/"${SELECTOR_CONFLICT_NVIDIA_DRIVER_NAME}" + kubectl delete nvidiadriver/"${SELECTOR_CONFLICT_OVERLAP_NVIDIA_DRIVER_NAME}" +} + test_arbitrary_name_default_nvidiadriver +test_selector_label_conflict_status_on_single_gpu_node create_nvidiadriver wait_for_nvidiadriver_owner "${NVIDIA_DRIVER_NAME}" wait_for_nvidiadriver_daemonsets "${NVIDIA_DRIVER_NAME}"