@@ -20,8 +20,10 @@ import (
2020 "context"
2121 "encoding/json"
2222 "fmt"
23+ "slices"
2324
2425 "sigs.k8s.io/controller-runtime/pkg/log"
26+
2527 "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/plugins"
2628 "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/framework"
2729 "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
@@ -36,53 +38,71 @@ const (
3638var _ framework.Picker = & MaxScorePicker {}
3739
3840// MaxScorePickerFactory defines the factory function for MaxScorePicker.
39- func MaxScorePickerFactory (name string , _ json.RawMessage , _ plugins.Handle ) (plugins.Plugin , error ) {
40- return NewMaxScorePicker ().WithName (name ), nil
41+ func MaxScorePickerFactory (name string , rawParameters json.RawMessage , _ plugins.Handle ) (plugins.Plugin , error ) {
42+ parameters := pickerParameters {MaxNumOfEndpoints : DefaultMaxNumOfEndpoints }
43+ if rawParameters != nil {
44+ if err := json .Unmarshal (rawParameters , & parameters ); err != nil {
45+ return nil , fmt .Errorf ("failed to parse the parameters of the '%s' picker - %w" , MaxScorePickerType , err )
46+ }
47+ }
48+
49+ return NewMaxScorePicker (parameters .MaxNumOfEndpoints ).WithName (name ), nil
4150}
4251
4352// NewMaxScorePicker initializes a new MaxScorePicker and returns its pointer.
44- func NewMaxScorePicker () * MaxScorePicker {
53+ func NewMaxScorePicker (maxNumOfEndpoints int ) * MaxScorePicker {
54+ if maxNumOfEndpoints <= 0 {
55+ maxNumOfEndpoints = DefaultMaxNumOfEndpoints // on invalid configuration value, fallback to default value
56+ }
57+
4558 return & MaxScorePicker {
46- tn : plugins.TypedName {Type : MaxScorePickerType , Name : MaxScorePickerType },
47- random : NewRandomPicker () ,
59+ typedName : plugins.TypedName {Type : MaxScorePickerType , Name : MaxScorePickerType },
60+ maxNumOfEndpoints : maxNumOfEndpoints ,
4861 }
4962}
5063
51- // MaxScorePicker picks the pod with the maximum score from the list of candidates.
64+ // MaxScorePicker picks pod(s) with the maximum score from the list of candidates.
5265type MaxScorePicker struct {
53- tn plugins.TypedName
54- random * RandomPicker
55- }
56-
57- // TypedName returns the type and name tuple of this plugin instance.
58- func (p * MaxScorePicker ) TypedName () plugins.TypedName {
59- return p .tn
66+ typedName plugins.TypedName
67+ maxNumOfEndpoints int // maximum number of endpoints to pick
6068}
6169
6270// WithName sets the picker's name
6371func (p * MaxScorePicker ) WithName (name string ) * MaxScorePicker {
64- p .tn .Name = name
72+ p .typedName .Name = name
6573 return p
6674}
6775
76+ // TypedName returns the type and name tuple of this plugin instance.
77+ func (p * MaxScorePicker ) TypedName () plugins.TypedName {
78+ return p .typedName
79+ }
80+
6881// Pick selects the pod with the maximum score from the list of candidates.
6982func (p * MaxScorePicker ) Pick (ctx context.Context , cycleState * types.CycleState , scoredPods []* types.ScoredPod ) * types.ProfileRunResult {
70- log .FromContext (ctx ).V (logutil .DEBUG ).Info (fmt .Sprintf ("Selecting a pod with the max score from %d candidates: %+v" , len (scoredPods ), scoredPods ))
71-
72- highestScorePods := []* types.ScoredPod {}
73- maxScore := - 1.0 // pods min score is 0, putting value lower than 0 in order to find at least one pod as highest
74- for _ , pod := range scoredPods {
75- if pod .Score > maxScore {
76- maxScore = pod .Score
77- highestScorePods = []* types.ScoredPod {pod }
78- } else if pod .Score == maxScore {
79- highestScorePods = append (highestScorePods , pod )
83+ log .FromContext (ctx ).V (logutil .DEBUG ).Info (fmt .Sprintf ("Selecting maximum '%d' pods from %d candidates sorted by max score: %+v" , p .maxNumOfEndpoints ,
84+ len (scoredPods ), scoredPods ))
85+
86+ slices .SortStableFunc (scoredPods , func (i , j * types.ScoredPod ) int { // highest score first
87+ if i .Score > j .Score {
88+ return - 1
89+ }
90+ if i .Score < j .Score {
91+ return 1
8092 }
93+ return 0
94+ })
95+
96+ // if we have enough pods to return keep only the "maxNumOfEndpoints" highest scored pods
97+ if p .maxNumOfEndpoints < len (scoredPods ) {
98+ scoredPods = scoredPods [:p .maxNumOfEndpoints ]
8199 }
82100
83- if len (highestScorePods ) > 1 {
84- return p .random .Pick (ctx , cycleState , highestScorePods ) // pick randomly from the highest score pods
101+ targetPods := make ([]types.Pod , len (scoredPods ))
102+ for i , scoredPod := range scoredPods {
103+ targetPods [i ] = scoredPod
85104 }
86105
87- return & types.ProfileRunResult {TargetPod : highestScorePods [0 ]}
106+ return & types.ProfileRunResult {TargetPods : targetPods }
107+
88108}
0 commit comments