Repository: kubernetes-sigs/kueue

Issue:
Add unit tests for `pkg/controller/jobframework/tas.go`
<!-- Please only use this template for submitting enhancement requests -->

**What would you like to be added**:
Add unit tests for `pkg/controller/jobframework/tas.go`

**Why is this needed**:

**Completion requirements**:

This enhancement requires the following artifacts:

- [ ] Design doc
- [ ] API change
- [ ] Docs update

The artifacts should be linked in subsequent comments.


Discussion:
@mbobrovskyi: /assign

Solution patch:
diff --git a/pkg/controller/jobframework/tas.go b/pkg/controller/jobframework/tas.go
index ea9634ed9..a7fa0569f 100644
--- a/pkg/controller/jobframework/tas.go
+++ b/pkg/controller/jobframework/tas.go
@@ -73,7 +73,7 @@ func (p *podSetTopologyRequestBuilder) Build() (*kueue.PodSetTopologyRequest, er
 		}
 		psTopologyReq.Unconstrained = &unconstrained
 	default:
-		if !sliceRequiredTopologyFound || !sliceSizeFound {
+		if (!sliceRequiredTopologyFound || !sliceSizeFound) && (p.podIndexLabel == nil && p.subGroupIndexLabel == nil && p.subGroupCount == nil) {
 			return nil, nil
 		}
 	}
diff --git a/pkg/scheduler/flavorassigner/tas_flavorassigner.go b/pkg/scheduler/flavorassigner/tas_flavorassigner.go
index 06fc657f6..1a3ffab75 100644
--- a/pkg/scheduler/flavorassigner/tas_flavorassigner.go
+++ b/pkg/scheduler/flavorassigner/tas_flavorassigner.go
@@ -136,8 +136,14 @@ func onlyFlavor(ra ResourceAssignment) (*kueue.ResourceFlavorReference, error) {
 }
 
 func checkPodSetAndFlavorMatchForTAS(cq *cache.ClusterQueueSnapshot, ps *kueue.PodSet, flavor *kueue.ResourceFlavor) *string {
-	// For PodSets which require TAS skip resource flavors which don't support it
-	if ps.TopologyRequest != nil {
+	if isTASRequested(ps, cq) {
+		if isTASImplied(ps, cq) {
+			// If this is a TAS-only CQ, then we don't need to check the flavor because
+			// all flavors in the ClusterQueue are TAS flavors, and all Workloads submitted
+			// to this ClusterQueue are expected to use TAS, and it's a match.
+			return nil
+		}
+		// PodSet explicitly requires TAS, so we need to check if the flavor supports it.
 		if flavor.Spec.TopologyName == nil {
 			return ptr.To(fmt.Sprintf("Flavor %q does not support TopologyAwareScheduling", flavor.Name))
 		}
@@ -152,26 +158,24 @@ func checkPodSetAndFlavorMatchForTAS(cq *cache.ClusterQueueSnapshot, ps *kueue.P
 			// Skip flavors which don't have the requested level
 			return ptr.To(fmt.Sprintf("Flavor %q does not contain the requested level", flavor.Name))
 		}
-	}
-	// If this is a TAS-only CQ, then no TopologyRequest is ok
-	if isTASImplied(ps, cq) {
+		// PodSet requires TAS and the flavor supports it, so it's a match.
 		return nil
 	}
-	// For PodSets which don't use TAS skip resource flavors which are only for TAS
-	if ps.TopologyRequest == nil && flavor.Spec.TopologyName != nil {
+	// PodSet doesn't require TAS, but the flavor supports it.
+	if flavor.Spec.TopologyName != nil {
 		return ptr.To(fmt.Sprintf("Flavor %q supports only TopologyAwareScheduling", flavor.Name))
 	}
+	// PodSet doesn't require TAS and the flavor doesn't support it, so it's a match.
 	return nil
 }
 
-// isTASImplied returns true if TAS is requested implicitly - there is no
-// explicit
+// isTASImplied returns true if TAS is requested implicitly
 func isTASImplied(ps *kueue.PodSet, cq *cache.ClusterQueueSnapshot) bool {
-	return ps.TopologyRequest == nil && cq.IsTASOnly()
+	return !workload.IsExplicitlyRequestingTAS(*ps) && cq.IsTASOnly()
 }
 
 // isTASRequested checks if TAS is requested for the input PodSet, either
 // explicitly or implicitly.
 func isTASRequested(ps *kueue.PodSet, cq *cache.ClusterQueueSnapshot) bool {
-	return ps.TopologyRequest != nil || cq.IsTASOnly()
+	return workload.IsExplicitlyRequestingTAS(*ps) || isTASImplied(ps, cq)
 }
diff --git a/pkg/scheduler/scheduler.go b/pkg/scheduler/scheduler.go
index 4aab6fa5c..c2a86a1fd 100644
--- a/pkg/scheduler/scheduler.go
+++ b/pkg/scheduler/scheduler.go
@@ -576,7 +576,8 @@ func (s *Scheduler) evictWorkloadAfterFailedTASReplacement(ctx context.Context,
 }
 
 func updateAssignmentForTAS(cq *cache.ClusterQueueSnapshot, wl *workload.Info, assignment *flavorassigner.Assignment, targets []*preemption.Target) {
-	if features.Enabled(features.TopologyAwareScheduling) && assignment.RepresentativeMode() == flavorassigner.Preempt && (wl.IsRequestingTAS() || cq.IsTASOnly()) && !workload.HasTopologyAssignmentWithNodeToReplace(wl.Obj) {
+	if features.Enabled(features.TopologyAwareScheduling) && assignment.RepresentativeMode() == flavorassigner.Preempt &&
+		(workload.IsExplicitlyRequestingTAS(wl.Obj.Spec.PodSets...) || cq.IsTASOnly()) && !workload.HasTopologyAssignmentWithNodeToReplace(wl.Obj) {
 		tasRequests := assignment.WorkloadsTopologyRequests(wl, cq)
 		var tasResult cache.TASAssignmentsResult
 		if len(targets) > 0 {
diff --git a/pkg/workload/workload.go b/pkg/workload/workload.go
index 646747c29..893c90020 100644
--- a/pkg/workload/workload.go
+++ b/pkg/workload/workload.go
@@ -348,11 +348,12 @@ func (i *Info) IsUsingTAS() bool {
 		})
 }
 
-// IsRequestingTAS returns information if the workload is requesting TAS
-func (i *Info) IsRequestingTAS() bool {
-	return slices.ContainsFunc(i.Obj.Spec.PodSets,
+// IsExplicitlyRequestingTAS returns information if the workload is requesting TAS
+func IsExplicitlyRequestingTAS(podSets ...kueue.PodSet) bool {
+	return slices.ContainsFunc(podSets,
 		func(ps kueue.PodSet) bool {
-			return ps.TopologyRequest != nil
+			tr := ps.TopologyRequest
+			return tr != nil && (tr.Unconstrained != nil || tr.Required != nil || tr.Preferred != nil || tr.PodSetSliceRequiredTopology != nil || tr.PodSetSliceSize != nil)
 		})
 }
 



Test patch:
diff --git a/pkg/controller/jobs/job/job_controller_test.go b/pkg/controller/jobs/job/job_controller_test.go
index ff2e105fa..d8fb4c253 100644
--- a/pkg/controller/jobs/job/job_controller_test.go
+++ b/pkg/controller/jobs/job/job_controller_test.go
@@ -447,6 +447,7 @@ func TestPodSets(t *testing.T) {
 					Annotations(map[string]string{
 						kueuealpha.PodSetSliceRequiredTopologyAnnotation: "cloud.com/block",
 					}).
+					PodIndexLabel(ptr.To(batchv1.JobCompletionIndexAnnotation)).
 					Obj(),
 			},
 			enableTopologyAwareScheduling: true,
@@ -464,6 +465,7 @@ func TestPodSets(t *testing.T) {
 					Annotations(map[string]string{
 						kueuealpha.PodSetSliceSizeAnnotation: "1",
 					}).
+					PodIndexLabel(ptr.To(batchv1.JobCompletionIndexAnnotation)).
 					Obj(),
 			},
 			enableTopologyAwareScheduling: true,
diff --git a/pkg/controller/jobs/jobset/jobset_controller_test.go b/pkg/controller/jobs/jobset/jobset_controller_test.go
index 9f8835846..98fd94df3 100644
--- a/pkg/controller/jobs/jobset/jobset_controller_test.go
+++ b/pkg/controller/jobs/jobset/jobset_controller_test.go
@@ -255,6 +255,9 @@ func TestPodSets(t *testing.T) {
 						Obj(),
 					*utiltesting.MakePodSet(kueue.NewPodSetReference(jobSet.Spec.ReplicatedJobs[1].Name), 6).
 						PodSpec(*jobSet.Spec.ReplicatedJobs[1].Template.Spec.Template.Spec.DeepCopy()).
+						PodIndexLabel(ptr.To(batchv1.JobCompletionIndexAnnotation)).
+						SubGroupIndexLabel(ptr.To(jobset.JobIndexKey)).
+						SubGroupCount(ptr.To[int32](3)).
 						Obj(),
 				}
 			},
@@ -279,6 +282,9 @@ func TestPodSets(t *testing.T) {
 				return []kueue.PodSet{
 					*utiltesting.MakePodSet(kueue.NewPodSetReference(jobSet.Spec.ReplicatedJobs[0].Name), 2).
 						PodSpec(*jobSet.Spec.ReplicatedJobs[0].Template.Spec.Template.Spec.DeepCopy()).
+						PodIndexLabel(ptr.To(batchv1.JobCompletionIndexAnnotation)).
+						SubGroupIndexLabel(ptr.To(jobset.JobIndexKey)).
+						SubGroupCount(ptr.To[int32](2)).
 						Obj(),
 					*utiltesting.MakePodSet(kueue.NewPodSetReference(jobSet.Spec.ReplicatedJobs[1].Name), 6).
 						PodSpec(*jobSet.Spec.ReplicatedJobs[1].Template.Spec.Template.Spec.DeepCopy()).
diff --git a/pkg/controller/jobs/kubeflow/jobs/tfjob/tfjob_controller_test.go b/pkg/controller/jobs/kubeflow/jobs/tfjob/tfjob_controller_test.go
index cd2abf488..7d92ba6ec 100644
--- a/pkg/controller/jobs/kubeflow/jobs/tfjob/tfjob_controller_test.go
+++ b/pkg/controller/jobs/kubeflow/jobs/tfjob/tfjob_controller_test.go
@@ -304,6 +304,7 @@ func TestPodSets(t *testing.T) {
 						Obj(),
 					*utiltesting.MakePodSet(kueue.NewPodSetReference(string(kftraining.TFJobReplicaTypeWorker)), 1).
 						PodSpec(job.Spec.TFReplicaSpecs[kftraining.TFJobReplicaTypeWorker].Template.Spec).
+						PodIndexLabel(ptr.To(kftraining.ReplicaIndexLabel)).
 						Obj(),
 				}
 			},
diff --git a/test/e2e/tas/job_test.go b/test/e2e/tas/job_test.go
index e9a7c9edb..6df7d5bee 100644
--- a/test/e2e/tas/job_test.go
+++ b/test/e2e/tas/job_test.go
@@ -303,5 +303,61 @@ var _ = ginkgo.Describe("TopologyAwareScheduling for Job", func() {
 				gomega.Expect(wantAssignment).Should(gomega.BeComparableTo(gotAssignment))
 			})
 		})
+
+		ginkgo.It("Should place pods based on the ranks-ordering even if the Job has no TAS annotation (Implicit TAS)", func() {
+			numPods := 4
+			sampleJob := testingjob.MakeJob("ranks-job", ns.Name).
+				Queue(kueue.LocalQueueName(localQueue.Name)).
+				Parallelism(int32(numPods)).
+				Completions(int32(numPods)).
+				Indexed(true).
+				RequestAndLimit(extraResource, "1").
+				Image(util.GetAgnHostImage(), util.BehaviorWaitForDeletion).
+				Obj()
+			util.MustCreate(ctx, k8sClient, sampleJob)
+
+			ginkgo.By("Job is unsuspended, and has all Pods active and ready", func() {
+				gomega.Eventually(func(g gomega.Gomega) {
+					g.Expect(k8sClient.Get(ctx, client.ObjectKeyFromObject(sampleJob), sampleJob)).To(gomega.Succeed())
+					g.Expect(sampleJob.Spec.Suspend).Should(gomega.Equal(ptr.To(false)))
+				}, util.LongTimeout, util.Interval).Should(gomega.Succeed())
+				gomega.Eventually(func(g gomega.Gomega) {
+					g.Expect(k8sClient.Get(ctx, client.ObjectKeyFromObject(sampleJob), sampleJob)).To(gomega.Succeed())
+					g.Expect(sampleJob.Status.Active).Should(gomega.Equal(int32(numPods)))
+				}, util.LongTimeout, util.Interval).Should(gomega.Succeed())
+				gomega.Eventually(func(g gomega.Gomega) {
+					g.Expect(k8sClient.Get(ctx, client.ObjectKeyFromObject(sampleJob), sampleJob)).To(gomega.Succeed())
+					g.Expect(sampleJob.Status.Ready).Should(gomega.Equal(ptr.To[int32](int32(numPods))))
+				}, util.LongTimeout, util.Interval).Should(gomega.Succeed())
+			})
+
+			pods := &corev1.PodList{}
+			ginkgo.By("ensure all pods are created and scheduled", func() {
+				listOpts := &client.ListOptions{
+					FieldSelector: fields.OneTermNotEqualSelector("spec.nodeName", ""),
+				}
+				gomega.Eventually(func(g gomega.Gomega) {
+					g.Expect(k8sClient.List(ctx, pods, client.InNamespace(ns.Name), listOpts)).To(gomega.Succeed())
+					g.Expect(pods.Items).Should(gomega.HaveLen(numPods))
+				}, util.LongTimeout, util.Interval).Should(gomega.Succeed())
+			})
+
+			ginkgo.By("verify the assignment of pods are as expected with rank-based ordering", func() {
+				gomega.Expect(k8sClient.List(ctx, pods, client.InNamespace(ns.Name),
+					client.MatchingLabels(sampleJob.Spec.Selector.MatchLabels))).To(gomega.Succeed())
+				gotAssignment := make(map[string]string, numPods)
+				for _, pod := range pods.Items {
+					index := pod.Labels[batchv1.JobCompletionIndexAnnotation]
+					gotAssignment[index] = pod.Spec.NodeName
+				}
+				wantAssignment := map[string]string{
+					"0": "kind-worker",
+					"1": "kind-worker2",
+					"2": "kind-worker3",
+					"3": "kind-worker4",
+				}
+				gomega.Expect(wantAssignment).Should(gomega.BeComparableTo(gotAssignment))
+			})
+		})
 	})
 })