Preemption while borrowing WIP

apis/kueue/v1beta1/clusterqueue_types.go

@@ -366,6 +366,10 @@ type ClusterQueuePreemption struct {
 	// +kubebuilder:validation:Enum=Never;LowerPriority;Any
 	ReclaimWithinCohort PreemptionPolicy `json:"reclaimWithinCohort,omitempty"`
 
+	// borrowWithinCohort provides configuration to allow preemption within
+	// cohort while borrowing.
+	BorrowWithinCohort *BorrowWithinCohort `json:"borrowFromCohortConfig,omitempty"`
+
 	// withinClusterQueue determines whether a pending Workload that doesn't fit
 	// within the nominal quota for its ClusterQueue, can preempt active Workloads in
 	// the ClusterQueue. The possible values are:
@@ -382,6 +386,34 @@ type ClusterQueuePreemption struct {
 	WithinClusterQueue PreemptionPolicy `json:"withinClusterQueue,omitempty"`
 }
 
+type BorrowWithinCohortPolicy string
+
+const (
+	BorrowWithinCohortPolicyNever         BorrowWithinCohortPolicy = "Never"
+	BorrowWithinCohortPolicyLowerPriority BorrowWithinCohortPolicy = "LowerPriority"
+)
+
+// BorrowWithinCohort contains configuration which allows to preempt workloads
+// within cohort while borrowing.
+type BorrowWithinCohort struct {
+	// policy determines the policy for preemption to reclaim quota within cohort while borrowing.
+	// Possible values are:
+	// - `Never` (default): do not reclaim quota within ClusterQueue while borrowing.
+	// - `LowerPriority`: allow preempting workloads to reclaim quota within the cohort,
+	//    but only by preempting lower-priority workloads.
+	//
+	// +kubebuilder:default=Never
+	// +kubebuilder:validation:Enum=Never;LowerPriority
+	Policy BorrowWithinCohortPolicy `json:"policy,omitempty"`
+
+	// maxPriorityThreshold allows to restrict the set of workloads for preemption
+	// to reclaim quota within the cohort, to only workloads with priority below
+	// or equal the specified level.
+	//
+	// +optional
+	MaxPriorityThreshold *int32 `json:"maxPriorityThreshold,omitempty"`
+}
+
 //+genclient
 //+genclient:nonNamespaced
 //+kubebuilder:object:root=true

charts/kueue/templates/crd/kueue.x-k8s.io_clusterqueues.yaml

@@ -176,6 +176,30 @@ spec:
                   of Workloads to preempt to accomomdate the pending Workload, preempting
                   Workloads with lower priority first."
                 properties:
+                  borrowFromCohortConfig:
+                    description: borrowWithinCohort provides configuration to allow
+                      preemption within cohort while borrowing.
+                    properties:
+                      maxPriorityThreshold:
+                        description: maxPriorityThreshold allows to restrict the set
+                          of workloads for preemption to reclaim quota within the
+                          cohort, to only workloads with priority below or equal the
+                          specified level.
+                        format: int32
+                        type: integer
+                      policy:
+                        default: Never
+                        description: 'policy determines the policy for preemption
+                          to reclaim quota within cohort while borrowing. Possible
+                          values are: - `Never` (default): do not reclaim quota within
+                          ClusterQueue while borrowing. - `LowerPriority`: allow preempting
+                          workloads to reclaim quota within the cohort, but only by
+                          preempting lower-priority workloads.'
+                        enum:
+                        - Never
+                        - LowerPriority
+                        type: string
+                    type: object
                   reclaimWithinCohort:
                     default: Never
                     description: "reclaimWithinCohort determines whether a pending

config/components/crd/bases/kueue.x-k8s.io_clusterqueues.yaml

@@ -163,6 +163,30 @@ spec:
                   of Workloads to preempt to accomomdate the pending Workload, preempting
                   Workloads with lower priority first."
                 properties:
+                  borrowFromCohortConfig:
+                    description: borrowWithinCohort provides configuration to allow
+                      preemption within cohort while borrowing.
+                    properties:
+                      maxPriorityThreshold:
+                        description: maxPriorityThreshold allows to restrict the set
+                          of workloads for preemption to reclaim quota within the
+                          cohort, to only workloads with priority below or equal the
+                          specified level.
+                        format: int32
+                        type: integer
+                      policy:
+                        default: Never
+                        description: 'policy determines the policy for preemption
+                          to reclaim quota within cohort while borrowing. Possible
+                          values are: - `Never` (default): do not reclaim quota within
+                          ClusterQueue while borrowing. - `LowerPriority`: allow preempting
+                          workloads to reclaim quota within the cohort, but only by
+                          preempting lower-priority workloads.'
+                        enum:
+                        - Never
+                        - LowerPriority
+                        type: string
+                    type: object
                   reclaimWithinCohort:
                     default: Never
                     description: "reclaimWithinCohort determines whether a pending

pkg/scheduler/flavorassigner/flavorassigner.go

@@ -428,7 +428,7 @@ func (a *Assignment) findFlavorForResourceGroup(
 			if mode < representativeMode {
 				representativeMode = mode
 			}
-			needsBorrowing = needsBorrowing || (mode == Fit && borrow > 0)
+			needsBorrowing = needsBorrowing || ((mode == Fit || mode == Preempt) && borrow > 0)
 			if representativeMode == NoFit {
 				// The flavor doesn't fit, no need to check other resources.
 				break
@@ -483,7 +483,9 @@ func shouldTryNextFlavor(representativeMode FlavorAssignmentMode, flavorFungibil
 	policyPreempt := flavorFungibility.WhenCanPreempt
 	policyBorrow := flavorFungibility.WhenCanBorrow
 	if representativeMode == Preempt && policyPreempt == kueue.Preempt {
-		return false
+		if !needsBorrowing || policyBorrow == kueue.Borrow {
+			return false
+		}
 	}
 
 	if representativeMode == Fit && needsBorrowing && policyBorrow == kueue.Borrow {
@@ -558,6 +560,13 @@ func fitsResourceQuota(fName kueue.ResourceFlavorReference, rName corev1.Resourc
 		// ClusterQueue are preempted.
 		mode = Preempt
 	}
+	if cq.Preemption.BorrowWithinCohort != nil && cq.Preemption.BorrowWithinCohort.Policy != kueue.BorrowWithinCohortPolicyNever {
+		// when preemption with borrowing is enabled, we can succeeded admitting the
+		// workload if preemption is used.
+		if rQuota.BorrowingLimit != nil && val <= rQuota.Nominal+*rQuota.BorrowingLimit {
+			mode = Preempt
+		}
+	}
 	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
@@ -589,6 +598,9 @@ func fitsResourceQuota(fName kueue.ResourceFlavorReference, rName corev1.Resourc
 		}
 	}
 	status.append(msg)
+	if borrow := val - rQuota.Nominal; borrow > 0 && mode == Preempt {
+		return mode, borrow, &status
+	}
 	return mode, 0, &status
 }