@@ -21,7 +21,6 @@ import (
21
21
"github.com/google/go-cmp/cmp"
22
22
common "github.com/kubeflow/common/pkg/apis/common/v1"
23
23
corev1 "k8s.io/api/core/v1"
24
- schedulingv1 "k8s.io/api/scheduling/v1"
25
24
"k8s.io/apimachinery/pkg/api/equality"
26
25
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
27
26
schedulinglisters "k8s.io/client-go/listers/scheduling/v1"
@@ -183,7 +182,7 @@ func (v *VolcanoCtrl) calculatePGMinResources(minMember *int32, mpiJob *kubeflow
183
182
}
184
183
185
184
// sort task by priorityClasses
186
- return calPgMinResource (minMember , mpiJob , v .PriorityClassLister . Get )
185
+ return calPgMinResource (minMember , mpiJob , v .PriorityClassLister )
187
186
}
188
187
189
188
func (v * VolcanoCtrl ) pgSpecsAreEqual (a , b metav1.Object ) bool {
@@ -323,7 +322,7 @@ func (s *SchedulerPluginsCtrl) calculatePGMinResources(minMember *int32, mpiJob
323
322
return nil
324
323
}
325
324
326
- return calPgMinResource (minMember , mpiJob , s .PriorityClassLister . Get )
325
+ return calPgMinResource (minMember , mpiJob , s .PriorityClassLister )
327
326
}
328
327
329
328
func (s * SchedulerPluginsCtrl ) pgSpecsAreEqual (a , b metav1.Object ) bool {
@@ -334,10 +333,8 @@ func (s *SchedulerPluginsCtrl) pgSpecsAreEqual(a, b metav1.Object) bool {
334
333
335
334
var _ PodGroupControl = & SchedulerPluginsCtrl {}
336
335
337
- type PriorityClassGetFunc func (string ) (* schedulingv1.PriorityClass , error )
338
-
339
336
// calPgMinResource returns the minimum resource for mpiJob with minMembers
340
- func calPgMinResource (minMember * int32 , mpiJob * kubeflow.MPIJob , pFunc PriorityClassGetFunc ) * corev1.ResourceList {
337
+ func calPgMinResource (minMember * int32 , mpiJob * kubeflow.MPIJob , pcLister schedulinglisters. PriorityClassLister ) * corev1.ResourceList {
341
338
var order replicasOrder
342
339
for rt , replica := range mpiJob .Spec .MPIReplicaSpecs {
343
340
rp := replicaPriority {
@@ -346,7 +343,7 @@ func calPgMinResource(minMember *int32, mpiJob *kubeflow.MPIJob, pFunc PriorityC
346
343
ReplicaSpec : * replica ,
347
344
}
348
345
349
- rp .priority = calculatePriority (& replica .Template .Spec , pFunc )
346
+ rp .priority = calculatePriority (& replica .Template .Spec , pcLister )
350
347
order = append (order , rp )
351
348
}
352
349
@@ -411,14 +408,14 @@ func calculatePriorityClassName(
411
408
// calculatePriority get priority
412
409
// 1. get priority with priorityClass's value
413
410
// 2. get priority with priority
414
- func calculatePriority (podSpec * corev1.PodSpec , pFunc PriorityClassGetFunc ) int32 {
411
+ func calculatePriority (podSpec * corev1.PodSpec , pcLister schedulinglisters. PriorityClassLister ) int32 {
415
412
if podSpec == nil {
416
413
return 0
417
414
}
418
415
var priority int32
419
416
pcName := podSpec .PriorityClassName
420
- if len (pcName ) != 0 && pFunc != nil {
421
- if priorityClass , err := pFunc (pcName ); err != nil {
417
+ if len (pcName ) != 0 && pcLister != nil {
418
+ if priorityClass , err := pcLister . Get (pcName ); err != nil {
422
419
klog .Warningf ("Ignore priority class %q: %v" , pcName , err )
423
420
if podSpec .Priority != nil {
424
421
priority = * podSpec .Priority
0 commit comments