use priorityclass lister

Signed-off-by: lowang_bh <lhui_wang@163.com>
kubeflow · Jun 13, 2023 · a66c818 · a66c818
1 parent a754b6f
commit a66c818
Showing 1 changed file with 7 additions and 10 deletions.
diff --git a/pkg/controller/podgroup.go b/pkg/controller/podgroup.go
@@ -21,7 +21,6 @@ import (
 	"github.com/google/go-cmp/cmp"
 	common "github.com/kubeflow/common/pkg/apis/common/v1"
 	corev1 "k8s.io/api/core/v1"
-	schedulingv1 "k8s.io/api/scheduling/v1"
 	"k8s.io/apimachinery/pkg/api/equality"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	schedulinglisters "k8s.io/client-go/listers/scheduling/v1"
@@ -183,7 +182,7 @@ func (v *VolcanoCtrl) calculatePGMinResources(minMember *int32, mpiJob *kubeflow
 	}
 
 	// sort task by priorityClasses
-	return calPgMinResource(minMember, mpiJob, v.PriorityClassLister.Get)
+	return calPgMinResource(minMember, mpiJob, v.PriorityClassLister)
 }
 
 func (v *VolcanoCtrl) pgSpecsAreEqual(a, b metav1.Object) bool {
@@ -323,7 +322,7 @@ func (s *SchedulerPluginsCtrl) calculatePGMinResources(minMember *int32, mpiJob
 		return nil
 	}
 
-	return calPgMinResource(minMember, mpiJob, s.PriorityClassLister.Get)
+	return calPgMinResource(minMember, mpiJob, s.PriorityClassLister)
 }
 
 func (s *SchedulerPluginsCtrl) pgSpecsAreEqual(a, b metav1.Object) bool {
@@ -334,10 +333,8 @@ func (s *SchedulerPluginsCtrl) pgSpecsAreEqual(a, b metav1.Object) bool {
 
 var _ PodGroupControl = &SchedulerPluginsCtrl{}
 
-type PriorityClassGetFunc func(string) (*schedulingv1.PriorityClass, error)
-
 // calPgMinResource returns the minimum resource for mpiJob with minMembers
-func calPgMinResource(minMember *int32, mpiJob *kubeflow.MPIJob, pFunc PriorityClassGetFunc) *corev1.ResourceList {
+func calPgMinResource(minMember *int32, mpiJob *kubeflow.MPIJob, pcLister schedulinglisters.PriorityClassLister) *corev1.ResourceList {
 	var order replicasOrder
 	for rt, replica := range mpiJob.Spec.MPIReplicaSpecs {
 		rp := replicaPriority{
@@ -346,7 +343,7 @@ func calPgMinResource(minMember *int32, mpiJob *kubeflow.MPIJob, pFunc PriorityC
 			ReplicaSpec: *replica,
 		}
 
-		rp.priority = calculatePriority(&replica.Template.Spec, pFunc)
+		rp.priority = calculatePriority(&replica.Template.Spec, pcLister)
 		order = append(order, rp)
 	}
 
@@ -411,14 +408,14 @@ func calculatePriorityClassName(
 // calculatePriority get priority
 // 1. get priority with priorityClass's value
 // 2. get priority with priority
-func calculatePriority(podSpec *corev1.PodSpec, pFunc PriorityClassGetFunc) int32 {
+func calculatePriority(podSpec *corev1.PodSpec, pcLister schedulinglisters.PriorityClassLister) int32 {
 	if podSpec == nil {
 		return 0
 	}
 	var priority int32
 	pcName := podSpec.PriorityClassName
-	if len(pcName) != 0 && pFunc != nil {
-		if priorityClass, err := pFunc(pcName); err != nil {
+	if len(pcName) != 0 && pcLister != nil {
+		if priorityClass, err := pcLister.Get(pcName); err != nil {
 			klog.Warningf("Ignore priority class %q: %v", pcName, err)
 			if podSpec.Priority != nil {
 				priority = *podSpec.Priority