Add AdmitRequest and PrepareData plugins

Author: rahulgurnani

pkg/epp/datalayer/attributemap.go

@@ -20,6 +20,8 @@ import (
 	"sync"
 )
 
+// TODO(rahulgurnani): Deprecate this AttributeMap in favor of AttributeMap added in scheduling layer.
+
 // Cloneable types support cloning of the value.
 type Cloneable interface {
 	Clone() Cloneable

pkg/epp/requestcontrol/director.go

@@ -36,6 +36,7 @@ import (
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/handlers"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/metadata"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/metrics"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
 	schedulingtypes "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
 	errutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/error"
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
@@ -167,8 +168,18 @@ func (d *Director) HandleRequest(ctx context.Context, reqCtx *handlers.RequestCo
 		logger.V(logutil.DEFAULT).Info("Request rejected by admission control", "error", err)
 		return reqCtx, err
 	}
+	copyOfCandidatePods := d.toSchedulerPodMetrics(candidatePods)
 
-	result, err := d.scheduler.Schedule(ctx, reqCtx.SchedulingRequest, d.toSchedulerPodMetrics(candidatePods))
+	// Prepare per request data
+	d.runPrepareDataPlugins(ctx, reqCtx.SchedulingRequest, copyOfCandidatePods)
+
+	// Run admit request plugins
+	if !d.runAdmitRequestPlugins(ctx, reqCtx.SchedulingRequest, copyOfCandidatePods) {
+		logger.V(logutil.DEFAULT).Info("Request cannot be admitted")
+		return reqCtx, errutil.Error{Code: errutil.Internal, Msg: "request cannot be admitted"}
+	}
+
+	result, err := d.scheduler.Schedule(ctx, reqCtx.SchedulingRequest, copyOfCandidatePods)
 	if err != nil {
 		return reqCtx, errutil.Error{Code: errutil.InferencePoolResourceExhausted, Msg: fmt.Errorf("failed to find target pod: %w", err).Error()}
 	}
@@ -333,6 +344,29 @@ func (d *Director) runPreRequestPlugins(ctx context.Context, request *scheduling
 	}
 }
 
+func (d *Director) runPrepareDataPlugins(ctx context.Context,
+	request *schedulingtypes.LLMRequest, pods []types.Pod) {
+	loggerDebug := log.FromContext(ctx).V(logutil.DEBUG)
+	for _, plugin := range d.requestControlPlugins.prepareDataPlugins {
+		loggerDebug.Info("Running PrepareData plugin", "plugin", plugin.TypedName())
+		plugin.PrepareData(ctx, request, pods)
+		loggerDebug.Info("Completed running PrepareData plugin successfully", "plugin", plugin.TypedName())
+	}
+}
+
+func (d *Director) runAdmitRequestPlugins(ctx context.Context,
+	request *schedulingtypes.LLMRequest, pods []types.Pod) bool {
+	loggerDebug := log.FromContext(ctx).V(logutil.DEBUG)
+	for _, plugin := range d.requestControlPlugins.admitRequestPlugins {
+		loggerDebug.Info("Running AdmitRequest plugin", "plugin", plugin.TypedName())
+		if !plugin.Admit(ctx, request, pods) {
+			return false
+		}
+		loggerDebug.Info("Completed running AdmitRequest plugin successfully", "plugin", plugin.TypedName())
+	}
+	return true
+}
+
 func (d *Director) runResponseReceivedPlugins(ctx context.Context, request *schedulingtypes.LLMRequest, response *Response, targetPod *backend.Pod) {
 	loggerDebug := log.FromContext(ctx).V(logutil.DEBUG)
 	for _, plugin := range d.requestControlPlugins.responseReceivedPlugins {

pkg/epp/requestcontrol/plugins.go

@@ -57,3 +57,15 @@ type ResponseComplete interface {
 	plugins.Plugin
 	ResponseComplete(ctx context.Context, request *types.LLMRequest, response *Response, targetPod *backend.Pod)
 }
+
+// PrepareData is called by the director before scheduling requests.
+type PrepareData interface {
+	plugins.Plugin
+	PrepareData(ctx context.Context, request *types.LLMRequest, pods []types.Pod)
+}
+
+// AdmitRequest is called by the director after the PrepareData plugins and before scheduling.
+type AdmitRequest interface {
+	plugins.Plugin
+	Admit(ctx context.Context, request *types.LLMRequest, pods []types.Pod) bool
+}

pkg/epp/requestcontrol/request_control_config.go

@@ -23,6 +23,8 @@ import (
 // NewConfig creates a new Config object and returns its pointer.
 func NewConfig() *Config {
 	return &Config{
+		admitRequestPlugins:      []AdmitRequest{},
+		prepareDataPlugins:       []PrepareData{},
 		preRequestPlugins:        []PreRequest{},
 		responseReceivedPlugins:  []ResponseReceived{},
 		responseStreamingPlugins: []ResponseStreaming{},
@@ -32,6 +34,8 @@ func NewConfig() *Config {
 
 // Config provides a configuration for the requestcontrol plugins.
 type Config struct {
+	admitRequestPlugins      []AdmitRequest
+	prepareDataPlugins       []PrepareData
 	preRequestPlugins        []PreRequest
 	responseReceivedPlugins  []ResponseReceived
 	responseStreamingPlugins []ResponseStreaming

pkg/epp/scheduling/types/types.go

@@ -21,6 +21,7 @@ import (
 	"errors"
 	"fmt"
 	"strings"
+	"sync"
 
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend"
 	backendmetrics "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend/metrics"
@@ -191,6 +192,9 @@ type Pod interface {
 	GetPod() *backend.Pod
 	GetMetrics() *backendmetrics.MetricsState
 	String() string
+	Put(key string, value Cloneable)
+	Get(key string) (Cloneable, bool)
+	Keys() []string
 }
 
 type ScoredPod struct {
@@ -217,6 +221,19 @@ func (pm *PodMetrics) GetMetrics() *backendmetrics.MetricsState {
 type PodMetrics struct {
 	*backend.Pod
 	*backendmetrics.MetricsState
+	AttributeMap
+}
+
+func (pm *PodMetrics) Put(key string, value Cloneable) {
+	pm.AttributeMap.Put(key, value)
+}
+
+func (pm *PodMetrics) Get(key string) (Cloneable, bool) {
+	return pm.AttributeMap.Get(key)
+}
+
+func (pm *PodMetrics) Keys() []string {
+	return pm.AttributeMap.Keys()
 }
 
 // ProfileRunResult captures the profile run result.
@@ -229,3 +246,72 @@ type SchedulingResult struct {
 	ProfileResults     map[string]*ProfileRunResult
 	PrimaryProfileName string
 }
+
+// Cloneable types support cloning of the value.
+type Cloneable interface {
+	Clone() Cloneable
+}
+
+// AttributeMap is used to store flexible metadata or traits
+// across different aspects of an inference server.
+// Stored values must be Cloneable.
+type AttributeMap interface {
+	Put(string, Cloneable)
+	Get(string) (Cloneable, bool)
+	Keys() []string
+}
+
+// Attributes provides a goroutine-safe implementation of AttributeMap.
+type Attributes struct {
+	data sync.Map // key: attribute name (string), value: attribute value (opaque, Cloneable)
+}
+
+// NewAttributes returns a new instance of Attributes.
+func NewAttributes() *Attributes {
+	return &Attributes{}
+}
+
+// Put adds or updates an attribute in the map.
+func (a *Attributes) Put(key string, value Cloneable) {
+	if value != nil {
+		a.data.Store(key, value) // TODO: Clone into map to ensure isolation
+	}
+}
+
+// Get retrieves an attribute by key, returning a cloned copy.
+func (a *Attributes) Get(key string) (Cloneable, bool) {
+	val, ok := a.data.Load(key)
+	if !ok {
+		return nil, false
+	}
+	if cloneable, ok := val.(Cloneable); ok {
+		return cloneable.Clone(), true
+	}
+	return nil, false
+}
+
+// Keys returns all keys in the attribute map.
+func (a *Attributes) Keys() []string {
+	var keys []string
+	a.data.Range(func(key, _ any) bool {
+		if sk, ok := key.(string); ok {
+			keys = append(keys, sk)
+		}
+		return true
+	})
+	return keys
+}
+
+// Clone creates a deep copy of the entire attribute map.
+func (a *Attributes) Clone() *Attributes {
+	clone := NewAttributes()
+	a.data.Range(func(key, value any) bool {
+		if sk, ok := key.(string); ok {
+			if v, ok := value.(Cloneable); ok {
+				clone.Put(sk, v)
+			}
+		}
+		return true
+	})
+	return clone
+}