Skip to content

Commit 4e9cf1e

Browse files
enoodleTensorRaya
enoodle
authored and
TensorRaya
committed
add test cases
Signed-off-by: Erez Freiberger <enoodle@gmail.com> Signed-off-by: Raya Solano <raya.solano@mbinf.de>
1 parent 8dc2871 commit 4e9cf1e

1 file changed

Lines changed: 260 additions & 0 deletions

File tree

‎pkg/scheduler/actions/integration_tests/allocate/allocate_shared_dra_test.go‎

Lines changed: 260 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,260 @@
1+
// Copyright 2026 NVIDIA CORPORATION
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package allocate
5+
6+
import (
7+
"testing"
8+
"time"
9+
10+
resourceapi "k8s.io/api/resource/v1"
11+
12+
commonconstants "github.com/kai-scheduler/KAI-scheduler/pkg/common/constants"
13+
featuregates "github.com/kai-scheduler/KAI-scheduler/pkg/common/feature_gates"
14+
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/actions/integration_tests/integration_tests_utils"
15+
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/pod_status"
16+
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/constants"
17+
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/test_utils"
18+
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/test_utils/dra_fake"
19+
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/test_utils/jobs_fake"
20+
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/test_utils/nodes_fake"
21+
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/test_utils/tasks_fake"
22+
)
23+
24+
func TestSharedDRADeviceDoesNotBlockCPUOnlyPod(t *testing.T) {
25+
featuregates.SetDynamicResourcesEnabledForTest(true)
26+
t.Cleanup(func() {
27+
featuregates.SetDynamicResourcesEnabledForTest(false)
28+
})
29+
30+
integration_tests_utils.RunTests(t, []integration_tests_utils.TestTopologyMetadata{
31+
{
32+
Name: "shared DRA device does not block CPU-only pod",
33+
TestTopologyBasic: test_utils.TestTopologyBasic{
34+
Name: "shared DRA device does not block CPU-only pod",
35+
Jobs: []*jobs_fake.TestJobBasic{
36+
{
37+
Name: "shared_dra_job0",
38+
Namespace: "test",
39+
Priority: constants.PriorityTrainNumber,
40+
QueueName: "queue1",
41+
Tasks: []*tasks_fake.TestTaskBasic{
42+
{
43+
NodeName: "node0",
44+
State: pod_status.Running,
45+
ResourceClaimNames: []string{"shared-claim"},
46+
},
47+
},
48+
},
49+
{
50+
Name: "shared_dra_job1",
51+
Namespace: "test",
52+
Priority: constants.PriorityTrainNumber,
53+
QueueName: "queue1",
54+
Tasks: []*tasks_fake.TestTaskBasic{
55+
{
56+
NodeName: "node0",
57+
State: pod_status.Running,
58+
ResourceClaimNames: []string{"shared-claim"},
59+
},
60+
},
61+
},
62+
{
63+
Name: "cpu_only_job",
64+
Namespace: "test",
65+
Priority: constants.PriorityTrainNumber,
66+
QueueName: "queue1",
67+
Tasks: []*tasks_fake.TestTaskBasic{
68+
{
69+
State: pod_status.Pending,
70+
NodeAffinityNames: []string{"node0"},
71+
},
72+
},
73+
},
74+
},
75+
TestDRAObjects: dra_fake.TestDRAObjects{
76+
DeviceClasses: []string{"nvidia.com/gpu"},
77+
ResourceSlices: []*dra_fake.TestResourceSlice{
78+
{
79+
Name: "node0-gpu",
80+
DeviceClassName: "nvidia.com/gpu",
81+
NodeName: "node0",
82+
Count: 1,
83+
},
84+
},
85+
ResourceClaims: []*dra_fake.TestResourceClaim{
86+
{
87+
Name: "shared-claim",
88+
Namespace: "test",
89+
DeviceClassName: "nvidia.com/gpu",
90+
Count: 1,
91+
Labels: map[string]string{
92+
commonconstants.DefaultQueueLabel: "queue1",
93+
},
94+
ClaimStatus: &resourceapi.ResourceClaimStatus{
95+
Allocation: &resourceapi.AllocationResult{
96+
Devices: resourceapi.DeviceAllocationResult{
97+
Results: []resourceapi.DeviceRequestAllocationResult{
98+
{
99+
Request: "request",
100+
Driver: "nvidia.com/gpu",
101+
Pool: "node0",
102+
Device: "0",
103+
},
104+
},
105+
},
106+
},
107+
ReservedFor: []resourceapi.ResourceClaimConsumerReference{
108+
{Resource: "pods", Name: "shared_dra_job0-0", UID: "shared_dra_job0-0"},
109+
{Resource: "pods", Name: "shared_dra_job1-0", UID: "shared_dra_job1-0"},
110+
},
111+
},
112+
},
113+
},
114+
},
115+
Nodes: map[string]nodes_fake.TestNodeBasic{
116+
"node0": {},
117+
},
118+
Queues: []test_utils.TestQueueBasic{
119+
{
120+
Name: "queue1",
121+
DeservedGPUs: 1,
122+
},
123+
},
124+
JobExpectedResults: map[string]test_utils.TestExpectedResultBasic{
125+
"shared_dra_job0": {
126+
NodeName: "node0",
127+
Status: pod_status.Running,
128+
},
129+
"shared_dra_job1": {
130+
NodeName: "node0",
131+
Status: pod_status.Running,
132+
},
133+
"cpu_only_job": {
134+
NodeName: "node0",
135+
Status: pod_status.Running,
136+
},
137+
},
138+
Mocks: &test_utils.TestMock{
139+
CacheRequirements: &test_utils.CacheMocking{
140+
NumberOfCacheBinds: 1,
141+
},
142+
},
143+
},
144+
RoundsUntilMatch: 1,
145+
RoundsAfterMatch: 1,
146+
SchedulingDuration: time.Millisecond,
147+
},
148+
})
149+
}
150+
151+
func TestPendingPodCanUseSharedDRADevice(t *testing.T) {
152+
featuregates.SetDynamicResourcesEnabledForTest(true)
153+
t.Cleanup(func() {
154+
featuregates.SetDynamicResourcesEnabledForTest(false)
155+
})
156+
157+
integration_tests_utils.RunTests(t, []integration_tests_utils.TestTopologyMetadata{
158+
{
159+
Name: "pending pod can use shared DRA device",
160+
TestTopologyBasic: test_utils.TestTopologyBasic{
161+
Name: "pending pod can use shared DRA device",
162+
Jobs: []*jobs_fake.TestJobBasic{
163+
{
164+
Name: "running_shared_dra_job",
165+
Namespace: "test",
166+
Priority: constants.PriorityTrainNumber,
167+
QueueName: "queue1",
168+
Tasks: []*tasks_fake.TestTaskBasic{
169+
{
170+
NodeName: "node0",
171+
State: pod_status.Running,
172+
ResourceClaimNames: []string{"shared-claim"},
173+
},
174+
},
175+
},
176+
{
177+
Name: "pending_shared_dra_job",
178+
Namespace: "test",
179+
Priority: constants.PriorityTrainNumber,
180+
QueueName: "queue1",
181+
Tasks: []*tasks_fake.TestTaskBasic{
182+
{
183+
State: pod_status.Pending,
184+
NodeAffinityNames: []string{"node0"},
185+
ResourceClaimNames: []string{"shared-claim"},
186+
},
187+
},
188+
},
189+
},
190+
TestDRAObjects: dra_fake.TestDRAObjects{
191+
DeviceClasses: []string{"nvidia.com/gpu"},
192+
ResourceSlices: []*dra_fake.TestResourceSlice{
193+
{
194+
Name: "node0-gpu",
195+
DeviceClassName: "nvidia.com/gpu",
196+
NodeName: "node0",
197+
Count: 1,
198+
},
199+
},
200+
ResourceClaims: []*dra_fake.TestResourceClaim{
201+
{
202+
Name: "shared-claim",
203+
Namespace: "test",
204+
DeviceClassName: "nvidia.com/gpu",
205+
Count: 1,
206+
Labels: map[string]string{
207+
commonconstants.DefaultQueueLabel: "queue1",
208+
},
209+
ClaimStatus: &resourceapi.ResourceClaimStatus{
210+
Allocation: &resourceapi.AllocationResult{
211+
Devices: resourceapi.DeviceAllocationResult{
212+
Results: []resourceapi.DeviceRequestAllocationResult{
213+
{
214+
Request: "request",
215+
Driver: "nvidia.com/gpu",
216+
Pool: "node0",
217+
Device: "0",
218+
},
219+
},
220+
},
221+
},
222+
ReservedFor: []resourceapi.ResourceClaimConsumerReference{
223+
{Resource: "pods", Name: "running_shared_dra_job-0", UID: "running_shared_dra_job-0"},
224+
{Resource: "pods", Name: "pending_shared_dra_job-0", UID: "pending_shared_dra_job-0"},
225+
},
226+
},
227+
},
228+
},
229+
},
230+
Nodes: map[string]nodes_fake.TestNodeBasic{
231+
"node0": {},
232+
},
233+
Queues: []test_utils.TestQueueBasic{
234+
{
235+
Name: "queue1",
236+
DeservedGPUs: 1,
237+
},
238+
},
239+
JobExpectedResults: map[string]test_utils.TestExpectedResultBasic{
240+
"running_shared_dra_job": {
241+
NodeName: "node0",
242+
Status: pod_status.Running,
243+
},
244+
"pending_shared_dra_job": {
245+
NodeName: "node0",
246+
Status: pod_status.Running,
247+
},
248+
},
249+
Mocks: &test_utils.TestMock{
250+
CacheRequirements: &test_utils.CacheMocking{
251+
NumberOfCacheBinds: 1,
252+
},
253+
},
254+
},
255+
RoundsUntilMatch: 1,
256+
RoundsAfterMatch: 1,
257+
SchedulingDuration: time.Millisecond,
258+
},
259+
})
260+
}

0 commit comments

Comments
 (0)