Fix restoring PodTemplate on JobSet suspend

kubernetes-sigs · Jul 25, 2024 · 24448c5 · 24448c5
1 parent 56c77da
commit 24448c5
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 2 deletions.
diff --git a/pkg/webhooks/jobset_webhook.go b/pkg/webhooks/jobset_webhook.go
@@ -309,9 +309,9 @@ func (j *jobSetWebhook) ValidateUpdate(ctx context.Context, old, newObj runtime.
 	}
 	mungedSpec := js.Spec.DeepCopy()
 
-	// Allow pod template to be mutated for suspended JobSets.
+	// Allow pod template to be mutated for suspended JobSets, or JobSets getting suspended.
 	// This is needed for integration with Kueue/DWS.
-	if ptr.Deref(oldJS.Spec.Suspend, false) {
+	if ptr.Deref(oldJS.Spec.Suspend, false) || ptr.Deref(js.Spec.Suspend, false) {
 		for index := range js.Spec.ReplicatedJobs {
 			// Pod values which must be mutable for Kueue are defined here: https://github.com/kubernetes-sigs/kueue/blob/a50d395c36a2cb3965be5232162cf1fded1bdb08/apis/kueue/v1beta1/workload_types.go#L256-L260
 			mungedSpec.ReplicatedJobs[index].Template.Spec.Template.Annotations = oldJS.Spec.ReplicatedJobs[index].Template.Spec.Template.Annotations

diff --git a/pkg/webhooks/jobset_webhook_test.go b/pkg/webhooks/jobset_webhook_test.go
@@ -1517,6 +1517,49 @@ func TestValidateUpdate(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "replicated job pod template can be updated for jobset getting suspended",
+			js: &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{
+								// Restoring the template by removing the annotation
+								Spec: batchv1.JobSpec{
+									Parallelism: ptr.To[int32](2),
+									Template:    corev1.PodTemplateSpec{},
+								},
+							},
+						},
+					},
+				},
+			},
+			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),
+									Template: corev1.PodTemplateSpec{
+										ObjectMeta: metav1.ObjectMeta{
+											Annotations: map[string]string{"key": "value"},
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
 		{
 			name: "replicated job pod template cannot be updated for running jobset",
 			js: &jobset.JobSet{