[PartialAdmission] Fix preemption while partially admitting.

Add an unit test for preemption while partially admitting.
kubernetes-sigs · Aug 12, 2024 · 8530e5a · 8530e5a
1 parent 57f8a72
commit 8530e5a
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 3 deletions.
diff --git a/pkg/scheduler/flavorassigner/flavorassigner.go b/pkg/scheduler/flavorassigner/flavorassigner.go
@@ -110,7 +110,12 @@ func (a *Assignment) TotalRequestsFor(wl *workload.Info) resources.FlavorResourc
 	usage := make(resources.FlavorResourceQuantities)
 	for i, ps := range wl.TotalRequests {
 		for res, q := range ps.Requests {
-			flv := a.PodSets[i].Flavors[res].Name
+			aps := a.PodSets[i]
+			flv := aps.Flavors[res].Name
+			// in case of partial admission scale down the quantity
+			if aps.Count != ps.Count {
+				q = q * int64(aps.Count) / int64(ps.Count)
+			}
 			usage[resources.FlavorResource{Flavor: flv, Resource: res}] += q
 		}
 	}

diff --git a/pkg/scheduler/scheduler.go b/pkg/scheduler/scheduler.go
@@ -471,9 +471,12 @@ func (s *Scheduler) getAssignments(log logr.Logger, wl *workload.Info, snap *cac
 	if wl.CanBePartiallyAdmitted() {
 		reducer := flavorassigner.NewPodSetReducer(wl.Obj.Spec.PodSets, func(nextCounts []int32) (*partialAssignment, bool) {
 			assignment := flvAssigner.Assign(log, nextCounts)
-			if assignment.RepresentativeMode() == flavorassigner.Fit {
+			mode := assignment.RepresentativeMode()
+			if mode == flavorassigner.Fit {
 				return &partialAssignment{assignment: assignment}, true
-			} else if assignment.RepresentativeMode() == flavorassigner.Preempt {
+			}
+
+			if mode == flavorassigner.Preempt {
 				preemptionTargets := s.preemptor.GetTargets(log, *wl, assignment, snap)
 				if len(preemptionTargets) > 0 {
 					return &partialAssignment{assignment: assignment, preemptionTargets: preemptionTargets}, true

diff --git a/pkg/scheduler/scheduler_test.go b/pkg/scheduler/scheduler_test.go
@@ -1231,6 +1231,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").