Skip to content

Commit 83a5be7

Browse files
romanbaronromanbaron
authored
Added allocation test with subgroups e2e test (#434)
1 parent 589a371 commit 83a5be7

File tree

3 files changed

+332
-9
lines changed

3 files changed

+332
-9
lines changed

‎pkg/scheduler/actions/allocate/allocate_subgroups_test.go‎

Lines changed: 135 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,20 +4,22 @@
44
package allocate_test
55

66
import (
7-
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/pod_status"
8-
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/podgroup_info"
9-
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/constants"
10-
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/test_utils/jobs_fake"
11-
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/test_utils/nodes_fake"
12-
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/test_utils/tasks_fake"
7+
"testing"
8+
139
"k8s.io/utils/pointer"
1410
"k8s.io/utils/ptr"
15-
"testing"
11+
12+
. "go.uber.org/mock/gomock"
1613

1714
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/actions/allocate"
1815
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/actions/integration_tests/integration_tests_utils"
16+
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/pod_status"
17+
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/podgroup_info"
18+
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/constants"
1919
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/test_utils"
20-
. "go.uber.org/mock/gomock"
20+
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/test_utils/jobs_fake"
21+
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/test_utils/nodes_fake"
22+
"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/test_utils/tasks_fake"
2123
)
2224

2325
func TestHandleSubGroupsAllocation(t *testing.T) {
@@ -533,5 +535,130 @@ func getAllocationSubGroupsTestsMetadata() []integration_tests_utils.TestTopolog
533535
},
534536
},
535537
},
538+
{
539+
TestTopologyBasic: test_utils.TestTopologyBasic{
540+
Name: "Allocate multiple jobs with SubGroups",
541+
Jobs: []*jobs_fake.TestJobBasic{
542+
{
543+
Name: "pending_job0",
544+
QueueName: "queue0",
545+
Priority: constants.PriorityTrainNumber,
546+
SubGroups: map[string]*podgroup_info.SubGroupInfo{
547+
"sub0": podgroup_info.NewSubGroupInfo("sub0", 1),
548+
"sub1": podgroup_info.NewSubGroupInfo("sub1", 1),
549+
},
550+
Tasks: []*tasks_fake.TestTaskBasic{
551+
{
552+
State: pod_status.Pending,
553+
SubGroupName: "sub0",
554+
RequiredGPUs: ptr.To(int64(1)),
555+
},
556+
{
557+
State: pod_status.Pending,
558+
SubGroupName: "sub0",
559+
RequiredGPUs: ptr.To(int64(1)),
560+
},
561+
{
562+
State: pod_status.Pending,
563+
SubGroupName: "sub1",
564+
RequiredGPUs: ptr.To(int64(1)),
565+
},
566+
{
567+
State: pod_status.Pending,
568+
SubGroupName: "sub1",
569+
RequiredGPUs: ptr.To(int64(1)),
570+
},
571+
},
572+
MinAvailable: pointer.Int32(2),
573+
},
574+
{
575+
Name: "pending_job1",
576+
QueueName: "queue0",
577+
Priority: constants.PriorityTrainNumber,
578+
SubGroups: map[string]*podgroup_info.SubGroupInfo{
579+
"sub0": podgroup_info.NewSubGroupInfo("sub0", 1),
580+
"sub1": podgroup_info.NewSubGroupInfo("sub1", 1),
581+
},
582+
Tasks: []*tasks_fake.TestTaskBasic{
583+
{
584+
State: pod_status.Pending,
585+
SubGroupName: "sub0",
586+
RequiredGPUs: ptr.To(int64(1)),
587+
},
588+
{
589+
State: pod_status.Pending,
590+
SubGroupName: "sub0",
591+
RequiredGPUs: ptr.To(int64(1)),
592+
},
593+
{
594+
State: pod_status.Pending,
595+
SubGroupName: "sub1",
596+
RequiredGPUs: ptr.To(int64(1)),
597+
},
598+
{
599+
State: pod_status.Pending,
600+
SubGroupName: "sub1",
601+
RequiredGPUs: ptr.To(int64(1)),
602+
},
603+
},
604+
MinAvailable: pointer.Int32(2),
605+
},
606+
},
607+
Nodes: map[string]nodes_fake.TestNodeBasic{
608+
"node0": {
609+
GPUs: 4,
610+
},
611+
},
612+
Queues: []test_utils.TestQueueBasic{
613+
{
614+
Name: "queue0",
615+
DeservedGPUs: 1,
616+
},
617+
},
618+
Mocks: &test_utils.TestMock{
619+
CacheRequirements: &test_utils.CacheMocking{
620+
NumberOfCacheBinds: 4,
621+
},
622+
},
623+
TaskExpectedResults: map[string]test_utils.TestExpectedResultBasic{
624+
"pending_job0-0": {
625+
GPUsRequired: 1,
626+
Status: pod_status.Binding,
627+
NodeName: "node0",
628+
},
629+
"pending_job0-1": {
630+
GPUsRequired: 1,
631+
Status: pod_status.Pending,
632+
},
633+
"pending_job0-2": {
634+
GPUsRequired: 1,
635+
Status: pod_status.Binding,
636+
NodeName: "node0",
637+
},
638+
"pending_job0-3": {
639+
GPUsRequired: 1,
640+
Status: pod_status.Pending,
641+
},
642+
"pending_job1-0": {
643+
GPUsRequired: 1,
644+
Status: pod_status.Binding,
645+
NodeName: "node0",
646+
},
647+
"pending_job1-1": {
648+
GPUsRequired: 1,
649+
Status: pod_status.Pending,
650+
},
651+
"pending_job1-2": {
652+
GPUsRequired: 1,
653+
Status: pod_status.Binding,
654+
NodeName: "node0",
655+
},
656+
"pending_job1-3": {
657+
GPUsRequired: 1,
658+
Status: pod_status.Pending,
659+
},
660+
},
661+
},
662+
},
536663
}
537664
}

‎test/e2e/suites/allocate/elastic/allocate_suite_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ import (
1313
. "github.com/onsi/gomega"
1414
)
1515

16-
func TestPriority(t *testing.T) {
16+
func TestAllocate(t *testing.T) {
1717
utils.SetLogger()
1818
RegisterFailHandler(Fail)
1919
RunSpecs(t, "Elastic allocation Suite")

‎test/e2e/suites/allocate/subgroups/subgroups_test.go‎

Lines changed: 196 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,196 @@
1+
/*
2+
Copyright 2025 NVIDIA CORPORATION
3+
SPDX-License-Identifier: Apache-2.0
4+
*/
5+
package subgroups
6+
7+
import (
8+
"context"
9+
"testing"
10+
11+
v1 "k8s.io/api/core/v1"
12+
"k8s.io/apimachinery/pkg/api/resource"
13+
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
14+
"k8s.io/client-go/kubernetes"
15+
16+
v2 "github.com/NVIDIA/KAI-scheduler/pkg/apis/scheduling/v2"
17+
schedulingv2alpha2 "github.com/NVIDIA/KAI-scheduler/pkg/apis/scheduling/v2alpha2"
18+
commonconsts "github.com/NVIDIA/KAI-scheduler/pkg/common/constants"
19+
testcontext "github.com/NVIDIA/KAI-scheduler/test/e2e/modules/context"
20+
"github.com/NVIDIA/KAI-scheduler/test/e2e/modules/resources/capacity"
21+
"github.com/NVIDIA/KAI-scheduler/test/e2e/modules/resources/rd"
22+
"github.com/NVIDIA/KAI-scheduler/test/e2e/modules/resources/rd/pod_group"
23+
"github.com/NVIDIA/KAI-scheduler/test/e2e/modules/resources/rd/queue"
24+
"github.com/NVIDIA/KAI-scheduler/test/e2e/modules/utils"
25+
"github.com/NVIDIA/KAI-scheduler/test/e2e/modules/wait"
26+
27+
. "github.com/onsi/ginkgo/v2"
28+
. "github.com/onsi/gomega"
29+
)
30+
31+
func TestSubGroups(t *testing.T) {
32+
utils.SetLogger()
33+
RegisterFailHandler(Fail)
34+
RunSpecs(t, "SubGroups Allocation Suite")
35+
}
36+
37+
var _ = Describe("Allocation scenario with subgroups", Ordered, func() {
38+
var (
39+
testCtx *testcontext.TestContext
40+
)
41+
42+
BeforeAll(func(ctx context.Context) {
43+
testCtx = testcontext.GetConnectivity(ctx, Default)
44+
45+
parentQueue := queue.CreateQueueObject(utils.GenerateRandomK8sName(10), "")
46+
childQueue := queue.CreateQueueObject(utils.GenerateRandomK8sName(10), parentQueue.Name)
47+
childQueue.Spec.Resources.CPU.Quota = 600
48+
childQueue.Spec.Resources.CPU.Limit = 600
49+
testCtx.InitQueues([]*v2.Queue{childQueue, parentQueue})
50+
51+
capacity.SkipIfInsufficientClusterTopologyResources(testCtx.KubeClientset, []capacity.ResourceList{
52+
{
53+
Cpu: resource.MustParse("600m"),
54+
PodCount: 6,
55+
},
56+
})
57+
})
58+
59+
AfterAll(func(ctx context.Context) {
60+
err := rd.DeleteAllE2EPriorityClasses(ctx, testCtx.ControllerClient)
61+
Expect(err).To(Succeed())
62+
testCtx.ClusterCleanup(ctx)
63+
})
64+
65+
AfterEach(func(ctx context.Context) {
66+
testCtx.TestContextCleanup(ctx)
67+
})
68+
69+
It("Partial allocation", func(ctx context.Context) {
70+
pgName := utils.GenerateRandomK8sName(10)
71+
subGroup1Pods := createSubGroupPods(ctx, testCtx.KubeClientset, testCtx.Queues[0], pgName, "sub-1", 5)
72+
subGroup2Pods := createSubGroupPods(ctx, testCtx.KubeClientset, testCtx.Queues[0], pgName, "sub-2", 5)
73+
74+
namespace := queue.GetConnectedNamespaceToQueue(testCtx.Queues[0])
75+
podGroup := pod_group.Create(namespace, pgName, testCtx.Queues[0].Name)
76+
podGroup.Spec.MinMember = 6
77+
podGroup.Spec.SubGroups = []schedulingv2alpha2.SubGroup{
78+
{Name: "sub-1", MinMember: 3},
79+
{Name: "sub-2", MinMember: 3},
80+
}
81+
_, err := testCtx.KubeAiSchedClientset.SchedulingV2alpha2().PodGroups(namespace).Create(ctx,
82+
podGroup, metav1.CreateOptions{})
83+
Expect(err).To(Succeed())
84+
85+
wait.ForAtLeastNPodsScheduled(ctx, testCtx.ControllerClient, podGroup.Namespace, subGroup1Pods, 3)
86+
wait.ForAtLeastNPodsUnschedulable(ctx, testCtx.ControllerClient, podGroup.Namespace, subGroup1Pods, 2)
87+
wait.ForAtLeastNPodsScheduled(ctx, testCtx.ControllerClient, podGroup.Namespace, subGroup2Pods, 3)
88+
wait.ForAtLeastNPodsUnschedulable(ctx, testCtx.ControllerClient, podGroup.Namespace, subGroup2Pods, 2)
89+
})
90+
91+
It("Balance 2 jobs with subgroups", func(ctx context.Context) {
92+
namespace := queue.GetConnectedNamespaceToQueue(testCtx.Queues[0])
93+
94+
pg1Name := utils.GenerateRandomK8sName(10)
95+
pg1SubGroup1Pods := createSubGroupPods(ctx, testCtx.KubeClientset, testCtx.Queues[0], pg1Name, "sub-1", 3)
96+
pg1SubGroup2Pods := createSubGroupPods(ctx, testCtx.KubeClientset, testCtx.Queues[0], pg1Name, "sub-2", 3)
97+
pg1 := pod_group.Create(namespace, pg1Name, testCtx.Queues[0].Name)
98+
pg1.Spec.MinMember = 2
99+
pg1.Spec.SubGroups = []schedulingv2alpha2.SubGroup{
100+
{Name: "sub-1", MinMember: 1},
101+
{Name: "sub-2", MinMember: 1},
102+
}
103+
104+
pg2Name := utils.GenerateRandomK8sName(10)
105+
pg2SubGroup1Pods := createSubGroupPods(ctx, testCtx.KubeClientset, testCtx.Queues[0], pg2Name, "sub-1", 3)
106+
pg2SubGroup2Pods := createSubGroupPods(ctx, testCtx.KubeClientset, testCtx.Queues[0], pg2Name, "sub-2", 3)
107+
pg2 := pod_group.Create(namespace, pg2Name, testCtx.Queues[0].Name)
108+
pg2.Spec.MinMember = 2
109+
pg2.Spec.SubGroups = []schedulingv2alpha2.SubGroup{
110+
{Name: "sub-1", MinMember: 1},
111+
{Name: "sub-2", MinMember: 1},
112+
}
113+
114+
_, err := testCtx.KubeAiSchedClientset.SchedulingV2alpha2().PodGroups(namespace).Create(ctx,
115+
pg1, metav1.CreateOptions{})
116+
Expect(err).To(Succeed())
117+
_, err = testCtx.KubeAiSchedClientset.SchedulingV2alpha2().PodGroups(namespace).Create(ctx,
118+
pg2, metav1.CreateOptions{})
119+
Expect(err).To(Succeed())
120+
121+
pg1Pods := append(pg1SubGroup1Pods, pg1SubGroup2Pods...)
122+
wait.ForAtLeastNPodsScheduled(ctx, testCtx.ControllerClient, namespace, pg1SubGroup1Pods, 1)
123+
wait.ForAtLeastNPodsScheduled(ctx, testCtx.ControllerClient, namespace, pg1SubGroup2Pods, 1)
124+
wait.ForAtLeastNPodsScheduled(ctx, testCtx.ControllerClient, namespace, pg1Pods, 3)
125+
wait.ForAtLeastNPodsUnschedulable(ctx, testCtx.ControllerClient, namespace, pg1Pods, 3)
126+
127+
pg2Pods := append(pg2SubGroup1Pods, pg2SubGroup2Pods...)
128+
wait.ForAtLeastNPodsScheduled(ctx, testCtx.ControllerClient, namespace, pg2SubGroup1Pods, 1)
129+
wait.ForAtLeastNPodsScheduled(ctx, testCtx.ControllerClient, namespace, pg2SubGroup2Pods, 1)
130+
wait.ForAtLeastNPodsScheduled(ctx, testCtx.ControllerClient, namespace, pg2Pods, 3)
131+
wait.ForAtLeastNPodsUnschedulable(ctx, testCtx.ControllerClient, namespace, pg2Pods, 3)
132+
})
133+
134+
It("Don't schedule job if subgroup gang is not satisfied", func(ctx context.Context) {
135+
namespace := queue.GetConnectedNamespaceToQueue(testCtx.Queues[0])
136+
137+
pg1Name := utils.GenerateRandomK8sName(10)
138+
pg1SubGroup1Pods := createSubGroupPods(ctx, testCtx.KubeClientset, testCtx.Queues[0], pg1Name, "sub-1", 3)
139+
pg1SubGroup2Pods := createSubGroupPods(ctx, testCtx.KubeClientset, testCtx.Queues[0], pg1Name, "sub-2", 3)
140+
pg1 := pod_group.Create(namespace, pg1Name, testCtx.Queues[0].Name)
141+
pg1.Spec.MinMember = 4
142+
pg1.Spec.SubGroups = []schedulingv2alpha2.SubGroup{
143+
{Name: "sub-1", MinMember: 2},
144+
{Name: "sub-2", MinMember: 2},
145+
}
146+
_, err := testCtx.KubeAiSchedClientset.SchedulingV2alpha2().PodGroups(namespace).Create(ctx,
147+
pg1, metav1.CreateOptions{})
148+
Expect(err).To(Succeed())
149+
150+
// wait until pg1 is scheduled to ensure that it will be the one running at the end of the test
151+
wait.ForAtLeastNPodsScheduled(ctx, testCtx.ControllerClient, namespace, pg1SubGroup1Pods, 1)
152+
153+
pg2Name := utils.GenerateRandomK8sName(10)
154+
pg2SubGroup1Pods := createSubGroupPods(ctx, testCtx.KubeClientset, testCtx.Queues[0], pg2Name, "sub-1", 3)
155+
pg2SubGroup2Pods := createSubGroupPods(ctx, testCtx.KubeClientset, testCtx.Queues[0], pg2Name, "sub-2", 3)
156+
pg2 := pod_group.Create(namespace, pg2Name, testCtx.Queues[0].Name)
157+
pg2.Spec.MinMember = 4
158+
pg2.Spec.SubGroups = []schedulingv2alpha2.SubGroup{
159+
{Name: "sub-1", MinMember: 2},
160+
{Name: "sub-2", MinMember: 2},
161+
}
162+
163+
_, err = testCtx.KubeAiSchedClientset.SchedulingV2alpha2().PodGroups(namespace).Create(ctx,
164+
pg2, metav1.CreateOptions{})
165+
Expect(err).To(Succeed())
166+
167+
pg1Pods := append(pg1SubGroup1Pods, pg1SubGroup2Pods...)
168+
wait.ForAtLeastNPodsScheduled(ctx, testCtx.ControllerClient, namespace, pg1Pods, 6)
169+
170+
pg2Pods := append(pg2SubGroup1Pods, pg2SubGroup2Pods...)
171+
wait.ForAtLeastNPodsUnschedulable(ctx, testCtx.ControllerClient, namespace, pg2Pods, 6)
172+
})
173+
})
174+
175+
func createSubGroupPods(ctx context.Context, client *kubernetes.Clientset, queue *v2.Queue,
176+
podGroupName string, subGroupName string, numPods int) []*v1.Pod {
177+
var pods []*v1.Pod
178+
for i := 0; i < numPods; i++ {
179+
pod := createPod(ctx, client, queue, podGroupName, subGroupName, "100m")
180+
pods = append(pods, pod)
181+
}
182+
return pods
183+
}
184+
185+
func createPod(ctx context.Context, client *kubernetes.Clientset, queue *v2.Queue, podGroupName string,
186+
subGroupName string, cpuPerPod string) *v1.Pod {
187+
pod := rd.CreatePodWithPodGroupReference(queue, podGroupName, v1.ResourceRequirements{
188+
Limits: map[v1.ResourceName]resource.Quantity{
189+
v1.ResourceCPU: resource.MustParse(cpuPerPod),
190+
},
191+
})
192+
pod.Labels[commonconsts.SubGroupLabelKey] = subGroupName
193+
pod, err := rd.CreatePod(ctx, client, pod)
194+
Expect(err).To(Succeed())
195+
return pod
196+
}

0 commit comments

Comments
 (0)