addressed code review comments

Signed-off-by: Nir Rozenbaum <[email protected]>

Author: nirrozenbaum

pkg/epp/scheduling/config.go

@@ -22,8 +22,8 @@ type SchedulerConfig struct {
 	preSchedulePlugins  []plugins.PreSchedule
 	filters             []plugins.Filter
 	scorers             map[plugins.Scorer]int // map from scorer to weight
-	postSchedulePlugins []plugins.PostSchedule
 	picker              plugins.Picker
+	postSchedulePlugins []plugins.PostSchedule
 }
 
 var defPlugin = &defaultPlugin{}
@@ -36,6 +36,6 @@ var defaultConfig = &SchedulerConfig{
 	preSchedulePlugins:  []plugins.PreSchedule{},
 	filters:             []plugins.Filter{defPlugin},
 	scorers:             map[plugins.Scorer]int{},
-	postSchedulePlugins: []plugins.PostSchedule{},
 	picker:              defPlugin,
+	postSchedulePlugins: []plugins.PostSchedule{},
 }

pkg/epp/scheduling/plugins/picker/random_picker.go

@@ -34,18 +34,8 @@ func (rp *RandomPicker) Name() string {
 	return "random"
 }
 
-func (rp *RandomPicker) Pick(ctx *types.SchedulingContext, scoredPods map[types.Pod]float64) *types.Result {
-	ctx.Logger.V(logutil.DEBUG).Info(fmt.Sprintf("Selecting a random pod from %d candidates", len(scoredPods)))
-	selectedIndex := rand.Intn(len(scoredPods))
-	i := 0
-	var randomPod types.Pod
-	for pod := range scoredPods {
-		if selectedIndex == i {
-			randomPod = pod
-			break
-
-		}
-		i++
-	}
-	return &types.Result{TargetPod: randomPod}
+func (rp *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}
 }

pkg/epp/scheduling/plugins/plugins.go

@@ -56,16 +56,16 @@ type Scorer interface {
 	Score(ctx *types.SchedulingContext, pods []types.Pod) map[types.Pod]float64
 }
 
-// PostSchedule is called by the scheduler after it selects a targetPod for the request.
-type PostSchedule interface {
+// Picker picks the final pod(s) to send the request to.
+type Picker interface {
 	Plugin
-	PostSchedule(ctx *types.SchedulingContext, res *types.Result)
+	Pick(ctx *types.SchedulingContext, scoredPods []*types.ScoredPod) *types.Result
 }
 
-// Picker picks the final pod(s) to send the request to.
-type Picker interface {
+// PostSchedule is called by the scheduler after it selects a targetPod for the request.
+type PostSchedule interface {
 	Plugin
-	Pick(ctx *types.SchedulingContext, scoredPods map[types.Pod]float64) *types.Result
+	PostSchedule(ctx *types.SchedulingContext, res *types.Result)
 }
 
 // PostResponse is called by the scheduler after a successful response was sent.

pkg/epp/scheduling/scheduler.go

@@ -72,31 +72,23 @@ func NewScheduler(datastore Datastore) *Scheduler {
 }
 
 func NewSchedulerWithConfig(datastore Datastore, config *SchedulerConfig) *Scheduler {
-	sumOfScorersWeights := 0
-	for _, weight := range config.scorers {
-		sumOfScorersWeights += weight
-	}
-	scheduler := &Scheduler{
+	return &Scheduler{
 		datastore:           datastore,
 		preSchedulePlugins:  config.preSchedulePlugins,
 		filters:             config.filters,
 		scorers:             config.scorers,
-		postSchedulePlugins: config.postSchedulePlugins,
 		picker:              config.picker,
-		sumOfScorersWeights: sumOfScorersWeights,
+		postSchedulePlugins: config.postSchedulePlugins,
 	}
-
-	return scheduler
 }
 
 type Scheduler struct {
 	datastore           Datastore
 	preSchedulePlugins  []plugins.PreSchedule
 	filters             []plugins.Filter
 	scorers             map[plugins.Scorer]int // map from scorer to its weight
-	postSchedulePlugins []plugins.PostSchedule
 	picker              plugins.Picker
-	sumOfScorersWeights int
+	postSchedulePlugins []plugins.PostSchedule
 }
 
 type Datastore interface {
@@ -123,14 +115,11 @@ func (s *Scheduler) Schedule(ctx context.Context, req *types.LLMRequest) (*types
 	// if we got here, there is at least one pod to score
 	weightedScorePerPod := s.runScorerPlugins(sCtx, pods)
 
-	before := time.Now()
-	res := s.picker.Pick(sCtx, weightedScorePerPod)
-	metrics.RecordSchedulerPluginProcessingLatency(plugins.PickerPluginType, s.picker.Name(), time.Since(before))
-	loggerDebug.Info("After running picker plugin", "result", res)
+	result := s.runPickerPlugin(sCtx, weightedScorePerPod)
 
-	s.runPostSchedulePlugins(sCtx, res)
+	s.runPostSchedulePlugins(sCtx, result)
 
-	return res, nil
+	return result, nil
 }
 
 func (s *Scheduler) runPreSchedulePlugins(ctx *types.SchedulingContext) {
@@ -157,6 +146,8 @@ func (s *Scheduler) runFilterPlugins(ctx *types.SchedulingContext) []types.Pod {
 			break
 		}
 	}
+	loggerDebug.Info("After running filter plugins")
+
 	return filteredPods
 }
 
@@ -175,15 +166,33 @@ func (s *Scheduler) runScorerPlugins(ctx *types.SchedulingContext, pods []types.
 		scores := scorer.Score(ctx, pods)
 		metrics.RecordSchedulerPluginProcessingLatency(plugins.ScorerPluginType, scorer.Name(), time.Since(before))
 		for pod, score := range scores { // weight is relative to the sum of weights
-			weightedScorePerPod[pod] += score * float64(weight) / float64(s.sumOfScorersWeights) // TODO normalize score before multiply with weight
+			weightedScorePerPod[pod] += score * float64(weight) // TODO normalize score before multiply with weight
 		}
 		loggerDebug.Info("After running scorer", "scorer", scorer.Name())
 	}
-	loggerDebug.Info("After running scorer plugins", "pods", pods)
+	loggerDebug.Info("After running scorer plugins")
 
 	return weightedScorePerPod
 }
 
+func (s *Scheduler) runPickerPlugin(ctx *types.SchedulingContext, weightedScorePerPod map[types.Pod]float64) *types.Result {
+	loggerDebug := ctx.Logger.V(logutil.DEBUG)
+	scoredPods := make([]*types.ScoredPod, len(weightedScorePerPod))
+	i := 0
+	for pod, score := range weightedScorePerPod {
+		scoredPods[i] = &types.ScoredPod{Pod: pod, Score: score}
+		i++
+	}
+
+	loggerDebug.Info("Before running picker plugin", "pods", weightedScorePerPod)
+	before := time.Now()
+	result := s.picker.Pick(ctx, scoredPods)
+	metrics.RecordSchedulerPluginProcessingLatency(plugins.PickerPluginType, s.picker.Name(), time.Since(before))
+	loggerDebug.Info("After running picker plugin", "result", result)
+
+	return result
+}
+
 func (s *Scheduler) runPostSchedulePlugins(ctx *types.SchedulingContext, res *types.Result) {
 	for _, plugin := range s.postSchedulePlugins {
 		ctx.Logger.V(logutil.DEBUG).Info("Running post-schedule plugin", "plugin", plugin.Name())

pkg/epp/scheduling/scheduler_test.go

@@ -274,16 +274,16 @@ func TestSchedulePlugins(t *testing.T) {
 					tp1: 1,
 					tp2: 1,
 				},
-				postSchedulePlugins: []plugins.PostSchedule{tp1, tp2},
 				picker:              pickerPlugin,
+				postSchedulePlugins: []plugins.PostSchedule{tp1, tp2},
 			},
 			input: []*backendmetrics.FakePodMetrics{
 				{Pod: &backendmetrics.Pod{NamespacedName: k8stypes.NamespacedName{Name: "pod1"}}},
 				{Pod: &backendmetrics.Pod{NamespacedName: k8stypes.NamespacedName{Name: "pod2"}}},
 				{Pod: &backendmetrics.Pod{NamespacedName: k8stypes.NamespacedName{Name: "pod3"}}},
 			},
 			wantTargetPod:  k8stypes.NamespacedName{Name: "pod1"},
-			targetPodScore: 0.55,
+			targetPodScore: 1.1,
 			numPodsToScore: 2,
 			err:            false,
 		},
@@ -296,16 +296,16 @@ func TestSchedulePlugins(t *testing.T) {
 					tp1: 60,
 					tp2: 40,
 				},
-				postSchedulePlugins: []plugins.PostSchedule{tp1, tp2},
 				picker:              pickerPlugin,
+				postSchedulePlugins: []plugins.PostSchedule{tp1, tp2},
 			},
 			input: []*backendmetrics.FakePodMetrics{
 				{Pod: &backendmetrics.Pod{NamespacedName: k8stypes.NamespacedName{Name: "pod1"}}},
 				{Pod: &backendmetrics.Pod{NamespacedName: k8stypes.NamespacedName{Name: "pod2"}}},
 				{Pod: &backendmetrics.Pod{NamespacedName: k8stypes.NamespacedName{Name: "pod3"}}},
 			},
 			wantTargetPod:  k8stypes.NamespacedName{Name: "pod1"},
-			targetPodScore: 0.5,
+			targetPodScore: 50,
 			numPodsToScore: 2,
 			err:            false,
 		},
@@ -318,8 +318,8 @@ func TestSchedulePlugins(t *testing.T) {
 					tp1: 1,
 					tp2: 1,
 				},
-				postSchedulePlugins: []plugins.PostSchedule{tp1, tp2},
 				picker:              pickerPlugin,
+				postSchedulePlugins: []plugins.PostSchedule{tp1, tp2},
 			},
 			input: []*backendmetrics.FakePodMetrics{
 				{Pod: &backendmetrics.Pod{NamespacedName: k8stypes.NamespacedName{Name: "pod1"}}},
@@ -337,16 +337,16 @@ func TestSchedulePlugins(t *testing.T) {
 			for _, plugin := range test.config.preSchedulePlugins {
 				plugin.(*TestPlugin).reset()
 			}
-			for _, plugin := range test.config.postSchedulePlugins {
-				plugin.(*TestPlugin).reset()
-			}
 			for _, plugin := range test.config.filters {
 				plugin.(*TestPlugin).reset()
 			}
 			for plugin := range test.config.scorers {
 				plugin.(*TestPlugin).reset()
 			}
 			test.config.picker.(*TestPlugin).reset()
+			for _, plugin := range test.config.postSchedulePlugins {
+				plugin.(*TestPlugin).reset()
+			}
 
 			// Initialize the scheduler
 			scheduler := NewSchedulerWithConfig(&fakeDataStore{pods: test.input}, &test.config)
@@ -397,13 +397,6 @@ func TestSchedulePlugins(t *testing.T) {
 				}
 			}
 
-			for _, plugin := range test.config.postSchedulePlugins {
-				tp, _ := plugin.(*TestPlugin)
-				if tp.PostScheduleCallCount != 1 {
-					t.Errorf("Plugin %s PostSchedule() called %d times, expected 1", plugin.Name(), tp.PostScheduleCallCount)
-				}
-			}
-
 			tp, _ := test.config.picker.(*TestPlugin)
 			if tp.NumOfPickerCandidates != test.numPodsToScore {
 				t.Errorf("Picker plugin %s Pick() called with %d candidates, expected %d", tp.Name(), tp.NumOfPickerCandidates, tp.NumOfScoredPods)
@@ -414,6 +407,13 @@ func TestSchedulePlugins(t *testing.T) {
 			if tp.WinnderPodScore != test.targetPodScore {
 				t.Errorf("winnder pod score %v, expected %v", tp.WinnderPodScore, test.targetPodScore)
 			}
+
+			for _, plugin := range test.config.postSchedulePlugins {
+				tp, _ := plugin.(*TestPlugin)
+				if tp.PostScheduleCallCount != 1 {
+					t.Errorf("Plugin %s PostSchedule() called %d times, expected 1", plugin.Name(), tp.PostScheduleCallCount)
+				}
+			}
 		})
 	}
 }
@@ -468,18 +468,18 @@ func (tp *TestPlugin) Score(ctx *types.SchedulingContext, pods []types.Pod) map[
 	return scoredPods
 }
 
-func (tp *TestPlugin) PostSchedule(ctx *types.SchedulingContext, res *types.Result) {
-	tp.PostScheduleCallCount++
-}
-
-func (tp *TestPlugin) Pick(ctx *types.SchedulingContext, scoredPods map[types.Pod]float64) *types.Result {
+func (tp *TestPlugin) Pick(ctx *types.SchedulingContext, scoredPods []*types.ScoredPod) *types.Result {
 	tp.PickCallCount++
 	tp.NumOfPickerCandidates = len(scoredPods)
 	pod := findPods(ctx, tp.PickRes)[0]
-	tp.WinnderPodScore = scoredPods[pod]
+	tp.WinnderPodScore = getPodScore(scoredPods, pod)
 	return &types.Result{TargetPod: pod}
 }
 
+func (tp *TestPlugin) PostSchedule(ctx *types.SchedulingContext, res *types.Result) {
+	tp.PostScheduleCallCount++
+}
+
 func (tp *TestPlugin) reset() {
 	tp.PreScheduleCallCount = 0
 	tp.FilterCallCount = 0
@@ -501,3 +501,14 @@ func findPods(ctx *types.SchedulingContext, names ...k8stypes.NamespacedName) []
 	}
 	return res
 }
+
+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() {
+			finalScore = scoredPod.Score
+			break
+		}
+	}
+	return finalScore
+}

pkg/epp/scheduling/types/types.go

@@ -46,6 +46,11 @@ type Pod interface {
 	String() string
 }
 
+type ScoredPod struct {
+	Pod   Pod
+	Score float64
+}
+
 // SchedulingContext holds contextual information during a scheduling operation.
 type SchedulingContext struct {
 	context.Context