diff --git a/pkg/epp/scheduling/plugins/picker/max_score_picker.go b/pkg/epp/scheduling/plugins/picker/max_score_picker.go new file mode 100644 index 000000000..1705b7dd1 --- /dev/null +++ b/pkg/epp/scheduling/plugins/picker/max_score_picker.go @@ -0,0 +1,49 @@ +package picker + +import ( + "fmt" + + "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/plugins" + "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types" + logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging" +) + +var _ plugins.Picker = &MaxScorePicker{} + +func NewMaxScorePicker() plugins.Picker { + return &MaxScorePicker{ + random: &RandomPicker{}, + } +} + +// MaxScorePicker picks the pod with the maximum score from the list of candidates. +type MaxScorePicker struct { + random *RandomPicker +} + +// Name returns the name of the picker. +func (p *MaxScorePicker) Name() string { + return "max_score" +} + +// Pick selects the pod with the maximum score from the list of candidates. +func (p *MaxScorePicker) Pick(ctx *types.SchedulingContext, scoredPods []*types.ScoredPod) *types.Result { + ctx.Logger.V(logutil.DEBUG).Info(fmt.Sprintf("Selecting a pod with the max score from %d candidates: %+v", len(scoredPods), scoredPods)) + + highestScorePods := []*types.ScoredPod{} + maxScore := -1.0 // pods min score is 0, putting value lower than 0 in order to find at least one pod as highest + for _, pod := range scoredPods { + if pod.Score > maxScore { + maxScore = pod.Score + highestScorePods = []*types.ScoredPod{pod} + } else if pod.Score == maxScore { + highestScorePods = append(highestScorePods, pod) + } + } + + if len(highestScorePods) > 1 { + return p.random.Pick(ctx, highestScorePods) // pick randomly from the highest score pods + } + + return &types.Result{TargetPod: highestScorePods[0]} +} diff --git a/pkg/epp/scheduling/plugins/picker/random_picker.go b/pkg/epp/scheduling/plugins/picker/random_picker.go index 6eecbb0da..fb9f9a295 100644 --- a/pkg/epp/scheduling/plugins/picker/random_picker.go +++ b/pkg/epp/scheduling/plugins/picker/random_picker.go @@ -30,12 +30,12 @@ var _ plugins.Picker = &RandomPicker{} // RandomPicker picks a random pod from the list of candidates. type RandomPicker struct{} -func (rp *RandomPicker) Name() string { +func (p *RandomPicker) Name() string { return "random" } -func (rp *RandomPicker) Pick(ctx *types.SchedulingContext, scoredPods []*types.ScoredPod) *types.Result { +func (p *RandomPicker) Pick(ctx *types.SchedulingContext, scoredPods []*types.ScoredPod) *types.Result { ctx.Logger.V(logutil.DEBUG).Info(fmt.Sprintf("Selecting a random pod from %d candidates: %+v", len(scoredPods), scoredPods)) i := rand.Intn(len(scoredPods)) - return &types.Result{TargetPod: scoredPods[i].Pod} + return &types.Result{TargetPod: scoredPods[i]} } diff --git a/pkg/epp/scheduling/scheduler_test.go b/pkg/epp/scheduling/scheduler_test.go index 559f53f8b..311f44e9f 100644 --- a/pkg/epp/scheduling/scheduler_test.go +++ b/pkg/epp/scheduling/scheduler_test.go @@ -93,17 +93,19 @@ func TestSchedule(t *testing.T) { }, }, wantRes: &types.Result{ - TargetPod: &types.PodMetrics{ - Pod: &backendmetrics.Pod{NamespacedName: k8stypes.NamespacedName{Name: "pod2"}}, - Metrics: &backendmetrics.Metrics{ - WaitingQueueSize: 3, - KVCacheUsagePercent: 0.1, - MaxActiveModels: 2, - ActiveModels: map[string]int{ - "foo": 1, - "critical": 1, + TargetPod: &types.ScoredPod{ + Pod: &types.PodMetrics{ + Pod: &backendmetrics.Pod{NamespacedName: k8stypes.NamespacedName{Name: "pod2"}}, + Metrics: &backendmetrics.Metrics{ + WaitingQueueSize: 3, + KVCacheUsagePercent: 0.1, + MaxActiveModels: 2, + ActiveModels: map[string]int{ + "foo": 1, + "critical": 1, + }, + WaitingModels: map[string]int{}, }, - WaitingModels: map[string]int{}, }, }, }, @@ -154,17 +156,19 @@ func TestSchedule(t *testing.T) { }, }, wantRes: &types.Result{ - TargetPod: &types.PodMetrics{ - Pod: &backendmetrics.Pod{NamespacedName: k8stypes.NamespacedName{Name: "pod1"}}, - Metrics: &backendmetrics.Metrics{ - WaitingQueueSize: 0, - KVCacheUsagePercent: 0.2, - MaxActiveModels: 2, - ActiveModels: map[string]int{ - "foo": 1, - "bar": 1, + TargetPod: &types.ScoredPod{ + Pod: &types.PodMetrics{ + Pod: &backendmetrics.Pod{NamespacedName: k8stypes.NamespacedName{Name: "pod1"}}, + Metrics: &backendmetrics.Metrics{ + WaitingQueueSize: 0, + KVCacheUsagePercent: 0.2, + MaxActiveModels: 2, + ActiveModels: map[string]int{ + "foo": 1, + "bar": 1, + }, + WaitingModels: map[string]int{}, }, - WaitingModels: map[string]int{}, }, }, }, @@ -505,7 +509,7 @@ func findPods(ctx *types.SchedulingContext, names ...k8stypes.NamespacedName) [] func getPodScore(scoredPods []*types.ScoredPod, selectedPod types.Pod) float64 { finalScore := 0.0 for _, scoredPod := range scoredPods { - if scoredPod.Pod.GetPod().NamespacedName.String() == selectedPod.GetPod().NamespacedName.String() { + if scoredPod.GetPod().NamespacedName.String() == selectedPod.GetPod().NamespacedName.String() { finalScore = scoredPod.Score break } diff --git a/pkg/epp/scheduling/types/types.go b/pkg/epp/scheduling/types/types.go index 5ccfbdcef..5198515be 100644 --- a/pkg/epp/scheduling/types/types.go +++ b/pkg/epp/scheduling/types/types.go @@ -47,7 +47,7 @@ type Pod interface { } type ScoredPod struct { - Pod Pod + Pod Score float64 }