Simplify flavorassigner to only return bool if borrowing is required

kubernetes-sigs · Dec 20, 2023 · 73d8e4e · 73d8e4e
1 parent 25630d0
commit 73d8e4e
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 51 deletions.
diff --git a/pkg/scheduler/flavorassigner/flavorassigner.go b/pkg/scheduler/flavorassigner/flavorassigner.go
@@ -39,9 +39,9 @@ import (
 )
 
 type Assignment struct {
-	PodSets     []PodSetAssignment
-	TotalBorrow cache.FlavorResourceQuantities
-	LastState   workload.AssigmentClusterQueueState
+	PodSets   []PodSetAssignment
+	Borrowing bool
+	LastState workload.AssigmentClusterQueueState
 
 	// Usage is the accumulated Usage of resources as pod sets get
 	// flavors assigned.
@@ -52,7 +52,7 @@ type Assignment struct {
 }
 
 func (a *Assignment) Borrows() bool {
-	return len(a.TotalBorrow) > 0
+	return a.Borrowing
 }
 
 // RepresentativeMode calculates the representative mode for the assigment as
@@ -223,7 +223,7 @@ type FlavorAssignment struct {
 	Name           kueue.ResourceFlavorReference
 	Mode           FlavorAssignmentMode
 	TriedFlavorIdx int
-	borrow         int64
+	borrow         bool
 }
 
 func lastAssignmentOutdated(wl *workload.Info, cq *cache.ClusterQueue) bool {
@@ -266,9 +266,8 @@ func AssignFlavors(log logr.Logger, wl *workload.Info, resourceFlavors map[kueue
 
 func assignFlavors(log logr.Logger, requests []workload.PodSetResources, podSets []kueue.PodSet, resourceFlavors map[kueue.ResourceFlavorReference]*kueue.ResourceFlavor, cq *cache.ClusterQueue, lastAssignment *workload.AssigmentClusterQueueState) Assignment {
 	assignment := Assignment{
-		TotalBorrow: make(cache.FlavorResourceQuantities),
-		PodSets:     make([]PodSetAssignment, 0, len(requests)),
-		Usage:       make(cache.FlavorResourceQuantities),
+		PodSets: make([]PodSetAssignment, 0, len(requests)),
+		Usage:   make(cache.FlavorResourceQuantities),
 		LastState: workload.AssigmentClusterQueueState{
 			LastTriedFlavorIdx:     make([]map[corev1.ResourceName]int, 0, len(podSets)),
 			CohortGeneration:       0,
@@ -323,14 +322,9 @@ func assignFlavors(log logr.Logger, requests []workload.PodSetResources, podSets
 
 		assignment.append(podSet.Requests, &psAssignment)
 		if psAssignment.Status.IsError() || (len(podSet.Requests) > 0 && len(psAssignment.Flavors) == 0) {
-			// This assignment failed, no need to continue tracking.
-			assignment.TotalBorrow = nil
 			return assignment
 		}
 	}
-	if len(assignment.TotalBorrow) == 0 {
-		assignment.TotalBorrow = nil
-	}
 	return assignment
 }
 
@@ -349,13 +343,8 @@ func (a *Assignment) append(requests workload.Requests, psAssignment *PodSetAssi
 	flavorIdx := make(map[corev1.ResourceName]int, len(psAssignment.Flavors))
 	a.PodSets = append(a.PodSets, *psAssignment)
 	for resource, flvAssignment := range psAssignment.Flavors {
-		if flvAssignment.borrow > 0 {
-			if a.TotalBorrow[flvAssignment.Name] == nil {
-				a.TotalBorrow[flvAssignment.Name] = make(map[corev1.ResourceName]int64)
-			}
-			// Don't accumulate borrowing. The returned `borrow` already considers
-			// usage from previous pod sets.
-			a.TotalBorrow[flvAssignment.Name][resource] = flvAssignment.borrow
+		if flvAssignment.borrow {
+			a.Borrowing = true
 		}
 		if a.Usage[flvAssignment.Name] == nil {
 			a.Usage[flvAssignment.Name] = make(map[corev1.ResourceName]int64)
@@ -428,7 +417,7 @@ func (a *Assignment) findFlavorForResourceGroup(
 			if mode < representativeMode {
 				representativeMode = mode
 			}
-			needsBorrowing = needsBorrowing || (mode == Fit && borrow > 0)
+			needsBorrowing = needsBorrowing || borrow
 			if representativeMode == NoFit {
 				// The flavor doesn't fit, no need to check other resources.
 				break
@@ -548,7 +537,7 @@ func flavorSelector(spec *corev1.PodSpec, allowedKeys sets.Set[string]) nodeaffi
 // If it fits, also returns any borrowing required.
 // If the flavor doesn't satisfy limits immediately (when waiting or preemption
 // could help), it returns a Status with reasons.
-func fitsResourceQuota(fName kueue.ResourceFlavorReference, rName corev1.ResourceName, val int64, cq *cache.ClusterQueue, rQuota *cache.ResourceQuota) (FlavorAssignmentMode, int64, *Status) {
+func fitsResourceQuota(fName kueue.ResourceFlavorReference, rName corev1.ResourceName, val int64, cq *cache.ClusterQueue, rQuota *cache.ResourceQuota) (FlavorAssignmentMode, bool, *Status) {
 	var status Status
 	used := cq.Usage[fName][rName]
 	mode := NoFit
@@ -560,7 +549,7 @@ func fitsResourceQuota(fName kueue.ResourceFlavorReference, rName corev1.Resourc
 	}
 	if rQuota.BorrowingLimit != nil && used+val > rQuota.Nominal+*rQuota.BorrowingLimit {
 		status.append(fmt.Sprintf("borrowing limit for %s in flavor %s exceeded", rName, fName))
-		return mode, 0, &status
+		return mode, false, &status
 	}
 
 	cohortUsed := used
@@ -572,11 +561,7 @@ func fitsResourceQuota(fName kueue.ResourceFlavorReference, rName corev1.Resourc
 
 	lack := cohortUsed + val - cohortAvailable
 	if lack <= 0 {
-		borrow := used + val - rQuota.Nominal
-		if borrow < 0 {
-			borrow = 0
-		}
-		return Fit, borrow, nil
+		return Fit, used+val > rQuota.Nominal, nil
 	}
 
 	lackQuantity := workload.ResourceQuantity(rName, lack)
@@ -589,7 +574,7 @@ func fitsResourceQuota(fName kueue.ResourceFlavorReference, rName corev1.Resourc
 		}
 	}
 	status.append(msg)
-	return mode, 0, &status
+	return mode, false, &status
 }
 
 func filterRequestedResources(req workload.Requests, allowList sets.Set[corev1.ResourceName]) workload.Requests {

diff --git a/pkg/scheduler/flavorassigner/flavorassigner_test.go b/pkg/scheduler/flavorassigner/flavorassigner_test.go
@@ -1049,12 +1049,7 @@ func TestAssignFlavors(t *testing.T) {
 						Count: 1,
 					},
 				},
-				TotalBorrow: cache.FlavorResourceQuantities{
-					"default": {
-						corev1.ResourceCPU:    8_000,
-						corev1.ResourceMemory: 3 * utiltesting.Gi,
-					},
-				},
+				Borrowing: true,
 				Usage: cache.FlavorResourceQuantities{
 					"default": map[corev1.ResourceName]int64{
 						corev1.ResourceCPU:    10000,
@@ -1713,7 +1708,7 @@ func TestAssignFlavors(t *testing.T) {
 			},
 			wantRepMode: Fit,
 			wantAssignment: Assignment{
-				TotalBorrow: cache.FlavorResourceQuantities{"one": {"cpu": 1000}},
+				Borrowing: true,
 				PodSets: []PodSetAssignment{{
 					Name: "main",
 					Flavors: ResourceAssignment{
@@ -1829,9 +1824,7 @@ func TestAssignFlavors(t *testing.T) {
 			},
 			wantRepMode: Fit,
 			wantAssignment: Assignment{
-				TotalBorrow: cache.FlavorResourceQuantities{
-					"one": {corev1.ResourceCPU: 1000},
-				},
+				Borrowing: true,
 				PodSets: []PodSetAssignment{{
 					Name: "main",
 					Flavors: ResourceAssignment{

diff --git a/pkg/scheduler/scheduler_test.go b/pkg/scheduler/scheduler_test.go
@@ -1230,9 +1230,7 @@ func TestEntryOrdering(t *testing.T) {
 				}},
 			},
 			assignment: flavorassigner.Assignment{
-				TotalBorrow: cache.FlavorResourceQuantities{
-					"flavor": {},
-				},
+				Borrowing: true,
 			},
 		},
 		{
@@ -1261,9 +1259,7 @@ func TestEntryOrdering(t *testing.T) {
 				}},
 			},
 			assignment: flavorassigner.Assignment{
-				TotalBorrow: cache.FlavorResourceQuantities{
-					"flavor": {},
-				},
+				Borrowing: true,
 			},
 		},
 		{
@@ -1284,9 +1280,7 @@ func TestEntryOrdering(t *testing.T) {
 				}},
 			},
 			assignment: flavorassigner.Assignment{
-				TotalBorrow: cache.FlavorResourceQuantities{
-					"flavor": {},
-				},
+				Borrowing: true,
 			},
 		},
 		{
@@ -1309,9 +1303,7 @@ func TestEntryOrdering(t *testing.T) {
 				},
 			},
 			assignment: flavorassigner.Assignment{
-				TotalBorrow: cache.FlavorResourceQuantities{
-					"flavor": {},
-				},
+				Borrowing: true,
 			},
 		},
 		{