implemented PrefixAwareScorer based on Ricardo's work

Author: vMaroon

pkg/epp/scheduling/local_config.go

@@ -30,10 +30,13 @@ import (
 const (
 	kvCacheScorerEnablementEnvVar   = "ENABLE_KVCACHE_AWARE_SCORER"
 	loadAwareScorerEnablementEnvVar = "ENABLE_LOAD_AWARE_SCORER"
-	pdFilterEnablementEnvVar        = "ENABLE_PD_FILTER"
+	prefixScorerEnablementEnvVar    = "ENABLE_PREFIX_AWARE_SCORER"
+
+	pdFilterEnablementEnvVar = "ENABLE_PD_FILTER"
 
 	kvCacheScorerWeightEnvVar   = "KVCACHE_AWARE_SCORER_WEIGHT"
 	loadAwareScorerWeightEnvVar = "LOAD_AWARE_SCORER_WEIGHT"
+	prefixScorerWeightEnvVar    = "PREFIX_AWARE_SCORER_WEIGHT"
 )
 
 func init() {
@@ -46,6 +49,7 @@ func setDefaultConfig() {
 	setLoadAwareScorer()
 	setKVCacheAwareScorer()
 	setPDFilter()
+	setPrefixScorer()
 
 	defaultConfig.picker = picker.NewMaxScorePicker()
 }
@@ -96,3 +100,20 @@ func setPDFilter() {
 	defaultConfig.filters = append(defaultConfig.filters, filter.PDFilter)
 	loggerDebug.Info("Initialized PDFilter")
 }
+
+func setPrefixScorer() {
+	ctx := context.Background()
+	loggerDebug := log.FromContext(ctx).WithName("scheduler_config").V(logutil.DEBUG)
+
+	if envutil.GetEnvString(prefixScorerEnablementEnvVar, "false", loggerDebug) != "true" {
+		loggerDebug.Info("Skipping PrefixScorer creation as it is not enabled")
+		return
+	}
+
+	prefixScorerWeight := envutil.GetEnvInt(prefixScorerWeightEnvVar, 1, loggerDebug)
+	prefixScorer := scorer.NewPrefixAwareScorer(nil)
+	defaultConfig.scorers[prefixScorer] = prefixScorerWeight // TODO: make configurable
+	defaultConfig.postResponsePlugins = append(defaultConfig.postResponsePlugins, prefixScorer)
+
+	loggerDebug.Info("Initialized PrefixAwareScorer", "weight", prefixScorerWeight)
+}

pkg/epp/scheduling/plugins/scorer/kvcache-aware-scorer.go

@@ -97,7 +97,20 @@ func (s *KVCacheAwareScorer) Score(ctx *types.SchedulingContext, pods []types.Po
 	}
 	loggerDebug.Info("Got pod scores", "scores", scores)
 
-	return indexerScoresToNormalizedScoredPods(pods, scores)
+	if len(scores) == 0 {
+		loggerDebug.Info("No scores found for pods")
+		return nil
+	}
+
+	podToKey := func(pod types.Pod) (string, bool) {
+		metricsPod := pod.GetPod()
+		if metricsPod == nil {
+			return "", false
+		}
+		return metricsPod.Address, true
+	}
+
+	return indexedScoresToNormalizedScoredPods(pods, podToKey, scores)
 }
 
 func getMinMax(scores map[string]int) (int, int) {
@@ -116,17 +129,21 @@ func getMinMax(scores map[string]int) (int, int) {
 	return minScore, maxScore
 }
 
-func indexerScoresToNormalizedScoredPods(pods []types.Pod, scores map[string]int) map[types.Pod]float64 {
+// podToKey is a function type that converts a Pod to a string key.
+// It returns the key and a boolean indicating success.
+type podToKeyFunc func(pod types.Pod) (string, bool)
+
+func indexedScoresToNormalizedScoredPods(pods []types.Pod, podToKey podToKeyFunc, scores map[string]int) map[types.Pod]float64 {
 	scoredPods := make(map[types.Pod]float64)
 	minScore, maxScore := getMinMax(scores)
 
 	for _, pod := range pods {
-		metricsPod := pod.GetPod()
-		if metricsPod == nil {
+		key, ok := podToKey(pod)
+		if !ok {
 			continue
 		}
 
-		if score, ok := scores[metricsPod.Address]; ok {
+		if score, ok := scores[key]; ok {
 			if minScore == maxScore {
 				scoredPods[pod] = 1.0
 				continue

pkg/epp/scheduling/plugins/scorer/prefix_aware_scorer.go

@@ -0,0 +1,100 @@
+/*
+Copyright 2025 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package scorer
+
+import (
+	"sigs.k8s.io/controller-runtime/pkg/log"
+	"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"
+)
+
+const prefixAwareScorerName = "prefix-aware-scorer"
+
+// PrefixAwareScorer is a routing scorer that scores pods based on the longest prefix match
+// between the request's prompt and stored prefixes. The score is normalized between 0 and 1,
+// where 1 represents the longest matching prefix.
+type PrefixAwareScorer struct {
+	prefixStore *PrefixStore
+}
+
+var _ plugins.Scorer = &PrefixAwareScorer{}
+
+// NewPrefixAwareScorer creates a new PrefixAwareScorer with the given
+// PrefixStoreConfig. If the config is nil, default is used.
+func NewPrefixAwareScorer(config *PrefixStoreConfig) *PrefixAwareScorer {
+	return &PrefixAwareScorer{
+		prefixStore: NewPrefixStore(config),
+	}
+}
+
+func (s *PrefixAwareScorer) Name() string {
+	return "prefix-aware-scorer"
+}
+
+// Score scores the target pods based on the longest prefix match.
+func (s *PrefixAwareScorer) Score(ctx *types.SchedulingContext, pods []types.Pod) map[types.Pod]float64 {
+	loggerDebug := log.FromContext(ctx).WithName(prefixAwareScorerName).V(logutil.DEBUG)
+	if ctx.Req == nil {
+		loggerDebug.Info("Request is nil, skipping scoring")
+		return nil
+	}
+
+	scores := s.prefixStore.FindMatchingPods(ctx.Req.Prompt, ctx.Req.Model)
+	loggerDebug.Info("Got pod scores", "scores", scores)
+
+	if len(scores) == 0 {
+		loggerDebug.Info("No scores found for pods")
+		return nil
+	}
+
+	podToKey := func(pod types.Pod) (string, bool) {
+		if pod.GetPod() == nil {
+			return "", false
+		}
+
+		return pod.GetPod().NamespacedName.String(), true
+	}
+
+	return indexedScoresToNormalizedScoredPods(pods, podToKey, scores)
+}
+
+// PostResponse implements the PostResponsePlugin interface.
+// It adds the prefix to the PrefixStore for the given pod.
+func (s *PrefixAwareScorer) PostResponse(ctx *types.SchedulingContext, pod types.Pod) {
+	debugLogger := log.FromContext(ctx).WithName(prefixAwareScorerName).V(logutil.DEBUG)
+
+	if ctx.Req == nil {
+		debugLogger.Info("Request is nil, skipping PostResponse")
+		return
+	}
+
+	if pod.GetPod() == nil {
+		debugLogger.Info("Pod is nil, skipping PostResponse", "req", ctx.Req, "pod", pod)
+		return
+	}
+
+	if err := s.prefixStore.AddEntry(ctx.Req.Model, ctx.Req.Prompt, &pod.GetPod().NamespacedName); err != nil {
+		debugLogger.Error(err, "Failed to add entry to prefix store", "req", ctx.Req, "pod", pod)
+		return
+	}
+}
+
+// GetPrefixStore returns the scorer's PrefixStore.
+func (s *PrefixAwareScorer) GetPrefixStore() *PrefixStore {
+	return s.prefixStore
+}

pkg/epp/scheduling/prefix_aware_scorer_test.go renamed to pkg/epp/scheduling/plugins/scorer/prefix_aware_scorer_test.go

@@ -14,32 +14,23 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package scheduling
+package scorer_test
 
 import (
 	"context"
-	"testing"
-	"time"
-
 	k8stypes "k8s.io/apimachinery/pkg/types"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 	backendmetrics "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend/metrics"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/plugins/scorer"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
+	"testing"
 )
 
 func TestPrefixAwareScorer(t *testing.T) {
 	ctx := context.Background()
 	logger := log.FromContext(ctx)
 	ctx = log.IntoContext(ctx, logger)
 
-	// Create a prefix store with test configuration
-	prefixStore := NewPrefixStore(PrefixStoreConfig{
-		MaxEntries:   100,
-		MinPrefixLen: 3,
-		MaxPrefixLen: 10,
-		EntryTTL:     1 * time.Hour,
-	})
-
 	// Create test pods
 	pod1 := &types.PodMetrics{
 		Pod: &backendmetrics.Pod{
@@ -68,7 +59,7 @@ func TestPrefixAwareScorer(t *testing.T) {
 		prefixToAdd    string
 		podToAdd       k8stypes.NamespacedName
 		prefixModel    string // Model name to use when adding the prefix
-		expectedScores []float64
+		expectedScores map[types.Pod]float64
 	}{
 		{
 			name:           "no prompt",
@@ -78,17 +69,20 @@ func TestPrefixAwareScorer(t *testing.T) {
 			prefixToAdd:    "hello",
 			podToAdd:       pod1.Pod.NamespacedName,
 			prefixModel:    "model1",
-			expectedScores: []float64{0, 0}, // No prompt means zero scores
+			expectedScores: map[types.Pod]float64{}, // No prompt means zero scores
 		},
 		{
-			name:           "exact prefix match",
-			weight:         1.0,
-			prompt:         "hello world",
-			modelName:      "model1",
-			prefixToAdd:    "hello",
-			podToAdd:       pod1.Pod.NamespacedName,
-			prefixModel:    "model1",
-			expectedScores: []float64{1.0, 0}, // pod1 matches, pod2 doesn't
+			name:        "exact prefix match",
+			weight:      1.0,
+			prompt:      "hello world",
+			modelName:   "model1",
+			prefixToAdd: "hello",
+			podToAdd:    pod1.Pod.NamespacedName,
+			prefixModel: "model1",
+			expectedScores: map[types.Pod]float64{
+				pod1: 1.0,
+				pod2: 0.0,
+			}, // pod1 matches, pod2 doesn't
 		},
 		{
 			name:           "no prefix match",
@@ -98,7 +92,7 @@ func TestPrefixAwareScorer(t *testing.T) {
 			prefixToAdd:    "hello",
 			podToAdd:       pod1.Pod.NamespacedName,
 			prefixModel:    "model1",
-			expectedScores: []float64{0, 0}, // No matching prefix
+			expectedScores: map[types.Pod]float64{}, // No matching prefix
 		},
 		{
 			name:           "different model name",
@@ -107,63 +101,54 @@ func TestPrefixAwareScorer(t *testing.T) {
 			modelName:      "model2", // Try to find with model2
 			prefixToAdd:    "hello",
 			podToAdd:       pod1.Pod.NamespacedName,
-			prefixModel:    "model1",        // But prefix was added with model1
-			expectedScores: []float64{0, 0}, // Model name mismatch should result in no match
+			prefixModel:    "model1",                // But prefix was added with model1
+			expectedScores: map[types.Pod]float64{}, // Model name mismatch should result in no match
 		},
 		{
-			name:           "custom weight",
-			weight:         0.5,
-			prompt:         "hello world",
-			modelName:      "model1",
-			prefixToAdd:    "hello",
-			podToAdd:       pod1.Pod.NamespacedName,
-			prefixModel:    "model1",
-			expectedScores: []float64{0.5, 0}, // Weight affects score
+			name:        "custom weight",
+			weight:      0.5,
+			prompt:      "hello world",
+			modelName:   "model1",
+			prefixToAdd: "hello",
+			podToAdd:    pod1.Pod.NamespacedName,
+			prefixModel: "model1",
+			expectedScores: map[types.Pod]float64{
+				pod1: 0.5, // Pod1 matches with weight
+				pod2: 0.0, // Pod2 doesn't match
+			}, // Weight affects score
 		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			// Reset prefix store for each test
-			prefixStore = NewPrefixStore(PrefixStoreConfig{
-				MaxEntries:   100,
-				MinPrefixLen: 3,
-				MaxPrefixLen: 10,
-				EntryTTL:     1 * time.Hour,
-			})
+			config := scorer.DefaultPrefixStoreConfig()
+			config.BlockSize = 5 // set small chunking for testing
+
+			s := scorer.NewPrefixAwareScorer(config)
 
 			// Add prefix if specified
 			if tt.prefixToAdd != "" {
-				err := prefixStore.AddPrefix(ctx, tt.prefixToAdd, tt.podToAdd, tt.prefixModel)
+				err := s.GetPrefixStore().AddEntry(tt.prefixModel,
+					tt.prefixToAdd, &tt.podToAdd)
 				if err != nil {
 					t.Fatalf("Failed to add prefix: %v", err)
 				}
 			}
 
-			// Create scorer with test weight
-			scorer := NewPrefixAwareScorer(tt.weight, prefixStore)
-
 			// Create test context
-			sCtx := types.NewContext(ctx, &types.LLMRequest{
+			sCtx := types.NewSchedulingContext(ctx, &types.LLMRequest{
 				Prompt:              tt.prompt,
 				ResolvedTargetModel: tt.modelName,
-			}, []*types.PodMetrics{})
+			}, []types.Pod{}, 0)
 
 			// Score pods
-			pods := []*types.PodMetrics{pod1, pod2}
-			scores, err := scorer.ScoreTargets(sCtx, pods)
-			if err != nil {
-				t.Fatalf("Unexpected error: %v", err)
-			}
-
-			// Verify scores
-			if len(scores) != len(tt.expectedScores) {
-				t.Fatalf("Expected %d scores, got %d", len(tt.expectedScores), len(scores))
-			}
+			pods := []types.Pod{pod1, pod2}
+			scores := s.Score(sCtx, pods)
 
-			for i, score := range scores {
-				if score.Score != tt.expectedScores[i] {
-					t.Errorf("Pod %d: expected score %v, got %v", i, tt.expectedScores[i], score.Score)
+			for p, score := range scores {
+				if score != tt.expectedScores[p] {
+					t.Errorf("Pod %v: expected score %v, got %v", p, tt.expectedScores[p], score)
 				}
 			}
 		})