Repository: kubernetes-sigs/kueue

Issue:
Make Topology-Aware Scheduling compatible with cohorts and preemption
**What would you like to be added**:

Make TAS compatible with cohorts and preemption.

**Why is this needed**:

Currently, TAS does not work when cohorts or preemption are used. In that case the CQ is marked as inactive (see https://github.com/kubernetes-sigs/kueue/tree/main/keps/2724-topology-aware-scheduling#alpha). 

**Completion requirements**:

This enhancement requires the following artifacts:

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

The artifacts should be linked in subsequent comments.


Hints from issue comments:
@mimowo: cc @mwielgus @mwysokin @tenzen-y @mimowo: Listing cleanups to do identified during https://github.com/kubernetes-sigs/kueue/pull/4200:
- introduce PodSetReference to represent PodSet names for static code analysis
- add a unit test for `if len(cq.TASFlavors) == 0 {` as requested in [comment](https://github.com/kubernetes-sigs/kueue/pull/4200#discussion_r1951691235)
- add a unit test for `if cq.TASFlavors[*tasFlvr] == nil {` as requested in [comment](https://github.com/kubernetes-sigs/kueue/pull/4200#discussion_r1952616625)
- attempt to commonize the structures for `PodSetResources` and `TASPodSetRequests` as requested in [comment](https://github.com/kubernetes-sigs/kueue/pull/4200#discussion_r1950788801)@tenzen-y: Could you add https://github.com/kubernetes-sigs/kueue/pull/4200#discussion_r1952616625 to your list?
@mimowo: Sure, updated the list, but I think long-term it is better to manage in the refreshed [spreadsheet](https://docs.google.com/spreadsheets/d/1MXCjKZtAfqBTb61bJo46u7jIUqRIlu1NrYj1L8Xz-UU/edit?resourcekey=0-gr1ML2A1Axi8s6Lxr-Zhlw&gid=0#gid=0).@tenzen-y: > spreadsheet

Thanks, SGTM

Solution patch:
diff --git a/pkg/cache/clusterqueue_snapshot.go b/pkg/cache/clusterqueue_snapshot.go
index c95a02a58..60aa8fd2c 100644
--- a/pkg/cache/clusterqueue_snapshot.go
+++ b/pkg/cache/clusterqueue_snapshot.go
@@ -149,3 +149,21 @@ func (c *ClusterQueueSnapshot) DominantResourceShareWithout(wlReq resources.Flav
 	}
 	return dominantResourceShare(c, without)
 }
+
+type WorkloadTASRequests map[kueue.ResourceFlavorReference]FlavorTASRequests
+
+func (c *ClusterQueueSnapshot) FindTopologyAssignmentsForWorkload(
+	tasRequestsByFlavor WorkloadTASRequests) TASAssignmentsResult {
+	result := make(TASAssignmentsResult)
+	for tasFlavor, flavorTASRequests := range tasRequestsByFlavor {
+		// We assume the `tasFlavor` is already in the snapshot as this was
+		// already checked earlier during flavor assignment, and the set of
+		// flavors is immutable in snapshot.
+		tasFlavorCache := c.TASFlavors[tasFlavor]
+		flvResult := tasFlavorCache.FindTopologyAssignmentsForFlavor(flavorTASRequests)
+		for psName, psAssignment := range flvResult {
+			result[psName] = psAssignment
+		}
+	}
+	return result
+}
diff --git a/pkg/cache/tas_flavor_snapshot.go b/pkg/cache/tas_flavor_snapshot.go
index c7b8ee638..092ffba42 100644
--- a/pkg/cache/tas_flavor_snapshot.go
+++ b/pkg/cache/tas_flavor_snapshot.go
@@ -17,6 +17,7 @@ limitations under the License.
 package cache
 
 import (
+	"cmp"
 	"errors"
 	"fmt"
 	"maps"
@@ -223,6 +224,63 @@ func (s *TASFlavorSnapshot) addTASUsage(domainID utiltas.TopologyDomainID, usage
 	s.leaves[domainID].tasUsage.Add(usage)
 }
 
+type TASPodSetRequests struct {
+	PodSet            *kueue.PodSet
+	SinglePodRequests resources.Requests
+	Count             int32
+	Flavor            kueue.ResourceFlavorReference
+}
+
+type FailureInfo struct {
+	// PodSetName indicates the name of the PodSet for which computing the
+	// TAS assignment failed.
+	PodSetName string
+
+	// Reason indicates the reason why computing the TAS assignment failed.
+	Reason string
+}
+
+// the key in this map is PodSet name
+type TASAssignmentsResult map[string]tasPodSetAssignmentResult
+
+func (r TASAssignmentsResult) Failure() *FailureInfo {
+	for psName, psAssignment := range r {
+		if psAssignment.FailureReason != "" {
+			return &FailureInfo{PodSetName: psName, Reason: psAssignment.FailureReason}
+		}
+	}
+	return nil
+}
+
+type tasPodSetAssignmentResult struct {
+	TopologyAssignment *kueue.TopologyAssignment
+	FailureReason      string
+}
+
+type FlavorTASRequests []TASPodSetRequests
+
+func (s *TASFlavorSnapshot) FindTopologyAssignmentsForFlavor(flavorTASRequests FlavorTASRequests) TASAssignmentsResult {
+	result := make(map[string]tasPodSetAssignmentResult)
+	assumedUsage := make(map[utiltas.TopologyDomainID]resources.Requests)
+	for _, tr := range flavorTASRequests {
+		assignment, reason := s.findTopologyAssignment(tr, assumedUsage)
+		result[tr.PodSet.Name] = tasPodSetAssignmentResult{TopologyAssignment: assignment, FailureReason: reason}
+		if reason != "" {
+			return result
+		}
+		for _, domain := range assignment.Domains {
+			domainID := utiltas.DomainID(domain.Values)
+			assumedDomainUsage := tr.SinglePodRequests.Clone()
+			assumedDomainUsage.Mul(int64(domain.Count))
+			if assumedUsage[domainID] == nil {
+				assumedUsage[domainID] = resources.Requests{}
+			}
+			assumedUsage[domainID].Add(assumedDomainUsage)
+		}
+	}
+	return result
+}
+
 // Algorithm overview:
 // Phase 1:
 //
@@ -235,11 +293,13 @@ func (s *TASFlavorSnapshot) addTASUsage(domainID utiltas.TopologyDomainID, usage
 //	b) traverse the structure down level-by-level optimizing the number of used
 //	  domains at each level
 //	c) build the assignment for the lowest level in the hierarchy
-func (s *TASFlavorSnapshot) FindTopologyAssignment(
-	topologyRequest *kueue.PodSetTopologyRequest,
-	requests resources.Requests,
-	count int32,
-	podSetTolerations []corev1.Toleration) (*kueue.TopologyAssignment, string) {
+func (s *TASFlavorSnapshot) findTopologyAssignment(
+	tasPodSetRequests TASPodSetRequests,
+	assumedUsage map[utiltas.TopologyDomainID]resources.Requests) (*kueue.TopologyAssignment, string) {
+	topologyRequest := tasPodSetRequests.PodSet.TopologyRequest
+	requests := tasPodSetRequests.SinglePodRequests
+	podSetTolerations := tasPodSetRequests.PodSet.Template.Spec.Tolerations
+	count := tasPodSetRequests.Count
 	required := topologyRequest.Required != nil
 	key := levelKey(topologyRequest)
 	if key == nil {
@@ -250,7 +310,7 @@ func (s *TASFlavorSnapshot) FindTopologyAssignment(
 		return nil, fmt.Sprintf("no requested topology level: %s", *key)
 	}
 	// phase 1 - determine the number of pods which can fit in each topology domain
-	s.fillInCounts(requests, append(podSetTolerations, s.tolerations...))
+	s.fillInCounts(requests, assumedUsage, append(podSetTolerations, s.tolerations...))
 
 	// phase 2a: determine the level at which the assignment is done along with
 	// the domains which can accommodate all pods
@@ -359,7 +419,6 @@ func (s *TASFlavorSnapshot) buildTopologyAssignmentForLevels(domains []*domain,
 		for resName, resValue := range singlePodRequest {
 			usage[resName] = resValue * int64(domain.state)
 		}
-		s.addTASUsage(domain.id, usage)
 	}
 	return assignment
 }
@@ -386,22 +445,20 @@ func (s *TASFlavorSnapshot) lowerLevelDomains(domains []*domain) []*domain {
 }
 
 func (s *TASFlavorSnapshot) sortedDomains(domains []*domain) []*domain {
-	result := make([]*domain, len(domains))
-	copy(result, domains)
+	result := slices.Clone(domains)
 	slices.SortFunc(result, func(a, b *domain) int {
-		switch {
-		case a.state == b.state:
+		if a.state == b.state {
 			return slices.Compare(a.levelValues, b.levelValues)
-		case a.state > b.state:
-			return -1
-		default:
-			return 1
 		}
+		// descending order.
+		return cmp.Compare(b.state, a.state)
 	})
 	return result
 }
 
-func (s *TASFlavorSnapshot) fillInCounts(requests resources.Requests, tolerations []corev1.Toleration) {
+func (s *TASFlavorSnapshot) fillInCounts(requests resources.Requests,
+	assumedUsage map[utiltas.TopologyDomainID]resources.Requests,
+	tolerations []corev1.Toleration) {
 	for _, domain := range s.domains {
 		// cleanup the state in case some remaining values are present from computing
 		// assignments for previous PodSets.
@@ -417,6 +474,9 @@ func (s *TASFlavorSnapshot) fillInCounts(requests resources.Requests, toleration
 		}
 		remainingCapacity := leaf.freeCapacity.Clone()
 		remainingCapacity.Sub(leaf.tasUsage)
+		if leafAssumedUsage, found := assumedUsage[leaf.domain.id]; found {
+			remainingCapacity.Sub(leafAssumedUsage)
+		}
 		leaf.state = requests.CountIn(remainingCapacity)
 	}
 	for _, root := range s.roots {
diff --git a/pkg/resources/requests.go b/pkg/resources/requests.go
index 2e5b7b42b..81b8172a1 100644
--- a/pkg/resources/requests.go
+++ b/pkg/resources/requests.go
@@ -49,6 +49,12 @@ func (r Requests) Divide(f int64) {
 	}
 }
 
+func (r Requests) Mul(f int64) {
+	for k := range r {
+		r[k] *= f
+	}
+}
+
 func (r Requests) Add(addRequests Requests) {
 	for k, v := range addRequests {
 		r[k] += v
diff --git a/pkg/scheduler/flavorassigner/flavorassigner.go b/pkg/scheduler/flavorassigner/flavorassigner.go
index 4bddd08bc..e97d63308 100644
--- a/pkg/scheduler/flavorassigner/flavorassigner.go
+++ b/pkg/scheduler/flavorassigner/flavorassigner.go
@@ -19,6 +19,7 @@ package flavorassigner
 import (
 	"errors"
 	"fmt"
+	"slices"
 	"sort"
 	"strings"
 
@@ -57,6 +58,13 @@ func (a *Assignment) Borrows() bool {
 	return a.Borrowing
 }
 
+func (a *Assignment) podSetAssignmentByName(psName string) *PodSetAssignment {
+	if idx := slices.IndexFunc(a.PodSets, func(ps PodSetAssignment) bool { return ps.Name == psName }); idx != -1 {
+		return &a.PodSets[idx]
+	}
+	return nil
+}
+
 // RepresentativeMode calculates the representative mode for the assignment as
 // the worst assignment mode among all the pod sets.
 func (a *Assignment) RepresentativeMode() FlavorAssignmentMode {
@@ -196,6 +204,20 @@ func (psa *PodSetAssignment) RepresentativeMode() FlavorAssignmentMode {
 	return mode
 }
 
+func (psa *PodSetAssignment) reason(reason string) {
+	if psa.Status == nil {
+		psa.Status = &Status{}
+	}
+	psa.Status.reasons = append(psa.Status.reasons, reason)
+}
+
+func (psa *PodSetAssignment) error(err error) {
+	if psa.Status == nil {
+		psa.Status = &Status{}
+	}
+	psa.Status.err = err
+}
+
 type ResourceAssignment map[corev1.ResourceName]*FlavorAssignment
 
 func (psa *PodSetAssignment) toAPI() kueue.PodSetAssignment {
@@ -363,17 +385,35 @@ func (a *FlavorAssigner) assignFlavors(log logr.Logger, counts []int32) Assignme
 			}
 			psAssignment.append(flavors, status)
 		}
-		if features.Enabled(features.TopologyAwareScheduling) {
-			if a.wl.Obj.Spec.PodSets[i].TopologyRequest != nil {
-				assignTopology(log, &psAssignment, a.cq, a.wl.TotalRequests[i], &a.wl.Obj.Spec.PodSets[i])
-			}
-		}
 
 		assignment.append(podSet.Requests, &psAssignment)
 		if psAssignment.Status.IsError() || (len(podSet.Requests) > 0 && len(psAssignment.Flavors) == 0) {
 			return assignment
 		}
 	}
+	if assignment.RepresentativeMode() == NoFit {
+		return assignment
+	}
+
+	if features.Enabled(features.TopologyAwareScheduling) && assignment.RepresentativeMode() == Fit {
+		tasRequests := assignment.WorkloadsTopologyRequests(a.wl, requests, a.cq)
+		result := a.cq.FindTopologyAssignmentsForWorkload(tasRequests)
+		if failure := result.Failure(); failure != nil {
+			// There is at least one PodSet which does not fit
+			psAssignment := assignment.podSetAssignmentByName(failure.PodSetName)
+			psAssignment.reason(failure.Reason)
+			// clear flavors to make the PodSet assignment as NoFit
+			psAssignment.Flavors = nil
+			// update the representative mode accordingly
+			assignment.representativeMode = ptr.To(NoFit)
+		} else {
+			// All PodSets fit, we just update the TopologyAssignments
+			for psName, psResult := range result {
+				psAssignment := assignment.podSetAssignmentByName(psName)
+				psAssignment.TopologyAssignment = psResult.TopologyAssignment
+			}
+		}
+	}
 	return assignment
 }
 
diff --git a/pkg/scheduler/flavorassigner/tas_flavorassigner.go b/pkg/scheduler/flavorassigner/tas_flavorassigner.go
index dc92756a4..7a80f52bf 100644
--- a/pkg/scheduler/flavorassigner/tas_flavorassigner.go
+++ b/pkg/scheduler/flavorassigner/tas_flavorassigner.go
@@ -20,7 +20,6 @@ import (
 	"errors"
 	"fmt"
 
-	"github.com/go-logr/logr"
 	"k8s.io/utils/ptr"
 
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
@@ -28,56 +27,54 @@ import (
 	"sigs.k8s.io/kueue/pkg/workload"
 )
 
-func assignTopology(log logr.Logger,
-	psAssignment *PodSetAssignment,
-	cq *cache.ClusterQueueSnapshot,
-	psResources workload.PodSetResources,
-	podSet *kueue.PodSet) {
-	switch {
-	case psAssignment.Status.IsError():
-		log.V(2).Info("There is no resource quota assignment for the workload. No need to check TAS.", "message", psAssignment.Status.Message())
-	case len(psAssignment.Flavors) == 0:
-		log.V(2).Info("There is no flavor assignment for the workload. No need to check TAS.", "message", psAssignment.Status.Message())
-	case len(cq.TASFlavors) == 0:
-		if psAssignment.Status == nil {
-			psAssignment.Status = &Status{}
-		}
-		psAssignment.Status.append("Workload requires Topology, but there is no TAS cache information")
-		psAssignment.Flavors = nil
-	default:
-		singlePodRequests := psResources.Requests.Clone()
-		singlePodRequests.Divide(int64(psResources.Count))
-		podCount := psAssignment.Count
-		tasFlvr, err := onlyFlavor(psAssignment.Flavors)
-		if err != nil {
-			if psAssignment.Status == nil {
-				psAssignment.Status = &Status{}
-			}
-			psAssignment.Status.err = err
-			psAssignment.Flavors = nil
-			return
-		}
-		snapshot := cq.TASFlavors[*tasFlvr]
-		if snapshot == nil {
-			if psAssignment.Status == nil {
-				psAssignment.Status = &Status{}
+// WorkloadsTopologyRequests - returns the TopologyRequests of the workload
+func (a *Assignment) WorkloadsTopologyRequests(wl *workload.Info,
+	requests []workload.PodSetResources,
+	cq *cache.ClusterQueueSnapshot) cache.WorkloadTASRequests {
+	tasRequests := make(cache.WorkloadTASRequests)
+	for i, podSet := range wl.Obj.Spec.PodSets {
+		if podSet.TopologyRequest != nil {
+			psAssignment := a.podSetAssignmentByName(podSet.Name)
+			if psAssignment.Status.IsError() {
+				// There is no resource quota assignment for the PodSet - no need to check TAS.
+				continue
 			}
-			psAssignment.Status.append("Workload requires Topology, but there is no TAS cache information for the assigned flavor")
-			psAssignment.Flavors = nil
-			return
-		}
-		var reason string
-		psAssignment.TopologyAssignment, reason = snapshot.FindTopologyAssignment(podSet.TopologyRequest,
-			singlePodRequests, podCount, podSet.Template.Spec.Tolerations)
-		if psAssignment.TopologyAssignment == nil {
-			if psAssignment.Status == nil {
-				psAssignment.Status = &Status{}
+			psTASRequest, err := podSetTopologyRequest(psAssignment, wl, cq, i)
+			if err != nil {
+				psAssignment.error(err)
+			} else {
+				tasRequests[psTASRequest.Flavor] = append(tasRequests[psTASRequest.Flavor], *psTASRequest)
 			}
-			psAssignment.Status.append(reason)
-			psAssignment.Flavors = nil
 		}
-		log.Info("TAS PodSet assignment", "tasAssignment", psAssignment.TopologyAssignment)
 	}
+	return tasRequests
+}
+
+func podSetTopologyRequest(psAssignment *PodSetAssignment,
+	wl *workload.Info,
+	cq *cache.ClusterQueueSnapshot,
+	podSetIndex int) (*cache.TASPodSetRequests, error) {
+	if len(cq.TASFlavors) == 0 {
+		return nil, errors.New("workload requires Topology, but there is no TAS cache information")
+	}
+	psResources := wl.TotalRequests[podSetIndex]
+	singlePodRequests := psResources.Requests.Clone()
+	singlePodRequests.Divide(int64(psResources.Count))
+	podCount := psAssignment.Count
+	tasFlvr, err := onlyFlavor(psAssignment.Flavors)
+	if err != nil {
+		return nil, err
+	}
+	if cq.TASFlavors[*tasFlvr] == nil {
+		return nil, errors.New("workload requires Topology, but there is no TAS cache information for the assigned flavor")
+	}
+	podSet := &wl.Obj.Spec.PodSets[podSetIndex]
+	return &cache.TASPodSetRequests{
+		Count:             podCount,
+		SinglePodRequests: singlePodRequests,
+		PodSet:            podSet,
+		Flavor:            *tasFlvr,
+	}, nil
 }
 
 func onlyFlavor(ra ResourceAssignment) (*kueue.ResourceFlavorReference, error) {



Test patch:
diff --git a/pkg/cache/tas_cache_test.go b/pkg/cache/tas_cache_test.go
index d65538999..160ecd5e1 100644
--- a/pkg/cache/tas_cache_test.go
+++ b/pkg/cache/tas_cache_test.go
@@ -1083,18 +1083,35 @@ func TestFindTopologyAssignment(t *testing.T) {
 			if err != nil {
 				t.Fatalf("failed to build the snapshot: %v", err)
 			}
-			gotAssignment, reason := snapshot.FindTopologyAssignment(&tc.request, tc.requests, tc.count, tc.tolerations)
-			if gotAssignment != nil {
+			tasInput := TASPodSetRequests{
+				PodSet: &kueue.PodSet{
+					Name:            "main",
+					TopologyRequest: &tc.request,
+					Template: corev1.PodTemplateSpec{
+						Spec: corev1.PodSpec{
+							Tolerations: tc.tolerations,
+						},
+					},
+				},
+				SinglePodRequests: tc.requests,
+				Count:             tc.count,
+			}
+			flavorTASRequests := []TASPodSetRequests{tasInput}
+			wantResult := make(TASAssignmentsResult)
+			wantMainPodSetResult := tasPodSetAssignmentResult{
+				FailureReason: tc.wantReason,
+			}
+			if tc.wantAssignment != nil {
 				sort.Slice(tc.wantAssignment.Domains, func(i, j int) bool {
 					return utiltas.DomainID(tc.wantAssignment.Domains[i].Values) < utiltas.DomainID(tc.wantAssignment.Domains[j].Values)
 				})
+				wantMainPodSetResult.TopologyAssignment = tc.wantAssignment
 			}
-			if diff := cmp.Diff(tc.wantAssignment, gotAssignment); diff != "" {
+			wantResult["main"] = wantMainPodSetResult
+			gotResult := snapshot.FindTopologyAssignmentsForFlavor(flavorTASRequests)
+			if diff := cmp.Diff(wantResult, gotResult); diff != "" {
 				t.Errorf("unexpected topology assignment (-want,+got): %s", diff)
 			}
-			if diff := cmp.Diff(tc.wantReason, reason); diff != "" {
-				t.Errorf("unexpected error (-want,+got): %s", diff)
-			}
 		})
 	}
 }