Allow manageJobsWithoutQueueName on StatefulSet.

kubernetes-sigs · Nov 18, 2024 · 752d4dc · 752d4dc
1 parent e0be159
commit 752d4dc
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 4 deletions.
diff --git a/pkg/controller/jobs/statefulset/statefulset_webhook.go b/pkg/controller/jobs/statefulset/statefulset_webhook.go
@@ -62,16 +62,24 @@ func (wh *Webhook) Default(ctx context.Context, obj runtime.Object) error {
 	log := ctrl.LoggerFrom(ctx).WithName("statefulset-webhook")
 	log.V(5).Info("Applying defaults")
 
-	cqLabel, ok := ss.Labels[constants.QueueLabel]
-	if !ok {
+	queueName := jobframework.QueueNameForObject(ss.Object())
+
+	// Skip not managed by Kueue.
+	if queueName == "" && !wh.manageJobsWithoutQueueName {
 		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
+	}
+
 	if ss.Spec.Template.Labels == nil {
-		ss.Spec.Template.Labels = make(map[string]string, 2)
+		ss.Spec.Template.Labels = make(map[string]string, 1)
 	}
 
-	ss.Spec.Template.Labels[constants.QueueLabel] = cqLabel
 	ss.Spec.Template.Labels[pod.GroupNameLabel] = GetWorkloadName(ss.Name)
 
 	if ss.Spec.Template.Annotations == nil {

diff --git a/pkg/controller/jobs/statefulset/statefulset_webhook_test.go b/pkg/controller/jobs/statefulset/statefulset_webhook_test.go
@@ -42,6 +42,28 @@ func TestDefault(t *testing.T) {
 		enableIntegrations         []string
 		want                       *appsv1.StatefulSet
 	}{
+		"statefulset not managed by kueue": {
+			enableIntegrations: []string{"pod"},
+			statefulset: testingstatefulset.MakeStatefulSet("test-pod", "").
+				Replicas(10).
+				Obj(),
+			want: testingstatefulset.MakeStatefulSet("test-pod", "").
+				Replicas(10).
+				Obj(),
+		},
+		"statefulset managed by kueue without queue name": {
+			enableIntegrations:         []string{"pod"},
+			manageJobsWithoutQueueName: true,
+			statefulset: testingstatefulset.MakeStatefulSet("test-pod", "").
+				Replicas(10).
+				Obj(),
+			want: testingstatefulset.MakeStatefulSet("test-pod", "").
+				Replicas(10).
+				PodTemplateSpecPodGroupNameLabel("test-pod", "", gvk).
+				PodTemplateSpecPodGroupTotalCountAnnotation(10).
+				PodTemplateSpecPodGroupFastAdmissionAnnotation(true).
+				Obj(),
+		},
 		"statefulset with queue": {
 			enableIntegrations: []string{"pod"},
 			statefulset: testingstatefulset.MakeStatefulSet("test-pod", "").