Add fallback logic

Author: rlakhtakia

pkg/epp/requestcontrol/director.go

@@ -239,6 +239,7 @@ func (d *Director) prepareRequest(ctx context.Context, reqCtx *handlers.RequestC
 	}
 	// primary profile is used to set destination
 	targetPod := result.ProfileResults[result.PrimaryProfileName].TargetPod.GetPod()
+	backupPods := result.ProfileResults[result.PrimaryProfileName].FallbackPods
 
 	pool, err := d.datastore.PoolGet()
 	if err != nil {
@@ -249,8 +250,20 @@ func (d *Director) prepareRequest(ctx context.Context, reqCtx *handlers.RequestC
 	endpoint := net.JoinHostPort(targetPod.Address, strconv.Itoa(targetPort))
 	logger.V(logutil.DEFAULT).Info("Request handled", "model", reqCtx.Model, "targetModel", reqCtx.ResolvedTargetModel, "endpoint", targetPod)
 
+	allEndpoints := []string{endpoint}
+
+	for _, pod := range backupPods {
+		curPod := pod.GetPod()
+		curEndpoint := net.JoinHostPort(curPod.Address, strconv.Itoa(targetPort))
+
+		allEndpoints = append(allEndpoints, curEndpoint)
+
+	}
+
+	combinedEndpointsString := strings.Join(allEndpoints, ",")
+
 	reqCtx.TargetPod = targetPod
-	reqCtx.TargetEndpoint = endpoint
+	reqCtx.TargetEndpoint = combinedEndpointsString
 
 	d.runPreRequestPlugins(ctx, reqCtx.SchedulingRequest, result, targetPort)
 

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"
@@ -80,9 +81,27 @@ func (p *MaxScorePicker) Pick(ctx context.Context, cycleState *types.CycleState,
 		}
 	}
 
+	var target types.Pod
+	fallbackPods := []*types.ScoredPod{}
 	if len(highestScorePods) > 1 {
-		return p.random.Pick(ctx, cycleState, highestScorePods) // pick randomly from the highest score pods
+		profile := p.random.Pick(ctx, cycleState, highestScorePods) // pick randomly from the highest score pods
+		target = profile.TargetPod
+		fallbackPods = profile.FallbackPods
+	} else {
+		target = highestScorePods[0]
 	}
 
-	return &types.ProfileRunResult{TargetPod: highestScorePods[0]}
+	// Add all pods that had a lower score than highestScorePods
+	for _, pod := range scoredPods {
+		if pod.Score < maxScore {
+			fallbackPods = append(fallbackPods, pod)
+		}
+	}
+
+	// Sort the final fallback pods by score in descending order
+	sort.Slice(fallbackPods, func(i, j int) bool {
+		return fallbackPods[i].Score > fallbackPods[j].Score
+	})
+
+	return &types.ProfileRunResult{TargetPod: target, FallbackPods: fallbackPods}
 }

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

@@ -20,18 +20,70 @@ import (
 	"context"
 	"testing"
 
+	"github.com/google/go-cmp/cmp"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/plugins"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/framework"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
 
 	k8stypes "k8s.io/apimachinery/pkg/types"
 )
 
+var _ framework.Picker = &TestRandomPicker{}
+
+// NewTestRandomPicker initializes a new NewTestRandomPicker and returns its pointer.
+func NewTestRandomPicker(pickRes string) *TestRandomPicker {
+	return &TestRandomPicker{
+		Picked: pickRes,
+		tn:     plugins.TypedName{Type: "random-test", Name: "random-test"},
+	}
+}
+
+// NewTestRandomPicker picks the selected pod from the list of candidates.
+type TestRandomPicker struct {
+	tn     plugins.TypedName
+	Picked string
+}
+
+// TypedName returns the type and name tuple of this plugin instance.
+func (p *TestRandomPicker) TypedName() plugins.TypedName {
+	return p.tn
+}
+
+// WithName sets the name of the picker.
+func (p *TestRandomPicker) WithName(name string) *TestRandomPicker {
+	p.tn.Name = name
+	return p
+}
+
+// Type returns the type of the picker.
+func (p *TestRandomPicker) Type() string {
+	return RandomPickerType
+}
+
+func (tp *TestRandomPicker) Pick(_ context.Context, _ *types.CycleState, scoredPods []*types.ScoredPod) *types.ProfileRunResult {
+	fallbackPods := []*types.ScoredPod{}
+
+	var winnerPod types.Pod
+	for _, scoredPod := range scoredPods {
+		if scoredPod.GetPod().NamespacedName.String() == tp.Picked {
+			winnerPod = scoredPod
+		} else {
+			fallbackPods = append(fallbackPods, scoredPod)
+		}
+	}
+
+	return &types.ProfileRunResult{TargetPod: winnerPod, FallbackPods: fallbackPods}
+}
+
 func TestPickMaxScorePicker(t *testing.T) {
 	tests := []struct {
-		name        string
-		scoredPods  []*types.ScoredPod
-		wantNames   []string
-		shouldPanic bool
+		name              string
+		scoredPods        []*types.ScoredPod
+		wantPodName       string
+		wantFallbackNames []string
+		shouldPanic       bool
+		picker            framework.Picker
 	}{
 		{
 			name: "Single max score",
@@ -40,7 +92,9 @@ func TestPickMaxScorePicker(t *testing.T) {
 				{Pod: &types.PodMetrics{Pod: &backend.Pod{NamespacedName: k8stypes.NamespacedName{Name: "pod2"}}}, Score: 25},
 				{Pod: &types.PodMetrics{Pod: &backend.Pod{NamespacedName: k8stypes.NamespacedName{Name: "pod3"}}}, Score: 15},
 			},
-			wantNames: []string{"pod2"},
+			wantPodName:       "pod2",
+			wantFallbackNames: []string{"pod3", "pod1"},
+			picker:            NewMaxScorePicker(),
 		},
 		{
 			name: "Multiple max scores",
@@ -49,13 +103,16 @@ func TestPickMaxScorePicker(t *testing.T) {
 				{Pod: &types.PodMetrics{Pod: &backend.Pod{NamespacedName: k8stypes.NamespacedName{Name: "podB"}}}, Score: 50},
 				{Pod: &types.PodMetrics{Pod: &backend.Pod{NamespacedName: k8stypes.NamespacedName{Name: "podC"}}}, Score: 30},
 			},
-			wantNames: []string{"podA", "podB"},
+			wantPodName:       "podA",
+			wantFallbackNames: []string{"podB", "podC"},
+			picker:            NewTestRandomPicker(k8stypes.NamespacedName{Name: "podA"}.String()),
 		},
 		{
-			name:        "Empty pod list",
-			scoredPods:  []*types.ScoredPod{},
-			wantNames:   nil,
-			shouldPanic: true,
+			name:              "Empty pod list",
+			scoredPods:        []*types.ScoredPod{},
+			wantFallbackNames: nil,
+			shouldPanic:       true,
+			picker:            NewMaxScorePicker(),
 		},
 	}
 
@@ -69,8 +126,7 @@ func TestPickMaxScorePicker(t *testing.T) {
 				}()
 			}
 
-			p := NewMaxScorePicker()
-			result := p.Pick(context.Background(), nil, tt.scoredPods)
+			result := tt.picker.Pick(context.Background(), nil, tt.scoredPods)
 
 			if len(tt.scoredPods) == 0 && result != nil {
 				t.Errorf("expected nil result for empty input, got %+v", result)
@@ -79,15 +135,14 @@ func TestPickMaxScorePicker(t *testing.T) {
 
 			if result != nil {
 				got := result.TargetPod.GetPod().NamespacedName.Name
-				match := false
-				for _, want := range tt.wantNames {
-					if got == want {
-						match = true
-						break
-					}
+				if diff := cmp.Diff(tt.wantPodName, got); diff != "" {
+					t.Errorf("Unexpected target pod name (-want +got): %v", diff)
 				}
-				if !match {
-					t.Errorf("got %q, want one of %v", got, tt.wantNames)
+				gotFallback := result.FallbackPods
+				for i, wantName := range tt.wantFallbackNames {
+					if diff := cmp.Diff(wantName, gotFallback[i].GetPod().NamespacedName.Name); diff != "" {
+						t.Errorf("Unexpected target pod name (-want +got): %v", diff)
+					}
 				}
 			}
 		})

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

@@ -69,5 +69,15 @@ func (p *RandomPicker) WithName(name string) *RandomPicker {
 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]}
+	pods := []*types.ScoredPod{}
+
+	for index, pod := range scoredPods {
+		fmt.Println(pod)
+		if index == i {
+			continue
+		}
+		pods = append(pods, pod)
+	}
+
+	return &types.ProfileRunResult{TargetPod: scoredPods[i], FallbackPods: pods}
 }

pkg/epp/scheduling/framework/scheduler_profile_test.go

@@ -42,11 +42,12 @@ func TestSchedulePlugins(t *testing.T) {
 		[]k8stypes.NamespacedName{}, k8stypes.NamespacedName{Name: "pod1"})
 
 	tests := []struct {
-		name           string
-		profile        *SchedulerProfile
-		input          []types.Pod
-		wantTargetPod  k8stypes.NamespacedName
-		targetPodScore float64
+		name             string
+		profile          *SchedulerProfile
+		input            []types.Pod
+		wantTargetPod    k8stypes.NamespacedName
+		wantFallbackPods []*types.ScoredPod
+		targetPodScore   float64
 		// Number of expected pods to score (after filter)
 		numPodsToScore int
 		err            bool
@@ -63,7 +64,16 @@ func TestSchedulePlugins(t *testing.T) {
 				&types.PodMetrics{Pod: &backend.Pod{NamespacedName: k8stypes.NamespacedName{Name: "pod2"}}},
 				&types.PodMetrics{Pod: &backend.Pod{NamespacedName: k8stypes.NamespacedName{Name: "pod3"}}},
 			},
-			wantTargetPod:  k8stypes.NamespacedName{Name: "pod1"},
+			wantTargetPod: k8stypes.NamespacedName{Name: "pod1"},
+			wantFallbackPods: []*types.ScoredPod{
+				{
+					Pod: &types.PodMetrics{Pod: &backend.Pod{
+						NamespacedName: k8stypes.NamespacedName{Name: "pod2"},
+						Labels:         map[string]string{},
+					}},
+					Score: 1.1,
+				},
+			},
 			targetPodScore: 1.1,
 			numPodsToScore: 2,
 			err:            false,
@@ -80,7 +90,16 @@ func TestSchedulePlugins(t *testing.T) {
 				&types.PodMetrics{Pod: &backend.Pod{NamespacedName: k8stypes.NamespacedName{Name: "pod2"}}},
 				&types.PodMetrics{Pod: &backend.Pod{NamespacedName: k8stypes.NamespacedName{Name: "pod3"}}},
 			},
-			wantTargetPod:  k8stypes.NamespacedName{Name: "pod1"},
+			wantTargetPod: k8stypes.NamespacedName{Name: "pod1"},
+			wantFallbackPods: []*types.ScoredPod{
+				{
+					Pod: &types.PodMetrics{Pod: &backend.Pod{
+						NamespacedName: k8stypes.NamespacedName{Name: "pod2"},
+						Labels:         map[string]string{},
+					}},
+					Score: 50,
+				},
+			},
 			targetPodScore: 50,
 			numPodsToScore: 2,
 			err:            false,
@@ -138,7 +157,8 @@ func TestSchedulePlugins(t *testing.T) {
 				Pod: &backend.Pod{NamespacedName: test.wantTargetPod},
 			}
 			wantRes := &types.ProfileRunResult{
-				TargetPod: wantPod,
+				TargetPod:    wantPod,
+				FallbackPods: test.wantFallbackPods,
 			}
 
 			if diff := cmp.Diff(wantRes, got); diff != "" {
@@ -235,16 +255,20 @@ func (tp *testPlugin) Score(_ context.Context, _ *types.CycleState, _ *types.LLM
 func (tp *testPlugin) Pick(_ context.Context, _ *types.CycleState, scoredPods []*types.ScoredPod) *types.ProfileRunResult {
 	tp.PickCallCount++
 	tp.NumOfPickerCandidates = len(scoredPods)
+	fallbackPods := []*types.ScoredPod{}
 
 	var winnerPod types.Pod
 	for _, scoredPod := range scoredPods {
 		if scoredPod.GetPod().NamespacedName.String() == tp.PickRes.String() {
 			winnerPod = scoredPod.Pod
 			tp.WinnerPodScore = scoredPod.Score
+		} else {
+			fallbackPods = append(fallbackPods, scoredPod)
 		}
+
 	}
 
-	return &types.ProfileRunResult{TargetPod: winnerPod}
+	return &types.ProfileRunResult{TargetPod: winnerPod, FallbackPods: fallbackPods}
 }
 
 func (tp *testPlugin) PostCycle(_ context.Context, _ *types.CycleState, res *types.ProfileRunResult) {

pkg/epp/scheduling/scheduler_test.go

@@ -109,6 +109,7 @@ func TestSchedule(t *testing.T) {
 								},
 							},
 						},
+						FallbackPods: []*types.ScoredPod{},
 					},
 				},
 				PrimaryProfileName: "default",

pkg/epp/scheduling/types/types.go

@@ -72,7 +72,8 @@ type PodMetrics struct {
 
 // ProfileRunResult captures the profile run result.
 type ProfileRunResult struct {
-	TargetPod Pod
+	TargetPod    Pod
+	FallbackPods []*ScoredPod
 }
 
 // SchedulingResult captures the result of the scheduling cycle.