[release-0.8] [PartialAdmission] Fix preemption while partially admitting.

pkg/scheduler/flavorassigner/flavorassigner.go

@@ -106,11 +106,18 @@ func (a *Assignment) ToAPI() []kueue.PodSetAssignment {
 	return psFlavors
 }
 
+// TotalRequestsFor - returns the total quota needs of the wl, taking into account the potential
+// scaling needed in case of partial admission.
 func (a *Assignment) TotalRequestsFor(wl *workload.Info) resources.FlavorResourceQuantities {
 	usage := make(resources.FlavorResourceQuantities)
 	for i, ps := range wl.TotalRequests {
+		// in case of partial admission scale down the quantity
+		aps := a.PodSets[i]
+		if aps.Count != ps.Count {
+			ps = *ps.ScaledTo(aps.Count)
+		}
 		for res, q := range ps.Requests {
-			flv := a.PodSets[i].Flavors[res].Name
+			flv := aps.Flavors[res].Name
 			usage[resources.FlavorResource{Flavor: flv, Resource: res}] += q
 		}
 	}

pkg/scheduler/preemption/preemption_test.go

@@ -910,6 +910,7 @@ func TestPreemption(t *testing.T) {
 								Mode: flavorassigner.Preempt,
 							},
 						},
+						Count: 1,
 					},
 					{
 						Name: "workers",
@@ -919,6 +920,7 @@ func TestPreemption(t *testing.T) {
 								Mode: flavorassigner.Preempt,
 							},
 						},
+						Count: 2,
 					},
 				},
 			},
@@ -2044,6 +2046,7 @@ func singlePodSetAssignment(assignments flavorassigner.ResourceAssignment) flavo
 		PodSets: []flavorassigner.PodSetAssignment{{
 			Name:    kueue.DefaultPodSetName,
 			Flavors: assignments,
+			Count:   1,
 		}},
 	}
 }

pkg/scheduler/scheduler_test.go

@@ -1245,6 +1245,46 @@ func TestSchedule(t *testing.T) {
 				"eng-beta": {"eng-beta/new"},
 			},
 		},
+		"partial admission single variable pod set, preempt with partial admission": {
+			workloads: []kueue.Workload{
+				*utiltesting.MakeWorkload("new", "eng-beta").
+					Queue("main").
+					Priority(4).
+					PodSets(*utiltesting.MakePodSet("one", 30).
+						SetMinimumCount(10).
+						Request("example.com/gpu", "1").
+						Obj()).
+					Obj(),
+				*utiltesting.MakeWorkload("old", "eng-beta").
+					Priority(-4).
+					PodSets(*utiltesting.MakePodSet("one", 10).
+						Request("example.com/gpu", "1").
+						Obj()).
+					ReserveQuota(utiltesting.MakeAdmission("eng-beta", "one").Assignment("example.com/gpu", "model-a", "10").AssignmentPodCount(10).Obj()).
+					Obj(),
+			},
+			wantAssignments: map[string]kueue.Admission{
+				"eng-beta/old": {
+					ClusterQueue: "eng-beta",
+					PodSetAssignments: []kueue.PodSetAssignment{
+						{
+							Name: "one",
+							Flavors: map[corev1.ResourceName]kueue.ResourceFlavorReference{
+								"example.com/gpu": "model-a",
+							},
+							ResourceUsage: corev1.ResourceList{
+								"example.com/gpu": resource.MustParse("10"),
+							},
+							Count: ptr.To[int32](10),
+						},
+					},
+				},
+			},
+			wantPreempted: sets.New("eng-beta/old"),
+			wantLeft: map[string][]string{
+				"eng-beta": {"eng-beta/new"},
+			},
+		},
 		"partial admission multiple variable pod sets": {
 			workloads: []kueue.Workload{
 				*utiltesting.MakeWorkload("new", "sales").

pkg/workload/workload.go

@@ -170,9 +170,11 @@ func (psr *PodSetResources) ScaledTo(newCount int32) *PodSetResources {
 		Flavors:  maps.Clone(psr.Flavors),
 	}
 
-	scaleDown(ret.Requests, int64(ret.Count))
-	scaleUp(ret.Requests, int64(newCount))
-	ret.Count = newCount
+	if psr.Count != 0 && psr.Count != newCount {
+		scaleDown(ret.Requests, int64(ret.Count))
+		scaleUp(ret.Requests, int64(newCount))
+		ret.Count = newCount
+	}
 	return ret
 }