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.

Additionally addresses review feedback:
- Rename handleModuleDeletion to deleteDRADaemonSets and
  deleteDevicePluginDaemonSets for clarity
- Add DevicePluginRoleLabelValue label to device plugin DaemonSets;
  the exclusion-based query is kept for backward compatibility with
  a TODO(3.0) to switch to direct role-label querying
- Add cleanup logic to DevicePluginReconciler to delete existing
  DaemonSets and clear status when spec.devicePlugin is removed

Author: TomerNewman

internal/constants/constants.go

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

internal/controllers/device_plugin_reconciler.go

@@ -100,12 +100,21 @@ func (r *DevicePluginReconciler) Reconcile(ctx context.Context, mod *kmmv1beta1.
 		if err = r.reconHelperAPI.removeDevicePluginTargetLabels(ctx, mod); err != nil {
 			return ctrl.Result{}, fmt.Errorf("could not remove device-plugin-target labels on deletion: %v", err)
 		}
-		err = r.reconHelperAPI.handleModuleDeletion(ctx, existingDevicePluginDS)
+		err = r.reconHelperAPI.deleteDevicePluginDaemonSets(ctx, existingDevicePluginDS)
 		return ctrl.Result{}, err
 	}
 
 	r.reconHelperAPI.setKMMOMetrics(ctx)
 
+	if mod.Spec.DevicePlugin == nil {
+		if len(existingDevicePluginDS) > 0 {
+			if err = r.reconHelperAPI.deleteDevicePluginDaemonSets(ctx, existingDevicePluginDS); err != nil {
+				return ctrl.Result{}, err
+			}
+		}
+		return ctrl.Result{}, r.reconHelperAPI.clearDevicePluginStatus(ctx, mod)
+	}
+
 	if err = r.reconHelperAPI.handleDevicePluginTargetLabels(ctx, mod); err != nil {
 		return res, fmt.Errorf("could not reconcile device-plugin-target labels: %v", err)
 	}
@@ -131,16 +140,17 @@ func (r *DevicePluginReconciler) Reconcile(ctx context.Context, mod *kmmv1beta1.
 	return res, nil
 }
 
-//go:generate mockgen -source=device_plugin_reconciler.go -package=controllers -destination=mock_device_plugin_reconciler.go devicePluginReconcilerHelperAPI
+//go:generate mockgen -source=device_plugin_reconciler.go -package=controllers -destination=mock_device_plugin_reconciler.go devicePluginReconcilerHelperAPI,daemonSetCreator
 
 type devicePluginReconcilerHelperAPI interface {
 	setKMMOMetrics(ctx context.Context)
 	handleDevicePlugin(ctx context.Context, mod *kmmv1beta1.Module, existingDevicePluginDS []appsv1.DaemonSet) error
 	handleDevicePluginTargetLabels(ctx context.Context, mod *kmmv1beta1.Module) error
 	removeDevicePluginTargetLabels(ctx context.Context, mod *kmmv1beta1.Module) error
 	garbageCollect(ctx context.Context, mod *kmmv1beta1.Module, existingDS []appsv1.DaemonSet) error
-	handleModuleDeletion(ctx context.Context, existingDevicePluginDS []appsv1.DaemonSet) error
+	deleteDevicePluginDaemonSets(ctx context.Context, existingDevicePluginDS []appsv1.DaemonSet) error
 	moduleUpdateDevicePluginStatus(ctx context.Context, mod *kmmv1beta1.Module, existingDevicePluginDS []appsv1.DaemonSet) error
+	clearDevicePluginStatus(ctx context.Context, mod *kmmv1beta1.Module) error
 	getModuleDevicePluginDaemonSets(ctx context.Context, name, namespace string) ([]appsv1.DaemonSet, error)
 }
 
@@ -176,9 +186,11 @@ func (dprh *devicePluginReconcilerHelper) getModuleDevicePluginDaemonSets(ctx co
 	}
 
 	devicePluginsList := make([]appsv1.DaemonSet, 0, len(dsList.Items))
-	// remove the older version module loader daemonsets
+	// TODO(3.0): switch to querying by DevicePluginRoleLabelValue directly once all
+	// clusters have been upgraded and existing DaemonSets have the role label.
 	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)
 		}
 	}
@@ -286,7 +298,7 @@ func (dprh *devicePluginReconcilerHelper) garbageCollect(ctx context.Context,
 	return nil
 }
 
-func (dprh *devicePluginReconcilerHelper) handleModuleDeletion(ctx context.Context, existingDevicePluginDS []appsv1.DaemonSet) error {
+func (dprh *devicePluginReconcilerHelper) deleteDevicePluginDaemonSets(ctx context.Context, existingDevicePluginDS []appsv1.DaemonSet) error {
 	// delete all the Device Plugin Daemonset, in order to allow worker pods to delete kernel modules
 	for _, ds := range existingDevicePluginDS {
 		err := dprh.client.Delete(ctx, &ds)
@@ -369,7 +381,18 @@ func (dprh *devicePluginReconcilerHelper) moduleUpdateDevicePluginStatus(ctx con
 	return dprh.client.Status().Patch(ctx, mod, client.MergeFrom(unmodifiedMod))
 }
 
-//go:generate mockgen -source=device_plugin_reconciler.go -package=controllers -destination=mock_device_plugin_reconciler.go daemonSetCreator
+func (dprh *devicePluginReconcilerHelper) clearDevicePluginStatus(ctx context.Context, mod *kmmv1beta1.Module) error {
+	emptyStatus := kmmv1beta1.DaemonSetStatus{}
+	if mod.Status.DevicePlugin == emptyStatus {
+		return nil
+	}
+
+	unmodifiedMod := mod.DeepCopy()
+
+	mod.Status.DevicePlugin = kmmv1beta1.DaemonSetStatus{}
+
+	return dprh.client.Status().Patch(ctx, mod, client.MergeFrom(unmodifiedMod))
+}
 
 type daemonSetCreator interface {
 	setDevicePluginAsDesired(ctx context.Context, ds *appsv1.DaemonSet, mod *kmmv1beta1.Module) error
@@ -467,7 +490,10 @@ func generatePodContainerSpec(containerSpec *kmmv1beta1.DevicePluginContainerSpe
 }
 
 func generateDevicePluginLabelsAndSelector(mod *kmmv1beta1.Module) (map[string]string, map[string]string) {
-	labels := map[string]string{constants.ModuleNameLabel: mod.Name}
+	labels := map[string]string{
+		constants.ModuleNameLabel: mod.Name,
+		constants.DaemonSetRole:   constants.DevicePluginRoleLabelValue,
+	}
 	nodeSelector := map[string]string{
 		utils.GetKernelModuleReadyNodeLabel(mod.Namespace, mod.Name):  "",
 		utils.GetDevicePluginTargetNodeLabel(mod.Namespace, mod.Name): "",

internal/controllers/device_plugin_reconciler_test.go

@@ -41,6 +41,9 @@ var _ = Describe("DevicePluginReconciler_Reconcile", func() {
 
 		mod = &kmmv1beta1.Module{
 			ObjectMeta: metav1.ObjectMeta{Namespace: namespace, Name: moduleName},
+			Spec: kmmv1beta1.ModuleSpec{
+				DevicePlugin: &kmmv1beta1.DevicePluginSpec{},
+			},
 		}
 
 		dpr = &DevicePluginReconciler{
@@ -115,7 +118,7 @@ var _ = Describe("DevicePluginReconciler_Reconcile", func() {
 		gomock.InOrder(
 			mockReconHelper.EXPECT().getModuleDevicePluginDaemonSets(ctx, mod.Name, mod.Namespace).Return(devicePluginDS, nil),
 			mockReconHelper.EXPECT().removeDevicePluginTargetLabels(ctx, mod).Return(nil),
-			mockReconHelper.EXPECT().handleModuleDeletion(ctx, devicePluginDS).Return(nil),
+			mockReconHelper.EXPECT().deleteDevicePluginDaemonSets(ctx, devicePluginDS).Return(nil),
 		)
 
 		res, err := dpr.Reconcile(ctx, mod)
@@ -132,18 +135,64 @@ var _ = Describe("DevicePluginReconciler_Reconcile", func() {
 		Expect(res).To(Equal(reconcile.Result{}))
 		Expect(err).To(HaveOccurred())
 
-		By("error flow - handleModuleDeletion fails")
+		By("error flow - deleteDevicePluginDaemonSets fails")
 		gomock.InOrder(
 			mockReconHelper.EXPECT().getModuleDevicePluginDaemonSets(ctx, mod.Name, mod.Namespace).Return(devicePluginDS, nil),
 			mockReconHelper.EXPECT().removeDevicePluginTargetLabels(ctx, mod).Return(nil),
-			mockReconHelper.EXPECT().handleModuleDeletion(ctx, devicePluginDS).Return(fmt.Errorf("some error")),
+			mockReconHelper.EXPECT().deleteDevicePluginDaemonSets(ctx, devicePluginDS).Return(fmt.Errorf("some error")),
 		)
 
 		res, err = dpr.Reconcile(ctx, mod)
 		Expect(res).To(Equal(reconcile.Result{}))
 		Expect(err).To(HaveOccurred())
 	})
 
+	It("no-op when spec.devicePlugin is nil and no existing DaemonSets", func() {
+		mod.Spec.DevicePlugin = nil
+		gomock.InOrder(
+			mockReconHelper.EXPECT().getModuleDevicePluginDaemonSets(ctx, mod.Name, mod.Namespace).Return(nil, nil),
+			mockReconHelper.EXPECT().setKMMOMetrics(ctx),
+			mockReconHelper.EXPECT().clearDevicePluginStatus(ctx, mod).Return(nil),
+		)
+
+		res, err := dpr.Reconcile(ctx, mod)
+
+		Expect(res).To(Equal(reconcile.Result{}))
+		Expect(err).NotTo(HaveOccurred())
+	})
+
+	It("cleanup when spec.devicePlugin is nil but existing DaemonSets present", func() {
+		mod.Spec.DevicePlugin = nil
+		devicePluginDS := []appsv1.DaemonSet{{}}
+
+		gomock.InOrder(
+			mockReconHelper.EXPECT().getModuleDevicePluginDaemonSets(ctx, mod.Name, mod.Namespace).Return(devicePluginDS, nil),
+			mockReconHelper.EXPECT().setKMMOMetrics(ctx),
+			mockReconHelper.EXPECT().deleteDevicePluginDaemonSets(ctx, devicePluginDS).Return(nil),
+			mockReconHelper.EXPECT().clearDevicePluginStatus(ctx, mod).Return(nil),
+		)
+
+		res, err := dpr.Reconcile(ctx, mod)
+
+		Expect(res).To(Equal(reconcile.Result{}))
+		Expect(err).NotTo(HaveOccurred())
+	})
+
+	It("cleanup when spec.devicePlugin is nil and clearDevicePluginStatus fails", func() {
+		mod.Spec.DevicePlugin = nil
+
+		gomock.InOrder(
+			mockReconHelper.EXPECT().getModuleDevicePluginDaemonSets(ctx, mod.Name, mod.Namespace).Return(nil, nil),
+			mockReconHelper.EXPECT().setKMMOMetrics(ctx),
+			mockReconHelper.EXPECT().clearDevicePluginStatus(ctx, mod).Return(fmt.Errorf("some error")),
+		)
+
+		res, err := dpr.Reconcile(ctx, mod)
+
+		Expect(res).To(Equal(reconcile.Result{}))
+		Expect(err).To(HaveOccurred())
+	})
+
 })
 
 var _ = Describe("DevicePluginReconciler_handleDevicePlugin", func() {
@@ -337,7 +386,7 @@ var _ = Describe("DevicePluginReconciler_garbageCollect", func() {
 	})
 })
 
-var _ = Describe("DevicePluginReconciler_handleModuleDeletion", func() {
+var _ = Describe("DevicePluginReconciler_deleteDevicePluginDaemonSets", func() {
 	var (
 		ctrl *gomock.Controller
 		clnt *client.MockClient
@@ -358,14 +407,14 @@ var _ = Describe("DevicePluginReconciler_handleModuleDeletion", func() {
 	It("good flow", func() {
 		clnt.EXPECT().Delete(ctx, &existingDevicePluginDS[0]).Return(nil)
 
-		err := dprh.handleModuleDeletion(ctx, existingDevicePluginDS)
+		err := dprh.deleteDevicePluginDaemonSets(ctx, existingDevicePluginDS)
 		Expect(err).NotTo(HaveOccurred())
 	})
 
 	It("error flow", func() {
 		clnt.EXPECT().Delete(ctx, &existingDevicePluginDS[0]).Return(fmt.Errorf("some error"))
 
-		err := dprh.handleModuleDeletion(ctx, existingDevicePluginDS)
+		err := dprh.deleteDevicePluginDaemonSets(ctx, existingDevicePluginDS)
 		Expect(err).To(HaveOccurred())
 	})
 })
@@ -556,6 +605,53 @@ var _ = Describe("DevicePluginReconciler_moduleUpdateDevicePluginStatus", func()
 	)
 })
 
+var _ = Describe("DevicePluginReconciler_clearDevicePluginStatus", func() {
+	var (
+		ctrl         *gomock.Controller
+		clnt         *client.MockClient
+		statusWriter *client.MockStatusWriter
+		dprh         devicePluginReconcilerHelper
+	)
+
+	BeforeEach(func() {
+		ctrl = gomock.NewController(GinkgoT())
+		clnt = client.NewMockClient(ctrl)
+		statusWriter = client.NewMockStatusWriter(ctrl)
+		dprh = devicePluginReconcilerHelper{
+			client: clnt,
+		}
+	})
+
+	ctx := context.Background()
+
+	It("should be a no-op when status.devicePlugin is already empty", func() {
+		mod := kmmv1beta1.Module{}
+		err := dprh.clearDevicePluginStatus(ctx, &mod)
+		Expect(err).NotTo(HaveOccurred())
+	})
+
+	It("should clear status.devicePlugin when it has values", func() {
+		mod := kmmv1beta1.Module{
+			Status: kmmv1beta1.ModuleStatus{
+				DevicePlugin: kmmv1beta1.DaemonSetStatus{
+					NodesMatchingSelectorNumber: 3,
+					DesiredNumber:               3,
+					AvailableNumber:             2,
+				},
+			},
+		}
+
+		expectedMod := mod.DeepCopy()
+		expectedMod.Status.DevicePlugin = kmmv1beta1.DaemonSetStatus{}
+
+		clnt.EXPECT().Status().Return(statusWriter)
+		statusWriter.EXPECT().Patch(ctx, expectedMod, gomock.Any())
+
+		err := dprh.clearDevicePluginStatus(ctx, &mod)
+		Expect(err).NotTo(HaveOccurred())
+	})
+})
+
 var _ = Describe("DevicePluginReconciler_setDevicePluginAsDesired", func() {
 	const (
 		devicePluginImage = "device-plugin-image"
@@ -747,7 +843,10 @@ var _ = Describe("DevicePluginReconciler_setDevicePluginAsDesired", func() {
 			err := dsc.setDevicePluginAsDesired(context.Background(), &ds, &mod)
 			Expect(err).NotTo(HaveOccurred())
 
-			podLabels := map[string]string{constants.ModuleNameLabel: moduleName}
+			podLabels := map[string]string{
+				constants.ModuleNameLabel: moduleName,
+				constants.DaemonSetRole:   constants.DevicePluginRoleLabelValue,
+			}
 
 			expectedInitContainer := []v1.Container{
 				{
@@ -927,7 +1026,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 +1050,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
 			},
 		)