Add Produces and Consumes methods to Plugin (#1754)

* Add consumes and produces method to plugin with stub implementations

* Refactor: Move metrics constants to a separate package

* Add a specific Consumer Plugin

* Add descriptive comments and cleanup code

* Remove empty Consumes implementations from tests

Author: rahulgurnani

pkg/epp/datalayer/metrics/extractor.go

@@ -28,6 +28,7 @@ import (
 	dto "github.com/prometheus/client_model/go"
 
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/datalayer"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/metrics"
 )
 
 const (
@@ -47,6 +48,17 @@ type Extractor struct {
 	mapping *Mapping
 }
 
+func Produces() map[string]any {
+	return map[string]any{
+		metrics.WaitingQueueSizeKey:    int(0),
+		metrics.KVCacheUsagePercentKey: float64(0),
+		metrics.ActiveModelsKey:        map[string]int{},
+		metrics.WaitingModelsKey:       map[string]int{},
+		metrics.MaxActiveModelsKey:     int(0),
+		metrics.UpdateTimeKey:          time.Time{},
+	}
+}
+
 // NewExtractor returns a new model server protocol (MSP) metrics extractor,
 // configured with the given metrics' specifications.
 // These are mandatory metrics per the MSP specification, and are used

pkg/epp/metrics/metrics.go

@@ -34,6 +34,13 @@ const (
 	InferenceObjectiveComponent = "inference_objective"
 	InferencePoolComponent      = "inference_pool"
 	InferenceExtension          = "inference_extension"
+
+	KVCacheUsagePercentKey = "KVCacheUsagePercent"
+	WaitingQueueSizeKey    = "WaitingQueueSize"
+	MaxActiveModelsKey     = "MaxActiveModels"
+	ActiveModelsKey        = "ActiveModels"
+	WaitingModelsKey       = "WaitingModels"
+	UpdateTimeKey          = "UpdateTime"
 )
 
 var (

pkg/epp/plugins/plugins.go

@@ -22,3 +22,21 @@ type Plugin interface {
 	// TypedName returns the type and name tuple of this plugin instance.
 	TypedName() TypedName
 }
+
+// ConsumerPlugin defines the interface for a consumer.
+type ConsumerPlugin interface {
+	Plugin
+	// Consumes returns data consumed by the plugin.
+	// This is a map from data key (string) produced to
+	// the data type of the key (represented as data with default value casted as any field).
+	Consumes() map[string]any
+}
+
+// ProducerPlugin defines the interface for a producer.
+type ProducerPlugin interface {
+	Plugin
+	// Produces returns data produced by the producer.
+	// This is a map from data key (string) produced to
+	// the data type of the key (represented as data with default value casted as any field).
+	Produces() map[string]any
+}

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

@@ -20,6 +20,7 @@ import (
 	"context"
 	"encoding/json"
 
+	"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/scheduling/framework"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
@@ -54,6 +55,13 @@ func (s *KVCacheUtilizationScorer) TypedName() plugins.TypedName {
 	return s.typedName
 }
 
+// Consumes returns the list of data that is consumed by the plugin.
+func (s *KVCacheUtilizationScorer) Consumes() map[string]any {
+	return map[string]any{
+		metrics.KVCacheUsagePercentKey: float64(0),
+	}
+}
+
 // WithName sets the name of the scorer.
 func (s *KVCacheUtilizationScorer) WithName(name string) *KVCacheUtilizationScorer {
 	s.typedName.Name = name

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

@@ -20,6 +20,7 @@ import (
 	"context"
 	"encoding/json"
 
+	"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/scheduling/framework"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
@@ -54,6 +55,14 @@ func (s *LoraAffinityScorer) TypedName() plugins.TypedName {
 	return s.tn
 }
 
+// Consumes returns the list of data that is consumed by the plugin.
+func (s *LoraAffinityScorer) Consumes() map[string]any {
+	return map[string]any{
+		metrics.ActiveModelsKey:  map[string]int{},
+		metrics.WaitingModelsKey: map[string]int{},
+	}
+}
+
 // WithName sets the name of the scorer.
 func (s *LoraAffinityScorer) WithName(name string) *LoraAffinityScorer {
 	s.tn.Name = name

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

@@ -21,6 +21,7 @@ import (
 	"encoding/json"
 	"math"
 
+	"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/scheduling/framework"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
@@ -56,6 +57,13 @@ func (s *QueueScorer) TypedName() plugins.TypedName {
 	return s.typedName
 }
 
+// Consumes returns the list of data that is consumed by the plugin.
+func (s *QueueScorer) Consumes() map[string]any {
+	return map[string]any{
+		metrics.WaitingQueueSizeKey: int(0),
+	}
+}
+
 // WithName sets the name of the scorer.
 func (s *QueueScorer) WithName(name string) *QueueScorer {
 	s.typedName.Name = name

pkg/epp/scheduling/types/cycle_state.go

@@ -33,6 +33,7 @@ func NewCycleState() *CycleState {
 // CycleState does not provide any data protection, as all plugins are assumed to be
 // trusted.
 // Note: CycleState uses a sync.Map to back the storage, because it is thread safe. It's aimed to optimize for the "write once and read many times" scenarios.
+// TODO: Perhaps, deprecate CycleState once datalayer producer-consumer changes are made.
 type CycleState struct {
 	// key: StateKey, value: StateData
 	storage sync.Map