Add DRA DaemonSet reconciler

Implement a DRAReconciler that manages DRA driver DaemonSets from
Module spec.dra, mirroring the DevicePluginReconciler pattern. The
reconciler creates DaemonSets with mandatory kubelet-plugins and
kubelet-plugins-registry host-path volumes, DRA-specific container
names, and a DaemonSetRole label to distinguish DRA DaemonSets.

Status patching reports targeted nodes and available pods in
status.dra, and clearing status.dra when spec.dra is removed.

Also fixes DevicePluginReconciler to exclude DRA-labeled DaemonSets
from its device-plugin list, preventing cross-reconciler interference.

Author: TomerNewman

internal/constants/constants.go

@@ -25,6 +25,7 @@ const (
 	PrivateSignDataKey             = "key"
 
 	ModuleLoaderRoleLabelValue = "module-loader"
+	DRARoleLabelValue          = "dra"
 
 	OperatorNamespaceEnvVar = "OPERATOR_NAMESPACE"
 )

internal/controllers/device_plugin_reconciler.go

@@ -176,9 +176,9 @@ func (dprh *devicePluginReconcilerHelper) getModuleDevicePluginDaemonSets(ctx co
 	}
 
 	devicePluginsList := make([]appsv1.DaemonSet, 0, len(dsList.Items))
-	// remove the older version module loader daemonsets
 	for _, ds := range dsList.Items {
-		if ds.GetLabels()[constants.DaemonSetRole] != constants.ModuleLoaderRoleLabelValue {
+		role := ds.GetLabels()[constants.DaemonSetRole]
+		if role != constants.ModuleLoaderRoleLabelValue && role != constants.DRARoleLabelValue {
 			devicePluginsList = append(devicePluginsList, ds)
 		}
 	}

internal/controllers/device_plugin_reconciler_test.go

@@ -927,7 +927,7 @@ var _ = Describe("DevicePluginReconciler_getModuleDevicePluginDaemonSets", func(
 		Expect(dsList).To(BeNil())
 	})
 
-	It("good flow, return only device plugin DSs, either v2 or v1", func() {
+	It("good flow, return only device plugin DSs, excluding module-loader and DRA roles", func() {
 		ds1 := appsv1.DaemonSet{
 			ObjectMeta: metav1.ObjectMeta{
 				Labels: map[string]string{
@@ -951,9 +951,17 @@ var _ = Describe("DevicePluginReconciler_getModuleDevicePluginDaemonSets", func(
 				},
 			},
 		}
+		ds4 := appsv1.DaemonSet{
+			ObjectMeta: metav1.ObjectMeta{
+				Labels: map[string]string{
+					constants.ModuleNameLabel: "some name",
+					constants.DaemonSetRole:   constants.DRARoleLabelValue,
+				},
+			},
+		}
 		clnt.EXPECT().List(ctx, gomock.Any(), gomock.Any()).DoAndReturn(
 			func(_ interface{}, list *appsv1.DaemonSetList, _ ...interface{}) error {
-				list.Items = []appsv1.DaemonSet{ds1, ds2, ds3}
+				list.Items = []appsv1.DaemonSet{ds1, ds2, ds3, ds4}
 				return nil
 			},
 		)

internal/controllers/dra_reconciler.go

@@ -0,0 +1,326 @@
+/*
+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"
+	"errors"
+	"fmt"
+
+	kmmv1beta1 "github.com/kubernetes-sigs/kernel-module-management/api/v1beta1"
+	"github.com/kubernetes-sigs/kernel-module-management/internal/constants"
+	"github.com/kubernetes-sigs/kernel-module-management/internal/node"
+	"github.com/kubernetes-sigs/kernel-module-management/internal/utils"
+	appsv1 "k8s.io/api/apps/v1"
+	v1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	ctrl "sigs.k8s.io/controller-runtime"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
+	"sigs.k8s.io/controller-runtime/pkg/log"
+	"sigs.k8s.io/controller-runtime/pkg/reconcile"
+)
+
+const (
+	DRAReconcilerName = "DRAReconciler"
+
+	kubeletPluginsVolumeName         = "kubelet-plugins"
+	kubeletPluginsPath               = "/var/lib/kubelet/plugins/"
+	kubeletPluginsRegistryVolumeName = "kubelet-plugins-registry"
+	kubeletPluginsRegistryPath       = "/var/lib/kubelet/plugins_registry/"
+)
+
+type DRAReconciler struct {
+	client         client.Client
+	reconHelperAPI draReconcilerHelperAPI
+}
+
+func NewDRAReconciler(
+	client client.Client,
+	nodeAPI node.Node,
+	scheme *runtime.Scheme,
+) *DRAReconciler {
+	reconHelperAPI := newDRAReconcilerHelper(client, nodeAPI, scheme)
+	return &DRAReconciler{
+		client:         client,
+		reconHelperAPI: reconHelperAPI,
+	}
+}
+
+func (r *DRAReconciler) SetupWithManager(mgr ctrl.Manager) error {
+	return ctrl.NewControllerManagedBy(mgr).
+		For(&kmmv1beta1.Module{}).
+		Owns(&appsv1.DaemonSet{}).
+		Named(DRAReconcilerName).
+		Complete(
+			reconcile.AsReconciler[*kmmv1beta1.Module](r.client, r),
+		)
+}
+
+func (r *DRAReconciler) Reconcile(ctx context.Context, mod *kmmv1beta1.Module) (ctrl.Result, error) {
+	res := ctrl.Result{}
+
+	logger := log.FromContext(ctx)
+
+	existingDRADS, err := r.reconHelperAPI.getModuleDRADaemonSets(ctx, mod.Name, mod.Namespace)
+	if err != nil {
+		return res, fmt.Errorf("could not get DRA DaemonSets for module %s, namespace %s: %v", mod.Name, mod.Namespace, err)
+	}
+
+	if mod.GetDeletionTimestamp() != nil {
+		err = r.reconHelperAPI.handleModuleDeletion(ctx, existingDRADS)
+		return ctrl.Result{}, err
+	}
+
+	if mod.Spec.DRA == nil {
+		if len(existingDRADS) > 0 {
+			if err = r.reconHelperAPI.handleModuleDeletion(ctx, existingDRADS); err != nil {
+				return ctrl.Result{}, err
+			}
+		}
+		return ctrl.Result{}, r.reconHelperAPI.clearDRAStatus(ctx, mod)
+	}
+
+	err = r.reconHelperAPI.handleDRA(ctx, mod, existingDRADS)
+	if err != nil {
+		return res, fmt.Errorf("could not handle DRA: %v", err)
+	}
+
+	err = r.reconHelperAPI.moduleUpdateDRAStatus(ctx, mod, existingDRADS)
+	if err != nil {
+		return res, fmt.Errorf("failed to update DRA status of the module: %v", err)
+	}
+
+	logger.Info("DRA reconcile loop finished successfully")
+
+	return res, nil
+}
+
+//go:generate mockgen -source=dra_reconciler.go -package=controllers -destination=mock_dra_reconciler.go draReconcilerHelperAPI,draDaemonSetCreator
+
+type draReconcilerHelperAPI interface {
+	getModuleDRADaemonSets(ctx context.Context, name, namespace string) ([]appsv1.DaemonSet, error)
+	handleDRA(ctx context.Context, mod *kmmv1beta1.Module, existingDRADS []appsv1.DaemonSet) error
+	handleModuleDeletion(ctx context.Context, existingDRADS []appsv1.DaemonSet) error
+	moduleUpdateDRAStatus(ctx context.Context, mod *kmmv1beta1.Module, existingDRADS []appsv1.DaemonSet) error
+	clearDRAStatus(ctx context.Context, mod *kmmv1beta1.Module) error
+}
+
+type draReconcilerHelper struct {
+	client          client.Client
+	daemonSetHelper draDaemonSetCreator
+	nodeAPI         node.Node
+}
+
+func newDRAReconcilerHelper(client client.Client,
+	nodeAPI node.Node,
+	scheme *runtime.Scheme,
+) draReconcilerHelperAPI {
+	daemonSetHelper := newDRADaemonSetCreator(scheme)
+	return &draReconcilerHelper{
+		client:          client,
+		daemonSetHelper: daemonSetHelper,
+		nodeAPI:         nodeAPI,
+	}
+}
+
+func (drh *draReconcilerHelper) getModuleDRADaemonSets(ctx context.Context, name, namespace string) ([]appsv1.DaemonSet, error) {
+	dsList := appsv1.DaemonSetList{}
+	opts := []client.ListOption{
+		client.MatchingLabels(map[string]string{
+			constants.ModuleNameLabel: name,
+			constants.DaemonSetRole:   constants.DRARoleLabelValue,
+		}),
+		client.InNamespace(namespace),
+	}
+	if err := drh.client.List(ctx, &dsList, opts...); err != nil {
+		return nil, fmt.Errorf("could not list DaemonSets: %v", err)
+	}
+
+	return dsList.Items, nil
+}
+
+func (drh *draReconcilerHelper) handleDRA(ctx context.Context, mod *kmmv1beta1.Module, existingDRADS []appsv1.DaemonSet) error {
+	if mod.Spec.DRA == nil {
+		return nil
+	}
+
+	logger := log.FromContext(ctx)
+
+	var ds *appsv1.DaemonSet
+	if len(existingDRADS) > 0 {
+		ds = &existingDRADS[0]
+	} else {
+		logger.Info("creating new DRA DaemonSet")
+		ds = &appsv1.DaemonSet{
+			ObjectMeta: metav1.ObjectMeta{Namespace: mod.Namespace, GenerateName: mod.Name + "-dra-"},
+		}
+	}
+
+	opRes, err := controllerutil.CreateOrPatch(ctx, drh.client, ds, func() error {
+		return drh.daemonSetHelper.setDRAAsDesired(ctx, ds, mod)
+	})
+
+	if err == nil {
+		logger.Info("Reconciled DRA", "name", ds.Name, "result", opRes)
+	}
+
+	return err
+}
+
+func (drh *draReconcilerHelper) handleModuleDeletion(ctx context.Context, existingDRADS []appsv1.DaemonSet) error {
+	for _, ds := range existingDRADS {
+		err := drh.client.Delete(ctx, &ds)
+		if err != nil {
+			return fmt.Errorf("failed to delete DRA DaemonSet %s/%s: %v", ds.Namespace, ds.Name, err)
+		}
+	}
+	return nil
+}
+
+func (drh *draReconcilerHelper) moduleUpdateDRAStatus(ctx context.Context,
+	mod *kmmv1beta1.Module,
+	existingDRADS []appsv1.DaemonSet) error {
+
+	if mod.Spec.DRA == nil {
+		return nil
+	}
+
+	numTargetedNodes, err := drh.nodeAPI.GetNumTargetedNodes(ctx, mod.Spec.Selector, mod.Spec.Tolerations)
+	if err != nil {
+		return fmt.Errorf("failed to determine the number of nodes targeted by Module %s/%s selector: %v", mod.Namespace, mod.Name, err)
+	}
+
+	numAvailable := int32(0)
+	for _, ds := range existingDRADS {
+		numAvailable += ds.Status.NumberAvailable
+	}
+
+	unmodifiedMod := mod.DeepCopy()
+
+	mod.Status.DRA.NodesMatchingSelectorNumber = int32(numTargetedNodes)
+	mod.Status.DRA.DesiredNumber = int32(numTargetedNodes)
+	mod.Status.DRA.AvailableNumber = numAvailable
+
+	return drh.client.Status().Patch(ctx, mod, client.MergeFrom(unmodifiedMod))
+}
+
+func (drh *draReconcilerHelper) clearDRAStatus(ctx context.Context, mod *kmmv1beta1.Module) error {
+	emptyStatus := kmmv1beta1.DaemonSetStatus{}
+	if mod.Status.DRA == emptyStatus {
+		return nil
+	}
+
+	unmodifiedMod := mod.DeepCopy()
+
+	mod.Status.DRA = kmmv1beta1.DaemonSetStatus{}
+
+	return drh.client.Status().Patch(ctx, mod, client.MergeFrom(unmodifiedMod))
+}
+
+type draDaemonSetCreator interface {
+	setDRAAsDesired(ctx context.Context, ds *appsv1.DaemonSet, mod *kmmv1beta1.Module) error
+}
+
+type draDaemonSetCreatorImpl struct {
+	scheme *runtime.Scheme
+}
+
+func newDRADaemonSetCreator(scheme *runtime.Scheme) draDaemonSetCreator {
+	return &draDaemonSetCreatorImpl{
+		scheme: scheme,
+	}
+}
+
+func (dsci *draDaemonSetCreatorImpl) setDRAAsDesired(
+	ctx context.Context,
+	ds *appsv1.DaemonSet,
+	mod *kmmv1beta1.Module,
+) error {
+	if ds == nil {
+		return errors.New("ds cannot be nil")
+	}
+
+	if mod.Spec.DRA == nil {
+		return errors.New("DRA spec in module should not be nil")
+	}
+
+	hostPathDirOrCreate := v1.HostPathDirectoryOrCreate
+	hostPathDir := v1.HostPathDirectory
+
+	pluginsVolume := v1.Volume{
+		Name: kubeletPluginsVolumeName,
+		VolumeSource: v1.VolumeSource{
+			HostPath: &v1.HostPathVolumeSource{
+				Path: kubeletPluginsPath,
+				Type: &hostPathDirOrCreate,
+			},
+		},
+	}
+
+	registryVolume := v1.Volume{
+		Name: kubeletPluginsRegistryVolumeName,
+		VolumeSource: v1.VolumeSource{
+			HostPath: &v1.HostPathVolumeSource{
+				Path: kubeletPluginsRegistryPath,
+				Type: &hostPathDir,
+			},
+		},
+	}
+
+	containerVolumeMounts := []v1.VolumeMount{
+		{Name: kubeletPluginsVolumeName, MountPath: kubeletPluginsPath},
+		{Name: kubeletPluginsRegistryVolumeName, MountPath: kubeletPluginsRegistryPath},
+	}
+
+	standardLabels := map[string]string{
+		constants.ModuleNameLabel: mod.Name,
+		constants.DaemonSetRole:   constants.DRARoleLabelValue,
+	}
+
+	nodeSelector := map[string]string{
+		utils.GetKernelModuleReadyNodeLabel(mod.Namespace, mod.Name): "",
+	}
+
+	ds.SetLabels(
+		overrideLabels(ds.GetLabels(), standardLabels),
+	)
+
+	ds.Spec = appsv1.DaemonSetSpec{
+		Selector: &metav1.LabelSelector{MatchLabels: standardLabels},
+		Template: v1.PodTemplateSpec{
+			ObjectMeta: metav1.ObjectMeta{
+				Labels:     standardLabels,
+				Finalizers: []string{constants.NodeLabelerFinalizer},
+			},
+			Spec: v1.PodSpec{
+				InitContainers:               generatePodContainerSpec(mod.Spec.DRA.InitContainer, "dra-init", nil),
+				Containers:                   generatePodContainerSpec(&mod.Spec.DRA.Container, "dra", containerVolumeMounts),
+				PriorityClassName:            "system-node-critical",
+				ImagePullSecrets:             getPodPullSecrets(mod.Spec.ImageRepoSecret),
+				NodeSelector:                 nodeSelector,
+				ServiceAccountName:           mod.Spec.DRA.ServiceAccountName,
+				Volumes:                      append([]v1.Volume{pluginsVolume, registryVolume}, mod.Spec.DRA.Volumes...),
+				Tolerations:                  mod.Spec.Tolerations,
+				AutomountServiceAccountToken: mod.Spec.DRA.AutomountServiceAccountToken,
+			},
+		},
+	}
+
+	return controllerutil.SetControllerReference(mod, ds, dsci.scheme)
+}