Fix the borrowing while preemption when no borrowingLimit

kubernetes-sigs · Jan 9, 2024 · 24bdf54 · 24bdf54
1 parent 1664221
commit 24bdf54
Show file tree

Hide file tree

Showing 2 changed files with 76 additions and 13 deletions.
diff --git a/pkg/scheduler/preemption/preemption.go b/pkg/scheduler/preemption/preemption.go
@@ -361,20 +361,28 @@ func workloadFits(wlReq cache.FlavorResourceQuantities, cq *cache.ClusterQueue,
 				cohortResRequestable = cq.Cohort.RequestableResources[flvQuotas.Name]
 			}
 			for rName, rReq := range flvReq {
-				limit := flvQuotas.Resources[rName].Nominal
-
-				// No need to check whether cohort is nil here because when borrowingLimit
-				// is non nil, cohort is always non nil.
-				if allowBorrowing && flvQuotas.Resources[rName].BorrowingLimit != nil {
-					limit += *flvQuotas.Resources[rName].BorrowingLimit
-				}
-
-				if cqResUsage[rName]+rReq > limit {
-					return false
+				resource := flvQuotas.Resources[rName]
+
+				if cq.Cohort != nil {
+					if allowBorrowing {
+						// When resource.BorrowingLimit == nil there is no borrowing
+						// limit. In that case we only need to compare the requested
+						// resources + current usage against the amount of requestable
+						// resources in the cohort.
+						if resource.BorrowingLimit != nil {
+							if cqResUsage[rName]+rReq > resource.Nominal+*resource.BorrowingLimit {
+								return false
+							}
+						}
+					}
+					if cohortResUsage[rName]+rReq > cohortResRequestable[rName] {
+						return false
+					}
 				}
-
-				if cq.Cohort != nil && cohortResUsage[rName]+rReq > cohortResRequestable[rName] {
-					return false
+				if cq.Cohort == nil || !allowBorrowing {
+					if cqResUsage[rName]+rReq > resource.Nominal {
+						return false
+					}
 				}
 			}
 		}

diff --git a/pkg/scheduler/preemption/preemption_test.go b/pkg/scheduler/preemption/preemption_test.go
@@ -102,6 +102,30 @@ func TestPreemption(t *testing.T) {
 				ReclaimWithinCohort: kueue.PreemptionPolicyAny,
 			}).
 			Obj(),
+		utiltesting.MakeClusterQueue("d1").
+			Cohort("cohort-no-limits").
+			ResourceGroup(*utiltesting.MakeFlavorQuotas("default").
+				Resource(corev1.ResourceCPU, "6").
+				Resource(corev1.ResourceMemory, "3Gi").
+				Obj(),
+			).
+			Preemption(kueue.ClusterQueuePreemption{
+				WithinClusterQueue:  kueue.PreemptionPolicyLowerPriority,
+				ReclaimWithinCohort: kueue.PreemptionPolicyLowerPriority,
+			}).
+			Obj(),
+		utiltesting.MakeClusterQueue("d2").
+			Cohort("cohort-no-limits").
+			ResourceGroup(*utiltesting.MakeFlavorQuotas("default").
+				Resource(corev1.ResourceCPU, "6").
+				Resource(corev1.ResourceMemory, "3Gi").
+				Obj(),
+			).
+			Preemption(kueue.ClusterQueuePreemption{
+				WithinClusterQueue:  kueue.PreemptionPolicyNever,
+				ReclaimWithinCohort: kueue.PreemptionPolicyAny,
+			}).
+			Obj(),
 		utiltesting.MakeClusterQueue("l1").
 			Cohort("legion").
 			ResourceGroup(*utiltesting.MakeFlavorQuotas("default").
@@ -540,6 +564,37 @@ func TestPreemption(t *testing.T) {
 			}),
 			wantPreempted: sets.New("/c1-low"),
 		},
+		"preempting locally and borrowing same resource in cohort; no borrowing limit in the cohort": {
+			admitted: []kueue.Workload{
+				*utiltesting.MakeWorkload("d1-med", "").
+					Priority(0).
+					Request(corev1.ResourceCPU, "4").
+					ReserveQuota(utiltesting.MakeAdmission("d1").Assignment(corev1.ResourceCPU, "default", "4000m").Obj()).
+					Obj(),
+				*utiltesting.MakeWorkload("d1-low", "").
+					Priority(-1).
+					Request(corev1.ResourceCPU, "4").
+					ReserveQuota(utiltesting.MakeAdmission("d1").Assignment(corev1.ResourceCPU, "default", "4000m").Obj()).
+					Obj(),
+				*utiltesting.MakeWorkload("d2-low-1", "").
+					Priority(-1).
+					Request(corev1.ResourceCPU, "4").
+					ReserveQuota(utiltesting.MakeAdmission("d2").Assignment(corev1.ResourceCPU, "default", "4000m").Obj()).
+					Obj(),
+			},
+			incoming: utiltesting.MakeWorkload("in", "").
+				Priority(1).
+				Request(corev1.ResourceCPU, "4").
+				Obj(),
+			targetCQ: "d1",
+			assignment: singlePodSetAssignment(flavorassigner.ResourceAssignment{
+				corev1.ResourceCPU: &flavorassigner.FlavorAssignment{
+					Name: "default",
+					Mode: flavorassigner.Preempt,
+				},
+			}),
+			wantPreempted: sets.New("/d1-low"),
+		},
 		"preempting locally and borrowing other resources in cohort, with cohort candidates": {
 			admitted: []kueue.Workload{
 				*utiltesting.MakeWorkload("c1-med", "").