Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 27 additions & 20 deletions controllers/nodelabeling_controller.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
}
Expand All @@ -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 {
Expand Down
20 changes: 10 additions & 10 deletions controllers/nodelabeling_controller_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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},
}
Expand Down Expand Up @@ -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")
Expand All @@ -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")
Expand All @@ -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")
Expand All @@ -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
}},
Expand Down
104 changes: 71 additions & 33 deletions controllers/nvidiadriver_controller.go
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@ import (
"context"
"errors"
"fmt"
"maps"
"time"

appsv1 "k8s.io/api/apps/v1"
Expand Down Expand Up @@ -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
Expand All @@ -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
}
Expand All @@ -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
}
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 {
Expand All @@ -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(),
Expand Down Expand Up @@ -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 {
Expand Down
Loading
Loading