[release-2.6] Add configurable StartupProbe and LivenessProbe to Module CR

api/v1beta1/module_types.go

@@ -298,6 +298,16 @@ type DevicePluginContainerSpec struct {
 	// VolumeMounts is a list of volume mounts that are appended to the default ones.
 	// +optional
 	VolumeMounts []v1.VolumeMount `json:"volumeMounts,omitempty"`
+
+	// StartupProbe allows a startup probe to be set for this container.
+	// When nil, no startup probes are configured.
+	// +optional
+	StartupProbe *v1.Probe `json:"startupProbe,omitempty"`
+
+	// LivenessProbe allows a liveness probe to be set for this container.
+	// When nil, no liveness probe is configured.
+	// +optional
+	LivenessProbe *v1.Probe `json:"livenessProbe,omitempty"`
 }
 
 type DevicePluginSpec struct {

internal/controllers/device_plugin_reconciler.go

@@ -443,8 +443,9 @@ func (dsci *daemonSetCreatorImpl) setDevicePluginAsDesired(
 				Finalizers: []string{constants.NodeLabelerFinalizer},
 			},
 			Spec: v1.PodSpec{
-				InitContainers:               generatePodContainerSpec(mod.Spec.DevicePlugin.InitContainer, "device-plugin-init", nil),
-				Containers:                   generatePodContainerSpec(&mod.Spec.DevicePlugin.Container, "device-plugin", containerVolumeMounts),
+				InitContainers: generatePodContainerSpec(mod.Spec.DevicePlugin.InitContainer, "device-plugin-init", nil, nil, nil),
+				// DevicePlugin has no default probe, so the CR value (possibly nil) is passed directly.
+				Containers:                   generatePodContainerSpec(&mod.Spec.DevicePlugin.Container, "device-plugin", containerVolumeMounts, mod.Spec.DevicePlugin.Container.LivenessProbe, mod.Spec.DevicePlugin.Container.StartupProbe),
 				PriorityClassName:            "system-node-critical",
 				ImagePullSecrets:             getPodPullSecrets(mod.Spec.ImageRepoSecret),
 				NodeSelector:                 nodeSelector,
@@ -459,7 +460,13 @@ func (dsci *daemonSetCreatorImpl) setDevicePluginAsDesired(
 	return controllerutil.SetControllerReference(mod, ds, dsci.scheme)
 }
 
-func generatePodContainerSpec(containerSpec *kmmv1beta1.DevicePluginContainerSpec, containerName string, presetVolumeMounts []v1.VolumeMount) []v1.Container {
+func generatePodContainerSpec(
+	containerSpec *kmmv1beta1.DevicePluginContainerSpec,
+	containerName string,
+	presetVolumeMounts []v1.VolumeMount,
+	livenessProbe *v1.Probe,
+	startupProbe *v1.Probe,
+) []v1.Container {
 	if containerSpec == nil {
 		return nil
 	}
@@ -474,6 +481,8 @@ func generatePodContainerSpec(containerSpec *kmmv1beta1.DevicePluginContainerSpe
 			Resources:       containerSpec.Resources,
 			SecurityContext: &v1.SecurityContext{Privileged: ptr.To(true)},
 			VolumeMounts:    append(containerSpec.VolumeMounts, presetVolumeMounts...),
+			LivenessProbe:   livenessProbe,
+			StartupProbe:    startupProbe,
 		},
 	}
 }

internal/controllers/device_plugin_reconciler_test.go

@@ -21,6 +21,7 @@ import (
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime/schema"
+	"k8s.io/apimachinery/pkg/util/intstr"
 	"k8s.io/utils/ptr"
 	ctrlclient "sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/reconcile"
@@ -672,7 +673,7 @@ var _ = Describe("DevicePluginReconciler_setDevicePluginAsDesired", func() {
 	})
 
 	DescribeTable("should work as expected",
-		func(moduleLoader *kmmv1beta1.ModuleLoaderSpec, expectedNodeSelector map[string]string, withInitContainer bool) {
+		func(moduleLoader *kmmv1beta1.ModuleLoaderSpec, expectedNodeSelector map[string]string, withInitContainer bool, customLiveness *v1.Probe, customStartup *v1.Probe) {
 			const (
 				dsName             = "ds-name"
 				serviceAccountName = "some-service-account"
@@ -752,6 +753,8 @@ var _ = Describe("DevicePluginReconciler_setDevicePluginAsDesired", func() {
 							ImagePullPolicy: ipp,
 							Resources:       resources,
 							VolumeMounts:    []v1.VolumeMount{dpVolMount},
+							LivenessProbe:   customLiveness,
+							StartupProbe:    customStartup,
 						},
 						ServiceAccountName:           serviceAccountName,
 						Volumes:                      []v1.Volume{dpVol},
@@ -839,6 +842,8 @@ var _ = Describe("DevicePluginReconciler_setDevicePluginAsDesired", func() {
 											MountPath: "/var/lib/kubelet/device-plugins",
 										},
 									},
+									LivenessProbe: customLiveness,
+									StartupProbe:  customStartup,
 								},
 							},
 							ImagePullSecrets:   []v1.LocalObjectReference{repoSecret},
@@ -873,6 +878,7 @@ var _ = Describe("DevicePluginReconciler_setDevicePluginAsDesired", func() {
 			nil,
 			map[string]string{"has-feature-x": "true"},
 			false,
+			nil, nil,
 		),
 		Entry("moduleLoader is defined",
 			&kmmv1beta1.ModuleLoaderSpec{},
@@ -881,6 +887,51 @@ var _ = Describe("DevicePluginReconciler_setDevicePluginAsDesired", func() {
 				utils.GetDevicePluginTargetNodeLabel(namespace, moduleName): "",
 			},
 			true,
+			nil, nil,
+		),
+		Entry("with custom liveness probe",
+			nil,
+			map[string]string{"has-feature-x": "true"},
+			false,
+			&v1.Probe{
+				ProbeHandler: v1.ProbeHandler{
+					HTTPGet: &v1.HTTPGetAction{Path: "/healthz", Port: intstr.FromInt32(8080)},
+				},
+				PeriodSeconds: 30,
+			}, nil,
+		),
+		Entry("with custom startup probe",
+			nil,
+			map[string]string{"has-feature-x": "true"},
+			false,
+			nil,
+			&v1.Probe{
+				ProbeHandler: v1.ProbeHandler{
+					HTTPGet: &v1.HTTPGetAction{Path: "/ready", Port: intstr.FromInt32(8080)},
+				},
+				InitialDelaySeconds: 10,
+				PeriodSeconds:       5,
+				FailureThreshold:    30,
+			},
+		),
+		Entry("with custom liveness and startup probes",
+			nil,
+			map[string]string{"has-feature-x": "true"},
+			false,
+			&v1.Probe{
+				ProbeHandler: v1.ProbeHandler{
+					HTTPGet: &v1.HTTPGetAction{Path: "/healthz", Port: intstr.FromInt32(8080)},
+				},
+				PeriodSeconds: 30,
+			},
+			&v1.Probe{
+				ProbeHandler: v1.ProbeHandler{
+					HTTPGet: &v1.HTTPGetAction{Path: "/ready", Port: intstr.FromInt32(8080)},
+				},
+				InitialDelaySeconds: 10,
+				PeriodSeconds:       5,
+				FailureThreshold:    30,
+			},
 		),
 	)
 })

internal/webhook/module.go

@@ -117,6 +117,15 @@ func validateModule(mod *kmmv1beta1.Module) (admission.Warnings, error) {
 		return nil, fmt.Errorf("failed to validate device plugin volumes: %v", err)
 	}
 
+	if mod.Spec.DevicePlugin != nil && mod.Spec.DevicePlugin.InitContainer != nil {
+		if mod.Spec.DevicePlugin.InitContainer.StartupProbe != nil {
+			return nil, fmt.Errorf("spec.devicePlugin.initContainer.startupProbe is not supported")
+		}
+		if mod.Spec.DevicePlugin.InitContainer.LivenessProbe != nil {
+			return nil, fmt.Errorf("spec.devicePlugin.initContainer.livenessProbe is not supported")
+		}
+	}
+
 	if mod.Spec.ModuleLoader == nil {
 		// If ModuleLoader is nil, there is no need to validate related fields
 		return nil, nil

internal/webhook/module_test.go

@@ -429,6 +429,36 @@ var _ = Describe("validateModule", func() {
 		_, err := validateModule(&mod)
 		Expect(err).NotTo(HaveOccurred())
 	})
+
+	It("should reject startupProbe on devicePlugin initContainer", func() {
+		mod := validModule
+		mod.Spec.DevicePlugin = &kmmv1beta1.DevicePluginSpec{
+			Container: kmmv1beta1.DevicePluginContainerSpec{
+				Image: "plugin:latest",
+			},
+			InitContainer: &kmmv1beta1.DevicePluginContainerSpec{
+				Image:        "init:latest",
+				StartupProbe: &v1.Probe{},
+			},
+		}
+		_, err := validateModule(&mod)
+		Expect(err).To(MatchError(ContainSubstring("spec.devicePlugin.initContainer.startupProbe is not supported")))
+	})
+
+	It("should reject livenessProbe on devicePlugin initContainer", func() {
+		mod := validModule
+		mod.Spec.DevicePlugin = &kmmv1beta1.DevicePluginSpec{
+			Container: kmmv1beta1.DevicePluginContainerSpec{
+				Image: "plugin:latest",
+			},
+			InitContainer: &kmmv1beta1.DevicePluginContainerSpec{
+				Image:         "init:latest",
+				LivenessProbe: &v1.Probe{},
+			},
+		}
+		_, err := validateModule(&mod)
+		Expect(err).To(MatchError(ContainSubstring("spec.devicePlugin.initContainer.livenessProbe is not supported")))
+	})
 })
 
 var _ = Describe("ValidateCreate", func() {