TAS: fix scenario when parallelism < completions (#3559)

kubernetes-sigs · Nov 15, 2024 · 30da293 · 30da293
1 parent e0be159
commit 30da293
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 6 deletions.
diff --git a/pkg/controller/tas/topology_ungater.go b/pkg/controller/tas/topology_ungater.go
@@ -262,9 +262,9 @@ func (r *topologyUngater) podsForPodSet(ctx context.Context, ns, wlName, psName
 	}
 	result := make([]*corev1.Pod, 0, len(pods.Items))
 	for i := range pods.Items {
-		if pods.Items[i].Status.Phase == corev1.PodFailed {
-			// ignore failed pods as they need to be replaced, and so we don't
-			// want to count them as already ungated Pods.
+		if phase := pods.Items[i].Status.Phase; phase == corev1.PodFailed || phase == corev1.PodSucceeded {
+			// ignore failed or succeeded pods as they need to be replaced, and
+			// so we don't want to count them as already ungated Pods.
 			continue
 		}
 		result = append(result, &pods.Items[i])

diff --git a/pkg/controller/tas/topology_ungater_test.go b/pkg/controller/tas/topology_ungater_test.go
@@ -87,7 +87,7 @@ func TestReconcile(t *testing.T) {
 			if utilpod.HasGate(&pod, kueuealpha.TopologySchedulingGate) {
 				continue
 			}
-			if pod.Status.Phase == corev1.PodFailed {
+			if pod.Status.Phase == corev1.PodFailed || pod.Status.Phase == corev1.PodSucceeded {
 				continue
 			}
 			key := mapToJSON(t, pod.Spec.NodeSelector)
@@ -744,7 +744,7 @@ func TestReconcile(t *testing.T) {
 				},
 			},
 		},
-		"expect single pod; one ungated pod succeeded - don't ungate second pod": {
+		"expect single pod; one ungated pod succeeded - ungate second pod": {
 			workloads: []kueue.Workload{
 				*utiltesting.MakeWorkload("unit-test", "ns").Finalizers(kueue.ResourceInUseFinalizerName).
 					PodSets(*utiltesting.MakePodSet(kueue.DefaultPodSetName, 1).Request(corev1.ResourceCPU, "1").Obj()).
@@ -794,7 +794,6 @@ func TestReconcile(t *testing.T) {
 				*testingpod.MakePod("pod-gated", "ns").
 					Annotation(kueuealpha.WorkloadAnnotation, "unit-test").
 					Label(kueuealpha.PodSetLabel, kueue.DefaultPodSetName).
-					TopologySchedulingGate().
 					Obj(),
 			},
 			wantCounts: []counts{

diff --git a/test/e2e/tas/tas_test.go b/test/e2e/tas/tas_test.go
@@ -233,6 +233,32 @@ var _ = ginkgo.Describe("TopologyAwareScheduling", func() {
 				}, util.LongTimeout, util.Interval).Should(gomega.Succeed())
 			})
 		})
+
+		ginkgo.It("Should allow to run a Job with parallelism < completions", func() {
+			sampleJob := testingjob.MakeJob("test-job", ns.Name).
+				Queue(localQueue.Name).
+				Parallelism(2).
+				Completions(3).
+				Request(extraResource, "1").
+				Limit(extraResource, "1").
+				Obj()
+			sampleJob = (&testingjob.JobWrapper{Job: *sampleJob}).
+				PodAnnotation(kueuealpha.PodSetRequiredTopologyAnnotation, topologyLevelBlock).
+				Image(util.E2eTestSleepImage, []string{"10ms"}).
+				Obj()
+			gomega.Expect(k8sClient.Create(ctx, sampleJob)).Should(gomega.Succeed())
+
+			wlLookupKey := types.NamespacedName{Name: workloadjob.GetWorkloadNameForJob(sampleJob.Name, sampleJob.UID), Namespace: ns.Name}
+			createdWorkload := &kueue.Workload{}
+
+			ginkgo.By(fmt.Sprintf("verify the workload %q gets TopologyAssignment becomes finished", wlLookupKey), func() {
+				gomega.Eventually(func(g gomega.Gomega) {
+					g.Expect(k8sClient.Get(ctx, wlLookupKey, createdWorkload)).Should(gomega.Succeed())
+					g.Expect(createdWorkload.Status.Admission.PodSetAssignments[0].TopologyAssignment).ShouldNot(gomega.BeNil())
+					g.Expect(createdWorkload.Status.Conditions).Should(testing.HaveConditionStatusTrue(kueue.WorkloadFinished))
+				}, util.LongTimeout, util.Interval).Should(gomega.Succeed())
+			})
+		})
 	})
 })