Repository: kubernetes-sigs/kueue

Issue:
TAS: Two-level JobSet scheduling
<!-- Please only use this template for submitting enhancement requests -->

**What would you like to be added**:
A possibility to define topology for each Job in ReplicatedJob to co-locate the whole ReplicatedJob in one domain and each Job in that ReplicatedJob in lower-level domain within that domain.

In more generic terms - an API to define required sub PodSet topology and allow for two-level scheduling for the whole PodSet and lower-level topology domain for each chunk within that PodSet.

An example is a JobSet in which user can define topology for all Jobs resulting from a single ReplicatedJob. We want to allow user to define topology for each such Job and possibly co-locate all such Jobs in higher-level topology.

**Why is this needed**:
Currently Topology-Aware Scheduling considers the whole PodSet as a single unit for which Kueue looks for a matching domain. For JobSet PodSet covers all pods resulting from the whole ReplicatedJob. However, there are use-cases in which user would like to co-locate each Job in a lower-level domains and co-locate all those Jobs in higher-level domain.

Additionally, user would like to define a topology only for each Job without any expectations for a topology of the whole ReplicatedJob.

**Completion requirements**:

This enhancement requires the following artifacts:

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

The artifacts should be linked in subsequent comments.


Hints from issue comments:
@mimowo: /retitle TAS: Two-level JobSet scheduling@lchrzaszcz: TAS KEP update to reflect two-level scheduling and podset chunk topology: https://github.com/kubernetes-sigs/kueue/pull/5449@kerthcet: One quick question: can this solve the https://github.com/kubernetes-sigs/kueue/pull/3941? E.g. co-locate the master and worker jobs in pytorchjob.
Thanks!@kerthcet: I'll read the KEP in detail later.@kerthcet: This focused on podSet chunks, not between multi-podsets. @lchrzaszcz: I will be looking into cross-podsets scheduling in few following weeks. I'll do it mostly for LeaderWorkerSet though, but I think the underlying mechanism will be universal, so probably it'll be easy to add pytorchjob support afterwards.@kerthcet: Thanks @lchrzaszcz for your response, I think leaderworkerset pytorchjob are just two different workload types, the mechanism should be universal to them all I believe. Glad to jump into the discussion.@tenzen-y: /assign @lchrzaszcz 

Solution patch:
diff --git a/pkg/controller/jobframework/tas_validation.go b/pkg/controller/jobframework/tas_validation.go
index 20ca35928..f52c2cc0c 100644
--- a/pkg/controller/jobframework/tas_validation.go
+++ b/pkg/controller/jobframework/tas_validation.go
@@ -18,6 +18,7 @@ package jobframework
 
 import (
 	"fmt"
+	"strconv"
 
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	metavalidation "k8s.io/apimachinery/pkg/apis/meta/v1/validation"
@@ -31,6 +32,8 @@ func ValidateTASPodSetRequest(replicaPath *field.Path, replicaMetadata *metav1.O
 	requiredValue, requiredFound := replicaMetadata.Annotations[kueuealpha.PodSetRequiredTopologyAnnotation]
 	preferredValue, preferredFound := replicaMetadata.Annotations[kueuealpha.PodSetPreferredTopologyAnnotation]
 	_, unconstrainedFound := replicaMetadata.Annotations[kueuealpha.PodSetUnconstrainedTopologyAnnotation]
+	sliceRequiredValue, sliceRequiredFound := replicaMetadata.Annotations[kueuealpha.PodSetSliceRequiredTopologyAnnotation]
+	sliceSizeValue, sliceSizeFound := replicaMetadata.Annotations[kueuealpha.PodSetSliceSizeAnnotation]
 
 	// validate no more than 1 annotation
 	asInt := func(b bool) int {
@@ -57,5 +60,33 @@ func ValidateTASPodSetRequest(replicaPath *field.Path, replicaMetadata *metav1.O
 	if preferredFound {
 		allErrs = append(allErrs, metavalidation.ValidateLabelName(preferredValue, annotationsPath.Key(kueuealpha.PodSetPreferredTopologyAnnotation))...)
 	}
+	if sliceRequiredFound {
+		allErrs = append(allErrs, metavalidation.ValidateLabelName(sliceRequiredValue, annotationsPath.Key(kueuealpha.PodSetSliceRequiredTopologyAnnotation))...)
+	}
+
+	// validate slice size annotation
+	if sliceSizeFound {
+		val, err := strconv.Atoi(sliceSizeValue)
+		if err != nil {
+			allErrs = append(allErrs, field.Invalid(annotationsPath.Key(kueuealpha.PodSetSliceSizeAnnotation), sliceSizeValue, "must be a numeric value"))
+		} else if val < 1 {
+			allErrs = append(allErrs, field.Invalid(annotationsPath.Key(kueuealpha.PodSetSliceSizeAnnotation), sliceSizeValue, "must be greater than or equal to 1"))
+		}
+	}
+
+	// validate slice annotations
+	if sliceRequiredFound {
+		if !requiredFound && !preferredFound {
+			allErrs = append(allErrs, field.Forbidden(annotationsPath.Key(kueuealpha.PodSetSliceRequiredTopologyAnnotation), "can be used with required or preferred topology only"))
+		}
+		if !sliceSizeFound {
+			allErrs = append(allErrs, field.Required(annotationsPath.Key(kueuealpha.PodSetSliceSizeAnnotation), "slice size is required if slice topology is requested"))
+		}
+	}
+
+	if !sliceRequiredFound && sliceSizeFound {
+		allErrs = append(allErrs, field.Forbidden(annotationsPath.Key(kueuealpha.PodSetSliceSizeAnnotation), fmt.Sprintf("cannot be set when '%s' is not present", kueuealpha.PodSetSliceRequiredTopologyAnnotation)))
+	}
+
 	return allErrs
 }
diff --git a/pkg/scheduler/preemption/classical/candidate_generator.go b/pkg/scheduler/preemption/classical/candidate_generator.go
index c782a5520..11ce99e4e 100644
--- a/pkg/scheduler/preemption/classical/candidate_generator.go
+++ b/pkg/scheduler/preemption/classical/candidate_generator.go
@@ -49,19 +49,6 @@ type candidateElem struct {
 	preemptionVariant preemptionVariant
 }
 
-// Need a separate function for candidateElem data type
-// in the future we will modify the ordering to take into
-// account the distance in the cohort hierarchy tree
-func candidateElemsOrdering(candidates []*candidateElem, cq kueue.ClusterQueueReference, now time.Time, ordering func([]*workload.Info, kueue.ClusterQueueReference, time.Time) func(int, int) bool) func(int, int) bool {
-	// Adapt the candidatesOrdering function to work with candidateElem
-	adaptedOrdering := func(i, j int) bool {
-		a := candidates[i].wl
-		b := candidates[j].wl
-		return ordering([]*workload.Info{a, b}, cq, now)(0, 1)
-	}
-	return adaptedOrdering
-}
-
 func WorkloadUsesResources(wl *workload.Info, frsNeedPreemption sets.Set[resources.FlavorResource]) bool {
 	for _, ps := range wl.TotalRequests {
 		for res, flv := range ps.Flavors {
@@ -86,12 +73,19 @@ func splitEvicted(workloads []*candidateElem) ([]*candidateElem, []*candidateEle
 // with and without borrowing. The runs are independent which means that the same candidates
 // might be returned for both, but note that the candidates with borrrowing are a subset of
 // candidates without borrowing.
-func NewCandidateIterator(hierarchicalReclaimCtx *HierarchicalPreemptionCtx, frsNeedPreemption sets.Set[resources.FlavorResource], snapshot *cache.Snapshot, clock clock.Clock, ordering func([]*workload.Info, kueue.ClusterQueueReference, time.Time) func(int, int) bool) *candidateIterator {
+func NewCandidateIterator(hierarchicalReclaimCtx *HierarchicalPreemptionCtx, frsNeedPreemption sets.Set[resources.FlavorResource], snapshot *cache.Snapshot, clock clock.Clock, ordering func(*workload.Info, *workload.Info, kueue.ClusterQueueReference, time.Time) bool) *candidateIterator {
 	sameQueueCandidates := collectSameQueueCandidates(hierarchicalReclaimCtx)
 	hierarchyCandidates, priorityCandidates := collectCandidatesForHierarchicalReclaim(hierarchicalReclaimCtx)
-	sort.Slice(sameQueueCandidates, candidateElemsOrdering(sameQueueCandidates, hierarchicalReclaimCtx.Cq.Name, clock.Now(), ordering))
-	sort.Slice(priorityCandidates, candidateElemsOrdering(priorityCandidates, hierarchicalReclaimCtx.Cq.Name, clock.Now(), ordering))
-	sort.Slice(hierarchyCandidates, candidateElemsOrdering(hierarchyCandidates, hierarchicalReclaimCtx.Cq.Name, clock.Now(), ordering))
+	sort.Slice(sameQueueCandidates, func(i, j int) bool {
+		return ordering(sameQueueCandidates[i].wl, sameQueueCandidates[j].wl, hierarchicalReclaimCtx.Cq.Name, clock.Now())
+	})
+	sort.Slice(priorityCandidates, func(i, j int) bool {
+		return ordering(priorityCandidates[i].wl, priorityCandidates[j].wl, hierarchicalReclaimCtx.Cq.Name, clock.Now())
+	})
+	sort.Slice(hierarchyCandidates, func(i, j int) bool {
+		return ordering(hierarchyCandidates[i].wl, hierarchyCandidates[j].wl, hierarchicalReclaimCtx.Cq.Name, clock.Now())
+	})
+
 	evictedHierarchicalReclaimCandidates, nonEvictedHierarchicalReclaimCandidates := splitEvicted(hierarchyCandidates)
 	evictedSTCandidates, nonEvictedSTCandidates := splitEvicted(priorityCandidates)
 	evictedSameQueueCandidates, nonEvictedSameQueueCandidates := splitEvicted(sameQueueCandidates)
diff --git a/pkg/scheduler/preemption/preemption.go b/pkg/scheduler/preemption/preemption.go
index d8f92272f..bcaa0c87a 100644
--- a/pkg/scheduler/preemption/preemption.go
+++ b/pkg/scheduler/preemption/preemption.go
@@ -379,7 +379,9 @@ func (p *Preemptor) fairPreemptions(preemptionCtx *preemptionCtx, strategies []f
 	if len(candidates) == 0 {
 		return nil
 	}
-	sort.Slice(candidates, CandidatesOrdering(candidates, preemptionCtx.preemptorCQ.Name, p.clock.Now()))
+	sort.Slice(candidates, func(i, j int) bool {
+		return CandidatesOrdering(candidates[i], candidates[j], preemptionCtx.preemptorCQ.Name, p.clock.Now())
+	})
 	if logV := preemptionCtx.log.V(5); logV.Enabled() {
 		logV.Info("Simulating fair preemption", "candidates", workload.References(candidates), "resourcesRequiringPreemption", preemptionCtx.frsNeedPreemption.UnsortedList(), "preemptingWorkload", klog.KObj(preemptionCtx.preemptor.Obj))
 	}
@@ -517,33 +519,30 @@ func queueUnderNominalInResourcesNeedingPreemption(preemptionCtx *preemptionCtx)
 // same ClusterQueue as the preemptor.
 // 2. Workloads with lower priority first.
 // 3. Workloads admitted more recently first.
-func CandidatesOrdering(candidates []*workload.Info, cq kueue.ClusterQueueReference, now time.Time) func(int, int) bool {
-	return func(i, j int) bool {
-		a := candidates[i]
-		b := candidates[j]
-		aEvicted := meta.IsStatusConditionTrue(a.Obj.Status.Conditions, kueue.WorkloadEvicted)
-		bEvicted := meta.IsStatusConditionTrue(b.Obj.Status.Conditions, kueue.WorkloadEvicted)
-		if aEvicted != bEvicted {
-			return aEvicted
-		}
-		aInCQ := a.ClusterQueue == cq
-		bInCQ := b.ClusterQueue == cq
-		if aInCQ != bInCQ {
-			return !aInCQ
-		}
-		pa := priority.Priority(a.Obj)
-		pb := priority.Priority(b.Obj)
-		if pa != pb {
-			return pa < pb
-		}
-		timeA := quotaReservationTime(a.Obj, now)
-		timeB := quotaReservationTime(b.Obj, now)
-		if !timeA.Equal(timeB) {
-			return timeA.After(timeB)
-		}
-		// Arbitrary comparison for deterministic sorting.
-		return a.Obj.UID < b.Obj.UID
+func CandidatesOrdering(a, b *workload.Info, cq kueue.ClusterQueueReference, now time.Time) bool {
+	aEvicted := meta.IsStatusConditionTrue(a.Obj.Status.Conditions, kueue.WorkloadEvicted)
+	bEvicted := meta.IsStatusConditionTrue(b.Obj.Status.Conditions, kueue.WorkloadEvicted)
+	if aEvicted != bEvicted {
+		return aEvicted
+	}
+	aInCQ := a.ClusterQueue == cq
+	bInCQ := b.ClusterQueue == cq
+	if aInCQ != bInCQ {
+		return !aInCQ
+	}
+
+	pa := priority.Priority(a.Obj)
+	pb := priority.Priority(b.Obj)
+	if pa != pb {
+		return pa < pb
+	}
+	timeA := quotaReservationTime(a.Obj, now)
+	timeB := quotaReservationTime(b.Obj, now)
+	if !timeA.Equal(timeB) {
+		return timeA.After(timeB)
 	}
+	// Arbitrary comparison for deterministic sorting.
+	return a.Obj.UID < b.Obj.UID
 }
 
 func quotaReservationTime(wl *kueue.Workload, now time.Time) time.Time {



Test patch:
diff --git a/pkg/controller/jobs/job/job_webhook_test.go b/pkg/controller/jobs/job/job_webhook_test.go
index 4a7e23848..0ee068168 100644
--- a/pkg/controller/jobs/job/job_webhook_test.go
+++ b/pkg/controller/jobs/job/job_webhook_test.go
@@ -277,6 +277,89 @@ func TestValidateCreate(t *testing.T) {
 					invalidLabelKeyMessage).WithOrigin("labelKey"),
 			},
 		},
+		{
+			name: "valid slice topology request",
+			job: testingutil.MakeJob("job", "default").
+				PodAnnotation(kueuealpha.PodSetRequiredTopologyAnnotation, "cloud.com/block").
+				PodAnnotation(kueuealpha.PodSetSliceRequiredTopologyAnnotation, "cloud.com/block").
+				PodAnnotation(kueuealpha.PodSetSliceSizeAnnotation, "1").
+				Obj(),
+			wantErr: nil,
+		},
+		{
+			name: "invalid topology request - unconstrained with slices defined",
+			job: testingutil.MakeJob("job", "default").
+				PodAnnotation(kueuealpha.PodSetUnconstrainedTopologyAnnotation, "true").
+				PodAnnotation(kueuealpha.PodSetSliceRequiredTopologyAnnotation, "cloud.com/block").
+				PodAnnotation(kueuealpha.PodSetSliceSizeAnnotation, "1").
+				Obj(),
+			wantErr: field.ErrorList{
+				field.Forbidden(replicaMetaPath.Child("annotations").Key("kueue.x-k8s.io/podset-slice-required-topology"), "can be used with required or preferred topology only"),
+			},
+		},
+		{
+			name: "invalid topology request - slice requested without slice size",
+			job: testingutil.MakeJob("job", "default").
+				PodAnnotation(kueuealpha.PodSetRequiredTopologyAnnotation, "cloud.com/block").
+				PodAnnotation(kueuealpha.PodSetSliceRequiredTopologyAnnotation, "cloud.com/block").
+				Obj(),
+			wantErr: field.ErrorList{
+				field.Required(replicaMetaPath.Child("annotations").Key("kueue.x-k8s.io/podset-slice-size"), "slice size is required if slice topology is requested"),
+			},
+		},
+		{
+			name: "invalid topology request - slice size is not a number",
+			job: testingutil.MakeJob("job", "default").
+				PodAnnotation(kueuealpha.PodSetRequiredTopologyAnnotation, "cloud.com/block").
+				PodAnnotation(kueuealpha.PodSetSliceRequiredTopologyAnnotation, "cloud.com/block").
+				PodAnnotation(kueuealpha.PodSetSliceSizeAnnotation, "not a number").
+				Obj(),
+			wantErr: field.ErrorList{
+				field.Invalid(replicaMetaPath.Child("annotations").Key("kueue.x-k8s.io/podset-slice-size"), "not a number", "must be a numeric value"),
+			},
+		},
+		{
+			name: "invalid topology request - slice size is negative",
+			job: testingutil.MakeJob("job", "default").
+				PodAnnotation(kueuealpha.PodSetRequiredTopologyAnnotation, "cloud.com/block").
+				PodAnnotation(kueuealpha.PodSetSliceRequiredTopologyAnnotation, "cloud.com/block").
+				PodAnnotation(kueuealpha.PodSetSliceSizeAnnotation, "-1").
+				Obj(),
+			wantErr: field.ErrorList{
+				field.Invalid(replicaMetaPath.Child("annotations").Key("kueue.x-k8s.io/podset-slice-size"), "-1", "must be greater than or equal to 1"),
+			},
+		},
+		{
+			name: "invalid topology request - slice size is zero",
+			job: testingutil.MakeJob("job", "default").
+				PodAnnotation(kueuealpha.PodSetRequiredTopologyAnnotation, "cloud.com/block").
+				PodAnnotation(kueuealpha.PodSetSliceRequiredTopologyAnnotation, "cloud.com/block").
+				PodAnnotation(kueuealpha.PodSetSliceSizeAnnotation, "0").
+				Obj(),
+			wantErr: field.ErrorList{
+				field.Invalid(replicaMetaPath.Child("annotations").Key("kueue.x-k8s.io/podset-slice-size"), "0", "must be greater than or equal to 1"),
+			},
+		},
+		{
+			name: "invalid topology request - slice size provided without slice topology",
+			job: testingutil.MakeJob("job", "default").
+				PodAnnotation(kueuealpha.PodSetRequiredTopologyAnnotation, "cloud.com/block").
+				PodAnnotation(kueuealpha.PodSetSliceSizeAnnotation, "1").
+				Obj(),
+			wantErr: field.ErrorList{
+				field.Forbidden(replicaMetaPath.Child("annotations").Key("kueue.x-k8s.io/podset-slice-size"), "cannot be set when 'kueue.x-k8s.io/podset-slice-required-topology' is not present"),
+			},
+		},
+		{
+			name: "invalid topology request - slice topology requested without podset topology",
+			job: testingutil.MakeJob("job", "default").
+				PodAnnotation(kueuealpha.PodSetSliceRequiredTopologyAnnotation, "cloud.com/block").
+				PodAnnotation(kueuealpha.PodSetSliceSizeAnnotation, "1").
+				Obj(),
+			wantErr: field.ErrorList{
+				field.Forbidden(replicaMetaPath.Child("annotations").Key("kueue.x-k8s.io/podset-slice-required-topology"), "can be used with required or preferred topology only"),
+			},
+		},
 	}
 
 	for _, tc := range testcases {
@@ -541,6 +624,29 @@ func TestValidateUpdate(t *testing.T) {
 					`must not contain more than one topology annotation: ["kueue.x-k8s.io/podset-required-topology", `+
 						`"kueue.x-k8s.io/podset-preferred-topology", "kueue.x-k8s.io/podset-unconstrained-topology"]`)},
 		},
+		{
+			name: "valid slice topology request",
+			oldJob: testingutil.MakeJob("job", "default").
+				Obj(),
+			newJob: testingutil.MakeJob("job", "default").
+				PodAnnotation(kueuealpha.PodSetRequiredTopologyAnnotation, "cloud.com/block").
+				PodAnnotation(kueuealpha.PodSetSliceRequiredTopologyAnnotation, "cloud.com/block").
+				PodAnnotation(kueuealpha.PodSetSliceSizeAnnotation, "1").
+				Obj(),
+		},
+		{
+			name: "attempt to set invalid slice topology request",
+			oldJob: testingutil.MakeJob("job", "default").
+				Obj(),
+			newJob: testingutil.MakeJob("job", "default").
+				PodAnnotation(kueuealpha.PodSetUnconstrainedTopologyAnnotation, "true").
+				PodAnnotation(kueuealpha.PodSetSliceRequiredTopologyAnnotation, "cloud.com/block").
+				PodAnnotation(kueuealpha.PodSetSliceSizeAnnotation, "1").
+				Obj(),
+			wantErr: field.ErrorList{
+				field.Forbidden(replicaMetaPath.Child("annotations").Key("kueue.x-k8s.io/podset-slice-required-topology"), "can be used with required or preferred topology only"),
+			},
+		},
 	}
 
 	for _, tc := range testcases {
diff --git a/pkg/scheduler/preemption/preemption_test.go b/pkg/scheduler/preemption/preemption_test.go
index 3347a659c..ab60b57e8 100644
--- a/pkg/scheduler/preemption/preemption_test.go
+++ b/pkg/scheduler/preemption/preemption_test.go
@@ -2824,7 +2824,9 @@ func TestCandidatesOrdering(t *testing.T) {
 			ReserveQuotaAt(utiltesting.MakeAdmission("self").Obj(), now.Add(time.Second)).
 			Obj()),
 	}
-	sort.Slice(candidates, CandidatesOrdering(candidates, "self", now))
+	sort.Slice(candidates, func(i int, j int) bool {
+		return CandidatesOrdering(candidates[i], candidates[j], "self", now)
+	})
 	gotNames := make([]workload.WorkloadReference, len(candidates))
 	for i, c := range candidates {
 		gotNames[i] = workload.Key(c.Obj)