add e2e test and fixes

# Conflicts: # pkg/webhooks/jobset_webhook.go
kubernetes-sigs · Jul 26, 2024 · 0974200 · 0974200
1 parent 001e6d6
commit 0974200
Show file tree

Hide file tree

Showing 5 changed files with 124 additions and 17 deletions.
diff --git a/pkg/controllers/jobset_controller.go b/pkg/controllers/jobset_controller.go
@@ -209,6 +209,10 @@ func (r *JobSetReconciler) reconcile(ctx context.Context, js *jobset.JobSet, upd
 			log.Error(err, "suspending jobset")
 			return ctrl.Result{}, err
 		}
+		if err := r.deleteJobs(ctx, ownedJobs.active); err != nil {
+			log.Error(err, "deleting jobs")
+			return ctrl.Result{}, err
+		}
 	} else {
 		if err := r.resumeJobsIfNecessary(ctx, js, ownedJobs.active, rjobStatuses, updateStatusOpts); err != nil {
 			log.Error(err, "resuming jobset")
@@ -493,6 +497,11 @@ func (r *JobSetReconciler) reconcileReplicatedJobs(ctx context.Context, js *jobs
 			return err
 		}
 
+		// Don't create child Jobs if the JobSet is suspended
+		if jobSetSuspended(js) {
+			continue
+		}
+
 		status := findReplicatedJobStatus(replicatedJobStatus, replicatedJob.Name)
 
 		// For startup policy, if the replicatedJob is started we can skip this loop.

diff --git a/pkg/webhooks/jobset_webhook.go b/pkg/webhooks/jobset_webhook.go
@@ -313,14 +313,27 @@ func (j *jobSetWebhook) ValidateUpdate(ctx context.Context, old, newObj runtime.
 	// This is needed for integration with Kueue/DWS.
 	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
-			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
+			munge := true
+			// Don't allow to unsuspend if there are still active Jobs with
+			// different PodTemplate. We do this to avoid a race condition when
+			// Jobs with an old PodTemplate (before suspension) are still
+			// runnable and would not be deleted if we unsuspend the JobSet.
+			if !ptr.Deref(js.Spec.Suspend, false) {
+				rStatus := js.Status.ReplicatedJobsStatus
+				if len(rStatus) > index && (rStatus[index].Active > 0 || rStatus[index].Suspended > 0) {
+					munge = false
+				}
+			}
+			if munge {
+				// 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
+				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/test/e2e/e2e_test.go b/test/e2e/e2e_test.go
@@ -24,6 +24,7 @@ import (
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/utils/ptr"
 
 	jobset "sigs.k8s.io/jobset/api/jobset/v1alpha2"
 	"sigs.k8s.io/jobset/pkg/util/testing"
@@ -131,6 +132,67 @@ var _ = ginkgo.Describe("JobSet", func() {
 		})
 	})
 
+	ginkgo.When("job is unsuspended and suspend", func() {
+		for rep := 0; rep < 100; rep++ {
+			ginkgo.It(fmt.Sprintf("should allow to update PodTemplate on unsuspend and restore the PodTemplate on suspend-%v", rep), func() {
+				ctx := context.Background()
+				js := shortSleepTestJobSet(ns).Obj()
+				jsKey := types.NamespacedName{Name: js.Name, Namespace: js.Namespace}
+
+				ginkgo.By("Create a suspended JobSet", func() {
+					js.Spec.Suspend = ptr.To(true)
+					js.Spec.TTLSecondsAfterFinished = ptr.To[int32](5)
+					gomega.Expect(k8sClient.Create(ctx, js)).Should(gomega.Succeed())
+				})
+
+				ginkgo.By("Unsuspend the JobSet setting nodeSelectors that prevent pods from being scheduled", func() {
+					gomega.Eventually(func() error {
+						gomega.Expect(k8sClient.Get(ctx, jsKey, js)).Should(gomega.Succeed())
+						js.Spec.Suspend = ptr.To(false)
+						podTemplate := &js.Spec.ReplicatedJobs[0].Template.Spec.Template
+						if podTemplate.Spec.NodeSelector == nil {
+							podTemplate.Spec.NodeSelector = make(map[string]string)
+						}
+						podTemplate.Spec.NodeSelector["kubernetes.io/hostname"] = "non-existing-node"
+						if podTemplate.Labels == nil {
+							podTemplate.Labels = make(map[string]string)
+						}
+						podTemplate.Labels["custom-label-key"] = "custom-label-value"
+						if podTemplate.Annotations == nil {
+							podTemplate.Annotations = make(map[string]string)
+						}
+						podTemplate.Annotations["custom-annotation-key"] = "custom-annotation-value"
+						return k8sClient.Update(ctx, js)
+					}, timeout, interval).Should(gomega.Succeed())
+				})
+
+				ginkgo.By("Suspend the JobSet restoring the PodTemplate properties", func() {
+					gomega.Eventually(func() error {
+						gomega.Expect(k8sClient.Get(ctx, jsKey, js)).Should(gomega.Succeed())
+						js.Spec.Suspend = ptr.To(true)
+						podTemplate := &js.Spec.ReplicatedJobs[0].Template.Spec.Template
+						delete(podTemplate.Spec.NodeSelector, "kubernetes.io/hostname")
+						delete(podTemplate.Labels, "custom-label-key")
+						delete(podTemplate.Annotations, "custom-annotation-key")
+						return k8sClient.Update(ctx, js)
+					}, timeout, interval).Should(gomega.Succeed())
+				})
+
+				ginkgo.By("Unsuspending the JobSet again with PodTemplate allowing completion", func() {
+					gomega.Eventually(func() error {
+						gomega.Expect(k8sClient.Get(ctx, jsKey, js)).Should(gomega.Succeed())
+						js.Spec.Suspend = ptr.To(false)
+						return k8sClient.Update(ctx, js)
+					}, timeout, interval).Should(gomega.Succeed())
+				})
+
+				ginkgo.By("Await for the JobSet to complete successfully", func() {
+					util.JobSetCompleted(ctx, k8sClient, js, timeout)
+				})
+			})
+		}
+	})
+
 }) // end of Describe
 
 // getPingCommand returns ping command for 4 hostnames
@@ -246,3 +308,25 @@ func sleepTestJobSet(ns *corev1.Namespace) *testing.JobSetWrapper {
 			Replicas(int32(replicas)).
 			Obj())
 }
+
+func shortSleepTestJobSet(ns *corev1.Namespace) *testing.JobSetWrapper {
+	jsName := "js"
+	rjobName := "rjob"
+	replicas := 3
+	return testing.MakeJobSet(jsName, ns.Name).
+		ReplicatedJob(testing.MakeReplicatedJob(rjobName).
+			Job(testing.MakeJobTemplate("job", ns.Name).
+				PodSpec(corev1.PodSpec{
+					RestartPolicy: "Never",
+					Containers: []corev1.Container{
+						{
+							Name:    "short-sleep-test-container",
+							Image:   "bash:latest",
+							Command: []string{"bash", "-c"},
+							Args:    []string{"sleep 1"},
+						},
+					},
+				}).Obj()).
+			Replicas(int32(replicas)).
+			Obj())
+}
diff --git a/test/integration/controller/jobset_controller_test.go b/test/integration/controller/jobset_controller_test.go
@@ -180,9 +180,10 @@ var _ = ginkgo.Describe("JobSet controller", func() {
 
 				if up.jobSetUpdateFn != nil {
 					up.jobSetUpdateFn(&jobSet)
+					gomega.Expect(k8sClient.Get(ctx, types.NamespacedName{Name: js.Name, Namespace: js.Namespace}, &jobSet)).To(gomega.Succeed())
 				} else if up.jobUpdateFn != nil {
 					if up.checkJobCreation == nil {
-						gomega.Eventually(testutil.NumJobs, timeout, interval).WithArguments(ctx, k8sClient, js).Should(gomega.Equal(testutil.NumExpectedJobs(js)))
+						gomega.Eventually(testutil.NumJobs, timeout, interval).WithArguments(ctx, k8sClient, js).Should(gomega.Equal(testutil.NumExpectedJobs(&jobSet)))
 					} else {
 						up.checkJobCreation(&jobSet)
 					}
@@ -898,12 +899,10 @@ var _ = ginkgo.Describe("JobSet controller", func() {
 						ginkgo.By("Check ReplicatedJobStatus for suspend")
 						matchJobSetReplicatedStatus(js, []jobset.ReplicatedJobStatus{
 							{
-								Name:      "replicated-job-b",
-								Suspended: 3,
+								Name: "replicated-job-b",
 							},
 							{
-								Name:      "replicated-job-a",
-								Suspended: 1,
+								Name: "replicated-job-a",
 							},
 						})
 					},
@@ -930,12 +929,10 @@ var _ = ginkgo.Describe("JobSet controller", func() {
 						ginkgo.By("Check ReplicatedJobStatus for suspend")
 						matchJobSetReplicatedStatus(js, []jobset.ReplicatedJobStatus{
 							{
-								Name:      "replicated-job-b",
-								Suspended: 3,
+								Name: "replicated-job-b",
 							},
 							{
-								Name:      "replicated-job-a",
-								Suspended: 1,
+								Name: "replicated-job-a",
 							},
 						})
 					},

diff --git a/test/util/util.go b/test/util/util.go
@@ -24,6 +24,7 @@ import (
 	batchv1 "k8s.io/api/batch/v1"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/utils/ptr"
 
 	apierrors "k8s.io/apimachinery/pkg/api/errors"
 
@@ -38,6 +39,9 @@ const interval = time.Millisecond * 250
 
 func NumExpectedJobs(js *jobset.JobSet) int {
 	expectedJobs := 0
+	if ptr.Deref(js.Spec.Suspend, false) {
+		return expectedJobs
+	}
 	for _, rjob := range js.Spec.ReplicatedJobs {
 		expectedJobs += int(rjob.Replicas)
 	}