Implement PrepareDataPlugin for prefix cache match plugin (#1942)

* Add PrepareRequestData method for the prefix cache plugin

* Add prefix cache match scorer

* Ensure prefix cache plugin implements all preprae data plugin methods

* Enable prepare data plugins behind a feature flag

* Add static check to ensure plugin implements PrepareDataPlugin interface

* Rename prefix cache match

* Add feature gate for prepare data plugin

* Update prefix cache match to have total and match length

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

* Move prefix cache scorer out of this PR

* Fix rebase errors

* Update directory structure

* Update TODOs

* Validate type as well in dag validation

* Fix rebase errors, add TODO

Author: rahulgurnani

cmd/epp/runner/runner.go

@@ -245,7 +245,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
 	}
@@ -387,7 +387,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")
@@ -479,8 +479,9 @@ 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)
 
 	r.registerInTreePlugins()
 
@@ -520,10 +521,16 @@ 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() != nil {
 		return nil, errors.New("failed to load the configuration - prepare data plugins have cyclic dependencies")
 	}
+	// TODO(#1970): Remove feature gate check once prepare data plugins are stable.
+	if !r.featureGates[datalayer.PrepareDataPluginsFeatureGate] {
+		// If the feature gate is disabled, clear any prepare data plugins so they are not used.
+		r.requestControlConfig.WithPrepareDataPlugins()
+	}
 
 	// Handler deprecated configuration options
 	r.deprecatedConfigurationHelper(cfg, logger)
@@ -545,7 +552,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.go

@@ -97,7 +97,7 @@ func InstantiateAndConfigure(
 	}
 
 	featureGates := loadFeatureConfig(rawConfig.FeatureGates)
-	dataConfig, err := buildDataLayerConfig(rawConfig.Data, featureGates[datalayer.FeatureGate], handle)
+	dataConfig, err := buildDataLayerConfig(rawConfig.Data, featureGates[datalayer.ExperimentalDatalayerFeatureGate], handle)
 	if err != nil {
 		return nil, fmt.Errorf("data layer config build failed: %w", err)
 	}

pkg/epp/config/loader/configloader_test.go

@@ -58,7 +58,7 @@ func TestLoadRawConfiguration(t *testing.T) {
 	t.Parallel()
 
 	// Register known feature gates for validation.
-	RegisterFeatureGate(datalayer.FeatureGate)
+	RegisterFeatureGate(datalayer.ExperimentalDatalayerFeatureGate)
 
 	tests := []struct {
 		name       string
@@ -90,7 +90,7 @@ func TestLoadRawConfiguration(t *testing.T) {
 						},
 					},
 				},
-				FeatureGates: configapi.FeatureGates{datalayer.FeatureGate},
+				FeatureGates: configapi.FeatureGates{datalayer.ExperimentalDatalayerFeatureGate},
 				SaturationDetector: &configapi.SaturationDetector{
 					QueueDepthThreshold:       10,
 					KVCacheUtilThreshold:      0.8,
@@ -150,7 +150,7 @@ func TestInstantiateAndConfigure(t *testing.T) {
 	// Not parallel because it modifies global plugin registry.
 	registerTestPlugins(t)
 
-	RegisterFeatureGate(datalayer.FeatureGate)
+	RegisterFeatureGate(datalayer.ExperimentalDatalayerFeatureGate)
 
 	tests := []struct {
 		name       string

pkg/epp/datalayer/factory.go

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

pkg/epp/datalayer/plugins/approximateprefix/data_types.go

@@ -0,0 +1,52 @@
+/*
+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 approximateprefix
+
+import (
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/datalayer"
+)
+
+const (
+	PrefixCacheMatchInfoKey = "PrefixCacheMatchInfoKey"
+)
+
+type PrefixCacheMatchInfo struct {
+	matchLength int
+	totalBlocks int
+}
+
+func NewPrefixCacheMatchInfo(matchLen int, blockHashLen int) *PrefixCacheMatchInfo {
+	return &PrefixCacheMatchInfo{
+		matchLength: matchLen,
+		totalBlocks: blockHashLen,
+	}
+}
+
+func (p *PrefixCacheMatchInfo) MatchLength() int {
+	return p.matchLength
+}
+
+func (p *PrefixCacheMatchInfo) TotalLength() int {
+	return p.totalBlocks
+}
+
+func (p *PrefixCacheMatchInfo) Clone() datalayer.Cloneable {
+	return &PrefixCacheMatchInfo{
+		matchLength: p.matchLength,
+		totalBlocks: p.totalBlocks,
+	}
+}

pkg/epp/requestcontrol/dag.go

@@ -23,7 +23,7 @@ import (
 
 // buildDAG builds a dependency graph among data preparation plugins based on their
 // produced and consumed data keys.
-func buildDAG(plugins []PrepareDataPlugin) map[string][]string {
+func buildDAG(plugins []PrepareDataPlugin) (map[string][]string, error) {
 	dag := make(map[string][]string)
 	for _, plugin := range plugins {
 		dag[plugin.TypedName().String()] = []string{}
@@ -36,11 +36,14 @@ func buildDAG(plugins []PrepareDataPlugin) map[string][]string {
 			}
 			// Check whether plugin[i] produces something consumed by plugin[j]. In that case, j depends on i.
 			if plugins[i].Produces() != nil && plugins[j].Consumes() != nil {
-				// For all the keys produced by plugin i, check if plugin j consumes any of them.
-				// If yes, then j depends on i.
-				for producedKey := range plugins[i].Produces() {
+				for producedKey, producedData := range plugins[i].Produces() {
 					// If plugin j consumes the produced key, then j depends on i. We can break after the first match.
-					if _, ok := plugins[j].Consumes()[producedKey]; ok {
+					if consumedData, ok := plugins[j].Consumes()[producedKey]; ok {
+						// Check types are same. Reflection is avoided here for simplicity.
+						// TODO(#1985): Document this detail in IGW docs.
+						if producedData != consumedData {
+							return nil, errors.New("data type mismatch between produced and consumed data for key: " + producedKey)
+						}
 						iPluginName := plugins[i].TypedName().String()
 						jPluginName := plugins[j].TypedName().String()
 						dag[jPluginName] = append(dag[jPluginName], iPluginName)
@@ -50,7 +53,7 @@ func buildDAG(plugins []PrepareDataPlugin) map[string][]string {
 			}
 		}
 	}
-	return dag
+	return dag, nil
 }
 
 // sortPlugins builds the dependency graph and returns the plugins ordered in topological order.

pkg/epp/requestcontrol/dag_test.go

@@ -61,6 +61,10 @@ func TestPrepareDataGraph(t *testing.T) {
 	pluginX := &mockPrepareRequestDataP{name: "X", produces: map[string]any{"keyX": nil}, consumes: map[string]any{"keyY": nil}}
 	pluginY := &mockPrepareRequestDataP{name: "Y", produces: map[string]any{"keyY": nil}, consumes: map[string]any{"keyX": nil}}
 
+	// Data type mismatch plugin.
+	pluginZ1 := &mockPrepareRequestDataP{name: "Z1", produces: map[string]any{"keyZ": int(0)}}
+	pluginZ2 := &mockPrepareRequestDataP{name: "Z2", consumes: map[string]any{"keyZ": string("")}}
+
 	testCases := []struct {
 		name        string
 		plugins     []PrepareDataPlugin
@@ -109,11 +113,23 @@ func TestPrepareDataGraph(t *testing.T) {
 			expectedDAG: nil,
 			expectError: true,
 		},
+		{
+			name:        "Data type mismatch between produced and consumed data",
+			plugins:     []PrepareDataPlugin{pluginZ1, pluginZ2},
+			expectedDAG: nil,
+			expectError: true,
+		},
 	}
 
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
-			dag := buildDAG(tc.plugins)
+			dag, err := buildDAG(tc.plugins)
+			if err != nil {
+				if tc.expectError {
+					assert.Error(t, err)
+					return
+				}
+			}
 			orderedPlugins, err := sortPlugins(dag, tc.plugins)
 
 			if tc.expectError {

pkg/epp/requestcontrol/director.go

@@ -347,6 +347,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

@@ -108,7 +108,14 @@ 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() error {
-	dag := buildDAG(c.prepareDataPlugins)
+	// TODO(#1988): Add all producer and consumer plugins to the graph.
+	if len(c.prepareDataPlugins) == 0 {
+		return nil
+	}
+	dag, err := buildDAG(c.prepareDataPlugins)
+	if err != nil {
+		return err
+	}
 	plugins, err := sortPlugins(dag, c.prepareDataPlugins)
 	if err != nil {
 		return err

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

@@ -28,6 +28,7 @@ import (
 	k8stypes "k8s.io/apimachinery/pkg/types"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/datalayer/plugins/approximateprefix"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/metrics"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/plugins"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/requestcontrol"
@@ -206,6 +207,30 @@ func (p *Plugin) WithName(name string) *Plugin {
 	return p
 }
 
+func (p *Plugin) Produces() map[string]any {
+	return map[string]any{approximateprefix.PrefixCacheMatchInfoKey: approximateprefix.PrefixCacheMatchInfo{}}
+}
+
+func (p *Plugin) Consumes() map[string]any {
+	return map[string]any{}
+}
+
+// PrepareRequestData hashes prompt, finds longest prefix match and stores it in pod as attribute.
+func (p *Plugin) PrepareRequestData(ctx context.Context, request *types.LLMRequest, pods []types.Pod) error {
+	hashes := hashPrompt(ctx, request, getBlockSize(pods, p.config), p.config.MaxPrefixBlocksToMatch)
+	state := &SchedulingContextState{
+		PrefixHashes:       hashes,
+		PrefixCacheServers: p.matchLongestPrefix(ctx, hashes),
+	}
+	total := len(state.PrefixHashes)
+
+	for _, pod := range pods {
+		matchLen := state.PrefixCacheServers[ServerID(pod.GetPod().NamespacedName)]
+		pod.Put(approximateprefix.PrefixCacheMatchInfoKey, approximateprefix.NewPrefixCacheMatchInfo(matchLen, total))
+	}
+	return nil
+}
+
 // Score returns the scoring result for the given list of pods based on context.
 func (p *Plugin) Score(ctx context.Context, cycleState *types.CycleState, request *types.LLMRequest, pods []types.Pod) map[types.Pod]float64 {
 	// pre score step, hashing prompt and find longest prefix match.