initial commit - multiple destinations

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

Author: nirrozenbaum

pkg/epp/scheduling/framework/plugins/filter/lora_affinity_filter.go

@@ -46,6 +46,7 @@ func LoraAffinityFilterFactory(name string, rawParameters json.RawMessage, _ plu
 	if err := json.Unmarshal(rawParameters, &parameters); err != nil {
 		return nil, fmt.Errorf("failed to parse the parameters of the '%s' filter - %w", LoraAffinityFilterType, err)
 	}
+
 	return NewLoraAffinityFilter(parameters.Threshold).WithName(name), nil
 }
 

pkg/epp/scheduling/framework/plugins/multi/prefix/plugin.go

@@ -198,7 +198,7 @@ func (m *Plugin) Score(ctx context.Context, cycleState *types.CycleState, reques
 
 // PostCycle records in the plugin cache the result of the scheduling selection.
 func (m *Plugin) PostCycle(ctx context.Context, cycleState *types.CycleState, res *types.ProfileRunResult) {
-	targetPod := res.TargetPod.GetPod()
+	targetPod := res.TargetPods[0].GetPod() // assumes single TargetPod. this should become PostResponse to remove this assumption
 	state, err := m.getPrefixState(cycleState)
 	if err != nil {
 		log.FromContext(ctx).Error(err, "failed to read prefix plugin cycle state")

pkg/epp/scheduling/framework/plugins/multi/prefix/plugin_test.go

@@ -61,7 +61,7 @@ func TestPrefixPlugin(t *testing.T) {
 	assert.Equal(t, float64(0), scores[pod2], "score for pod2")
 
 	// Simulate pod1 was picked.
-	plugin.PostCycle(context.Background(), cycleState1, &types.ProfileRunResult{TargetPod: pod1})
+	plugin.PostCycle(context.Background(), cycleState1, &types.ProfileRunResult{TargetPods: []types.Pod{pod1}})
 
 	// Second request doesn't share any prefix with first one. It should be added to the cache but
 	// the pod score should be 0.
@@ -82,7 +82,7 @@ func TestPrefixPlugin(t *testing.T) {
 	assert.Equal(t, float64(0), scores[pod2], "score for pod2")
 
 	// Simulate pod2 was picked.
-	plugin.PostCycle(context.Background(), cycleState2, &types.ProfileRunResult{TargetPod: pod2})
+	plugin.PostCycle(context.Background(), cycleState2, &types.ProfileRunResult{TargetPods: []types.Pod{pod2}})
 
 	// Third request shares partial prefix with first one.
 	req3 := &types.LLMRequest{
@@ -101,7 +101,7 @@ func TestPrefixPlugin(t *testing.T) {
 	assert.Equal(t, float64(2)/float64(3), scores[pod1], "score should be 2/3 - the model and the first prefix block match")
 	assert.Equal(t, float64(0), scores[pod2], "score for pod2")
 
-	plugin.PostCycle(context.Background(), cycleState3, &types.ProfileRunResult{TargetPod: pod1})
+	plugin.PostCycle(context.Background(), cycleState3, &types.ProfileRunResult{TargetPods: []types.Pod{pod1}})
 
 	// 4th request is same as req3 except the model is different, still no match.
 	req4 := &types.LLMRequest{
@@ -120,7 +120,7 @@ func TestPrefixPlugin(t *testing.T) {
 	assert.Equal(t, float64(0), scores[pod1], "score for pod1")
 	assert.Equal(t, float64(0), scores[pod2], "score for pod2")
 
-	plugin.PostCycle(context.Background(), cycleState4, &types.ProfileRunResult{TargetPod: pod1})
+	plugin.PostCycle(context.Background(), cycleState4, &types.ProfileRunResult{TargetPods: []types.Pod{pod1}})
 
 	// 5th request shares partial prefix with 3rd one.
 	req5 := &types.LLMRequest{
@@ -139,7 +139,7 @@ func TestPrefixPlugin(t *testing.T) {
 	assert.Equal(t, 0.75, scores[pod1], "score should be 0.75 - the model and the first 2 prefix blocks match")
 	assert.Equal(t, float64(0), scores[pod2], "score for pod2")
 
-	plugin.PostCycle(context.Background(), cycleState5, &types.ProfileRunResult{TargetPod: pod1})
+	plugin.PostCycle(context.Background(), cycleState5, &types.ProfileRunResult{TargetPods: []types.Pod{pod1}})
 }
 
 // TestPrefixPluginStress is a stress test for the prefix scoring plugin, using prompts of increasing length.
@@ -180,7 +180,7 @@ func BenchmarkPrefixPluginStress(b *testing.B) {
 		// First cycle: simulate scheduling and insert prefix info into the cache
 		cycleState := types.NewCycleState()
 		plugin.Score(context.Background(), cycleState, req, pods)
-		plugin.PostCycle(context.Background(), cycleState, &types.ProfileRunResult{TargetPod: pod})
+		plugin.PostCycle(context.Background(), cycleState, &types.ProfileRunResult{TargetPods: []types.Pod{pod}})
 
 		// Second cycle: validate internal state
 		state, err := plugin.getPrefixState(cycleState)

pkg/epp/scheduling/framework/plugins/picker/max_score_picker.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"sort"
 
 	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/plugins"
@@ -32,26 +33,39 @@ const (
 	MaxScorePickerType = "max-score"
 )
 
+type maxScorePickerParameters struct {
+	MaxNumOfEndpoints int `json:"maxEndpoints"`
+}
+
 // compile-time type validation
 var _ framework.Picker = &MaxScorePicker{}
 
 // MaxScorePickerFactory defines the factory function for MaxScorePicker.
-func MaxScorePickerFactory(name string, _ json.RawMessage, _ plugins.Handle) (plugins.Plugin, error) {
-	return NewMaxScorePicker().WithName(name), nil
+func MaxScorePickerFactory(name string, rawParameters json.RawMessage, _ plugins.Handle) (plugins.Plugin, error) {
+	parameters := maxScorePickerParameters{}
+	if err := json.Unmarshal(rawParameters, &parameters); err != nil {
+		return nil, fmt.Errorf("failed to parse the parameters of the '%s' picker - %w", MaxScorePickerType, err)
+	}
+
+	return NewMaxScorePicker(parameters.MaxNumOfEndpoints).WithName(name), nil
 }
 
 // NewMaxScorePicker initializes a new MaxScorePicker and returns its pointer.
-func NewMaxScorePicker() *MaxScorePicker {
+func NewMaxScorePicker(maxNumOfEndpoints int) *MaxScorePicker {
+	if maxNumOfEndpoints <= 0 {
+		maxNumOfEndpoints = 1 // on invalid configruation value, fallback to 1
+	}
+
 	return &MaxScorePicker{
-		name:   MaxScorePickerType,
-		random: NewRandomPicker(),
+		name:              MaxScorePickerType,
+		maxNumOfEndpoints: maxNumOfEndpoints,
 	}
 }
 
 // MaxScorePicker picks the pod with the maximum score from the list of candidates.
 type MaxScorePicker struct {
-	name   string
-	random *RandomPicker
+	name              string
+	maxNumOfEndpoints int // maximum number of endpoints to pick
 }
 
 // Type returns the type of the picker.
@@ -72,22 +86,22 @@ func (p *MaxScorePicker) WithName(name string) *MaxScorePicker {
 
 // Pick selects the pod with the maximum score from the list of candidates.
 func (p *MaxScorePicker) Pick(ctx context.Context, cycleState *types.CycleState, scoredPods []*types.ScoredPod) *types.ProfileRunResult {
-	log.FromContext(ctx).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)
-		}
+	log.FromContext(ctx).V(logutil.DEBUG).Info(fmt.Sprintf("Selecting maximum '%d' pods from %d candidates sorted by max score: %+v", p.maxNumOfEndpoints,
+		len(scoredPods), scoredPods))
+
+	sort.Slice(scoredPods, func(i, j int) bool { // highest score first
+		return scoredPods[i].Score > scoredPods[j].Score
+	})
+
+	// if we have enough pods to return keep only the "maxNumOfEndpoints" highest scored pods
+	if p.maxNumOfEndpoints < len(scoredPods) {
+		scoredPods = scoredPods[:p.maxNumOfEndpoints]
 	}
 
-	if len(highestScorePods) > 1 {
-		return p.random.Pick(ctx, cycleState, highestScorePods) // pick randomly from the highest score pods
+	targetPods := make([]types.Pod, len(scoredPods))
+	for i, scoredPod := range scoredPods {
+		targetPods[i] = scoredPod
 	}
 
-	return &types.ProfileRunResult{TargetPod: highestScorePods[0]}
+	return &types.ProfileRunResult{TargetPods: targetPods}
 }

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

@@ -33,24 +33,39 @@ const (
 	RandomPickerType = "random"
 )
 
+type randomPickerParameters struct {
+	MaxNumOfEndpoints int `json:"maxEndpoints"`
+}
+
 // compile-time type validation
 var _ framework.Picker = &RandomPicker{}
 
 // RandomPickerFactory defines the factory function for RandomPicker.
-func RandomPickerFactory(name string, _ json.RawMessage, _ plugins.Handle) (plugins.Plugin, error) {
-	return NewRandomPicker().WithName(name), nil
+func RandomPickerFactory(name string, rawParameters json.RawMessage, _ plugins.Handle) (plugins.Plugin, error) {
+	parameters := randomPickerParameters{}
+	if err := json.Unmarshal(rawParameters, &parameters); err != nil {
+		return nil, fmt.Errorf("failed to parse the parameters of the '%s' picker - %w", RandomPickerType, err)
+	}
+
+	return NewRandomPicker(parameters.MaxNumOfEndpoints).WithName(name), nil
 }
 
 // NewRandomPicker initializes a new RandomPicker and returns its pointer.
-func NewRandomPicker() *RandomPicker {
+func NewRandomPicker(maxNumOfEndpoints int) *RandomPicker {
+	if maxNumOfEndpoints <= 0 {
+		maxNumOfEndpoints = 1 // on invalid configruation value, fallback to 1
+	}
+
 	return &RandomPicker{
-		name: RandomPickerType,
+		name:              RandomPickerType,
+		maxNumOfEndpoints: maxNumOfEndpoints,
 	}
 }
 
 // RandomPicker picks a random pod from the list of candidates.
 type RandomPicker struct {
-	name string
+	name              string
+	maxNumOfEndpoints int // maximum number of endpoints to pick
 }
 
 // Type returns the type of the picker.
@@ -71,7 +86,23 @@ func (p *RandomPicker) WithName(name string) *RandomPicker {
 
 // Pick selects a random pod from the list of candidates.
 func (p *RandomPicker) Pick(ctx context.Context, _ *types.CycleState, scoredPods []*types.ScoredPod) *types.ProfileRunResult {
-	log.FromContext(ctx).V(logutil.DEBUG).Info(fmt.Sprintf("Selecting a random pod from %d candidates: %+v", len(scoredPods), scoredPods))
-	i := rand.Intn(len(scoredPods))
-	return &types.ProfileRunResult{TargetPod: scoredPods[i]}
+	log.FromContext(ctx).V(logutil.DEBUG).Info(fmt.Sprintf("Selecting maximum '%d' pods from %d candidates randomly: %+v", p.maxNumOfEndpoints,
+		len(scoredPods), scoredPods))
+
+	// Shuffle in-place
+	rand.Shuffle(len(scoredPods), func(i, j int) {
+		scoredPods[i], scoredPods[j] = scoredPods[j], scoredPods[i]
+	})
+
+	// if we have enough pods to return keep only the "maxNumOfEndpoints" highest scored pods
+	if p.maxNumOfEndpoints < len(scoredPods) {
+		scoredPods = scoredPods[:p.maxNumOfEndpoints]
+	}
+
+	targetPods := make([]types.Pod, len(scoredPods))
+	for i, scoredPod := range scoredPods {
+		targetPods[i] = scoredPod
+	}
+
+	return &types.ProfileRunResult{TargetPods: targetPods}
 }

pkg/epp/scheduling/framework/scheduler_profile_test.go

@@ -141,11 +141,12 @@ func TestSchedulePlugins(t *testing.T) {
 			}
 
 			// Validate output
-			wantPod := &types.PodMetrics{
-				Pod: &backend.Pod{NamespacedName: test.wantTargetPod, Labels: make(map[string]string)},
-			}
 			wantRes := &types.ProfileRunResult{
-				TargetPod: wantPod,
+				TargetPods: []types.Pod{
+					&types.PodMetrics{
+						Pod: &backend.Pod{NamespacedName: test.wantTargetPod, Labels: make(map[string]string)},
+					},
+				},
 			}
 
 			if diff := cmp.Diff(wantRes, got); diff != "" {
@@ -231,15 +232,15 @@ func (tp *testPlugin) Pick(_ context.Context, _ *types.CycleState, scoredPods []
 	tp.PickCallCount++
 	tp.NumOfPickerCandidates = len(scoredPods)
 
-	var winnerPod types.Pod
+	winnerPods := []types.Pod{}
 	for _, scoredPod := range scoredPods {
 		if scoredPod.GetPod().NamespacedName.String() == tp.PickRes.String() {
-			winnerPod = scoredPod.Pod
+			winnerPods = append(winnerPods, scoredPod.Pod)
 			tp.WinnderPodScore = scoredPod.Score
 		}
 	}
 
-	return &types.ProfileRunResult{TargetPod: winnerPod}
+	return &types.ProfileRunResult{TargetPods: winnerPods}
 }
 
 func (tp *testPlugin) PostCycle(_ context.Context, _ *types.CycleState, res *types.ProfileRunResult) {

pkg/epp/scheduling/scheduler_test.go

@@ -95,22 +95,23 @@ func TestSchedule(t *testing.T) {
 			wantRes: &types.SchedulingResult{
 				ProfileResults: map[string]*types.ProfileRunResult{
 					"default": {
-						TargetPod: &types.ScoredPod{
-							Pod: &types.PodMetrics{
-								Pod: &backend.Pod{NamespacedName: k8stypes.NamespacedName{Name: "pod2"}, Labels: make(map[string]string)},
-								MetricsState: &backendmetrics.MetricsState{
-									WaitingQueueSize:    3,
-									KVCacheUsagePercent: 0.1,
-									MaxActiveModels:     2,
-									ActiveModels: map[string]int{
-										"foo":      1,
-										"critical": 1,
+						TargetPods: []types.Pod{
+							&types.ScoredPod{
+								Pod: &types.PodMetrics{
+									Pod: &backend.Pod{NamespacedName: k8stypes.NamespacedName{Name: "pod2"}, Labels: make(map[string]string)},
+									MetricsState: &backendmetrics.MetricsState{
+										WaitingQueueSize:    3,
+										KVCacheUsagePercent: 0.1,
+										MaxActiveModels:     2,
+										ActiveModels: map[string]int{
+											"foo":      1,
+											"critical": 1,
+										},
+										WaitingModels: map[string]int{},
 									},
-									WaitingModels: map[string]int{},
 								},
 							},
-						},
-					},
+						}},
 				},
 				PrimaryProfileName: "default",
 			},

pkg/epp/scheduling/types/types.go

@@ -80,7 +80,7 @@ func ToSchedulerPodMetrics(pods []backendmetrics.PodMetrics) []Pod {
 
 // ProfileRunResult captures the profile run result.
 type ProfileRunResult struct {
-	TargetPod Pod
+	TargetPods []Pod
 }
 
 // SchedulingResult captures the result of the scheduling cycle.