Add more tests, register prefix cache scorer and address other review comments

Author: rahulgurnani

cmd/epp/runner/runner.go

@@ -240,7 +240,7 @@ func (r *Runner) Run(ctx context.Context) error {
 	}
 
 	// --- Setup Datastore ---
-	epf, err := r.setupMetricsCollection(setupLog, r.featureGates[datalayer.FeatureGate])
+	epf, err := r.setupMetricsCollection(setupLog, r.featureGates[datalayer.ExperimentalDatalayerFeatureGate])
 	if err != nil {
 		return err
 	}
@@ -372,7 +372,7 @@ func (r *Runner) Run(ctx context.Context) error {
 		MetricsStalenessThreshold:        *metricsStalenessThreshold,
 		Director:                         director,
 		SaturationDetector:               saturationDetector,
-		UseExperimentalDatalayerV2:       r.featureGates[datalayer.FeatureGate], // pluggable data layer feature flag
+		UseExperimentalDatalayerV2:       r.featureGates[datalayer.ExperimentalDatalayerFeatureGate], // pluggable data layer feature flag
 	}
 	if err := serverRunner.SetupWithManager(ctx, mgr); err != nil {
 		setupLog.Error(err, "Failed to setup EPP controllers")
@@ -463,7 +463,7 @@ func (r *Runner) parseConfigurationPhaseOne(ctx context.Context) (*configapi.End
 		}
 	}
 
-	loader.RegisterFeatureGate(datalayer.FeatureGate)
+	loader.RegisterFeatureGate(datalayer.ExperimentalDatalayerFeatureGate)
 	loader.RegisterFeatureGate(flowcontrol.FeatureGate)
 	loader.RegisterFeatureGate(datalayer.PrepareDataPluginsFeatureGate)
 
@@ -506,9 +506,15 @@ func (r *Runner) parseConfigurationPhaseTwo(ctx context.Context, rawConfig *conf
 	// Add requestControl plugins
 	r.requestControlConfig.AddPlugins(handle.GetAllPlugins()...)
 
-	// Sort prepare data plugins in DAG order (topological sort). Also check prepare data plugins for cycles.
-	if r.requestControlConfig.PrepareDataPluginGraph(r.featureGates[datalayer.PrepareDataPluginsFeatureGate]) != nil {
-		return nil, errors.New("failed to load the configuration - prepare data plugins have cyclic dependencies")
+	// TODO(rahulgurnani): Remove feature gate check once prepare data plugins are stable.
+	if r.featureGates[datalayer.PrepareDataPluginsFeatureGate] {
+		// Sort prepare data plugins in DAG order (topological sort). Also check prepare data plugins for cycles.
+		if r.requestControlConfig.PrepareDataPluginGraph() != nil {
+			return nil, errors.New("failed to load the configuration - prepare data plugins have cyclic dependencies")
+		}
+		plugins.Register(scorer.PrefixCacheMatchScorerType, scorer.PrefixCacheScorerFactory) // register PrefixCacheMatchScorer used by prepare data plugins
+	} else {
+		r.requestControlConfig.WithPrepareDataPlugins()
 	}
 
 	// Handler deprecated configuration options
@@ -531,7 +537,7 @@ func (r *Runner) deprecatedConfigurationHelper(cfg *config.Config, logger logr.L
 
 	if _, ok := os.LookupEnv(enableExperimentalDatalayerV2); ok {
 		logger.Info("Enabling the experimental Data Layer V2 using environment variables is deprecated and will be removed in next version")
-		r.featureGates[datalayer.FeatureGate] = env.GetEnvBool(enableExperimentalDatalayerV2, false, logger)
+		r.featureGates[datalayer.ExperimentalDatalayerFeatureGate] = env.GetEnvBool(enableExperimentalDatalayerV2, false, logger)
 	}
 	if _, ok := os.LookupEnv(enableExperimentalFlowControlLayer); ok {
 		logger.Info("Enabling the experimental Flow Control layer using environment variables is deprecated and will be removed in next version")

pkg/epp/config/loader/configloader_test.go

@@ -100,7 +100,7 @@ func TestLoadRawConfiguration(t *testing.T) {
 				},
 			},
 		},
-		FeatureGates: configapi.FeatureGates{datalayer.FeatureGate},
+		FeatureGates: configapi.FeatureGates{datalayer.ExperimentalDatalayerFeatureGate},
 		SaturationDetector: &configapi.SaturationDetector{
 			MetricsStalenessThreshold: metav1.Duration{Duration: 150 * time.Millisecond},
 		},
@@ -206,7 +206,7 @@ func TestLoadRawConfigurationWithDefaults(t *testing.T) {
 				},
 			},
 		},
-		FeatureGates: configapi.FeatureGates{datalayer.FeatureGate},
+		FeatureGates: configapi.FeatureGates{datalayer.ExperimentalDatalayerFeatureGate},
 		SaturationDetector: &configapi.SaturationDetector{
 			QueueDepthThreshold:       saturationdetector.DefaultQueueDepthThreshold,
 			KVCacheUtilThreshold:      saturationdetector.DefaultKVCacheUtilThreshold,
@@ -488,7 +488,7 @@ func TestLoadConfig(t *testing.T) {
 }
 
 func registerNeededFeatureGates() {
-	RegisterFeatureGate(datalayer.FeatureGate)
+	RegisterFeatureGate(datalayer.ExperimentalDatalayerFeatureGate)
 }
 
 func registerNeededPlgugins() {

pkg/epp/datalayer/factory.go

@@ -26,8 +26,8 @@ import (
 )
 
 const (
-	FeatureGate                   = "dataLayer"
-	PrepareDataPluginsFeatureGate = "prepareDataPlugins"
+	ExperimentalDatalayerFeatureGate = "dataLayer"
+	PrepareDataPluginsFeatureGate    = "prepareDataPlugins"
 )
 
 // PoolInfo represents the DataStore information needed for endpoints.

pkg/epp/requestcontrol/director.go

@@ -388,6 +388,9 @@ func (d *Director) runPreRequestPlugins(ctx context.Context, request *scheduling
 
 func (d *Director) runPrepareDataPlugins(ctx context.Context,
 	request *schedulingtypes.LLMRequest, pods []schedulingtypes.Pod) error {
+	if len(d.requestControlPlugins.prepareDataPlugins) == 0 {
+		return nil
+	}
 	return prepareDataPluginsWithTimeout(prepareDataTimeout, d.requestControlPlugins.prepareDataPlugins, ctx, request, pods)
 }
 

pkg/epp/requestcontrol/request_control_config.go

@@ -107,9 +107,8 @@ func (c *Config) AddPlugins(pluginObjects ...plugins.Plugin) {
 
 // PrepareDataPluginGraph creates data dependency graph and sorts the plugins in topological order.
 // If a cycle is detected, it returns an error.
-func (c *Config) PrepareDataPluginGraph(enablePrepareDataPlugins bool) error {
-	if !enablePrepareDataPlugins {
-		c.prepareDataPlugins = []PrepareDataPlugin{}
+func (c *Config) PrepareDataPluginGraph() error {
+	if len(c.prepareDataPlugins) == 0 {
 		return nil
 	}
 	dag := buildDAG(c.prepareDataPlugins)

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

@@ -222,11 +222,6 @@ func (p *Plugin) PrepareRequestData(ctx context.Context, request *types.LLMReque
 		PrefixHashes:       hashes,
 		PrefixCacheServers: p.matchLongestPrefix(ctx, hashes),
 	}
-	for server, matchLen := range state.PrefixCacheServers {
-		log.FromContext(ctx).V(logutil.TRACE).Info("prefix cached state", "server", server, "longest-prefix-match", matchLen)
-
-	}
-
 	total := len(state.PrefixHashes)
 
 	for _, pod := range pods {

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

@@ -29,6 +29,8 @@ import (
 
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend"
 	backendmetrics "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend/metrics"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/datalayer"
+	dplugins "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/datalayer/plugins"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/plugins"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/requestcontrol"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
@@ -575,6 +577,67 @@ func randomPrompt(n int) string {
 	return sb.String()
 }
 
+func TestPrepareRequestData(t *testing.T) {
+	config := Config{
+		BlockSize:              4,
+		MaxPrefixBlocksToMatch: DefaultMaxPrefixBlocks,
+		LRUCapacityPerServer:   DefaultLRUCapacityPerServer,
+	}
+	plugin := New(context.Background(), config)
+
+	pod1 := &types.PodMetrics{Pod: &backend.Pod{NamespacedName: k8stypes.NamespacedName{Name: "pod1"}}, MetricsState: backendmetrics.NewMetricsState(), AttributeMap: datalayer.NewAttributes()}
+	pod2 := &types.PodMetrics{Pod: &backend.Pod{NamespacedName: k8stypes.NamespacedName{Name: "pod2"}}, MetricsState: backendmetrics.NewMetricsState(), AttributeMap: datalayer.NewAttributes()}
+	pods := []types.Pod{pod1, pod2}
+
+	// First request to populate cache.
+	req1 := &types.LLMRequest{
+		RequestId:   uuid.NewString(),
+		TargetModel: "test-model1",
+		Body: &types.LLMRequestBody{
+			Completions: &types.CompletionsRequest{
+				Prompt: "aaaabbbb",
+			},
+		},
+	}
+	_ = plugin.Score(context.Background(), types.NewCycleState(), req1, pods)
+	schedulingResult := &types.SchedulingResult{
+		PrimaryProfileName: "default",
+		ProfileResults: map[string]*types.ProfileRunResult{
+			"default": {TargetPods: []types.Pod{pod1}},
+		},
+	}
+	plugin.PreRequest(context.Background(), req1, schedulingResult)
+	plugin.wg.Wait()
+
+	// Second request that shares a prefix.
+	req2 := &types.LLMRequest{
+		RequestId:   uuid.NewString(),
+		TargetModel: "test-model1",
+		Body: &types.LLMRequestBody{
+			Completions: &types.CompletionsRequest{
+				Prompt: "aaaacccc",
+			},
+		},
+	}
+
+	err := plugin.PrepareRequestData(context.Background(), req2, pods)
+	assert.NoError(t, err)
+
+	// Verify pod1 has the correct prefix match info
+	info1, ok := pod1.Get(dplugins.PrefixCacheMatchInfoKey)
+	assert.True(t, ok)
+	prefixInfo1 := info1.(*dplugins.PrefixCacheMatchInfo)
+	assert.Equal(t, 1, prefixInfo1.MatchLength()) // "aaaa" matches
+	assert.Equal(t, 2, prefixInfo1.TotalLength()) // "aaaacccc" -> 2 blocks
+
+	// Verify pod2 has no match info
+	info2, ok := pod2.Get(dplugins.PrefixCacheMatchInfoKey)
+	assert.True(t, ok)
+	prefixInfo2 := info2.(*dplugins.PrefixCacheMatchInfo)
+	assert.Equal(t, 0, prefixInfo2.MatchLength()) // No match for pod2
+	assert.Equal(t, 2, prefixInfo2.TotalLength())
+}
+
 // BenchmarkPrefixPluginChatCompletionsStress is a stress test for chat completions with varying message counts and lengths
 func BenchmarkPrefixPluginChatCompletionsStress(b *testing.B) {
 	blockSize := 8

pkg/epp/scheduling/framework/plugins/scorer/prefix_cache_match_scorer.go

@@ -60,7 +60,9 @@ func (s *PrefixCacheScorer) TypedName() plugins.TypedName {
 
 // Consumes returns the list of data that is consumed by the plugin.
 func (s *PrefixCacheScorer) Consumes() map[string]any {
-	return map[string]any{}
+	return map[string]any{
+		dplugins.PrefixCacheMatchInfoKey: &dplugins.PrefixCacheMatchInfo{},
+	}
 }
 
 // WithName sets the name of the scorer.
@@ -74,12 +76,12 @@ func (s *PrefixCacheScorer) Score(_ context.Context, cycleState *types.CycleStat
 	scores := make(map[types.Pod]float64, len(pods))
 
 	for _, pod := range pods {
-		matchPercent, ok := pod.Get(dplugins.PrefixCacheMatchInfoKey)
+		matchInfo, ok := pod.Get(dplugins.PrefixCacheMatchInfoKey)
 		if !ok {
 			scores[pod] = 0.0
 			continue
 		}
-		scores[pod] = float64(matchPercent.(*dplugins.PrefixCacheMatchInfo).MatchLength()) / float64(matchPercent.(*dplugins.PrefixCacheMatchInfo).TotalLength()) * 100
+		scores[pod] = float64(matchInfo.(*dplugins.PrefixCacheMatchInfo).MatchLength()) / float64(matchInfo.(*dplugins.PrefixCacheMatchInfo).TotalLength())
 	}
 	return scores
 }

pkg/epp/scheduling/framework/plugins/scorer/prefix_cache_match_scorer_test.go

@@ -86,8 +86,8 @@ func TestPrefixCacheScorer_Score(t *testing.T) {
 			name: "pods with prefix cache match percent",
 			pods: []types.Pod{pod1, pod2},
 			expected: map[types.Pod]float64{
-				pod1: 50.0,
-				pod2: 100.0,
+				pod1: 0.5,
+				pod2: 1.0,
 			},
 		},
 		{
@@ -101,7 +101,7 @@ func TestPrefixCacheScorer_Score(t *testing.T) {
 			name: "mixed pods",
 			pods: []types.Pod{pod1, pod3},
 			expected: map[types.Pod]float64{
-				pod1: 50.0,
+				pod1: 0.5,
 				pod3: 0.0,
 			},
 		},
@@ -146,5 +146,5 @@ func TestPrefixCacheScorer_TypedName(t *testing.T) {
 func TestPrefixCacheScorer_Consumes(t *testing.T) {
 	scorer := NewPrefixCacheScorer()
 	consumes := scorer.Consumes()
-	assert.Empty(t, consumes)
+	assert.Contains(t, consumes, dplugins.PrefixCacheMatchInfoKey)
 }