Replace OptionalConsumes with ConsumesResult struct

Signed-off-by: Alok Behera <alokbeherak061@gmail.com>

Author: sudoalok

pkg/epp/datalayer/data_graph.go

@@ -84,7 +84,8 @@ func CreateMissingDataProducers(ctx context.Context, defaultProducerRegistry map
 	// Warn about optional keys with no producer — no error, just a warning.
 	for _, p := range handle.GetAllPlugins() {
 		if consumer, ok := p.(plugin.ConsumerPlugin); ok {
-			for key := range consumer.OptionalConsumes() {
+			result := consumer.Consumes()
+			for key := range result.Optional {
 				if !producedKeys[key.String()] {
 					logger.Info("Warning: optional data key has no producer, plugin will use fallback",
 						"plugin", p.TypedName().Name, "dataKey", key.String())
@@ -97,7 +98,8 @@ func CreateMissingDataProducers(ctx context.Context, defaultProducerRegistry map
 	missingKeys := make(map[string]string)
 	for _, p := range handle.GetAllPlugins() {
 		if consumer, ok := p.(plugin.ConsumerPlugin); ok {
-			for key := range consumer.Consumes() {
+			result := consumer.Consumes()
+			for key := range result.Required {
 				if !producedKeys[key.String()] {
 					missingKeys[key.String()] = consumer.TypedName().Name
 				}
@@ -203,9 +205,10 @@ func buildDAG(producers map[string]plugin.ProducerPlugin, consumers map[string]p
 			if pName == cName {
 				continue
 			}
-			if producer.Produces() != nil && consumer.Consumes() != nil {
+			consumed := consumer.Consumes()
+			if producer.Produces() != nil && consumed.Required != nil {
 				for producedKey, producedData := range producer.Produces() {
-					if consumedData, ok := consumer.Consumes()[producedKey]; ok {
+					if consumedData, ok := consumed.Required[producedKey]; ok {
 						// Check types are same.
 						if reflect.TypeOf(producedData) != reflect.TypeOf(consumedData) {
 							return nil, errors.New("data type mismatch between produced and consumed data for key: " + producedKey.String())

pkg/epp/datalayer/data_graph_test.go

@@ -58,12 +58,10 @@ func (m *mockDataProducerP) Produces() map[fwkplugin.DataKey]any {
 	return m.produces
 }
 
-func (m *mockDataProducerP) Consumes() map[fwkplugin.DataKey]any {
-	return m.consumes
+func (m *mockDataProducerP) Consumes() fwkplugin.ConsumesResult {
+	return fwkplugin.ConsumesResult{Required: m.consumes}
 }
 
-func (m *mockDataProducerP) OptionalConsumes() map[fwkplugin.DataKey]any { return nil }
-
 func (m *mockDataProducerP) Produce(ctx context.Context, request *fwksched.InferenceRequest, endpoints []fwksched.Endpoint) error {
 	endpoints[0].Put(mockProducedDataKey, &mockProducedDataType{value: 42})
 	return nil
@@ -90,12 +88,10 @@ type MockConsumerFairnessPolicy struct {
 	consumes map[fwkplugin.DataKey]any
 }
 
-func (m *MockConsumerFairnessPolicy) Consumes() map[fwkplugin.DataKey]any {
-	return m.consumes
+func (m *MockConsumerFairnessPolicy) Consumes() fwkplugin.ConsumesResult {
+	return fwkplugin.ConsumesResult{Required: m.consumes}
 }
 
-func (m *MockConsumerFairnessPolicy) OptionalConsumes() map[fwkplugin.DataKey]any { return nil }
-
 type MockSchedulingPlugin struct {
 	fwksched.Scorer
 	consumes map[fwkplugin.DataKey]any
@@ -105,12 +101,10 @@ func (m *MockSchedulingPlugin) TypedName() fwkplugin.TypedName {
 	return fwkplugin.TypedName{Name: "MockSchedulingPlugin", Type: "mock"}
 }
 
-func (m *MockSchedulingPlugin) Consumes() map[fwkplugin.DataKey]any {
-	return m.consumes
+func (m *MockSchedulingPlugin) Consumes() fwkplugin.ConsumesResult {
+	return fwkplugin.ConsumesResult{Required: m.consumes}
 }
 
-func (m *MockSchedulingPlugin) OptionalConsumes() map[fwkplugin.DataKey]any { return nil }
-
 func TestValidatePluginExecutionOrder(t *testing.T) {
 	dkA := fwkplugin.NewDataKey("keyA", "mock")
 	// Request control plugin that produces data.
@@ -445,10 +439,8 @@ func (m *mockMayConsumerPlugin) TypedName() fwkplugin.TypedName {
 	return fwkplugin.TypedName{Name: m.name, Type: "mock"}
 }
 
-func (m *mockMayConsumerPlugin) Consumes() map[fwkplugin.DataKey]any { return nil }
-
-func (m *mockMayConsumerPlugin) OptionalConsumes() map[fwkplugin.DataKey]any {
-	return m.optionalConsumes
+func (m *mockMayConsumerPlugin) Consumes() fwkplugin.ConsumesResult {
+	return fwkplugin.ConsumesResult{Optional: m.optionalConsumes}
 }
 
 // mockMixedConsumerPlugin is a plugin that has both required Consumes and optional OptionalConsumes.
@@ -464,12 +456,8 @@ func (m *mockMixedConsumerPlugin) TypedName() fwkplugin.TypedName {
 	return fwkplugin.TypedName{Name: m.name, Type: "mock"}
 }
 
-func (m *mockMixedConsumerPlugin) Consumes() map[fwkplugin.DataKey]any {
-	return m.consumes
-}
-
-func (m *mockMixedConsumerPlugin) OptionalConsumes() map[fwkplugin.DataKey]any {
-	return m.optionalConsumes
+func (m *mockMixedConsumerPlugin) Consumes() fwkplugin.ConsumesResult {
+	return fwkplugin.ConsumesResult{Required: m.consumes, Optional: m.optionalConsumes}
 }
 
 func TestCreateMissingDataProducers_MayConsume(t *testing.T) {

pkg/epp/framework/interface/plugin/plugins.go

@@ -23,18 +23,24 @@ type Plugin interface {
 	TypedName() TypedName
 }
 
+// ConsumesResult holds the data keys a plugin consumes, split by whether they
+// are required (framework errors if no producer exists) or optional (framework
+// logs a warning but continues if no producer exists).
+type ConsumesResult struct {
+	// Required keys — the framework will error at init time if no producer exists for any of these.
+	Required map[DataKey]any
+	// Optional keys — the framework logs a warning at init time but does NOT error if no producer exists.
+	// The plugin must handle the case where this data is absent at runtime.
+	Optional map[DataKey]any
+}
+
 // ConsumerPlugin defines the interface for a consumer.
 type ConsumerPlugin interface {
 	Plugin
-	// Consumes returns data consumed by the plugin.
-	// This is a map from DataKey produced to
-	// the data type of the key (represented as data with default value casted as any field).
-	Consumes() map[DataKey]any
-	// OptionalConsumes returns data keys the plugin can optionally consume.
-	// Unlike Consumes(), the framework will NOT error if no producer exists for these keys.
-	// Instead it logs a warning at init time. The plugin must handle the case where this data is absent.
-	// Return nil if there are no optional keys.
-	OptionalConsumes() map[DataKey]any
+	// Consumes returns the data keys consumed by this plugin, split into Required and Optional.
+	// Required keys: the framework errors at init time if no producer exists.
+	// Optional keys: the framework logs a warning but does not error; the plugin must handle absence.
+	Consumes() ConsumesResult
 }
 
 // ProducerPlugin defines the interface for a producer.

pkg/epp/framework/interface/plugin/plugins_test.go

@@ -33,8 +33,7 @@ type consumerImpl struct {
 	consumes map[DataKey]any
 }
 
-func (c *consumerImpl) Consumes() map[DataKey]any         { return c.consumes }
-func (c *consumerImpl) OptionalConsumes() map[DataKey]any { return nil }
+func (c *consumerImpl) Consumes() ConsumesResult { return ConsumesResult{Required: c.consumes} }
 
 type producerImpl struct {
 	basePlugin
@@ -59,8 +58,8 @@ func TestConsumerPlugin_Contract(t *testing.T) {
 	var cp ConsumerPlugin = c
 
 	got := cp.Consumes()
-	assert.Len(t, got, 1)
-	_, ok := got[key]
+	assert.Len(t, got.Required, 1)
+	_, ok := got.Required[key]
 	assert.True(t, ok)
 	assert.Equal(t, "queue-filter/filter", cp.TypedName().String())
 }
@@ -87,7 +86,7 @@ func TestConsumerPlugin_EmptyConsumes(t *testing.T) {
 		basePlugin: basePlugin{name: TypedName{Type: "filter", Name: "noop"}},
 		consumes:   nil,
 	}
-	assert.Nil(t, c.Consumes())
+	assert.Nil(t, c.Consumes().Required)
 }
 
 func TestProducerPlugin_EmptyProduces(t *testing.T) {

pkg/epp/framework/plugins/flowcontrol/saturationdetector/concurrency/detector.go

@@ -103,9 +103,9 @@ func (d *detector) TypedName() fwkplugin.TypedName {
 	return d.typedName
 }
 
-func (d *detector) Consumes() map[fwkplugin.DataKey]any {
-	return map[fwkplugin.DataKey]any{
-		d.inFlightLoadDataKey: attrconcurrency.InFlightLoad{},
+func (d *detector) Consumes() fwkplugin.ConsumesResult {
+	return fwkplugin.ConsumesResult{
+		Required: map[fwkplugin.DataKey]any{d.inFlightLoadDataKey: attrconcurrency.InFlightLoad{}},
 	}
 }
 

pkg/epp/framework/plugins/requestcontrol/admitter/latencyslo/plugin.go

@@ -89,9 +89,9 @@ func (p *LatencyAdmission) TypedName() fwkplugin.TypedName {
 }
 
 // Consumes declares that this plugin reads latency prediction data from endpoints.
-func (p *LatencyAdmission) Consumes() map[fwkplugin.DataKey]any {
-	return map[fwkplugin.DataKey]any{
-		p.latencyPredictionInfoDataKey: attrlatency.LatencyPredictionInfo{},
+func (p *LatencyAdmission) Consumes() fwkplugin.ConsumesResult {
+	return fwkplugin.ConsumesResult{
+		Required: map[fwkplugin.DataKey]any{p.latencyPredictionInfoDataKey: attrlatency.LatencyPredictionInfo{}},
 	}
 }
 

pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer.go

@@ -209,8 +209,10 @@ func (p *Producer) Produces() map[plugin.DataKey]any {
 
 // Consumes declares the TokenizedPrompt dependency from token-producer so
 // the data-layer DAG orders tokenization before this producer runs.
-func (p *Producer) Consumes() map[plugin.DataKey]any {
-	return map[plugin.DataKey]any{tokenproducer.TokenizedPromptDataKey: scheduling.TokenizedPrompt{}}
+func (p *Producer) Consumes() plugin.ConsumesResult {
+	return plugin.ConsumesResult{
+		Required: map[plugin.DataKey]any{tokenproducer.TokenizedPromptDataKey: scheduling.TokenizedPrompt{}},
+	}
 }
 
 // Produce hashes the request's TokenizedPrompt into KV-block keys, looks

pkg/epp/framework/plugins/requestcontrol/dataproducer/preciseprefixcache/producer_test.go

@@ -306,7 +306,7 @@ func TestProduces_DeclaresPrefixCacheMatchInfo(t *testing.T) {
 func TestConsumes_DeclaresTokenizedPrompt(t *testing.T) {
 	p := &Producer{typedName: plugin.TypedName{Type: PluginType, Name: "x"}}
 	expected := plugin.NewDataKey("TokenizedPrompt", "token-producer")
-	_, ok := p.Consumes()[expected]
+	_, ok := p.Consumes().Required[expected]
 	require.True(t, ok)
 }
 

pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/dataproducer_hooks.go

@@ -112,6 +112,8 @@ func (pl *PredictedLatency) Produces() map[plugin.DataKey]any {
 	}
 }
 
-func (pl *PredictedLatency) Consumes() map[plugin.DataKey]any {
-	return map[plugin.DataKey]any{pl.prefixMatchDataKey: attrprefix.PrefixCacheMatchInfo{}}
+func (pl *PredictedLatency) Consumes() plugin.ConsumesResult {
+	return plugin.ConsumesResult{
+		Required: map[plugin.DataKey]any{pl.prefixMatchDataKey: attrprefix.PrefixCacheMatchInfo{}},
+	}
 }

pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/dataproducer_hooks_test.go

@@ -35,7 +35,7 @@ func TestProducesConsumes(t *testing.T) {
 	assert.Contains(t, produces, expectedProduceKey)
 
 	consumes := pl.Consumes()
-	assert.Contains(t, consumes, attrprefix.PrefixCacheMatchInfoDataKey)
+	assert.Contains(t, consumes.Required, attrprefix.PrefixCacheMatchInfoDataKey)
 }
 
 // TestProduce_CancelledContextDoesNotPublish verifies that when the