managedJobsNamespaceSelector for Deployment, StatefulSets, and Pods (#…

…3765)

* managedJobsNamespaceSelector for deployment and statefulset

* normalize manageJobsWithoutQueueName logic in pod webhook

* refactor to use common helper function

* unit tests for WorkloadShouldBeSuspended

* linter fixes

* add test for job owned by kueue managed parent

* bug fix in deployment/statefulset webhooks

pkg/controller/jobframework/defaults.go

@@ -30,27 +30,44 @@ import (
 
 func ApplyDefaultForSuspend(ctx context.Context, job GenericJob, k8sClient client.Client,
 	manageJobsWithoutQueueName bool, managedJobsNamespaceSelector labels.Selector) error {
+	suspend, err := WorkloadShouldBeSuspended(ctx, job.Object(), k8sClient, manageJobsWithoutQueueName, managedJobsNamespaceSelector)
+	if err != nil {
+		return err
+	}
+	if suspend && !job.IsSuspended() {
+		job.Suspend()
+	}
+	return nil
+}
+
+// WorkloadShouldBeSuspended determines whether jobObj should be default suspended on creation
+func WorkloadShouldBeSuspended(ctx context.Context, jobObj client.Object, k8sClient client.Client,
+	manageJobsWithoutQueueName bool, managedJobsNamespaceSelector labels.Selector) (bool, error) {
 	// Do not default suspend a job whose owner is already managed by Kueue
-	if owner := metav1.GetControllerOf(job.Object()); owner != nil && IsOwnerManagedByKueue(owner) {
-		return nil
+	if owner := metav1.GetControllerOf(jobObj); owner != nil && IsOwnerManagedByKueue(owner) {
+		return false, nil
 	}
 
-	// Do not default suspend a job without a queue name unless the namespace selector also matches
-	if features.Enabled(features.ManagedJobsNamespaceSelector) && manageJobsWithoutQueueName && QueueName(job) == "" {
-		ns := corev1.Namespace{}
-		err := k8sClient.Get(ctx, client.ObjectKey{Name: job.Object().GetNamespace()}, &ns)
-		if err != nil {
-			return fmt.Errorf("failed to get namespace: %w", err)
-		}
-		if !managedJobsNamespaceSelector.Matches(labels.Set(ns.GetLabels())) {
-			return nil
-		}
+	// Jobs with queue names whose parents are not managed by Kueue are default suspended
+	if QueueNameForObject(jobObj) != "" {
+		return true, nil
 	}
 
-	if QueueName(job) != "" || manageJobsWithoutQueueName {
-		if !job.IsSuspended() {
-			job.Suspend()
+	// Logic for managing jobs without queue names.
+	if manageJobsWithoutQueueName {
+		if features.Enabled(features.ManagedJobsNamespaceSelector) {
+			// Default suspend the job if the namespace selector matches
+			ns := corev1.Namespace{}
+			err := k8sClient.Get(ctx, client.ObjectKey{Name: jobObj.GetNamespace()}, &ns)
+			if err != nil {
+				return false, fmt.Errorf("failed to get namespace: %w", err)
+			}
+			return managedJobsNamespaceSelector.Matches(labels.Set(ns.GetLabels())), nil
+		} else {
+			// Namespace filtering is disabled; unconditionally default suspend
+			return true, nil
 		}
 	}
-	return nil
+
+	return false, nil
 }

pkg/controller/jobframework/defaults_test.go

@@ -0,0 +1,133 @@
+/*
+Copyright 2024 The Kubernetes Authors.
+
+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 jobframework
+
+import (
+	"testing"
+
+	batchv1 "k8s.io/api/batch/v1"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+
+	"sigs.k8s.io/kueue/pkg/features"
+	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
+	utiltestingjob "sigs.k8s.io/kueue/pkg/util/testingjobs/job"
+)
+
+func TestWorkloadShouldBeSuspended(t *testing.T) {
+	t.Cleanup(EnableIntegrationsForTest(t, "batch/job"))
+	managedNamespace := &corev1.Namespace{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:   "managed-ns",
+			Labels: map[string]string{"kubernetes.io/metadata.name": "managed-ns"},
+		},
+	}
+	unmanagedNamespace := &corev1.Namespace{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:   "unmanaged-ns",
+			Labels: map[string]string{"kubernetes.io/metadata.name": "unmanaged-ns"},
+		},
+	}
+	parent := utiltestingjob.MakeJob("parent", managedNamespace.Name).Queue("default").Obj()
+	ls := &metav1.LabelSelector{
+		MatchExpressions: []metav1.LabelSelectorRequirement{
+			{
+				Key:      "kubernetes.io/metadata.name",
+				Operator: metav1.LabelSelectorOpNotIn,
+				Values:   []string{unmanagedNamespace.Name},
+			},
+		},
+	}
+	namespaceSelector, _ := metav1.LabelSelectorAsSelector(ls)
+
+	cases := map[string]struct {
+		obj                        client.Object
+		manageJobsWithoutQueueName bool
+		featureGateEnabled         bool
+		wantSuspend                bool
+	}{
+		"job with queue name ": {
+			obj:                        utiltestingjob.MakeJob("test-job", managedNamespace.Name).Queue("default").Obj(),
+			manageJobsWithoutQueueName: false,
+			featureGateEnabled:         true,
+			wantSuspend:                true,
+		},
+		"job with queue name manageJobs": {
+			obj:                        utiltestingjob.MakeJob("test-job", managedNamespace.Name).Queue("default").Obj(),
+			manageJobsWithoutQueueName: true,
+			featureGateEnabled:         true,
+			wantSuspend:                true,
+		},
+		"job without queue name": {
+			obj:                        utiltestingjob.MakeJob("test-job", managedNamespace.Name).Obj(),
+			manageJobsWithoutQueueName: false,
+			featureGateEnabled:         true,
+			wantSuspend:                false,
+		},
+		"job without queue name with manageJobs": {
+			obj:                        utiltestingjob.MakeJob("test-job", managedNamespace.Name).Obj(),
+			manageJobsWithoutQueueName: true,
+			featureGateEnabled:         true,
+			wantSuspend:                true,
+		},
+		"job without queue name but with managed parent with manageJobs": {
+			obj: utiltestingjob.MakeJob("test-job", managedNamespace.Name).
+				OwnerReference(parent.Name, batchv1.SchemeGroupVersion.WithKind("Job")).
+				Obj(),
+			manageJobsWithoutQueueName: true,
+			featureGateEnabled:         true,
+			wantSuspend:                false,
+		},
+		"job without queue name with manageJobs with feature disabled": {
+			obj:                        utiltestingjob.MakeJob("test-job", managedNamespace.Name).Obj(),
+			manageJobsWithoutQueueName: true,
+			featureGateEnabled:         false,
+			wantSuspend:                true,
+		},
+		"job without queue name with manageJobs in unmanaged ns": {
+			obj:                        utiltestingjob.MakeJob("test-job", unmanagedNamespace.Name).Obj(),
+			manageJobsWithoutQueueName: true,
+			featureGateEnabled:         true,
+			wantSuspend:                false,
+		},
+		"job without queue name with manageJobs in unmanaged ns with feature disabled": {
+			obj:                        utiltestingjob.MakeJob("test-job", unmanagedNamespace.Name).Obj(),
+			manageJobsWithoutQueueName: true,
+			featureGateEnabled:         false,
+			wantSuspend:                true,
+		},
+	}
+
+	for tcName, tc := range cases {
+		t.Run(tcName, func(t *testing.T) {
+			builder := utiltesting.NewClientBuilder()
+			builder.WithObjects(managedNamespace, unmanagedNamespace, tc.obj, parent)
+			client := builder.Build()
+			ctx, _ := utiltesting.ContextWithLog(t)
+
+			features.SetFeatureGateDuringTest(t, features.ManagedJobsNamespaceSelector, tc.featureGateEnabled)
+			suspend, err := WorkloadShouldBeSuspended(ctx, tc.obj, client, tc.manageJobsWithoutQueueName, namespaceSelector)
+			if err != nil {
+				t.Errorf("Got error: %v", err)
+			}
+			if suspend != tc.wantSuspend {
+				t.Errorf("Unexpected result: got %v wanted %v", suspend, tc.wantSuspend)
+			}
+		})
+	}
+}

pkg/controller/jobs/deployment/deployment_webhook.go

@@ -36,7 +36,7 @@ type Webhook struct {
 	client client.Client
 }
 
-func SetupWebhook(mgr ctrl.Manager, _ ...jobframework.Option) error {
+func SetupWebhook(mgr ctrl.Manager, opts ...jobframework.Option) error {
 	wh := &Webhook{
 		client: mgr.GetClient(),
 	}
@@ -56,9 +56,11 @@ func (wh *Webhook) Default(ctx context.Context, obj runtime.Object) error {
 	deployment := fromObject(obj)
 
 	log := ctrl.LoggerFrom(ctx).WithName("deployment-webhook")
-	log.V(5).Info("Applying defaults")
+	log.V(5).Info("Propagating queue-name")
 
-	if queueName := jobframework.QueueNameForObject(deployment.Object()); queueName != "" {
+	// Because Deployment is built using a NoOpReconciler handling of jobs without queue names is delegating to the Pod webhook.
+	queueName := jobframework.QueueNameForObject(deployment.Object())
+	if queueName != "" {
 		if deployment.Spec.Template.Labels == nil {
 			deployment.Spec.Template.Labels = make(map[string]string, 1)
 		}

pkg/controller/jobs/pod/pod_webhook.go

@@ -146,44 +146,44 @@ func (w *PodWebhook) Default(ctx context.Context, obj runtime.Object) error {
 	log := ctrl.LoggerFrom(ctx).WithName("pod-webhook")
 	log.V(5).Info("Applying defaults")
 
-	if IsPodOwnerManagedByKueue(pod) {
-		log.V(5).Info("Pod owner is managed by kueue, skipping")
-		return nil
-	}
-
-	// Check for pod label selector match
-	podSelector, err := metav1.LabelSelectorAsSelector(w.podSelector)
+	suspend, err := jobframework.WorkloadShouldBeSuspended(ctx, pod.Object(), w.client, w.manageJobsWithoutQueueName, w.managedJobsNamespaceSelector)
 	if err != nil {
-		return fmt.Errorf("failed to parse pod selector: %w", err)
-	}
-	if !podSelector.Matches(labels.Set(pod.pod.GetLabels())) {
-		return nil
+		return err
 	}
 
-	// Get pod namespace and check for namespace label selector match
-	ns := corev1.Namespace{}
-	err = w.client.Get(ctx, client.ObjectKey{Name: pod.pod.GetNamespace()}, &ns)
-	if err != nil {
-		return fmt.Errorf("failed to run mutating webhook on pod %s, error while getting namespace: %w",
-			pod.pod.GetName(),
-			err,
-		)
-	}
-	log.V(5).Info("Found pod namespace", "Namespace.Name", ns.GetName())
-	if features.Enabled(features.ManagedJobsNamespaceSelector) {
-		if !w.managedJobsNamespaceSelector.Matches(labels.Set(ns.GetLabels())) {
+	// Backwards compatibility support until podOptions.podSelector and podOptions.namespaceSelector are deprecated.
+	// When WorkloadShouldBeSuspend determines that suspend is true, also run the podOptions based checks
+	// and if either of them exempts the Pod from suspension, we return early.
+	if suspend {
+		// podOptions.podSelector
+		podSelector, err := metav1.LabelSelectorAsSelector(w.podSelector)
+		if err != nil {
+			return fmt.Errorf("failed to parse pod selector: %w", err)
+		}
+		if !podSelector.Matches(labels.Set(pod.pod.GetLabels())) {
+			return nil
+		}
+
+		// podOptions.namespaceSelector
+		ns := corev1.Namespace{}
+		err = w.client.Get(ctx, client.ObjectKey{Name: pod.pod.GetNamespace()}, &ns)
+		if err != nil {
+			return fmt.Errorf("failed to run mutating webhook on pod %s, error while getting namespace: %w",
+				pod.pod.GetName(),
+				err,
+			)
+		}
+		log.V(5).Info("Found pod namespace", "Namespace.Name", ns.GetName())
+		nsSelector, err := metav1.LabelSelectorAsSelector(w.namespaceSelector)
+		if err != nil {
+			return fmt.Errorf("failed to parse namespace selector: %w", err)
+		}
+		if !nsSelector.Matches(labels.Set(ns.GetLabels())) {
 			return nil
 		}
-	}
-	nsSelector, err := metav1.LabelSelectorAsSelector(w.namespaceSelector)
-	if err != nil {
-		return fmt.Errorf("failed to parse namespace selector: %w", err)
-	}
-	if !nsSelector.Matches(labels.Set(ns.GetLabels())) {
-		return nil
 	}
 
-	if jobframework.QueueName(pod) != "" || w.manageJobsWithoutQueueName {
+	if suspend {
 		controllerutil.AddFinalizer(pod.Object(), PodFinalizer)
 
 		if pod.pod.Labels == nil {

pkg/controller/jobs/statefulset/statefulset_webhook.go

@@ -37,15 +37,12 @@ import (
 )
 
 type Webhook struct {
-	client                     client.Client
-	manageJobsWithoutQueueName bool
+	client client.Client
 }
 
 func SetupWebhook(mgr ctrl.Manager, opts ...jobframework.Option) error {
-	options := jobframework.ProcessOptions(opts...)
 	wh := &Webhook{
-		client:                     mgr.GetClient(),
-		manageJobsWithoutQueueName: options.ManageJobsWithoutQueueName,
+		client: mgr.GetClient(),
 	}
 	return ctrl.NewWebhookManagedBy(mgr).
 		For(&appsv1.StatefulSet{}).
@@ -61,27 +58,26 @@ var _ webhook.CustomDefaulter = &Webhook{}
 func (wh *Webhook) Default(ctx context.Context, obj runtime.Object) error {
 	ss := fromObject(obj)
 	log := ctrl.LoggerFrom(ctx).WithName("statefulset-webhook")
-	log.V(5).Info("Applying defaults")
+	log.V(5).Info("Propagating queue-name")
 
+	// Because StatefuleSet is built using a NoOpReconciler handling of jobs without queue names is delegating to the Pod webhook.
 	queueName := jobframework.QueueNameForObject(ss.Object())
-	if queueName == "" {
-		return nil
+	if queueName != "" {
+		if ss.Spec.Template.Labels == nil {
+			ss.Spec.Template.Labels = make(map[string]string, 2)
+		}
+		ss.Spec.Template.Labels[constants.QueueLabel] = queueName
+		ss.Spec.Template.Labels[pod.GroupNameLabel] = GetWorkloadName(ss.Name)
+
+		if ss.Spec.Template.Annotations == nil {
+			ss.Spec.Template.Annotations = make(map[string]string, 4)
+		}
+		ss.Spec.Template.Annotations[pod.GroupTotalCountAnnotation] = fmt.Sprint(ptr.Deref(ss.Spec.Replicas, 1))
+		ss.Spec.Template.Annotations[pod.GroupFastAdmissionAnnotation] = "true"
+		ss.Spec.Template.Annotations[pod.GroupServingAnnotation] = "true"
+		ss.Spec.Template.Annotations[kueuealpha.PodGroupPodIndexLabelAnnotation] = appsv1.PodIndexLabel
 	}
 
-	if ss.Spec.Template.Labels == nil {
-		ss.Spec.Template.Labels = make(map[string]string, 2)
-	}
-	ss.Spec.Template.Labels[constants.QueueLabel] = queueName
-	ss.Spec.Template.Labels[pod.GroupNameLabel] = GetWorkloadName(ss.Name)
-
-	if ss.Spec.Template.Annotations == nil {
-		ss.Spec.Template.Annotations = make(map[string]string, 4)
-	}
-	ss.Spec.Template.Annotations[pod.GroupTotalCountAnnotation] = fmt.Sprint(ptr.Deref(ss.Spec.Replicas, 1))
-	ss.Spec.Template.Annotations[pod.GroupFastAdmissionAnnotation] = "true"
-	ss.Spec.Template.Annotations[pod.GroupServingAnnotation] = "true"
-	ss.Spec.Template.Annotations[kueuealpha.PodGroupPodIndexLabelAnnotation] = appsv1.PodIndexLabel
-
 	return nil
 }
 

pkg/controller/jobs/statefulset/statefulset_webhook_test.go

@@ -38,10 +38,9 @@ import (
 
 func TestDefault(t *testing.T) {
 	testCases := map[string]struct {
-		statefulset                *appsv1.StatefulSet
-		manageJobsWithoutQueueName bool
-		enableIntegrations         []string
-		want                       *appsv1.StatefulSet
+		statefulset        *appsv1.StatefulSet
+		enableIntegrations []string
+		want               *appsv1.StatefulSet
 	}{
 		"statefulset with queue": {
 			enableIntegrations: []string{"pod"},
@@ -84,8 +83,7 @@ func TestDefault(t *testing.T) {
 			cli := builder.Build()
 
 			w := &Webhook{
-				client:                     cli,
-				manageJobsWithoutQueueName: tc.manageJobsWithoutQueueName,
+				client: cli,
 			}
 
 			ctx, _ := utiltesting.ContextWithLog(t)