Add test

Signed-off-by: kerthcet <[email protected]>
kubernetes-sigs · Jan 13, 2025 · 21fa434 · 21fa434
1 parent 4f61106
commit 21fa434
Show file tree

Hide file tree

Showing 2 changed files with 113 additions and 4 deletions.
diff --git a/pkg/scheduler/scheduler_test.go b/pkg/scheduler/scheduler_test.go
@@ -4486,6 +4486,100 @@ func TestScheduleForTAS(t *testing.T) {
 				},
 			},
 		},
+		"workload with multiple PodSets requesting the same TAS flavor": {
+			nodes: []corev1.Node{
+				*testingnode.MakeNode("x1").
+					Label("tas-node", "true").
+					Label(tasRackLabel, "r1").
+					Label(corev1.LabelHostname, "x1").
+					StatusAllocatable(corev1.ResourceList{
+						corev1.ResourceCPU: resource.MustParse("8"),
+					}).
+					Ready().
+					Obj(),
+				*testingnode.MakeNode("y1").
+					Label("tas-node", "true").
+					Label(tasRackLabel, "r1").
+					Label(corev1.LabelHostname, "y1").
+					StatusAllocatable(corev1.ResourceList{
+						corev1.ResourceCPU: resource.MustParse("8"),
+					}).
+					Ready().
+					Obj(),
+			},
+			topologies:      []kueuealpha.Topology{defaultTwoLevelTopology},
+			resourceFlavors: []kueue.ResourceFlavor{defaultTASTwoLevelFlavor},
+			clusterQueues: []kueue.ClusterQueue{
+				*utiltesting.MakeClusterQueue("tas-main").
+					ResourceGroup(
+						*utiltesting.MakeFlavorQuotas("tas-default").
+							Resource(corev1.ResourceCPU, "16").Obj()).
+					Obj(),
+			},
+			workloads: []kueue.Workload{
+				*utiltesting.MakeWorkload("foo", "default").
+					Queue("tas-main").
+					PodSets(
+						*utiltesting.MakePodSet("launcher", 1).
+							RequiredTopologyRequest(corev1.LabelHostname).
+							Request(corev1.ResourceCPU, "1").
+							Obj(),
+						*utiltesting.MakePodSet("worker", 15).
+							PreferredTopologyRequest(corev1.LabelHostname).
+							Request(corev1.ResourceCPU, "1").
+							Obj()).
+					Obj(),
+			},
+			wantNewAssignments: map[string]kueue.Admission{
+				"default/foo": *utiltesting.MakeAdmission("tas-main", "launcher", "worker").
+					AssignmentWithIndex(0, corev1.ResourceCPU, "tas-default", "1000m").
+					AssignmentPodCountWithIndex(0, 1).
+					TopologyAssignmentWithIndex(0, &kueue.TopologyAssignment{
+						Levels: []string{"kubernetes.io/hostname"},
+						Domains: []kueue.TopologyDomainAssignment{
+							{
+								Count: 1,
+								Values: []string{
+									"x1",
+								},
+							},
+						},
+					}).
+					AssignmentWithIndex(1, corev1.ResourceCPU, "tas-default", "15000m").
+					AssignmentPodCountWithIndex(1, 15).
+					TopologyAssignmentWithIndex(1, &kueue.TopologyAssignment{
+						Levels: []string{"kubernetes.io/hostname"},
+						Domains: []kueue.TopologyDomainAssignment{
+							{
+								Count: 7,
+								Values: []string{
+									"x1",
+								},
+							},
+							{
+								Count: 8,
+								Values: []string{
+									"y1",
+								},
+							},
+						},
+					}).
+					Obj(),
+			},
+			eventCmpOpts: []cmp.Option{eventIgnoreMessage},
+			wantEvents: []utiltesting.EventRecord{
+				{
+					Key:       types.NamespacedName{Namespace: "default", Name: "foo"},
+					Reason:    "QuotaReserved",
+					EventType: corev1.EventTypeNormal,
+				},
+				{
+					Key:       types.NamespacedName{Namespace: "default", Name: "foo"},
+					Reason:    "Admitted",
+					EventType: corev1.EventTypeNormal,
+				},
+			},
+		},
 		"scheduling workload when the node for another admitted workload is deleted": {
 			// Here we have the "bar-admitted" workload, which is admitted and
 			// is using the "x1" node, which is deleted. Still, we have the y1

diff --git a/pkg/util/testing/wrappers.go b/pkg/util/testing/wrappers.go
@@ -548,18 +548,33 @@ func (w *AdmissionWrapper) Obj() *kueue.Admission {
 }
 
 func (w *AdmissionWrapper) Assignment(r corev1.ResourceName, f kueue.ResourceFlavorReference, value string) *AdmissionWrapper {
-	w.PodSetAssignments[0].Flavors[r] = f
-	w.PodSetAssignments[0].ResourceUsage[r] = resource.MustParse(value)
+	w.AssignmentWithIndex(0, r, f, value)
 	return w
 }
 
 func (w *AdmissionWrapper) AssignmentPodCount(value int32) *AdmissionWrapper {
-	w.PodSetAssignments[0].Count = ptr.To(value)
+	w.AssignmentPodCountWithIndex(0, value)
 	return w
 }
 
 func (w *AdmissionWrapper) TopologyAssignment(ts *kueue.TopologyAssignment) *AdmissionWrapper {
-	w.PodSetAssignments[0].TopologyAssignment = ts
+	w.TopologyAssignmentWithIndex(0, ts)
+	return w
+}
+
+func (w *AdmissionWrapper) AssignmentWithIndex(index int32, r corev1.ResourceName, f kueue.ResourceFlavorReference, value string) *AdmissionWrapper {
+	w.PodSetAssignments[index].Flavors[r] = f
+	w.PodSetAssignments[index].ResourceUsage[r] = resource.MustParse(value)
+	return w
+}
+
+func (w *AdmissionWrapper) AssignmentPodCountWithIndex(index, value int32) *AdmissionWrapper {
+	w.PodSetAssignments[index].Count = ptr.To(value)
+	return w
+}
+
+func (w *AdmissionWrapper) TopologyAssignmentWithIndex(index int32, ts *kueue.TopologyAssignment) *AdmissionWrapper {
+	w.PodSetAssignments[index].TopologyAssignment = ts
 	return w
 }