refactor: replace MayConsumePlugin with OptionalConsumes in ConsumerPlugin

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

Author: sudoalok

pkg/epp/datalayer/data_graph.go

@@ -81,10 +81,10 @@ func CreateMissingDataProducers(ctx context.Context, defaultProducerRegistry map
 		}
 	}
 
-	// Warn about optional keys (MayConsume) with no producer — no error, just a warning.
+	// Warn about optional keys with no producer — no error, just a warning.
 	for _, p := range handle.GetAllPlugins() {
-		if mayConsumer, ok := p.(plugin.MayConsumePlugin); ok {
-			for key := range mayConsumer.MayConsume() {
+		if consumer, ok := p.(plugin.ConsumerPlugin); ok {
+			for key := range consumer.OptionalConsumes() {
 				if !producedKeys[key.String()] {
 					logger.Info("Warning: optional data key has no producer, plugin will use fallback",
 						"plugin", p.TypedName().Name, "dataKey", key.String())

pkg/epp/datalayer/data_graph_test.go

@@ -62,6 +62,8 @@ func (m *mockDataProducerP) Consumes() map[fwkplugin.DataKey]any {
 	return 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
@@ -92,6 +94,8 @@ func (m *MockConsumerFairnessPolicy) Consumes() map[fwkplugin.DataKey]any {
 	return m.consumes
 }
 
+func (m *MockConsumerFairnessPolicy) OptionalConsumes() map[fwkplugin.DataKey]any { return nil }
+
 type MockSchedulingPlugin struct {
 	fwksched.Scorer
 	consumes map[fwkplugin.DataKey]any
@@ -105,6 +109,8 @@ func (m *MockSchedulingPlugin) Consumes() map[fwkplugin.DataKey]any {
 	return 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.
@@ -429,27 +435,29 @@ func TestCreateMissingDataProducers(t *testing.T) {
 	}
 }
 
-// mockMayConsumerPlugin is a plugin that optionally consumes certain data keys.
+// mockMayConsumerPlugin is a plugin that only optionally consumes certain data keys.
 type mockMayConsumerPlugin struct {
-	name       string
-	mayConsume map[fwkplugin.DataKey]any
+	name            string
+	optionalConsumes map[fwkplugin.DataKey]any
 }
 
 func (m *mockMayConsumerPlugin) TypedName() fwkplugin.TypedName {
 	return fwkplugin.TypedName{Name: m.name, Type: "mock"}
 }
 
-func (m *mockMayConsumerPlugin) MayConsume() map[fwkplugin.DataKey]any {
-	return m.mayConsume
+func (m *mockMayConsumerPlugin) Consumes() map[fwkplugin.DataKey]any { return nil }
+
+func (m *mockMayConsumerPlugin) OptionalConsumes() map[fwkplugin.DataKey]any {
+	return m.optionalConsumes
 }
 
-// mockMixedConsumerPlugin is a plugin that has both required Consumes and optional MayConsume.
+// mockMixedConsumerPlugin is a plugin that has both required Consumes and optional OptionalConsumes.
 // This models a real plugin like prefix cache scorer — requires prefix-match data,
 // but optionally uses tokenized input and falls back to raw text if unavailable.
 type mockMixedConsumerPlugin struct {
-	name       string
-	consumes   map[fwkplugin.DataKey]any
-	mayConsume map[fwkplugin.DataKey]any
+	name            string
+	consumes        map[fwkplugin.DataKey]any
+	optionalConsumes map[fwkplugin.DataKey]any
 }
 
 func (m *mockMixedConsumerPlugin) TypedName() fwkplugin.TypedName {
@@ -460,15 +468,15 @@ func (m *mockMixedConsumerPlugin) Consumes() map[fwkplugin.DataKey]any {
 	return m.consumes
 }
 
-func (m *mockMixedConsumerPlugin) MayConsume() map[fwkplugin.DataKey]any {
-	return m.mayConsume
+func (m *mockMixedConsumerPlugin) OptionalConsumes() map[fwkplugin.DataKey]any {
+	return m.optionalConsumes
 }
 
 func TestCreateMissingDataProducers_MayConsume(t *testing.T) {
 	producerTypeA := "producer-a"
 	keyA := fwkplugin.NewDataKey("keyA", producerTypeA)
 
-	producerAFactory := fwkplugin.FactoryFunc(func(name string, _ json.RawMessage, handle fwkplugin.Handle) (fwkplugin.Plugin, error) {
+	producerAFactory := fwkplugin.FactoryFunc(func(name string, _ *json.Decoder, handle fwkplugin.Handle) (fwkplugin.Plugin, error) {
 		return &mockDataProducerP{name: name, produces: map[fwkplugin.DataKey]any{keyA: nil}}, nil
 	})
 
@@ -479,40 +487,40 @@ func TestCreateMissingDataProducers_MayConsume(t *testing.T) {
 		wantErr         bool
 	}{
 		{
-			name: "MayConsume key with no producer — warning only, no error",
+			name: "OptionalConsumes key with no producer — warning only, no error",
 			existingPlugins: []fwkplugin.Plugin{
 				&mockMayConsumerPlugin{
-					name:       "optional-consumer",
-					mayConsume: map[fwkplugin.DataKey]any{keyA: nil},
+					name:            "optional-consumer",
+					optionalConsumes: map[fwkplugin.DataKey]any{keyA: nil},
 				},
 			},
 			factoryRegistry: map[string]fwkplugin.FactoryFunc{},
 			wantErr:         false, // must NOT error
 		},
 		{
-			name: "MayConsume key with a producer present — no warning, no error",
+			name: "OptionalConsumes key with a producer present — no warning, no error",
 			existingPlugins: []fwkplugin.Plugin{
 				&mockDataProducerP{name: "producer", produces: map[fwkplugin.DataKey]any{keyA: nil}},
 				&mockMayConsumerPlugin{
-					name:       "optional-consumer",
-					mayConsume: map[fwkplugin.DataKey]any{keyA: nil},
+					name:            "optional-consumer",
+					optionalConsumes: map[fwkplugin.DataKey]any{keyA: nil},
 				},
 			},
 			factoryRegistry: map[string]fwkplugin.FactoryFunc{producerTypeA: producerAFactory},
 			wantErr:         false,
 		},
 		{
 			// Models the real prefix cache scorer — it requires prefix-match data (Consumes)
-			// but optionally uses tokenized input (MayConsume), falling back to raw text.
+			// but optionally uses tokenized input (OptionalConsumes), falling back to raw text.
 			// The required key has a producer. The optional key does not.
 			// Result: no error. Warning logged for the missing optional key.
-			name: "plugin with both Consumes and MayConsume — required key has producer, optional does not",
+			name: "plugin with both Consumes and OptionalConsumes — required key has producer, optional does not",
 			existingPlugins: []fwkplugin.Plugin{
 				&mockDataProducerP{name: "required-producer", produces: map[fwkplugin.DataKey]any{keyA: nil}},
 				&mockMixedConsumerPlugin{
-					name:       "prefix-cache-scorer",
-					consumes:   map[fwkplugin.DataKey]any{keyA: nil},
-					mayConsume: map[fwkplugin.DataKey]any{fwkplugin.NewDataKey("tokenized-input", "tokenizer"): nil},
+					name:            "prefix-cache-scorer",
+					consumes:        map[fwkplugin.DataKey]any{keyA: nil},
+					optionalConsumes: map[fwkplugin.DataKey]any{fwkplugin.NewDataKey("tokenized-input", "tokenizer"): nil},
 				},
 			},
 			factoryRegistry: map[string]fwkplugin.FactoryFunc{},

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

@@ -30,16 +30,11 @@ type ConsumerPlugin interface {
 	// 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
-}
-
-// MayConsumePlugin defines the interface for a plugin that optionally consumes data.
-// Unlike ConsumerPlugin, 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.
-type MayConsumePlugin interface {
-	Plugin
-	// MayConsume returns data keys the plugin can optionally consume.
-	// The plugin must handle the case where this data is absent.
-	MayConsume() 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
 }
 
 // ProducerPlugin defines the interface for a producer.

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

@@ -33,7 +33,8 @@ type consumerImpl struct {
 	consumes map[DataKey]any
 }
 
-func (c *consumerImpl) Consumes() map[DataKey]any { return c.consumes }
+func (c *consumerImpl) Consumes() map[DataKey]any         { return c.consumes }
+func (c *consumerImpl) OptionalConsumes() map[DataKey]any { return nil }
 
 type producerImpl struct {
 	basePlugin

pkg/epp/framework/plugins/scheduling/scorer/activerequest/active_request.go

@@ -149,6 +149,9 @@ func (s *ActiveRequest) Consumes() map[plugin.DataKey]any {
 	}
 }
 
+// OptionalConsumes returns nil as this plugin has no optional data dependencies.
+func (s *ActiveRequest) OptionalConsumes() map[plugin.DataKey]any { return nil }
+
 // Score scores the given endpoints based on the number of active requests
 // being served by each endpoint. The score is normalized to a range of 0-1.
 func (s *ActiveRequest) Score(ctx context.Context, _ *scheduling.InferenceRequest,

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

@@ -89,6 +89,9 @@ func (s *Scorer) Consumes() map[plugin.DataKey]any {
 	return map[plugin.DataKey]any{s.mmMatchDataKey: attrmm.EncoderCacheMatchInfo{}}
 }
 
+// OptionalConsumes returns nil as this plugin has no optional data dependencies.
+func (s *Scorer) OptionalConsumes() map[plugin.DataKey]any { return nil }
+
 // Score scores endpoints by matched multimodal encoder-cache item size divided
 // by total multimodal request item size.
 func (s *Scorer) Score(ctx context.Context, req *scheduling.InferenceRequest, endpoints []scheduling.Endpoint) map[scheduling.Endpoint]float64 {

pkg/epp/framework/plugins/scheduling/scorer/nohitlru/no_hit_lru.go

@@ -128,6 +128,9 @@ func (s *NoHitLRU) Consumes() map[plugin.DataKey]any {
 	}
 }
 
+// OptionalConsumes returns nil as this plugin has no optional data dependencies.
+func (s *NoHitLRU) OptionalConsumes() map[plugin.DataKey]any { return nil }
+
 // isColdRequest determines if a request is cold by checking endpoint prefix-cache attributes.
 // Returns true when no endpoint reports any cache-hit blocks, or when no attribute is present.
 func (s *NoHitLRU) isColdRequest(ctx context.Context, endpoints []scheduling.Endpoint) bool {