migrate plugin keys to explicit scopes

Migrates all EPP in-tree dynamic attribute keys (predicted-latency,
prefix-cache, concurrency, multimodal, and tokenizer) to explicitly
declare RequestScope and EndpointScope boundaries.

Strictly demotes legacy UnspecifiedScope out-of-tree plugins to
Post-Admission phase at startup, ogging actionable warning logs and
utilizing DAG validations to fail-fast on pre-admission legacy
dependencies.

Signed-off-by: Luke Van Drie <lukevandrie@google.com>

Author: LukeAVanDrie

pkg/epp/datalayer/data_graph.go

@@ -45,6 +45,8 @@ func ValidateAndOrderDataDependencies(plugins []plugin.Plugin) ([]string, error)
 // CompilePipeline compiles the active plugins into two topologically sorted slices:
 // Pre-Admission and Post-Admission, applying pruning and scope validations.
 func CompilePipeline(plugins []plugin.Plugin) ([]string, []string, error) {
+	coercedScopes := auditAndCoerceScopes(plugins)
+
 	pluginMap := make(map[string]plugin.Plugin)
 	for _, p := range plugins {
 		pluginMap[p.TypedName().String()] = p
@@ -55,19 +57,19 @@ func CompilePipeline(plugins []plugin.Plugin) ([]string, []string, error) {
 		return nil, nil, err
 	}
 
-	dag, err := buildDAG(producers, consumers)
+	dag, err := buildDAG(producers, consumers, coercedScopes)
 	if err != nil {
 		return nil, nil, err
 	}
 
 	V_kept, filteredDag := pruneUnusedProducers(pluginMap, dag)
 
-	preAdmission, postAdmission, err := slicePipelineByAdmissionGating(pluginMap, V_kept, filteredDag)
+	preAdmission, postAdmission, err := slicePipelineByAdmissionGating(pluginMap, V_kept, filteredDag, coercedScopes)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	if err := validatePreAdmissionScopes(pluginMap, preAdmission); err != nil {
+	if err := validatePreAdmissionScopes(pluginMap, preAdmission, coercedScopes); err != nil {
 		return nil, nil, err
 	}
 
@@ -206,6 +208,7 @@ func slicePipelineByAdmissionGating(
 	pluginMap map[string]plugin.Plugin,
 	V_kept map[string]bool,
 	filteredDag map[string][]string,
+	coercedScopes map[string]plugin.DataScope,
 ) (map[string]bool, map[string]bool, error) {
 	preAdmissionRoots := make(map[string]bool)
 	for name := range V_kept {
@@ -218,10 +221,12 @@ func slicePipelineByAdmissionGating(
 		} else if _, ok := p.(fwkfc.OrderingPolicy); ok {
 			isRoot = true
 		} else if eagerP, ok := p.(plugin.EagerProducerPlugin); ok && eagerP.Eager() {
-			for key := range eagerP.Produces() {
-				if key.Scope() == plugin.RequestScope {
-					isRoot = true
-					break
+			if eagerP.Produces() != nil {
+				for key := range eagerP.Produces() {
+					if resolveScope(key, coercedScopes) == plugin.RequestScope {
+						isRoot = true
+						break
+					}
 				}
 			}
 		}
@@ -259,22 +264,22 @@ func slicePipelineByAdmissionGating(
 
 // validatePreAdmissionScopes validates that Pre-Admission plugins do not consume or produce EndpointScope or
 // UnspecifiedScope keys.
-func validatePreAdmissionScopes(pluginMap map[string]plugin.Plugin, preAdmission map[string]bool) error {
+func validatePreAdmissionScopes(pluginMap map[string]plugin.Plugin, preAdmission map[string]bool, coercedScopes map[string]plugin.DataScope) error {
 	for name := range preAdmission {
 		p := pluginMap[name]
 		if consumer, ok := p.(plugin.ConsumerPlugin); ok {
 			for key := range consumer.Consumes() {
-				if key.Scope() == plugin.EndpointScope || key.Scope() == plugin.UnspecifiedScope {
+				if resolveScope(key, coercedScopes) == plugin.EndpointScope {
 					return fmt.Errorf("scope mismatch: plugin %q in Pre-Admission consumes key %q with invalid scope %q",
-						name, key.String(), key.Scope())
+						name, key.String(), resolveScope(key, coercedScopes))
 				}
 			}
 		}
 		if producer, ok := p.(plugin.ProducerPlugin); ok {
 			for key := range producer.Produces() {
-				if key.Scope() == plugin.EndpointScope || key.Scope() == plugin.UnspecifiedScope {
+				if resolveScope(key, coercedScopes) == plugin.EndpointScope {
 					return fmt.Errorf("scope mismatch: plugin %q in Pre-Admission produces key %q with invalid scope %q",
-						name, key.String(), key.Scope())
+						name, key.String(), resolveScope(key, coercedScopes))
 				}
 			}
 		}
@@ -355,7 +360,7 @@ const (
 	DefaultLayer        = -1
 )
 
-func pluginToLayerExecutionOrder(p plugin.Plugin) int {
+func pluginToLayerExecutionOrder(p plugin.Plugin, coercedScopes map[string]plugin.DataScope) int {
 	// PreAdmitter
 	if _, ok := p.(fwkrc.PreAdmitter); ok {
 		return PreAdmissionLayer
@@ -376,7 +381,7 @@ func pluginToLayerExecutionOrder(p plugin.Plugin) int {
 		if producer.Produces() != nil && len(producer.Produces()) > 0 {
 			hasProduces = true
 			for key := range producer.Produces() {
-				if key.Scope() != plugin.RequestScope {
+				if resolveScope(key, coercedScopes) != plugin.RequestScope {
 					isRequestScoped = false
 					break
 				}
@@ -418,7 +423,7 @@ func pluginToLayerExecutionOrder(p plugin.Plugin) int {
 
 // buildDAG builds a dependency graph among data preparation plugins based on their
 // produced and consumed data keys.
-func buildDAG(producers map[string]plugin.ProducerPlugin, consumers map[string]plugin.ConsumerPlugin) (map[string][]string, error) {
+func buildDAG(producers map[string]plugin.ProducerPlugin, consumers map[string]plugin.ConsumerPlugin, coercedScopes map[string]plugin.DataScope) (map[string][]string, error) {
 	dag := make(map[string][]string)
 	// Create dependency graph as a DAG.
 	for _, producer := range producers {
@@ -432,25 +437,39 @@ func buildDAG(producers map[string]plugin.ProducerPlugin, consumers map[string]p
 			if pName == cName {
 				continue
 			}
-			if producer.Produces() != nil && consumer.Consumes() != nil {
-				hasDependency := false
-				for producedKey, producedData := range producer.Produces() {
-					if consumedData, ok := consumer.Consumes()[producedKey]; ok {
-						// Check types are same.
-						if reflect.TypeOf(producedData) != reflect.TypeOf(consumedData) {
-							return nil, fmt.Errorf("data type mismatch between produced and consumed data for key: %s", producedKey.String())
-						}
-						if pluginToLayerExecutionOrder(producer) > pluginToLayerExecutionOrder(consumer) {
-							return nil, fmt.Errorf("invalid plugin layer execution order: producer %s needs to be executed before consumer %s", pName, cName)
+			hasDependency := false
+			if producer.Produces() != nil {
+				// Required consumes check
+				if consumer.Consumes() != nil {
+					for producedKey, producedData := range producer.Produces() {
+						if consumedData, ok := consumer.Consumes()[producedKey]; ok {
+							// Check types are same.
+							if reflect.TypeOf(producedData) != reflect.TypeOf(consumedData) {
+								return nil, fmt.Errorf("data type mismatch between produced and consumed data for key: %s", producedKey.String())
+							}
+							if pluginToLayerExecutionOrder(producer, coercedScopes) > pluginToLayerExecutionOrder(consumer, coercedScopes) {
+								return nil, fmt.Errorf("invalid plugin layer execution order: producer %s needs to be executed before consumer %s", pName, cName)
+							}
+							hasDependency = true
 						}
-						hasDependency = true
 					}
 				}
-				if hasDependency {
-					// Consumer depends on producer, so add an edge from consumer to producer.
-					dag[cName] = append(dag[cName], pName)
+				// Optional consumes check
+				if optConsumer, ok := consumer.(plugin.OptionalConsumerPlugin); ok && optConsumer.OptionalConsumes() != nil {
+					for _, producedKey := range optConsumer.OptionalConsumes() {
+						if _, ok := producer.Produces()[producedKey]; ok {
+							if pluginToLayerExecutionOrder(producer, coercedScopes) > pluginToLayerExecutionOrder(consumer, coercedScopes) {
+								return nil, fmt.Errorf("invalid plugin layer execution order: producer %s needs to be executed before consumer %s (via optional dependency)", pName, cName)
+							}
+							hasDependency = true
+						}
+					}
 				}
 			}
+			if hasDependency {
+				// Consumer depends on producer, so add an edge from consumer to producer.
+				dag[cName] = append(dag[cName], pName)
+			}
 		}
 	}
 	return dag, nil
@@ -563,3 +582,65 @@ func findAndFormatCycle(graph map[string][]string, inDegree map[string]int) stri
 	}
 	return "unknown cycle"
 }
+
+// resolveScope resolves a key's scope, falling back to its coerced scope if it is UnspecifiedScope.
+func resolveScope(key plugin.DataKey, coercedScopes map[string]plugin.DataScope) plugin.DataScope {
+	if key.Scope() == plugin.UnspecifiedScope {
+		if scope, ok := coercedScopes[key.String()]; ok {
+			return scope
+		}
+	}
+	return key.Scope()
+}
+
+// auditAndCoerceScopes audits all active plugins and determines dynamic scopes for UnspecifiedScope keys.
+func auditAndCoerceScopes(plugins []plugin.Plugin) map[string]plugin.DataScope {
+	logger := log.Log.WithName("epp-graph-compiler")
+	coercedScopes := make(map[string]plugin.DataScope)
+	unspecifiedKeysLogged := make(map[string]bool)
+
+	for _, p := range plugins {
+		var unspecifiedKeys []plugin.DataKey
+
+		if consumer, ok := p.(plugin.ConsumerPlugin); ok && consumer.Consumes() != nil {
+			for key := range consumer.Consumes() {
+				if key.Scope() == plugin.UnspecifiedScope {
+					unspecifiedKeys = append(unspecifiedKeys, key)
+				}
+			}
+		}
+		if optConsumer, ok := p.(plugin.OptionalConsumerPlugin); ok && optConsumer.OptionalConsumes() != nil {
+			for _, key := range optConsumer.OptionalConsumes() {
+				if key.Scope() == plugin.UnspecifiedScope {
+					unspecifiedKeys = append(unspecifiedKeys, key)
+				}
+			}
+		}
+		if producer, ok := p.(plugin.ProducerPlugin); ok && producer.Produces() != nil {
+			for key := range producer.Produces() {
+				if key.Scope() == plugin.UnspecifiedScope {
+					unspecifiedKeys = append(unspecifiedKeys, key)
+				}
+			}
+		}
+
+		for _, key := range unspecifiedKeys {
+			keyStr := key.String()
+			if unspecifiedKeysLogged[keyStr] {
+				continue
+			}
+			unspecifiedKeysLogged[keyStr] = true
+
+			coercedScopes[keyStr] = plugin.EndpointScope
+
+			logger.Info(
+				"WARNING: Legacy/un-upgraded plugin registered using UnspecifiedScope data key. Coercing to EndpointScope. Please upgrade the plugin to declare an explicit scope.",
+				"plugin", p.TypedName().String(),
+				"dataKey", keyStr,
+				"coercedScope", plugin.EndpointScope,
+			)
+		}
+	}
+
+	return coercedScopes
+}

pkg/epp/datalayer/data_graph_test.go

@@ -153,7 +153,7 @@ func TestDAGAndTopologicalOrder(t *testing.T) {
 					consumers[p.TypedName().String()] = cp
 				}
 			}
-			dag, err := buildDAG(producers, consumers)
+			dag, err := buildDAG(producers, consumers, nil)
 			if err != nil {
 				if tc.expectedErr != "" {
 					assert.Error(t, err)
@@ -552,13 +552,13 @@ func TestCompilePipeline_ScopeConstraints(t *testing.T) {
 	})
 
 	t.Run("PreAdmissionConsumesUnspecifiedKey", func(t *testing.T) {
+		// UnspecifiedScope is strictly coerced to EndpointScope, so Pre-Admission consuming it throws a scope mismatch.
 		dkUnspec := fwkplugin.NewDataKey("unspecKey", "mock")
 		preAdmitterUnspec := &MockPreAdmitter{name: "PreAdmitterUnspec", consumes: map[fwkplugin.DataKey]any{dkUnspec: nil}}
 
 		_, _, err := CompilePipeline([]fwkplugin.Plugin{preAdmitterUnspec})
 		assert.Error(t, err)
 		assert.Contains(t, err.Error(), "scope mismatch")
-		assert.Contains(t, err.Error(), "consumes key")
 	})
 
 	t.Run("PreAdmissionProducesEndpointKey", func(t *testing.T) {
@@ -579,6 +579,7 @@ func TestCompilePipeline_ScopeConstraints(t *testing.T) {
 	})
 
 	t.Run("PreAdmissionProducesUnspecifiedKey", func(t *testing.T) {
+		// UnspecifiedScope is strictly coerced to EndpointScope, so Pre-Admission producing it throws a scope mismatch.
 		dkReq := fwkplugin.NewRequestDataKey("reqKey", "mock")
 		dkUnspec := fwkplugin.NewDataKey("unspecKey", "mock")
 		pEagerUnspec := &mockEagerProducer{
@@ -592,7 +593,6 @@ func TestCompilePipeline_ScopeConstraints(t *testing.T) {
 		_, _, err := CompilePipeline([]fwkplugin.Plugin{pEagerUnspec})
 		assert.Error(t, err)
 		assert.Contains(t, err.Error(), "scope mismatch")
-		assert.Contains(t, err.Error(), "produces key")
 	})
 }
 
@@ -666,6 +666,95 @@ func TestCompilePipeline_BoundaryEdgeCases(t *testing.T) {
 	})
 }
 
+func TestCompilePipeline_LegacyFallback(t *testing.T) {
+	t.Run("LegacyUnspecifiedKey_PreAdmissionCoercion", func(t *testing.T) {
+		// Key with UnspecifiedScope
+		dkUnspec := fwkplugin.NewDataKey("legacyKey", "mock")
+
+		// Producer producing the unspecified key
+		pLegacy := &mockDataProducerP{name: "PLegacy", produces: map[fwkplugin.DataKey]any{dkUnspec: nil}}
+
+		// Pre-admission plugin consuming the unspecified key
+		admitterReq := &MockPreAdmitter{name: "PreAdmitter", consumes: map[fwkplugin.DataKey]any{dkUnspec: nil}}
+
+		_, _, err := CompilePipeline([]fwkplugin.Plugin{pLegacy, admitterReq})
+		assert.Error(t, err)
+		// Expected invalid layer execution order error because legacy key is coerced to EndpointScope.
+		assert.Contains(t, err.Error(), "invalid plugin layer execution order")
+	})
+
+	t.Run("LegacyUnspecifiedKey_PostAdmissionCoercion", func(t *testing.T) {
+		// Key with UnspecifiedScope
+		dkUnspec := fwkplugin.NewDataKey("legacyKeyPost", "mock")
+
+		// Producer producing the unspecified key
+		pLegacy := &mockDataProducerP{name: "PLegacyPost", produces: map[fwkplugin.DataKey]any{dkUnspec: nil}}
+
+		// Post-admission (admitter) plugin consuming the unspecified key
+		admitterEp := &MockAdmitter{name: "Admitter", consumes: map[fwkplugin.DataKey]any{dkUnspec: nil}}
+
+		preAdmission, postAdmission, err := CompilePipeline([]fwkplugin.Plugin{pLegacy, admitterEp})
+		assert.NoError(t, err)
+
+		// Verify they compile and are correctly sliced into Post-Admission.
+		assert.Empty(t, preAdmission)
+		assert.Contains(t, postAdmission, "PLegacyPost/mock")
+		assert.Contains(t, postAdmission, "Admitter/mock")
+	})
+
+	t.Run("LegacyUnspecifiedKey_OptionalPreAdmissionCoercion", func(t *testing.T) {
+		// Key with UnspecifiedScope
+		dkUnspec := fwkplugin.NewDataKey("legacyOptionalKey", "mock")
+
+		// Producer producing the unspecified key
+		pLegacy := &mockDataProducerP{name: "PLegacyOpt", produces: map[fwkplugin.DataKey]any{dkUnspec: nil}}
+
+		// Pre-admission plugin optionally consuming the unspecified key
+		admitterReq := &MockOptionalPreAdmitter{
+			MockPreAdmitter:  MockPreAdmitter{name: "PreAdmitter", consumes: nil},
+			optionalConsumes: []fwkplugin.DataKey{dkUnspec},
+		}
+
+		_, _, err := CompilePipeline([]fwkplugin.Plugin{pLegacy, admitterReq})
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "invalid plugin layer execution order")
+	})
+	t.Run("OptionalConsumer_PruningAndTopologicalOrdering", func(t *testing.T) {
+		// Key with RequestScope
+		dkOpt := fwkplugin.NewRequestDataKey("optionalKey", "mock")
+
+		// Producer producing the optional key (non-eager)
+		pOpt := &mockDataProducerP{name: "POpt", produces: map[fwkplugin.DataKey]any{dkOpt: nil}}
+
+		// Consumer optionally consuming the key
+		admitterReq := &MockOptionalPreAdmitter{
+			MockPreAdmitter:  MockPreAdmitter{name: "PreAdmitter", consumes: nil},
+			optionalConsumes: []fwkplugin.DataKey{dkOpt},
+		}
+
+		preAdmission, _, err := CompilePipeline([]fwkplugin.Plugin{pOpt, admitterReq})
+		assert.NoError(t, err)
+
+		// Verify that the producer is NOT pruned (it is kept because it's optionally consumed by active PreAdmitter).
+		assert.Contains(t, preAdmission, "POpt/mock")
+		assert.Contains(t, preAdmission, "PreAdmitter/mock")
+
+		// Verify correct topological sorted order: producer must execute before consumer.
+		indexOfProducer := -1
+		indexOfConsumer := -1
+		for i, name := range preAdmission {
+			switch name {
+			case "POpt/mock":
+				indexOfProducer = i
+			case "PreAdmitter/mock":
+				indexOfConsumer = i
+			}
+		}
+		assert.True(t, indexOfProducer != -1 && indexOfConsumer != -1)
+		assert.Less(t, indexOfProducer, indexOfConsumer, "Producer POpt should execute before optional consumer PreAdmitter")
+	})
+}
+
 func TestCreateMissingDataProducers(t *testing.T) {
 	producerTypeA := "producer-a"
 	producerTypeB := "producer-b"
@@ -921,6 +1010,15 @@ func (m *MockPreAdmitter) PreAdmit(ctx context.Context, request *fwksched.Infere
 	return nil
 }
 
+type MockOptionalPreAdmitter struct {
+	MockPreAdmitter
+	optionalConsumes []fwkplugin.DataKey
+}
+
+func (m *MockOptionalPreAdmitter) OptionalConsumes() []fwkplugin.DataKey {
+	return m.optionalConsumes
+}
+
 func assertTopologicalOrder(t *testing.T, dag map[string][]string, ordered []string) {
 	t.Helper()
 	positions := make(map[string]int)

pkg/epp/framework/plugins/datalayer/attribute/concurrency/data_types.go

@@ -22,7 +22,7 @@ import (
 	inflightloadconstants "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol/dataproducer/inflightload/constants"
 )
 
-var InFlightLoadDataKey = plugin.NewDataKey("InFlightLoadDataKey", inflightloadconstants.InFlightLoadProducerType)
+var InFlightLoadDataKey = plugin.NewEndpointDataKey("InFlightLoadDataKey", inflightloadconstants.InFlightLoadProducerType)
 
 // InFlightLoad captures the current real-time load of an endpoint as tracked by the EPP.
 type InFlightLoad struct {