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 950320bd8..093409d8a 100644
--- a/pkg/cache/clusterqueue_snapshot.go
+++ b/pkg/cache/clusterqueue_snapshot.go
@@ -68,21 +68,23 @@ func (c *ClusterQueueSnapshot) RGByResource(resource corev1.ResourceName) *Resou
 	return nil
 }
 
-func (c *ClusterQueueSnapshot) AddUsage(frq resources.FlavorResourceQuantities) {
-	for fr, q := range frq {
+func (c *ClusterQueueSnapshot) AddUsage(usage workload.Usage) {
+	for fr, q := range usage.Quota {
 		addUsage(c, fr, q)
 	}
+	c.updateTASUsage(usage.TAS, add)
 }
 
-func (c *ClusterQueueSnapshot) removeUsage(frq resources.FlavorResourceQuantities) {
-	for fr, q := range frq {
+func (c *ClusterQueueSnapshot) removeUsage(usage workload.Usage) {
+	for fr, q := range usage.Quota {
 		removeUsage(c, fr, q)
 	}
+	c.updateTASUsage(usage.TAS, subtract)
 }
 
-func (c *ClusterQueueSnapshot) updateTASUsage(wi *workload.Info, op usageOp) {
-	if features.Enabled(features.TopologyAwareScheduling) && wi.IsUsingTAS() {
-		for tasFlavor, tasUsage := range wi.TASUsage() {
+func (c *ClusterQueueSnapshot) updateTASUsage(usage workload.TASUsage, op usageOp) {
+	if features.Enabled(features.TopologyAwareScheduling) {
+		for tasFlavor, tasUsage := range usage {
 			if tasFlvCache := c.TASFlavors[tasFlavor]; tasFlvCache != nil {
 				for _, tr := range tasUsage {
 					domainID := utiltas.DomainID(tr.Values)
@@ -93,8 +95,8 @@ func (c *ClusterQueueSnapshot) updateTASUsage(wi *workload.Info, op usageOp) {
 	}
 }
 
-func (c *ClusterQueueSnapshot) Fits(frq resources.FlavorResourceQuantities) bool {
-	for fr, q := range frq {
+func (c *ClusterQueueSnapshot) Fits(usage workload.Usage) bool {
+	for fr, q := range usage.Quota {
 		if c.Available(fr) < q {
 			return false
 		}
diff --git a/pkg/cache/snapshot.go b/pkg/cache/snapshot.go
index 053d7037f..d1ee360a7 100644
--- a/pkg/cache/snapshot.go
+++ b/pkg/cache/snapshot.go
@@ -44,8 +44,7 @@ type Snapshot struct {
 func (s *Snapshot) RemoveWorkload(wl *workload.Info) {
 	cq := s.ClusterQueues[wl.ClusterQueue]
 	delete(cq.Workloads, workload.Key(wl.Obj))
-	cq.removeUsage(wl.FlavorResourceUsage())
-	cq.updateTASUsage(wl, subtract)
+	cq.removeUsage(wl.Usage())
 }
 
 // AddWorkload adds a workload from its corresponding ClusterQueue and
@@ -53,8 +52,7 @@ func (s *Snapshot) RemoveWorkload(wl *workload.Info) {
 func (s *Snapshot) AddWorkload(wl *workload.Info) {
 	cq := s.ClusterQueues[wl.ClusterQueue]
 	cq.Workloads[workload.Key(wl.Obj)] = wl
-	cq.AddUsage(wl.FlavorResourceUsage())
-	cq.updateTASUsage(wl, add)
+	cq.AddUsage(wl.Usage())
 }
 
 func (s *Snapshot) Log(log logr.Logger) {
diff --git a/pkg/scheduler/scheduler.go b/pkg/scheduler/scheduler.go
index 758dd95d6..68ed47598 100644
--- a/pkg/scheduler/scheduler.go
+++ b/pkg/scheduler/scheduler.go
@@ -330,10 +330,16 @@ type entry struct {
 	preemptionTargets     []*preemption.Target
 }
 
-// netUsage returns how much capacity this entry will require from the ClusterQueue/Cohort.
+func (e *entry) netUsage() workload.Usage {
+	return workload.Usage{
+		Quota: e.netQuotaUsage(),
+	}
+}
+
+// netQuotaUsage returns how much capacity this entry will require from the ClusterQueue/Cohort.
 // When a workload is preempting, it subtracts the preempted resources from the resources
 // required, as the remaining quota is all we need from the CQ/Cohort.
-func (e *entry) netUsage() resources.FlavorResourceQuantities {
+func (e *entry) netQuotaUsage() resources.FlavorResourceQuantities {
 	if e.assignment.RepresentativeMode() == flavorassigner.Fit {
 		return e.assignment.Usage
 	}
@@ -392,7 +398,13 @@ func (s *Scheduler) nominate(ctx context.Context, workloads []workload.Info, sna
 }
 
 // resourcesToReserve calculates how much of the available resources in cq/cohort assignment should be reserved.
-func resourcesToReserve(e *entry, cq *cache.ClusterQueueSnapshot) resources.FlavorResourceQuantities {
+func resourcesToReserve(e *entry, cq *cache.ClusterQueueSnapshot) workload.Usage {
+	return workload.Usage{
+		Quota: quotaResourcesToReserve(e, cq),
+	}
+}
+
+func quotaResourcesToReserve(e *entry, cq *cache.ClusterQueueSnapshot) resources.FlavorResourceQuantities {
 	if e.assignment.RepresentativeMode() != flavorassigner.Preempt {
 		return e.assignment.Usage
 	}
diff --git a/pkg/workload/usage.go b/pkg/workload/usage.go
new file mode 100644
index 000000000..38f430471
--- /dev/null
+++ b/pkg/workload/usage.go
@@ -0,0 +1,30 @@
+/*
+Copyright The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package workload
+
+import (
+	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
+	"sigs.k8s.io/kueue/pkg/resources"
+)
+
+type TASFlavorUsage []TopologyDomainRequests
+type TASUsage map[kueue.ResourceFlavorReference]TASFlavorUsage
+
+type Usage struct {
+	Quota resources.FlavorResourceQuantities
+	TAS   TASUsage
+}
diff --git a/pkg/workload/workload.go b/pkg/workload/workload.go
index ae718439c..1e03c986d 100644
--- a/pkg/workload/workload.go
+++ b/pkg/workload/workload.go
@@ -238,6 +238,15 @@ func (i *Info) CanBePartiallyAdmitted() bool {
 	return CanBePartiallyAdmitted(i.Obj)
 }
 
+// Usage returns the total resource usage for the workload, including regular
+// quota and TAS usage.
+func (i *Info) Usage() Usage {
+	return Usage{
+		Quota: i.FlavorResourceUsage(),
+		TAS:   i.TASUsage(),
+	}
+}
+
 // FlavorResourceUsage returns the total resource usage for the workload,
 // per flavor (if assigned, otherwise flavor shows as empty string), per resource.
 func (i *Info) FlavorResourceUsage() resources.FlavorResourceQuantities {
@@ -280,8 +289,11 @@ func (i *Info) IsUsingTAS() bool {
 }
 
 // TASUsage returns topology usage requested by the Workload
-func (i *Info) TASUsage() map[kueue.ResourceFlavorReference][]TopologyDomainRequests {
-	result := make(map[kueue.ResourceFlavorReference][]TopologyDomainRequests, 0)
+func (i *Info) TASUsage() TASUsage {
+	if !features.Enabled(features.TopologyAwareScheduling) || !i.IsUsingTAS() {
+		return nil
+	}
+	result := make(TASUsage, 0)
 	for _, ps := range i.TotalRequests {
 		if ps.TopologyRequest != nil {
 			psFlavors := sets.New[kueue.ResourceFlavorReference]()



Test patch:
diff --git a/pkg/cache/resource_test.go b/pkg/cache/resource_test.go
index bae7cff18..9babe43c9 100644
--- a/pkg/cache/resource_test.go
+++ b/pkg/cache/resource_test.go
@@ -26,6 +26,7 @@ import (
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
 	"sigs.k8s.io/kueue/pkg/resources"
 	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
+	"sigs.k8s.io/kueue/pkg/workload"
 )
 
 func TestAvailable(t *testing.T) {
@@ -401,7 +402,7 @@ func TestAvailable(t *testing.T) {
 			// add usage
 			{
 				for cqName, usage := range tc.usage {
-					snapshot.ClusterQueues[cqName].AddUsage(usage)
+					snapshot.ClusterQueues[cqName].AddUsage(workload.Usage{Quota: usage})
 				}
 				gotAvailable := make(map[string]resources.FlavorResourceQuantities, len(snapshot.ClusterQueues))
 				gotPotentiallyAvailable := make(map[string]resources.FlavorResourceQuantities, len(snapshot.ClusterQueues))
@@ -425,7 +426,7 @@ func TestAvailable(t *testing.T) {
 			// remove usage
 			{
 				for cqName, usage := range tc.usage {
-					snapshot.ClusterQueues[cqName].removeUsage(usage)
+					snapshot.ClusterQueues[cqName].removeUsage(workload.Usage{Quota: usage})
 				}
 				gotAvailable := make(map[string]resources.FlavorResourceQuantities, len(snapshot.ClusterQueues))
 				gotPotentiallyAvailable := make(map[string]resources.FlavorResourceQuantities, len(snapshot.ClusterQueues))
diff --git a/pkg/scheduler/flavorassigner/flavorassigner_test.go b/pkg/scheduler/flavorassigner/flavorassigner_test.go
index e0a88a338..4ebf32cef 100644
--- a/pkg/scheduler/flavorassigner/flavorassigner_test.go
+++ b/pkg/scheduler/flavorassigner/flavorassigner_test.go
@@ -1967,7 +1967,7 @@ func TestAssignFlavors(t *testing.T) {
 				t.Fatalf("Failed to create CQ snapshot")
 			}
 			if tc.clusterQueueUsage != nil {
-				clusterQueue.AddUsage(tc.clusterQueueUsage)
+				clusterQueue.AddUsage(workload.Usage{Quota: tc.clusterQueueUsage})
 			}
 
 			if tc.secondaryClusterQueue != nil {
@@ -1975,7 +1975,7 @@ func TestAssignFlavors(t *testing.T) {
 				if secondaryClusterQueue == nil {
 					t.Fatalf("Failed to create secondary CQ snapshot")
 				}
-				secondaryClusterQueue.AddUsage(tc.secondaryClusterQueueUsage)
+				secondaryClusterQueue.AddUsage(workload.Usage{Quota: tc.secondaryClusterQueueUsage})
 			}
 
 			flvAssigner := New(wlInfo, clusterQueue, resourceFlavors, tc.enableFairSharing, &testOracle{})
@@ -2125,10 +2125,10 @@ func TestReclaimBeforePriorityPreemption(t *testing.T) {
 				t.Fatalf("unexpected error while building snapshot: %v", err)
 			}
 			otherClusterQueue := snapshot.ClusterQueues["other-clusterqueue"]
-			otherClusterQueue.AddUsage(tc.otherClusterQueueUsage)
+			otherClusterQueue.AddUsage(workload.Usage{Quota: tc.otherClusterQueueUsage})
 
 			testClusterQueue := snapshot.ClusterQueues["test-clusterqueue"]
-			testClusterQueue.AddUsage(tc.testClusterQueueUsage)
+			testClusterQueue.AddUsage(workload.Usage{Quota: tc.testClusterQueueUsage})
 
 			flvAssigner := New(wlInfo, testClusterQueue, resourceFlavors, false, &testOracle{})
 			log := testr.NewWithOptions(t, testr.Options{Verbosity: 2})
diff --git a/pkg/scheduler/scheduler_test.go b/pkg/scheduler/scheduler_test.go
index 36d62c9f2..d23e791fd 100644
--- a/pkg/scheduler/scheduler_test.go
+++ b/pkg/scheduler/scheduler_test.go
@@ -3818,7 +3818,7 @@ func TestResourcesToReserve(t *testing.T) {
 			cqSnapshot := snapshot.ClusterQueues["cq"]
 
 			got := resourcesToReserve(e, cqSnapshot)
-			if !reflect.DeepEqual(tc.wantReserved, got) {
+			if !reflect.DeepEqual(tc.wantReserved, got.Quota) {
 				t.Errorf("%s failed\n: Want reservedMem: %v, got: %v", tc.name, tc.wantReserved, got)
 			}
 		})