From f9aa14c299d9c242d4d2e7642bbd948926f051d8 Mon Sep 17 00:00:00 2001
From: Mykhailo Bobrovskyi <mikhail.bobrovsky@gmail.com>
Date: Wed, 13 Nov 2024 11:18:16 +0200
Subject: [PATCH] Allow manageJobsWithoutQueueName on StatefulSet.

---
 .../jobs/deployment/deployment_webhook.go     | 16 +++++---------
 .../deployment/deployment_webhook_test.go     | 19 +++++++++-------
 .../jobs/statefulset/statefulset_webhook.go   | 16 ++++++++++----
 .../statefulset/statefulset_webhook_test.go   | 22 +++++++++++++++++++
 4 files changed, 50 insertions(+), 23 deletions(-)

diff --git a/pkg/controller/jobs/deployment/deployment_webhook.go b/pkg/controller/jobs/deployment/deployment_webhook.go
index cbab9b6e62..c1f9c98b3a 100644
--- a/pkg/controller/jobs/deployment/deployment_webhook.go
+++ b/pkg/controller/jobs/deployment/deployment_webhook.go
@@ -34,16 +34,11 @@ 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,
-	}
+func SetupWebhook(mgr ctrl.Manager, _ ...jobframework.Option) error {
+	wh := &Webhook{client: mgr.GetClient()}
 	obj := &appsv1.Deployment{}
 	return webhook.WebhookManagedBy(mgr).
 		For(obj).
@@ -61,12 +56,11 @@ func (wh *Webhook) Default(ctx context.Context, obj runtime.Object) error {
 	log := ctrl.LoggerFrom(ctx).WithName("deployment-webhook").WithValues("deployment", klog.KObj(d))
 	log.V(5).Info("Applying defaults")
 
-	cqLabel, ok := d.Labels[constants.QueueLabel]
-	if ok {
+	if queueName := jobframework.QueueNameForObject(d.Object()); queueName != "" {
 		if d.Spec.Template.Labels == nil {
 			d.Spec.Template.Labels = make(map[string]string, 1)
 		}
-		d.Spec.Template.Labels[constants.QueueLabel] = cqLabel
+		d.Spec.Template.Labels[constants.QueueLabel] = queueName
 	}
 
 	return nil
diff --git a/pkg/controller/jobs/deployment/deployment_webhook_test.go b/pkg/controller/jobs/deployment/deployment_webhook_test.go
index f2ee99f46a..e6c1535617 100644
--- a/pkg/controller/jobs/deployment/deployment_webhook_test.go
+++ b/pkg/controller/jobs/deployment/deployment_webhook_test.go
@@ -29,11 +29,17 @@ import (
 
 func TestDefault(t *testing.T) {
 	testCases := map[string]struct {
-		deployment                 *appsv1.Deployment
-		manageJobsWithoutQueueName bool
-		enableIntegrations         []string
-		want                       *appsv1.Deployment
+		deployment         *appsv1.Deployment
+		enableIntegrations []string
+		want               *appsv1.Deployment
 	}{
+		"pod without queue": {
+			enableIntegrations: []string{"pod"},
+			deployment: testingdeployment.MakeDeployment("test-pod", "").
+				Obj(),
+			want: testingdeployment.MakeDeployment("test-pod", "").
+				Obj(),
+		},
 		"pod with queue": {
 			enableIntegrations: []string{"pod"},
 			deployment: testingdeployment.MakeDeployment("test-pod", "").
@@ -52,10 +58,7 @@ func TestDefault(t *testing.T) {
 			builder := utiltesting.NewClientBuilder()
 			cli := builder.Build()
 
-			w := &Webhook{
-				client:                     cli,
-				manageJobsWithoutQueueName: tc.manageJobsWithoutQueueName,
-			}
+			w := &Webhook{client: cli}
 
 			ctx, _ := utiltesting.ContextWithLog(t)
 
diff --git a/pkg/controller/jobs/statefulset/statefulset_webhook.go b/pkg/controller/jobs/statefulset/statefulset_webhook.go
index 6a5bfaa268..ca53e92132 100644
--- 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
index c40bc698d7..26565037c6 100644
--- 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", "").