From dd2dd5948b40dd465a92b66905f0d7cba5d206ba Mon Sep 17 00:00:00 2001
From: Michal Wozniak <michalwozniak@google.com>
Date: Wed, 24 Jul 2024 14:10:06 +0200
Subject: [PATCH] Allow mutating schedulingGates when the Jobset is suspended

---
 pkg/webhooks/jobset_webhook.go                |  3 +
 pkg/webhooks/jobset_webhook_test.go           | 94 +++++++++++++++++++
 .../webhook/jobset_webhook_test.go            |  5 +
 3 files changed, 102 insertions(+)

diff --git a/pkg/webhooks/jobset_webhook.go b/pkg/webhooks/jobset_webhook.go
index 85805eaf..d7f080b1 100644
--- a/pkg/webhooks/jobset_webhook.go
+++ b/pkg/webhooks/jobset_webhook.go
@@ -318,6 +318,9 @@ func (j *jobSetWebhook) ValidateUpdate(ctx context.Context, old, newObj runtime.
 			mungedSpec.ReplicatedJobs[index].Template.Spec.Template.Labels = oldJS.Spec.ReplicatedJobs[index].Template.Spec.Template.Labels
 			mungedSpec.ReplicatedJobs[index].Template.Spec.Template.Spec.NodeSelector = oldJS.Spec.ReplicatedJobs[index].Template.Spec.Template.Spec.NodeSelector
 			mungedSpec.ReplicatedJobs[index].Template.Spec.Template.Spec.Tolerations = oldJS.Spec.ReplicatedJobs[index].Template.Spec.Template.Spec.Tolerations
+
+			// Pod Scheduling Gates can be updated for batch/v1 Job: https://github.com/kubernetes/kubernetes/blob/ceb58a4dbc671b9d0a2de6d73a1616bc0c299863/pkg/apis/batch/validation/validation.go#L662
+			mungedSpec.ReplicatedJobs[index].Template.Spec.Template.Spec.SchedulingGates = oldJS.Spec.ReplicatedJobs[index].Template.Spec.Template.Spec.SchedulingGates
 		}
 	}
 
diff --git a/pkg/webhooks/jobset_webhook_test.go b/pkg/webhooks/jobset_webhook_test.go
index e1744c2e..8597a16d 100644
--- a/pkg/webhooks/jobset_webhook_test.go
+++ b/pkg/webhooks/jobset_webhook_test.go
@@ -1561,6 +1561,100 @@ func TestValidateUpdate(t *testing.T) {
 				field.Invalid(field.NewPath("spec").Child("replicatedJobs"), "", "field is immutable"),
 			}.ToAggregate(),
 		},
+		{
+			name: "schedulingGates for pod template can be updated for suspended JobSet",
+			js: &jobset.JobSet{
+				ObjectMeta: validObjectMeta,
+				Spec: jobset.JobSetSpec{
+					ReplicatedJobs: []jobset.ReplicatedJob{
+						{
+							Name:     "test-jobset-replicated-job-0",
+							Replicas: 2,
+							Template: batchv1.JobTemplateSpec{
+								Spec: batchv1.JobSpec{
+									Parallelism: ptr.To[int32](2),
+									Template: corev1.PodTemplateSpec{
+										Spec: corev1.PodSpec{
+											// Adding a scheduling gate
+											SchedulingGates: []corev1.PodSchedulingGate{
+												{
+													Name: "example.com/gate",
+												},
+											},
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			oldJs: &jobset.JobSet{
+				ObjectMeta: validObjectMeta,
+				Spec: jobset.JobSetSpec{
+					Suspend: ptr.To(true),
+					ReplicatedJobs: []jobset.ReplicatedJob{
+						{
+							Name:     "test-jobset-replicated-job-0",
+							Replicas: 2,
+							Template: batchv1.JobTemplateSpec{
+								Spec: batchv1.JobSpec{
+									Parallelism: ptr.To[int32](2),
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+		{
+			name: "schedulingGates for pod template cannot be updated for unsuspended JobSet",
+			js: &jobset.JobSet{
+				ObjectMeta: validObjectMeta,
+				Spec: jobset.JobSetSpec{
+					ReplicatedJobs: []jobset.ReplicatedJob{
+						{
+							Name:     "test-jobset-replicated-job-0",
+							Replicas: 2,
+							Template: batchv1.JobTemplateSpec{
+								Spec: batchv1.JobSpec{
+									Parallelism: ptr.To[int32](2),
+									Template: corev1.PodTemplateSpec{
+										Spec: corev1.PodSpec{
+											// Adding a scheduling gate
+											SchedulingGates: []corev1.PodSchedulingGate{
+												{
+													Name: "example.com/gate",
+												},
+											},
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			oldJs: &jobset.JobSet{
+				ObjectMeta: validObjectMeta,
+				Spec: jobset.JobSetSpec{
+					ReplicatedJobs: []jobset.ReplicatedJob{
+						{
+							Name:     "test-jobset-replicated-job-0",
+							Replicas: 2,
+							Template: batchv1.JobTemplateSpec{
+								Spec: batchv1.JobSpec{
+									Parallelism: ptr.To[int32](2),
+								},
+							},
+						},
+					},
+				},
+			},
+			want: field.ErrorList{
+				field.Invalid(field.NewPath("spec").Child("replicatedJobs"), "", "field is immutable"),
+			}.ToAggregate(),
+		},
 		{
 			name: "replicated job name cannot be updated",
 			js: &jobset.JobSet{
diff --git a/test/integration/webhook/jobset_webhook_test.go b/test/integration/webhook/jobset_webhook_test.go
index 58bf5a03..764542ce 100644
--- a/test/integration/webhook/jobset_webhook_test.go
+++ b/test/integration/webhook/jobset_webhook_test.go
@@ -372,6 +372,11 @@ var _ = ginkgo.Describe("jobset webhook defaulting", func() {
 			},
 			updateJobSet: func(js *jobset.JobSet) {
 				js.Spec.ReplicatedJobs[0].Template.Spec.Template.Annotations["new"] = "annotation"
+				js.Spec.ReplicatedJobs[0].Template.Spec.Template.Spec.SchedulingGates = []corev1.PodSchedulingGate{
+					{
+						Name: "example.com/gate",
+					},
+				}
 			},
 		}),
 		ginkgo.Entry("updating pod template in running jobset is not allowed", &testCase{