[BestEffortFIFO] Implement Sticky ClusterQueue Head Policy

Author: gabesaba

pkg/cache/queue/cluster_queue.go

@@ -56,6 +56,22 @@ var (
 	realClock = clock.RealClock{}
 )
 
+type stickyWorkload struct {
+	workloadName workload.Reference
+}
+
+func (s *stickyWorkload) matches(workload workload.Reference) bool {
+	return s.workloadName == workload
+}
+
+func (s *stickyWorkload) clear() {
+	s.workloadName = ""
+}
+
+func (s *stickyWorkload) set(workload workload.Reference) {
+	s.workloadName = workload
+}
+
 type ClusterQueue struct {
 	hierarchy.ClusterQueue[*cohort]
 	name              kueue.ClusterQueueReference
@@ -90,6 +106,8 @@ type ClusterQueue struct {
 
 	afsEntryPenalties         *utilmaps.SyncMap[utilqueue.LocalQueueReference, corev1.ResourceList]
 	localQueuesInClusterQueue map[utilqueue.LocalQueueReference]bool
+
+	sw *stickyWorkload
 }
 
 func (c *ClusterQueue) GetName() kueue.ClusterQueueReference {
@@ -111,7 +129,8 @@ func newClusterQueue(ctx context.Context, client client.Client, cq *kueue.Cluste
 }
 
 func newClusterQueueImpl(ctx context.Context, client client.Client, wo workload.Ordering, clock clock.Clock, fsResWeights map[corev1.ResourceName]float64, enableAdmissionFs bool, afsEntryPenalties *utilmaps.SyncMap[utilqueue.LocalQueueReference, corev1.ResourceList]) *ClusterQueue {
-	lessFunc := queueOrderingFunc(ctx, client, wo, fsResWeights, enableAdmissionFs, afsEntryPenalties)
+	sw := stickyWorkload{}
+	lessFunc := queueOrderingFunc(ctx, client, wo, fsResWeights, enableAdmissionFs, afsEntryPenalties, &sw)
 	return &ClusterQueue{
 		heap:                      *heap.New(workloadKey, lessFunc),
 		inadmissibleWorkloads:     make(map[workload.Reference]*workload.Info),
@@ -121,6 +140,7 @@ func newClusterQueueImpl(ctx context.Context, client client.Client, wo workload.
 		clock:                     clock,
 		afsEntryPenalties:         afsEntryPenalties,
 		localQueuesInClusterQueue: make(map[utilqueue.LocalQueueReference]bool),
+		sw:                        &sw,
 	}
 }
 
@@ -342,6 +362,7 @@ func (c *ClusterQueue) Pop() *workload.Info {
 	}
 
 	c.popCycle++
+	c.sw.clear()
 	if c.heap.Len() == 0 {
 		c.inflight = nil
 		return nil
@@ -453,6 +474,10 @@ func (c *ClusterQueue) Active() bool {
 // The workload should not be reinserted if it's already in the ClusterQueue.
 // Returns true if the workload was inserted.
 func (c *ClusterQueue) RequeueIfNotPresent(wInfo *workload.Info, reason RequeueReason) bool {
+	if reason == RequeueReasonPendingPreemption && c.queueingStrategy == kueue.BestEffortFIFO {
+		c.sw.set(workload.Key(wInfo.Obj))
+	}
+
 	if c.queueingStrategy == kueue.StrictFIFO {
 		return c.requeueIfNotPresent(wInfo, reason != RequeueReasonNamespaceMismatch)
 	}
@@ -463,7 +488,7 @@ func (c *ClusterQueue) RequeueIfNotPresent(wInfo *workload.Info, reason RequeueR
 // to sort workloads. The function sorts workloads based on their priority.
 // When priorities are equal, it uses the workload's creation or eviction
 // time.
-func queueOrderingFunc(ctx context.Context, c client.Client, wo workload.Ordering, fsResWeights map[corev1.ResourceName]float64, enableAdmissionFs bool, afsEntryPenalties *utilmaps.SyncMap[utilqueue.LocalQueueReference, corev1.ResourceList]) func(a, b *workload.Info) bool {
+func queueOrderingFunc(ctx context.Context, c client.Client, wo workload.Ordering, fsResWeights map[corev1.ResourceName]float64, enableAdmissionFs bool, afsEntryPenalties *utilmaps.SyncMap[utilqueue.LocalQueueReference, corev1.ResourceList], sticky *stickyWorkload) func(a, b *workload.Info) bool {
 	log := ctrl.LoggerFrom(ctx)
 	return func(a, b *workload.Info) bool {
 		if enableAdmissionFs {
@@ -482,6 +507,14 @@ func queueOrderingFunc(ctx context.Context, c client.Client, wo workload.Orderin
 				}
 			}
 		}
+
+		if sticky.matches(workload.Key(a.Obj)) {
+			return true
+		}
+		if sticky.matches(workload.Key(b.Obj)) {
+			return false
+		}
+
 		p1 := utilpriority.Priority(a.Obj)
 		p2 := utilpriority.Priority(b.Obj)
 

test/integration/singlecluster/scheduler/fairsharing/fair_sharing_test.go

@@ -649,6 +649,218 @@ var _ = ginkgo.Describe("Scheduler", ginkgo.Ordered, ginkgo.ContinueOnFailure, f
 		})
 	})
 
+	// kueue#7101
+	ginkgo.When("ClusterQueue head is ineligible for admission due to DominantResourceShare", func() {
+		var (
+			cqp1 *kueue.ClusterQueue
+			cqp2 *kueue.ClusterQueue
+		)
+		ginkgo.BeforeEach(func() {
+			createCohort(testing.MakeCohort("cohort-a").
+				Parent("root-cohort").
+				FairWeight(resource.MustParse("1")).
+				ResourceGroup(
+					*testing.MakeFlavorQuotas("flavor1").Resource(corev1.ResourceCPU, "9").Obj(),
+				).
+				Obj())
+
+			fungibility := kueue.FlavorFungibility{
+				WhenCanBorrow:  kueue.TryNextFlavor,
+				WhenCanPreempt: kueue.TryNextFlavor,
+			}
+			preemption := kueue.ClusterQueuePreemption{
+				ReclaimWithinCohort: kueue.PreemptionPolicyAny,
+				WithinClusterQueue:  kueue.PreemptionPolicyLowerPriority,
+			}
+
+			cqp1 = createQueue(testing.MakeClusterQueue("cq-p1").
+				Cohort("cohort-a").
+				FairWeight(resource.MustParse("1")).
+				ResourceGroup(
+					*testing.MakeFlavorQuotas("flavor1").Resource(corev1.ResourceCPU, "0").Obj(),
+				).
+				FlavorFungibility(fungibility).
+				Preemption(preemption).
+				Obj())
+
+			cqp2 = createQueue(testing.MakeClusterQueue("cq-p2").
+				Cohort("cohort-a").
+				FairWeight(resource.MustParse("1")).
+				ResourceGroup(
+					*testing.MakeFlavorQuotas("flavor1").Resource(corev1.ResourceCPU, "0").Obj(),
+				).
+				FlavorFungibility(fungibility).
+				Preemption(preemption).
+				Obj())
+			_ = features.SetEnable(features.FlavorFungibilityImplicitPreferenceDefault, true)
+		})
+		ginkgo.AfterEach(func() {
+			_ = features.SetEnable(features.FlavorFungibilityImplicitPreferenceDefault, false)
+		})
+
+		ginkgo.It("workload of size 5 preempts using LessThanInitialShare policy and admits", func() {
+			ginkgo.By("Create workloads in queue1")
+			for range 4 {
+				createWorkload("cq-p1", "2")
+			}
+			util.ExpectReservingActiveWorkloadsMetric(cqp1, 4)
+
+			ginkgo.By("Create workload in queue2")
+			createWorkload("cq-p2", "5")
+
+			ginkgo.By("Complete preemption")
+			util.FinishEvictionOfWorkloadsInCQ(ctx, k8sClient, cqp1, 2)
+
+			ginkgo.By("Expect active workloads")
+			util.ExpectReservingActiveWorkloadsMetric(cqp1, 2)
+			util.ExpectReservingActiveWorkloadsMetric(cqp2, 1)
+			util.ExpectClusterQueueWeightedShareMetric(cqp1, 445)
+			util.ExpectClusterQueueWeightedShareMetric(cqp2, 556)
+		})
+
+		ginkgo.It("workload of size 5 admits with inadmissible higher priority workload at ClusterQueue head", func() {
+			ginkgo.By("Create workloads in queue1")
+			for range 4 {
+				createWorkload("cq-p1", "2")
+			}
+			util.ExpectReservingActiveWorkloadsMetric(cqp1, 4)
+
+			ginkgo.By("Create workloads in queue2")
+			createWorkloadWithPriority("cq-p2", "6", 999)
+
+			ginkgo.By("Verify doesn't admit")
+			util.ExpectReservingActiveWorkloadsMetric(cqp2, 0)
+
+			ginkgo.By("Create admissible workload in queue2")
+			createWorkloadWithPriority("cq-p2", "5", 0)
+
+			ginkgo.By("Complete preemption")
+			util.FinishEvictionOfWorkloadsInCQ(ctx, k8sClient, cqp1, 2)
+
+			ginkgo.By("Expect active workloads")
+			util.ExpectReservingActiveWorkloadsMetric(cqp1, 2)
+			util.ExpectReservingActiveWorkloadsMetric(cqp2, 1)
+			util.ExpectClusterQueueWeightedShareMetric(cqp1, 445)
+			util.ExpectClusterQueueWeightedShareMetric(cqp2, 556)
+		})
+
+		ginkgo.It("workload of size 4 admits with inadmissible higher priority workload at ClusterQueue head", func() {
+			ginkgo.By("Create workloads in queue1")
+			for range 4 {
+				createWorkload("cq-p1", "2")
+			}
+			util.ExpectReservingActiveWorkloadsMetric(cqp1, 4)
+
+			ginkgo.By("Create workload in queue2")
+			createWorkloadWithPriority("cq-p2", "6", 999)
+
+			ginkgo.By("Verify doesn't admit")
+			util.ExpectReservingActiveWorkloadsMetric(cqp2, 0)
+
+			ginkgo.By("Create admissible workload in queue2")
+			createWorkloadWithPriority("cq-p2", "4", 0)
+
+			ginkgo.By("Complete preemption")
+			util.FinishEvictionOfWorkloadsInCQ(ctx, k8sClient, cqp1, 2)
+
+			ginkgo.By("Expect active workloads")
+			util.ExpectReservingActiveWorkloadsMetric(cqp1, 2)
+			util.ExpectReservingActiveWorkloadsMetric(cqp2, 1)
+			util.ExpectClusterQueueWeightedShareMetric(cqp1, 445)
+			util.ExpectClusterQueueWeightedShareMetric(cqp2, 445)
+		})
+
+		ginkgo.It("workload admits when several higher priority blocking workloads in front", func() {
+			ginkgo.By("Create workloads in queue1")
+			for range 4 {
+				createWorkload("cq-p1", "2")
+			}
+			util.ExpectReservingActiveWorkloadsMetric(cqp1, 4)
+
+			ginkgo.By("Create workloads in queue2")
+			createWorkloadWithPriority("cq-p2", "7", 999)
+			createWorkloadWithPriority("cq-p2", "6", 999)
+
+			ginkgo.By("Verify don't admit")
+			util.ExpectReservingActiveWorkloadsMetric(cqp2, 0)
+
+			ginkgo.By("Create admissible workload in queue2")
+			createWorkloadWithPriority("cq-p2", "5", 0)
+
+			ginkgo.By("Complete preemption")
+			util.FinishEvictionOfWorkloadsInCQ(ctx, k8sClient, cqp1, 2)
+
+			ginkgo.By("Expect active workloads")
+			util.ExpectReservingActiveWorkloadsMetric(cqp1, 2)
+			util.ExpectReservingActiveWorkloadsMetric(cqp2, 1)
+			util.ExpectClusterQueueWeightedShareMetric(cqp1, 445)
+			util.ExpectClusterQueueWeightedShareMetric(cqp2, 556)
+		})
+
+		// kueue#6929
+		ginkgo.When("ClusterQueue head has inadmissible workload", func() {
+			var (
+				cq1 *kueue.ClusterQueue
+				cq2 *kueue.ClusterQueue
+			)
+			ginkgo.BeforeEach(func() {
+				fungibility := kueue.FlavorFungibility{
+					WhenCanBorrow:  kueue.TryNextFlavor,
+					WhenCanPreempt: kueue.TryNextFlavor,
+				}
+				preemption := kueue.ClusterQueuePreemption{
+					ReclaimWithinCohort: kueue.PreemptionPolicyAny,
+					WithinClusterQueue:  kueue.PreemptionPolicyLowerPriority,
+				}
+
+				cq1 = createQueue(testing.MakeClusterQueue("cq1").
+					Cohort("root").
+					FairWeight(resource.MustParse("1")).
+					ResourceGroup(
+						*testing.MakeFlavorQuotas("flavor1").Resource(corev1.ResourceCPU, "3").Obj(),
+					).
+					FlavorFungibility(fungibility).
+					Preemption(preemption).
+					Obj())
+
+				cq2 = createQueue(testing.MakeClusterQueue("cq2").
+					Cohort("root").
+					FairWeight(resource.MustParse("1")).
+					ResourceGroup(
+						*testing.MakeFlavorQuotas("flavor1").Resource(corev1.ResourceCPU, "0").Obj(),
+					).
+					FlavorFungibility(fungibility).
+					Preemption(preemption).
+					Obj())
+			})
+
+			ginkgo.It("workload which fits behind ClusterQueue head is able to admit", func() {
+				ginkgo.By("Creating borrowing workloads in queue2")
+				createWorkload("cq2", "1")
+				createWorkload("cq2", "1")
+				util.ExpectReservingActiveWorkloadsMetric(cq2, 2)
+
+				ginkgo.By("Create inadmissible workload in queue2")
+				createWorkloadWithPriority("cq1", "4", 999)
+
+				ginkgo.By("Verify doesn't admit")
+				util.ExpectReservingActiveWorkloadsMetric(cq1, 0)
+
+				ginkgo.By("Create admissible workload in queue2")
+				createWorkloadWithPriority("cq1", "3", 0)
+
+				ginkgo.By("Complete preemption")
+				util.FinishEvictionOfWorkloadsInCQ(ctx, k8sClient, cq2, 2)
+
+				ginkgo.By("Expect active workloads")
+				util.ExpectReservingActiveWorkloadsMetric(cq1, 1)
+				util.ExpectReservingActiveWorkloadsMetric(cq2, 0)
+				util.ExpectClusterQueueWeightedShareMetric(cq1, 0)
+				util.ExpectClusterQueueWeightedShareMetric(cq2, 0)
+			})
+		})
+	})
+
 	ginkgo.When("Using AdmissionFairSharing at ClusterQueue level", func() {
 		var (
 			cq1 *kueue.ClusterQueue