@@ -5,10 +5,14 @@ import (
5
5
"testing"
6
6
7
7
"github.com/stretchr/testify/assert"
8
+ "github.com/stretchr/testify/require"
8
9
corev1 "k8s.io/api/core/v1"
9
10
"k8s.io/apimachinery/pkg/api/resource"
10
11
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
12
+ "k8s.io/apimachinery/pkg/runtime"
11
13
"k8s.io/utils/ptr"
14
+ "sigs.k8s.io/controller-runtime/pkg/client"
15
+ "sigs.k8s.io/controller-runtime/pkg/client/fake"
12
16
volcanobatchv1alpha1 "volcano.sh/apis/pkg/apis/batch/v1alpha1"
13
17
volcanoschedulingv1beta1 "volcano.sh/apis/pkg/apis/scheduling/v1beta1"
14
18
@@ -124,7 +128,7 @@ func createTestRayJob(numOfHosts int32) rayv1.RayJob {
124
128
Namespace : "default" ,
125
129
Labels : map [string ]string {
126
130
QueueNameLabelKey : "test-queue" ,
127
- utils .RayPriorityClassName : "high -priority" ,
131
+ utils .RayPriorityClassName : "test -priority" ,
128
132
},
129
133
},
130
134
Spec : rayv1.RayJobSpec {
@@ -204,32 +208,120 @@ func TestCreatePodGroupForRayCluster_NumOfHosts2(t *testing.T) {
204
208
205
209
func TestCreatePodGroupForRayJob (t * testing.T ) {
206
210
a := assert .New (t )
211
+ ctx := context .Background ()
207
212
208
- rayJob := createTestRayJob (1 )
213
+ scheme := runtime .NewScheme ()
214
+ a .NoError (rayv1 .AddToScheme (scheme ))
215
+ a .NoError (volcanoschedulingv1beta1 .AddToScheme (scheme ))
216
+ fakeCli := fake .NewClientBuilder ().WithScheme (scheme ).Build ()
217
+ scheduler := & VolcanoBatchScheduler {cli : fakeCli }
209
218
210
- // Create RayCluster from RayJob spec for calculation
211
- rayCluster := & rayv1.RayCluster {
212
- Spec : * rayJob .Spec .RayClusterSpec ,
213
- }
219
+ t .Run ("No submitter pod resources" , func (_ * testing.T ) {
220
+ rayJob := createTestRayJob (1 )
221
+ rayJob .Spec .SubmissionMode = rayv1 .HTTPMode
214
222
215
- minMember := utils .CalculateDesiredReplicas (context .Background (), rayCluster ) + 1
216
- totalResource := utils .CalculateDesiredResources (rayCluster )
217
- pg := createPodGroup (& rayJob , getAppPodGroupName (& rayJob ), minMember , totalResource )
223
+ err := scheduler .handleRayJob (ctx , & rayJob )
224
+ require .NoError (t , err )
218
225
219
- a .Equal (rayJob .Namespace , pg .Namespace )
220
- a .Equal ("ray-rayjob-sample-pg" , pg .Name )
226
+ var pg volcanoschedulingv1beta1.PodGroup
227
+ err = fakeCli .Get (ctx , client.ObjectKey {Namespace : rayJob .Namespace , Name : getAppPodGroupName (& rayJob )}, & pg )
228
+ require .NoError (t , err )
221
229
222
- // Verify owner reference is set to RayJob
223
- a .Len (pg .OwnerReferences , 1 )
224
- a .Equal ("RayJob" , pg .OwnerReferences [0 ].Kind )
225
- a .Equal (rayJob .Name , pg .OwnerReferences [0 ].Name )
230
+ // 1 head + 2 workers (desired, not min replicas)
231
+ a .Equal (int32 (3 ), pg .Spec .MinMember )
232
+ // 256m * 3 (requests, not limits)
233
+ a .Equal ("768m" , pg .Spec .MinResources .Cpu ().String ())
234
+ // 256m * 3 (requests, not limits)
235
+ a .Equal ("768Mi" , pg .Spec .MinResources .Memory ().String ())
236
+ a .Equal ("test-queue" , pg .Spec .Queue )
237
+ a .Equal ("test-priority" , pg .Spec .PriorityClassName )
238
+ a .Len (pg .OwnerReferences , 1 )
239
+ a .Equal ("RayJob" , pg .OwnerReferences [0 ].Kind )
240
+ })
226
241
227
- // Verify queue and priority class are set from RayJob labels
228
- a .Equal ("test-queue" , pg .Spec .Queue )
229
- a .Equal ("high-priority" , pg .Spec .PriorityClassName )
242
+ t .Run ("K8sJobMode includes submitter pod resources" , func (_ * testing.T ) {
243
+ rayJob := createTestRayJob (1 )
244
+ rayJob .Spec .SubmissionMode = rayv1 .K8sJobMode
245
+
246
+ err := scheduler .handleRayJob (ctx , & rayJob )
247
+ require .NoError (t , err )
248
+
249
+ var pg volcanoschedulingv1beta1.PodGroup
250
+ err = fakeCli .Get (ctx , client.ObjectKey {Namespace : rayJob .Namespace , Name : getAppPodGroupName (& rayJob )}, & pg )
251
+ require .NoError (t , err )
252
+
253
+ // 1 head + 2 workers (desired, not min replicas)
254
+ a .Equal (int32 (3 ), pg .Spec .MinMember )
255
+ // 768m + 500m = 1268m
256
+ a .Equal ("1268m" , pg .Spec .MinResources .Cpu ().String ())
257
+ // 768Mi + 200Mi = 968Mi
258
+ a .Equal ("968Mi" , pg .Spec .MinResources .Memory ().String ())
259
+ a .Equal ("test-queue" , pg .Spec .Queue )
260
+ a .Equal ("test-priority" , pg .Spec .PriorityClassName )
261
+ a .Len (pg .OwnerReferences , 1 )
262
+ a .Equal ("RayJob" , pg .OwnerReferences [0 ].Kind )
263
+ })
264
+ }
230
265
231
- // 1 head + 2 workers (desired, not min replicas)
232
- a .Equal (int32 (3 ), pg .Spec .MinMember )
266
+ func TestCreatePodGroupForRayJob_NumOfHosts2 (t * testing.T ) {
267
+ a := assert .New (t )
268
+ ctx := context .Background ()
269
+
270
+ scheme := runtime .NewScheme ()
271
+ a .NoError (rayv1 .AddToScheme (scheme ))
272
+ a .NoError (volcanoschedulingv1beta1 .AddToScheme (scheme ))
273
+ fakeCli := fake .NewClientBuilder ().WithScheme (scheme ).Build ()
274
+ scheduler := & VolcanoBatchScheduler {cli : fakeCli }
275
+
276
+ t .Run ("No submitter pod resources" , func (_ * testing.T ) {
277
+ rayJob := createTestRayJob (2 )
278
+ rayJob .Spec .SubmissionMode = rayv1 .HTTPMode
279
+
280
+ err := scheduler .handleRayJob (ctx , & rayJob )
281
+ require .NoError (t , err )
282
+
283
+ var pg volcanoschedulingv1beta1.PodGroup
284
+ err = fakeCli .Get (ctx , client.ObjectKey {Namespace : rayJob .Namespace , Name : getAppPodGroupName (& rayJob )}, & pg )
285
+ require .NoError (t , err )
286
+
287
+ // 2 workers (desired, not min replicas) * 2 (num of hosts) + 1 head
288
+ // 2 * 2 + 1 = 5
289
+ a .Equal (int32 (5 ), pg .Spec .MinMember )
290
+ // 256m * (2 (requests, not limits) * 2 (num of hosts) + 1 head)
291
+ // 256m * 5 = 1280m
292
+ a .Equal ("1280m" , pg .Spec .MinResources .Cpu ().String ())
293
+ // 256Mi * (2 (requests, not limits) * 2 (num of hosts) + 1 head)
294
+ // 256Mi * 5 = 1280Mi
295
+ a .Equal ("1280Mi" , pg .Spec .MinResources .Memory ().String ())
296
+ a .Equal ("test-queue" , pg .Spec .Queue )
297
+ a .Equal ("test-priority" , pg .Spec .PriorityClassName )
298
+ a .Len (pg .OwnerReferences , 1 )
299
+ a .Equal ("RayJob" , pg .OwnerReferences [0 ].Kind )
300
+ })
301
+
302
+ t .Run ("K8sJobMode includes submitter pod resources" , func (_ * testing.T ) {
303
+ rayJob := createTestRayJob (2 )
304
+ rayJob .Spec .SubmissionMode = rayv1 .K8sJobMode
305
+
306
+ err := scheduler .handleRayJob (ctx , & rayJob )
307
+ require .NoError (t , err )
308
+
309
+ var pg volcanoschedulingv1beta1.PodGroup
310
+ err = fakeCli .Get (ctx , client.ObjectKey {Namespace : rayJob .Namespace , Name : getAppPodGroupName (& rayJob )}, & pg )
311
+ require .NoError (t , err )
312
+
313
+ // 2 workers (desired, not min replicas) * 2 (num of hosts) + 1 head
314
+ // 2 * 2 + 1 = 5
315
+ a .Equal (int32 (5 ), pg .Spec .MinMember )
316
+ // 1280m + 500m = 1780m
317
+ a .Equal ("1780m" , pg .Spec .MinResources .Cpu ().String ())
318
+ // 1280Mi + 200Mi = 1480Mi
319
+ a .Equal ("1480Mi" , pg .Spec .MinResources .Memory ().String ())
320
+ a .Equal ("test-queue" , pg .Spec .Queue )
321
+ a .Equal ("test-priority" , pg .Spec .PriorityClassName )
322
+ a .Len (pg .OwnerReferences , 1 )
323
+ a .Equal ("RayJob" , pg .OwnerReferences [0 ].Kind )
324
+ })
233
325
}
234
326
235
327
func TestAddMetadataToSubmitterPod (t * testing.T ) {
@@ -253,7 +345,7 @@ func TestAddMetadataToSubmitterPod(t *testing.T) {
253
345
254
346
// Check labels
255
347
a .Equal ("test-queue" , submitterTemplate .Labels [QueueNameLabelKey ])
256
- a .Equal ("high -priority" , submitterTemplate .Labels [utils .RayPriorityClassName ])
348
+ a .Equal ("test -priority" , submitterTemplate .Labels [utils .RayPriorityClassName ])
257
349
258
350
// Check scheduler name
259
351
a .Equal (pluginName , submitterTemplate .Spec .SchedulerName )
0 commit comments