[release-0.13] Fix SimulateWorkloadRemoval bug (#7084)

* Fix SimulateWorkloadRemoval

* Test case for 7015

* introduce cqUsage type

---------

Co-authored-by: Gabe <15304068+gabesaba@users.noreply.github.com>

Author: k8s-infra-cherrypick-robot

pkg/cache/clusterqueue_snapshot.go

@@ -72,24 +72,6 @@ func (c *ClusterQueueSnapshot) RGByResource(resource corev1.ResourceName) *Resou
 	return nil
 }
 
-// SimulateWorkloadRemoval modifies the snapshot by removing the usage
-// corresponding to the list of workloads. It returns a function which
-// can be used to restore the usage.
-func (c *ClusterQueueSnapshot) SimulateWorkloadRemoval(workloads []*workload.Info) func() {
-	usage := make([]workload.Usage, 0, len(workloads))
-	for _, w := range workloads {
-		usage = append(usage, w.Usage())
-	}
-	for _, u := range usage {
-		c.RemoveUsage(u)
-	}
-	return func() {
-		for _, u := range usage {
-			c.AddUsage(u)
-		}
-	}
-}
-
 // SimulateUsageAddition modifies the snapshot by adding usage, and
 // returns a function used to restore the usage.
 func (c *ClusterQueueSnapshot) SimulateUsageAddition(usage workload.Usage) func() {

pkg/cache/snapshot.go

@@ -52,14 +52,37 @@ func (s *Snapshot) RemoveWorkload(wl *workload.Info) {
 	cq.RemoveUsage(wl.Usage())
 }
 
-// AddWorkload adds a workload from its corresponding ClusterQueue and
+// AddWorkload adds a workload to its corresponding ClusterQueue and
 // updates resource usage.
 func (s *Snapshot) AddWorkload(wl *workload.Info) {
 	cq := s.ClusterQueue(wl.ClusterQueue)
 	cq.Workloads[workload.Key(wl.Obj)] = wl
 	cq.AddUsage(wl.Usage())
 }
 
+// SimulateWorkloadRemoval modifies the snapshot by removing the usage
+// corresponding to the list of workloads from workloads' respective
+// ClusterQueues. It returns a function which can be used to restore
+// this usage.
+func (s *Snapshot) SimulateWorkloadRemoval(workloads []*workload.Info) func() {
+	type cqUsage struct {
+		cq    kueue.ClusterQueueReference
+		usage workload.Usage
+	}
+	cqUsages := make([]cqUsage, 0, len(workloads))
+	for _, w := range workloads {
+		cqUsages = append(cqUsages, cqUsage{cq: w.ClusterQueue, usage: w.Usage()})
+	}
+	for _, cqUsage := range cqUsages {
+		s.ClusterQueue(cqUsage.cq).RemoveUsage(cqUsage.usage)
+	}
+	return func() {
+		for _, cqUsage := range cqUsages {
+			s.ClusterQueue(cqUsage.cq).AddUsage(cqUsage.usage)
+		}
+	}
+}
+
 func (s *Snapshot) Log(log logr.Logger) {
 	for name, cq := range s.ClusterQueues() {
 		cohortName := "<none>"

pkg/scheduler/preemption/preemption_oracle.go

@@ -57,7 +57,7 @@ func (p *PreemptionOracle) SimulatePreemption(log logr.Logger, cq *schdcache.Clu
 	for i, c := range candidates {
 		workloadsToPreempt[i] = c.WorkloadInfo
 	}
-	revertRemoval := cq.SimulateWorkloadRemoval(workloadsToPreempt)
+	revertRemoval := p.snapshot.SimulateWorkloadRemoval(workloadsToPreempt)
 	borrowAfterPreemptions, _ := classical.FindHeightOfLowestSubtreeThatFits(cq, fr, quantity)
 	revertRemoval()
 

pkg/scheduler/scheduler.go

@@ -283,7 +283,7 @@ func (s *Scheduler) schedule(ctx context.Context) wait.SpeedSignal {
 		}
 
 		usage := e.assignmentUsage()
-		if !fits(cq, &usage, preemptedWorkloads, e.preemptionTargets) {
+		if !fits(snapshot, cq, &usage, preemptedWorkloads, e.preemptionTargets) {
 			setSkipped(e, "Workload no longer fits after processing another workload")
 			if mode == flavorassigner.Preempt {
 				skippedPreemptions[cq.Name]++
@@ -439,12 +439,12 @@ func (s *Scheduler) nominate(ctx context.Context, workloads []workload.Info, sna
 	return entries, inadmissibleEntries
 }
 
-func fits(cq *schdcache.ClusterQueueSnapshot, usage *workload.Usage, preemptedWorkloads preemption.PreemptedWorkloads, newTargets []*preemption.Target) bool {
+func fits(snapshot *schdcache.Snapshot, cq *schdcache.ClusterQueueSnapshot, usage *workload.Usage, preemptedWorkloads preemption.PreemptedWorkloads, newTargets []*preemption.Target) bool {
 	workloads := slices.Collect(maps.Values(preemptedWorkloads))
 	for _, target := range newTargets {
 		workloads = append(workloads, target.WorkloadInfo)
 	}
-	revertUsage := cq.SimulateWorkloadRemoval(workloads)
+	revertUsage := snapshot.SimulateWorkloadRemoval(workloads)
 	defer revertUsage()
 	return cq.Fits(*usage)
 }
@@ -494,7 +494,7 @@ type partialAssignment struct {
 func (s *Scheduler) getAssignments(log logr.Logger, wl *workload.Info, snap *schdcache.Snapshot) (flavorassigner.Assignment, []*preemption.Target) {
 	assignment, targets := s.getInitialAssignments(log, wl, snap)
 	cq := snap.ClusterQueue(wl.ClusterQueue)
-	updateAssignmentForTAS(cq, wl, &assignment, targets)
+	updateAssignmentForTAS(snap, cq, wl, &assignment, targets)
 	return assignment, targets
 }
 
@@ -575,7 +575,7 @@ func (s *Scheduler) evictWorkloadAfterFailedTASReplacement(ctx context.Context,
 	return nil
 }
 
-func updateAssignmentForTAS(cq *schdcache.ClusterQueueSnapshot, wl *workload.Info, assignment *flavorassigner.Assignment, targets []*preemption.Target) {
+func updateAssignmentForTAS(snapshot *schdcache.Snapshot, cq *schdcache.ClusterQueueSnapshot, wl *workload.Info, assignment *flavorassigner.Assignment, targets []*preemption.Target) {
 	if features.Enabled(features.TopologyAwareScheduling) && assignment.RepresentativeMode() == flavorassigner.Preempt &&
 		(workload.IsExplicitlyRequestingTAS(wl.Obj.Spec.PodSets...) || cq.IsTASOnly()) && !workload.HasTopologyAssignmentWithNodeToReplace(wl.Obj) {
 		tasRequests := assignment.WorkloadsTopologyRequests(wl, cq)
@@ -585,7 +585,7 @@ func updateAssignmentForTAS(cq *schdcache.ClusterQueueSnapshot, wl *workload.Inf
 			for _, target := range targets {
 				targetWorkloads = append(targetWorkloads, target.WorkloadInfo)
 			}
-			revertUsage := cq.SimulateWorkloadRemoval(targetWorkloads)
+			revertUsage := snapshot.SimulateWorkloadRemoval(targetWorkloads)
 			tasResult = cq.FindTopologyAssignmentsForWorkload(tasRequests)
 			revertUsage()
 		} else {

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

@@ -493,6 +493,7 @@ var _ = ginkgo.Describe("Scheduler", ginkgo.Ordered, ginkgo.ContinueOnFailure, f
 	ginkgo.When("using hierarchical cohorts with several flavors", func() {
 		var (
 			cqp1 *kueue.ClusterQueue
+			cqp5 *kueue.ClusterQueue
 		)
 		ginkgo.BeforeEach(func() {
 			createCohort(testing.MakeCohort("root-cohort").Obj())
@@ -566,7 +567,7 @@ var _ = ginkgo.Describe("Scheduler", ginkgo.Ordered, ginkgo.ContinueOnFailure, f
 				Preemption(preemption).
 				Obj())
 
-			createQueue(testing.MakeClusterQueue("cq-p5").
+			cqp5 = createQueue(testing.MakeClusterQueue("cq-p5").
 				Cohort("cohort-b").
 				FairWeight(resource.MustParse("1")).
 				ResourceGroup(
@@ -614,6 +615,38 @@ var _ = ginkgo.Describe("Scheduler", ginkgo.Ordered, ginkgo.ContinueOnFailure, f
 			util.ExpectClusterQueueWeightedShareMetric(cqp1, 600)
 			expectCohortWeightedShare("cohort-a", 0)
 		})
+
+		// scenario from Kueue#7015
+		// WeightedShare(CohortA) = 0/20 * 1000 = 0
+		// WeightedShare(cq-p1)  = 12/20 * 1000 = 600
+		// WeightedShare(CohortB) = 0/20 * 1000 = 0
+		// WeightedShare(cq-p5)   = 2/20 * 1000 = 100
+		ginkgo.It("CohortB preempts and schedules in flavor which has guarantees", func() {
+			_ = features.SetEnable(features.FlavorFungibilityImplicitPreferenceDefault, true)
+			ginkgo.By("Create workload which saturate all cohort resources")
+			for range 20 {
+				createWorkload("cq-p1", "1")
+			}
+			util.ExpectReservingActiveWorkloadsMetric(cqp1, 20)
+			expectCohortWeightedShare("cohort-a", 100)
+
+			ginkgo.By("Create workloads in CohortB which will preempt CohortA")
+			createWorkload("cq-p5", "1")
+			createWorkload("cq-p5", "1")
+
+			ginkgo.By("Finish Preemption of Workloads")
+			util.FinishEvictionOfWorkloadsInCQ(ctx, k8sClient, cqp1, 2)
+
+			ginkgo.By("Check expected workloads active")
+			util.ExpectReservingActiveWorkloadsMetric(cqp1, 18)
+			util.ExpectReservingActiveWorkloadsMetric(cqp5, 2)
+
+			ginkgo.By("Expected Weighted Shares")
+			util.ExpectClusterQueueWeightedShareMetric(cqp1, 600)
+			util.ExpectClusterQueueWeightedShareMetric(cqp5, 100)
+			expectCohortWeightedShare("cohort-a", 0)
+			expectCohortWeightedShare("cohort-b", 0)
+		})
 	})
 
 	ginkgo.When("Using AdmissionFairSharing at ClusterQueue level", func() {