Skip to content

Commit a1e1fb5

Browse files
0x54570x5457
authored
test: add comprehensive unit tests for GPU scheduling strategies (#75)
* test: add comprehensive unit tests for GPU scheduling strategies * chore: format code
1 parent b45f6f6 commit a1e1fb5

File tree

1 file changed

+317
-0
lines changed

1 file changed

+317
-0
lines changed

‎internal/scheduler/strategy_test.go

+317
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,317 @@
1+
package scheduler
2+
3+
import (
4+
"testing"
5+
6+
tfv1 "github.com/NexusGPU/tensor-fusion/api/v1"
7+
"github.com/stretchr/testify/assert"
8+
"k8s.io/apimachinery/pkg/api/resource"
9+
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
10+
)
11+
12+
func TestLowLoadFirst(t *testing.T) {
13+
// Create test cases
14+
testCases := []struct {
15+
name string
16+
gpus []tfv1.GPU
17+
expected string
18+
errorMsg string
19+
}{
20+
{
21+
name: "select highest available resources",
22+
gpus: []tfv1.GPU{
23+
{
24+
ObjectMeta: metav1.ObjectMeta{Name: "gpu-1"},
25+
Status: tfv1.GPUStatus{
26+
Available: &tfv1.Resource{
27+
Tflops: resource.MustParse("10"),
28+
Vram: resource.MustParse("40Gi"),
29+
},
30+
},
31+
},
32+
{
33+
ObjectMeta: metav1.ObjectMeta{Name: "gpu-2"},
34+
Status: tfv1.GPUStatus{
35+
Available: &tfv1.Resource{
36+
Tflops: resource.MustParse("15"),
37+
Vram: resource.MustParse("80Gi"),
38+
},
39+
},
40+
},
41+
{
42+
ObjectMeta: metav1.ObjectMeta{Name: "gpu-3"},
43+
Status: tfv1.GPUStatus{
44+
Available: &tfv1.Resource{
45+
Tflops: resource.MustParse("20"),
46+
Vram: resource.MustParse("60Gi"),
47+
},
48+
},
49+
},
50+
},
51+
expected: "gpu-2", // Has highest VRAM
52+
},
53+
{
54+
name: "select higher TFLOPS when VRAM is equal",
55+
gpus: []tfv1.GPU{
56+
{
57+
ObjectMeta: metav1.ObjectMeta{Name: "gpu-1"},
58+
Status: tfv1.GPUStatus{
59+
Available: &tfv1.Resource{
60+
Tflops: resource.MustParse("10"),
61+
Vram: resource.MustParse("40Gi"),
62+
},
63+
},
64+
},
65+
{
66+
ObjectMeta: metav1.ObjectMeta{Name: "gpu-2"},
67+
Status: tfv1.GPUStatus{
68+
Available: &tfv1.Resource{
69+
Tflops: resource.MustParse("15"),
70+
Vram: resource.MustParse("40Gi"),
71+
},
72+
},
73+
},
74+
},
75+
expected: "gpu-2", // Equal VRAM but higher TFLOPS
76+
},
77+
{
78+
name: "no GPUs available",
79+
gpus: []tfv1.GPU{},
80+
expected: "",
81+
errorMsg: "no GPUs available",
82+
},
83+
{
84+
name: "single GPU available",
85+
gpus: []tfv1.GPU{
86+
{
87+
ObjectMeta: metav1.ObjectMeta{Name: "gpu-1"},
88+
Status: tfv1.GPUStatus{
89+
Available: &tfv1.Resource{
90+
Tflops: resource.MustParse("10"),
91+
Vram: resource.MustParse("40Gi"),
92+
},
93+
},
94+
},
95+
},
96+
expected: "gpu-1",
97+
},
98+
}
99+
100+
// Run test cases
101+
for _, tc := range testCases {
102+
t.Run(tc.name, func(t *testing.T) {
103+
strategy := LowLoadFirst{}
104+
selected, err := strategy.SelectGPU(tc.gpus)
105+
106+
if tc.errorMsg != "" {
107+
assert.Error(t, err)
108+
assert.Contains(t, err.Error(), tc.errorMsg)
109+
assert.Nil(t, selected)
110+
} else {
111+
assert.NoError(t, err)
112+
assert.NotNil(t, selected)
113+
assert.Equal(t, tc.expected, selected.Name)
114+
}
115+
})
116+
}
117+
}
118+
119+
func TestCompactFirst(t *testing.T) {
120+
// Create test cases
121+
testCases := []struct {
122+
name string
123+
gpus []tfv1.GPU
124+
expected string
125+
errorMsg string
126+
}{
127+
{
128+
name: "select lowest available resources",
129+
gpus: []tfv1.GPU{
130+
{
131+
ObjectMeta: metav1.ObjectMeta{Name: "gpu-1"},
132+
Status: tfv1.GPUStatus{
133+
Available: &tfv1.Resource{
134+
Tflops: resource.MustParse("10"),
135+
Vram: resource.MustParse("40Gi"),
136+
},
137+
},
138+
},
139+
{
140+
ObjectMeta: metav1.ObjectMeta{Name: "gpu-2"},
141+
Status: tfv1.GPUStatus{
142+
Available: &tfv1.Resource{
143+
Tflops: resource.MustParse("15"),
144+
Vram: resource.MustParse("20Gi"),
145+
},
146+
},
147+
},
148+
{
149+
ObjectMeta: metav1.ObjectMeta{Name: "gpu-3"},
150+
Status: tfv1.GPUStatus{
151+
Available: &tfv1.Resource{
152+
Tflops: resource.MustParse("20"),
153+
Vram: resource.MustParse("60Gi"),
154+
},
155+
},
156+
},
157+
},
158+
expected: "gpu-2", // Has lowest VRAM
159+
},
160+
{
161+
name: "select lower TFLOPS when VRAM is equal",
162+
gpus: []tfv1.GPU{
163+
{
164+
ObjectMeta: metav1.ObjectMeta{Name: "gpu-1"},
165+
Status: tfv1.GPUStatus{
166+
Available: &tfv1.Resource{
167+
Tflops: resource.MustParse("10"),
168+
Vram: resource.MustParse("40Gi"),
169+
},
170+
},
171+
},
172+
{
173+
ObjectMeta: metav1.ObjectMeta{Name: "gpu-2"},
174+
Status: tfv1.GPUStatus{
175+
Available: &tfv1.Resource{
176+
Tflops: resource.MustParse("15"),
177+
Vram: resource.MustParse("40Gi"),
178+
},
179+
},
180+
},
181+
},
182+
expected: "gpu-1", // Equal VRAM but lower TFLOPS
183+
},
184+
{
185+
name: "no GPUs available",
186+
gpus: []tfv1.GPU{},
187+
expected: "",
188+
errorMsg: "no GPUs available",
189+
},
190+
{
191+
name: "single GPU available",
192+
gpus: []tfv1.GPU{
193+
{
194+
ObjectMeta: metav1.ObjectMeta{Name: "gpu-1"},
195+
Status: tfv1.GPUStatus{
196+
Available: &tfv1.Resource{
197+
Tflops: resource.MustParse("10"),
198+
Vram: resource.MustParse("40Gi"),
199+
},
200+
},
201+
},
202+
},
203+
expected: "gpu-1",
204+
},
205+
}
206+
207+
// Run test cases
208+
for _, tc := range testCases {
209+
t.Run(tc.name, func(t *testing.T) {
210+
strategy := CompactFirst{}
211+
selected, err := strategy.SelectGPU(tc.gpus)
212+
213+
if tc.errorMsg != "" {
214+
assert.Error(t, err)
215+
assert.Contains(t, err.Error(), tc.errorMsg)
216+
assert.Nil(t, selected)
217+
} else {
218+
assert.NoError(t, err)
219+
assert.NotNil(t, selected)
220+
assert.Equal(t, tc.expected, selected.Name)
221+
}
222+
})
223+
}
224+
}
225+
226+
func TestStrategyEdgeCases(t *testing.T) {
227+
// Test both strategies with different edge cases
228+
edgeCases := []struct {
229+
name string
230+
gpus []tfv1.GPU
231+
lowLoadName string // expected result for LowLoadFirst
232+
compactName string // expected result for CompactFirst
233+
errorExpected bool
234+
}{
235+
{
236+
name: "identical resources",
237+
gpus: []tfv1.GPU{
238+
{
239+
ObjectMeta: metav1.ObjectMeta{Name: "gpu-1"},
240+
Status: tfv1.GPUStatus{
241+
Available: &tfv1.Resource{
242+
Tflops: resource.MustParse("10"),
243+
Vram: resource.MustParse("40Gi"),
244+
},
245+
},
246+
},
247+
{
248+
ObjectMeta: metav1.ObjectMeta{Name: "gpu-2"},
249+
Status: tfv1.GPUStatus{
250+
Available: &tfv1.Resource{
251+
Tflops: resource.MustParse("10"),
252+
Vram: resource.MustParse("40Gi"),
253+
},
254+
},
255+
},
256+
},
257+
lowLoadName: "gpu-1", // Both strategies should pick the first one when values are equal
258+
compactName: "gpu-1",
259+
errorExpected: false,
260+
},
261+
{
262+
name: "extreme value differences",
263+
gpus: []tfv1.GPU{
264+
{
265+
ObjectMeta: metav1.ObjectMeta{Name: "gpu-1"},
266+
Status: tfv1.GPUStatus{
267+
Available: &tfv1.Resource{
268+
Tflops: resource.MustParse("1"),
269+
Vram: resource.MustParse("1Gi"),
270+
},
271+
},
272+
},
273+
{
274+
ObjectMeta: metav1.ObjectMeta{Name: "gpu-2"},
275+
Status: tfv1.GPUStatus{
276+
Available: &tfv1.Resource{
277+
Tflops: resource.MustParse("100"),
278+
Vram: resource.MustParse("1000Gi"),
279+
},
280+
},
281+
},
282+
},
283+
lowLoadName: "gpu-2", // Highest resources
284+
compactName: "gpu-1", // Lowest resources
285+
errorExpected: false,
286+
},
287+
}
288+
289+
// Test both strategies
290+
for _, tc := range edgeCases {
291+
t.Run("LowLoadFirst_"+tc.name, func(t *testing.T) {
292+
strategy := LowLoadFirst{}
293+
selected, err := strategy.SelectGPU(tc.gpus)
294+
295+
if tc.errorExpected {
296+
assert.Error(t, err)
297+
} else {
298+
assert.NoError(t, err)
299+
assert.NotNil(t, selected)
300+
assert.Equal(t, tc.lowLoadName, selected.Name)
301+
}
302+
})
303+
304+
t.Run("CompactFirst_"+tc.name, func(t *testing.T) {
305+
strategy := CompactFirst{}
306+
selected, err := strategy.SelectGPU(tc.gpus)
307+
308+
if tc.errorExpected {
309+
assert.Error(t, err)
310+
} else {
311+
assert.NoError(t, err)
312+
assert.NotNil(t, selected)
313+
assert.Equal(t, tc.compactName, selected.Name)
314+
}
315+
})
316+
}
317+
}

0 commit comments

Comments
 (0)