Add DRA pod reconciler and node labeling

Add DRAPodReconciler that watches DRA-role pods and labels/unlabels
nodes based on pod readiness, mirroring DevicePluginPodReconciler.

- Add DRAVersionLabelPrefix constant and DRA label utilities
  (GetDRANodeLabel, GetDRAVersionLabelName, IsDRAVersionLabel)
- Extend IsVersionLabel and GetNodesVersionLabels for DRA labels
- Create DRAPodReconciler with NodeLabelerFinalizer support
- Add isNotDRAPod predicate to DevicePluginPodReconciler to exclude
  DRA pods from event delivery
- Exclude DRA pods from DevicePluginPodReconciler replacement check
  to prevent a ready DRA pod from blocking device-plugin-ready label
  removal
- Register DRAPodReconciler in manager

Author: TomerNewman

cmd/manager/main.go

@@ -170,6 +170,10 @@ func main() {
 		cmd.FatalError(setupLogger, err, "unable to create controller", "name", controllers.DevicePluginPodReconcilerName)
 	}
 
+	if err = controllers.NewDRAPodReconciler(client).SetupWithManager(mgr); err != nil {
+		cmd.FatalError(setupLogger, err, "unable to create controller", "name", controllers.DRAPodReconcilerName)
+	}
+
 	if err = controllers.NewDRAReconciler(client, nodeAPI, scheme).SetupWithManager(mgr); err != nil {
 		cmd.FatalError(setupLogger, err, "unable to create controller", "name", controllers.DRAReconcilerName)
 	}

internal/constants/constants.go

@@ -13,6 +13,7 @@ const (
 
 	WorkerPodVersionLabelPrefix    = "beta.kmm.node.kubernetes.io/version-worker-pod"
 	DevicePluginVersionLabelPrefix = "beta.kmm.node.kubernetes.io/version-device-plugin"
+	DRAVersionLabelPrefix          = "beta.kmm.node.kubernetes.io/version-dra"
 	ModuleVersionLabelPrefix       = "kmm.node.kubernetes.io/version-module"
 
 	GCDelayFinalizer  = "kmm.node.kubernetes.io/gc-delay"

internal/controllers/device_plugin_pod_reconciler.go

@@ -8,14 +8,11 @@ import (
 	"github.com/kubernetes-sigs/kernel-module-management/internal/filter"
 	"github.com/kubernetes-sigs/kernel-module-management/internal/utils"
 	v1 "k8s.io/api/core/v1"
-	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/kubectl/pkg/util/podutils"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
-	"sigs.k8s.io/controller-runtime/pkg/reconcile"
-
-	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
 	"sigs.k8s.io/controller-runtime/pkg/predicate"
+	"sigs.k8s.io/controller-runtime/pkg/reconcile"
 )
 
 const DevicePluginPodReconcilerName = "DevicePluginPod"
@@ -59,7 +56,7 @@ func (dppr *DevicePluginPodReconciler) Reconcile(ctx context.Context, pod *v1.Po
 		if nodeName != "" {
 			logger.Info("Unlabeling node")
 
-			// Make sure we don't already have a new running pod before unlabeling the node
+			// Make sure we don't already have a new running device-plugin pod before unlabeling the node
 			labelSelector := client.MatchingLabels{constants.ModuleNameLabel: moduleName}
 			fieldSelector := client.MatchingFields{"spec.nodeName": nodeName}
 			var modulePodsList v1.PodList
@@ -69,13 +66,16 @@ func (dppr *DevicePluginPodReconciler) Reconcile(ctx context.Context, pod *v1.Po
 			}
 			var foundRunningPod bool
 			for _, p := range modulePodsList.Items {
+				if p.Labels[constants.DaemonSetRole] == constants.DRARoleLabelValue {
+					continue
+				}
 				if podutils.IsPodReady(&p) && p.DeletionTimestamp.IsZero() {
 					foundRunningPod = true
 					break
 				}
 			}
 			if !foundRunningPod {
-				if err := dppr.deleteLabel(ctx, nodeName, labelName); err != nil {
+				if err := utils.DeleteNodeLabel(ctx, dppr.client, nodeName, labelName); err != nil {
 					return ctrl.Result{}, fmt.Errorf("could not unlabel node %s with label %s: %v",
 						nodeName, labelName, err)
 				}
@@ -89,7 +89,7 @@ func (dppr *DevicePluginPodReconciler) Reconcile(ctx context.Context, pod *v1.Po
 			// the specified Pod has already been deleted. By ignoring NotFound errors we ensure
 			// that no additional, unnecessary reconciliation request will be queued (since a
 			// reconciliation result with a non-nil error will be requeued).
-			if err := dppr.deleteFinalizer(ctx, pod); client.IgnoreNotFound(err) != nil {
+			if err := utils.DeletePodFinalizer(ctx, dppr.client, pod); client.IgnoreNotFound(err) != nil {
 				return ctrl.Result{}, fmt.Errorf("could not delete the pod finalizer: %v", err)
 			}
 		}
@@ -99,7 +99,7 @@ func (dppr *DevicePluginPodReconciler) Reconcile(ctx context.Context, pod *v1.Po
 
 	logger.Info("Labeling node")
 
-	if err := dppr.addLabel(ctx, nodeName, labelName); err != nil {
+	if err := utils.AddNodeLabel(ctx, dppr.client, nodeName, labelName); err != nil {
 		return ctrl.Result{}, fmt.Errorf("could not label node %s with %s: %v", nodeName, labelName, err)
 	}
 
@@ -126,6 +126,10 @@ func (dppr *DevicePluginPodReconciler) SetupWithManager(mgr ctrl.Manager) error
 		return false
 	})
 
+	isNotDRAPod := predicate.NewPredicateFuncs(func(o client.Object) bool {
+		return o.GetLabels()[constants.DaemonSetRole] != constants.DRARoleLabelValue
+	})
+
 	p := predicate.And(
 		predicate.Or(
 			filter.PodReadinessChangedPredicate(
@@ -135,6 +139,7 @@ func (dppr *DevicePluginPodReconciler) SetupWithManager(mgr ctrl.Manager) error
 		),
 		filter.HasLabel(constants.ModuleNameLabel),
 		isDaemonSetPod,
+		isNotDRAPod,
 	)
 
 	return ctrl.
@@ -147,42 +152,3 @@ func (dppr *DevicePluginPodReconciler) SetupWithManager(mgr ctrl.Manager) error
 		)
 }
 
-func (dppr *DevicePluginPodReconciler) addLabel(ctx context.Context, nodeName string, labelName string) error {
-	node := v1.Node{}
-
-	if err := dppr.client.Get(ctx, types.NamespacedName{Name: nodeName}, &node); err != nil {
-		return fmt.Errorf("could not get node %s: %v", nodeName, err)
-	}
-
-	nodeCopy := node.DeepCopy()
-
-	if node.Labels == nil {
-		node.Labels = make(map[string]string)
-	}
-
-	node.Labels[labelName] = ""
-
-	return dppr.client.Patch(ctx, &node, client.MergeFrom(nodeCopy))
-}
-
-func (dppr *DevicePluginPodReconciler) deleteFinalizer(ctx context.Context, pod *v1.Pod) error {
-	podCopy := pod.DeepCopy()
-
-	controllerutil.RemoveFinalizer(pod, constants.NodeLabelerFinalizer)
-
-	return dppr.client.Patch(ctx, pod, client.MergeFrom(podCopy))
-}
-
-func (dppr *DevicePluginPodReconciler) deleteLabel(ctx context.Context, nodeName string, labelName string) error {
-	node := v1.Node{}
-
-	if err := dppr.client.Get(ctx, types.NamespacedName{Name: nodeName}, &node); err != nil {
-		return fmt.Errorf("could not get node %s: %v", nodeName, err)
-	}
-
-	nodeCopy := node.DeepCopy()
-
-	delete(node.Labels, labelName)
-
-	return dppr.client.Patch(ctx, &node, client.MergeFrom(nodeCopy))
-}

internal/controllers/device_plugin_pod_reconciler_test.go

@@ -126,6 +126,61 @@ var _ = Describe("DevicePluginPodReconciler_Reconcile", func() {
 		Expect(err).NotTo(HaveOccurred())
 	})
 
+	It("should unlabel the node when the only ready pod is a DRA pod", func() {
+		var (
+			labelSelector = client.MatchingLabels{constants.ModuleNameLabel: moduleName}
+			fieldSelector = client.MatchingFields{"spec.nodeName": nodeName}
+		)
+
+		gomock.InOrder(
+			kubeClient.EXPECT().List(ctx, gomock.Any(), labelSelector, fieldSelector).Do(
+				func(_ interface{}, modulePodsList *v1.PodList, _ ...client.ListOption) {
+					modulePodsList.Items = []v1.Pod{
+						{
+							ObjectMeta: metav1.ObjectMeta{
+								Labels: map[string]string{
+									constants.ModuleNameLabel: moduleName,
+									constants.DaemonSetRole:   constants.DRARoleLabelValue,
+								},
+							},
+							Status: v1.PodStatus{
+								Conditions: []v1.PodCondition{
+									{
+										Type:   v1.PodReady,
+										Status: v1.ConditionTrue,
+									},
+								},
+							},
+						},
+					}
+				},
+			),
+			kubeClient.EXPECT().Get(ctx, gomock.Any(), gomock.Any()).Do(
+				func(_ interface{}, _ interface{}, node *v1.Node, _ ...client.GetOption) {
+					node.SetLabels(map[string]string{utils.GetDevicePluginNodeLabel(podNamespace, moduleName): ""})
+				},
+			),
+			kubeClient.EXPECT().Patch(ctx, gomock.Any(), gomock.Any()).Do(
+				func(_ interface{}, node *v1.Node, p client.Patch, _ ...client.GetOption) {
+					Expect(p.Type()).To(Equal(types.MergePatchType))
+					Expect(p.Data(node)).To(Equal([]byte(`{"metadata":{"labels":null}}`)))
+				},
+			),
+		)
+
+		pod := &v1.Pod{
+			ObjectMeta: metav1.ObjectMeta{
+				Labels:    map[string]string{constants.ModuleNameLabel: moduleName},
+				Name:      podName,
+				Namespace: podNamespace,
+			},
+			Spec: v1.PodSpec{NodeName: nodeName},
+		}
+
+		_, err := r.Reconcile(ctx, pod)
+		Expect(err).NotTo(HaveOccurred())
+	})
+
 	now := metav1.Now()
 
 	patchRemoveFinalizerFunc := func(_ interface{}, pod *v1.Pod, p client.Patch, _ ...client.GetOption) {

internal/controllers/dra_pod_reconciler.go

@@ -0,0 +1,147 @@
+/*
+Copyright 2022.
+
+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"
+	"fmt"
+
+	"github.com/kubernetes-sigs/kernel-module-management/internal/constants"
+	"github.com/kubernetes-sigs/kernel-module-management/internal/filter"
+	"github.com/kubernetes-sigs/kernel-module-management/internal/utils"
+	v1 "k8s.io/api/core/v1"
+	"k8s.io/kubectl/pkg/util/podutils"
+	ctrl "sigs.k8s.io/controller-runtime"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/predicate"
+	"sigs.k8s.io/controller-runtime/pkg/reconcile"
+)
+
+const DRAPodReconcilerName = "DRAPod"
+
+type DRAPodReconciler struct {
+	client client.Client
+}
+
+func NewDRAPodReconciler(client client.Client) *DRAPodReconciler {
+	return &DRAPodReconciler{client: client}
+}
+
+func (r *DRAPodReconciler) Reconcile(ctx context.Context, pod *v1.Pod) (ctrl.Result, error) {
+	logger := ctrl.LoggerFrom(ctx)
+
+	nodeName := pod.Spec.NodeName
+
+	moduleName, ok := pod.Labels[constants.ModuleNameLabel]
+	if !ok {
+		return ctrl.Result{}, fmt.Errorf("pod %s/%s has no %q label", pod.Namespace, pod.Name, constants.ModuleNameLabel)
+	}
+
+	labelName := utils.GetDRANodeLabel(pod.Namespace, moduleName)
+
+	logger = logger.WithValues(
+		"node name", nodeName,
+		"module name", moduleName,
+		"label name", labelName,
+	)
+
+	if !podutils.IsPodReady(pod) || !pod.DeletionTimestamp.IsZero() {
+		if nodeName != "" {
+			logger.Info("Unlabeling node")
+
+			labelSelector := client.MatchingLabels{
+				constants.ModuleNameLabel: moduleName,
+				constants.DaemonSetRole:   constants.DRARoleLabelValue,
+			}
+			fieldSelector := client.MatchingFields{"spec.nodeName": nodeName}
+			var draPodsList v1.PodList
+			err := r.client.List(ctx, &draPodsList, labelSelector, fieldSelector)
+			if err != nil {
+				return ctrl.Result{}, fmt.Errorf("failed to get list of all DRA pods for module %s on node %s: %v", moduleName, nodeName, err)
+			}
+			var foundRunningPod bool
+			for _, p := range draPodsList.Items {
+				if podutils.IsPodReady(&p) && p.DeletionTimestamp.IsZero() {
+					foundRunningPod = true
+					break
+				}
+			}
+			if !foundRunningPod {
+				if err := utils.DeleteNodeLabel(ctx, r.client, nodeName, labelName); err != nil {
+					return ctrl.Result{}, fmt.Errorf("could not unlabel node %s with label %s: %v",
+						nodeName, labelName, err)
+				}
+			}
+		}
+
+		if !pod.DeletionTimestamp.IsZero() {
+			logger.Info("Pod deletion requested; removing finalizer")
+
+			if err := utils.DeletePodFinalizer(ctx, r.client, pod); client.IgnoreNotFound(err) != nil {
+				return ctrl.Result{}, fmt.Errorf("could not delete the pod finalizer: %v", err)
+			}
+		}
+
+		return ctrl.Result{}, nil
+	}
+
+	logger.Info("Labeling node")
+
+	if err := utils.AddNodeLabel(ctx, r.client, nodeName, labelName); err != nil {
+		return ctrl.Result{}, fmt.Errorf("could not label node %s with %s: %v", nodeName, labelName, err)
+	}
+
+	return ctrl.Result{}, nil
+}
+
+func (r *DRAPodReconciler) SetupWithManager(mgr ctrl.Manager) error {
+	isDaemonSetPod := predicate.NewPredicateFuncs(func(o client.Object) bool {
+		ownerReferences := o.GetOwnerReferences()
+		for _, ownerReference := range ownerReferences {
+			if ownerReference.Kind == "DaemonSet" {
+				return true
+			}
+		}
+		return false
+	})
+
+	isDRAPod := predicate.NewPredicateFuncs(func(o client.Object) bool {
+		return o.GetLabels()[constants.DaemonSetRole] == constants.DRARoleLabelValue
+	})
+
+	p := predicate.And(
+		predicate.Or(
+			filter.PodReadinessChangedPredicate(
+				mgr.GetLogger().WithName("dra-pod-readiness-changed"),
+			),
+			filter.DeletingPredicate(),
+		),
+		filter.HasLabel(constants.ModuleNameLabel),
+		isDaemonSetPod,
+		isDRAPod,
+	)
+
+	return ctrl.
+		NewControllerManagedBy(mgr).
+		Named(DRAPodReconcilerName).
+		For(&v1.Pod{}).
+		WithEventFilter(p).
+		Complete(
+			reconcile.AsReconciler[*v1.Pod](r.client, r),
+		)
+}
+