removed datastore dependency from saturation detector

Signed-off-by: Nir Rozenbaum <[email protected]>

Author: nirrozenbaum

cmd/epp/runner/runner.go

@@ -333,7 +333,7 @@ func (r *Runner) Run(ctx context.Context) error {
 
 	scheduler := scheduling.NewSchedulerWithConfig(r.schedulerConfig)
 
-	saturationDetector := saturationdetector.NewDetector(sdConfig, datastore, setupLog)
+	saturationDetector := saturationdetector.NewDetector(sdConfig, setupLog)
 
 	director := requestcontrol.NewDirectorWithConfig(datastore, scheduler, saturationDetector, r.requestControlConfig)
 

pkg/epp/requestcontrol/director.go

@@ -30,10 +30,10 @@ import (
 	"github.com/go-logr/logr"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 
+	v1 "sigs.k8s.io/gateway-api-inference-extension/api/v1"
 	"sigs.k8s.io/gateway-api-inference-extension/apix/v1alpha2"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend"
 	backendmetrics "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend/metrics"
-	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/datastore"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/handlers"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/metrics"
 	schedulingtypes "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
@@ -47,18 +47,25 @@ const (
 	subsetHintKey       = "x-gateway-destination-endpoint-subset"
 )
 
+// Datastore defines the interface required by the Director.
+type Datastore interface {
+	PoolGet() (*v1.InferencePool, error)
+	ObjectiveGet(modelName string) *v1alpha2.InferenceObjective
+	PodList(predicate func(backendmetrics.PodMetrics) bool) []backendmetrics.PodMetrics
+}
+
 // Scheduler defines the interface required by the Director for scheduling.
 type Scheduler interface {
 	Schedule(ctx context.Context, request *schedulingtypes.LLMRequest, candidatePods []schedulingtypes.Pod) (result *schedulingtypes.SchedulingResult, err error)
 }
 
 // SaturationDetector provides a signal indicating whether the backends are considered saturated.
 type SaturationDetector interface {
-	IsSaturated(ctx context.Context) bool
+	IsSaturated(ctx context.Context, candidatePods []backendmetrics.PodMetrics) bool
 }
 
 // NewDirectorWithConfig creates a new Director instance with all dependencies.
-func NewDirectorWithConfig(datastore datastore.Datastore, scheduler Scheduler, saturationDetector SaturationDetector, config *Config) *Director {
+func NewDirectorWithConfig(datastore Datastore, scheduler Scheduler, saturationDetector SaturationDetector, config *Config) *Director {
 	return &Director{
 		datastore:           datastore,
 		scheduler:           scheduler,
@@ -70,24 +77,19 @@ func NewDirectorWithConfig(datastore datastore.Datastore, scheduler Scheduler, s
 
 // Director orchestrates the request handling flow, including scheduling.
 type Director struct {
-	datastore           datastore.Datastore
+	datastore           Datastore
 	scheduler           Scheduler
 	saturationDetector  SaturationDetector
 	preRequestPlugins   []PreRequest
 	postResponsePlugins []PostResponse
 }
 
-// HandleRequest orchestrates the request lifecycle:
-//  1. Parses request details.
-//  2. Calls admitRequest for admission control.
-//  3. Calls Scheduler.Schedule if request is approved.
-//  4. Calls prepareRequest to populate RequestContext with result and call PreRequest plugins.
-//
+// HandleRequest orchestrates the request lifecycle.
 // It always returns the requestContext even in the error case, as the request context is used in error handling.
 func (d *Director) HandleRequest(ctx context.Context, reqCtx *handlers.RequestContext) (*handlers.RequestContext, error) {
 	logger := log.FromContext(ctx)
 
-	// --- 1. Parse Request, Resolve Target Models, and Determine Parameters ---
+	// Parse Request, Resolve Target Models, and Determine Parameters
 	var ok bool
 	requestBodyMap := reqCtx.Request.Body
 	reqCtx.Model, ok = requestBodyMap["model"].(string)
@@ -138,17 +140,18 @@ func (d *Director) HandleRequest(ctx context.Context, reqCtx *handlers.RequestCo
 	ctx = log.IntoContext(ctx, logger)
 	logger.V(logutil.DEBUG).Info("LLM request assembled")
 
-	// --- 2. Admission Control check --
-	if err := d.admitRequest(ctx, requestCriticality, reqCtx.FairnessID); err != nil {
-		return reqCtx, err
-	}
-
-	// --- 3. Call Scheduler (with the relevant candidate pods) ---
+	// Get candidate pods for scheduling
 	candidatePods := d.getCandidatePodsForScheduling(ctx, reqCtx.Request.Metadata)
 	if len(candidatePods) == 0 {
 		return reqCtx, errutil.Error{Code: errutil.ServiceUnavailable, Msg: "failed to find candidate pods for serving the request"}
 	}
-	result, err := d.scheduler.Schedule(ctx, reqCtx.SchedulingRequest, candidatePods)
+
+	// Admission Control check
+	if err := d.admitRequest(ctx, candidatePods, requestCriticality, reqCtx.FairnessID); err != nil {
+		return reqCtx, err
+	}
+
+	result, err := d.scheduler.Schedule(ctx, reqCtx.SchedulingRequest, d.toSchedulerPodMetrics(candidatePods))
 	if err != nil {
 		return reqCtx, errutil.Error{Code: errutil.InferencePoolResourceExhausted, Msg: fmt.Errorf("failed to find target pod: %w", err).Error()}
 	}
@@ -165,8 +168,9 @@ func (d *Director) HandleRequest(ctx context.Context, reqCtx *handlers.RequestCo
 }
 
 // admitRequest handles admission control to decide whether or not to accept the request
-// based on the request criticality and system saturation state.
-func (d *Director) admitRequest(ctx context.Context, requestCriticality v1alpha2.Criticality, fairnessID string) error {
+// based on the request criticality and the saturation state of the candidate pods.
+func (d *Director) admitRequest(ctx context.Context, candidatePods []backendmetrics.PodMetrics,
+	requestCriticality v1alpha2.Criticality, fairnessID string) error {
 	logger := log.FromContext(ctx)
 
 	logger.V(logutil.TRACE).Info("Entering Flow Control", "criticality", requestCriticality, "fairnessID", fairnessID)
@@ -177,7 +181,7 @@ func (d *Director) admitRequest(ctx context.Context, requestCriticality v1alpha2
 	}
 
 	logger.V(logutil.DEBUG).Info("Performing saturation check for non-critical request.")
-	if d.saturationDetector.IsSaturated(ctx) { // Assuming non-nil Saturation Detector
+	if d.saturationDetector.IsSaturated(ctx, candidatePods) {
 		return errutil.Error{
 			Code: errutil.InferencePoolResourceExhausted,
 			Msg:  "system saturated, non-critical request dropped",
@@ -193,21 +197,21 @@ func (d *Director) admitRequest(ctx context.Context, requestCriticality v1alpha2
 // Snapshot pod metrics from the datastore to:
 // 1. Reduce concurrent access to the datastore.
 // 2. Ensure consistent data during the scheduling operation of a request between all scheduling cycles.
-func (d *Director) getCandidatePodsForScheduling(ctx context.Context, requestMetadata map[string]any) []schedulingtypes.Pod {
+func (d *Director) getCandidatePodsForScheduling(ctx context.Context, requestMetadata map[string]any) []backendmetrics.PodMetrics {
 	loggerTrace := log.FromContext(ctx).V(logutil.TRACE)
 
 	subsetMap, found := requestMetadata[subsetHintNamespace].(map[string]any)
 	if !found {
-		return d.toSchedulerPodMetrics(d.datastore.PodList(backendmetrics.AllPodPredicate))
+		return d.datastore.PodList(backendmetrics.AllPodPredicate)
 	}
 
 	// Check if endpoint key is present in the subset map and ensure there is at least one value
 	endpointSubsetList, found := subsetMap[subsetHintKey].([]any)
 	if !found {
-		return d.toSchedulerPodMetrics(d.datastore.PodList(backendmetrics.AllPodPredicate))
+		return d.datastore.PodList(backendmetrics.AllPodPredicate)
 	} else if len(endpointSubsetList) == 0 {
 		loggerTrace.Info("found empty subset filter in request metadata, filtering all pods")
-		return []schedulingtypes.Pod{}
+		return []backendmetrics.PodMetrics{}
 	}
 
 	// Create a map of endpoint addresses for easy lookup
@@ -230,7 +234,7 @@ func (d *Director) getCandidatePodsForScheduling(ctx context.Context, requestMet
 
 	loggerTrace.Info("filtered candidate pods by subset filtering", "podTotalCount", podTotalCount, "filteredCount", len(podFitleredList))
 
-	return d.toSchedulerPodMetrics(podFitleredList)
+	return podFitleredList
 }
 
 // prepareRequest populates the RequestContext and calls the registered PreRequest plugins

pkg/epp/requestcontrol/director_test.go

@@ -31,6 +31,7 @@ import (
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/types"
 	clientgoscheme "k8s.io/client-go/kubernetes/scheme"
+	"k8s.io/utils/ptr"
 	"sigs.k8s.io/controller-runtime/pkg/client/fake"
 
 	v1 "sigs.k8s.io/gateway-api-inference-extension/api/v1"
@@ -53,7 +54,7 @@ type mockSaturationDetector struct {
 	isSaturated bool
 }
 
-func (m *mockSaturationDetector) IsSaturated(_ context.Context) bool {
+func (m *mockSaturationDetector) IsSaturated(_ context.Context, _ []backendmetrics.PodMetrics) bool {
 	return m.isSaturated
 }
 
@@ -66,6 +67,23 @@ func (m *mockScheduler) Schedule(_ context.Context, _ *schedulingtypes.LLMReques
 	return m.scheduleResults, m.scheduleErr
 }
 
+type mockDatastore struct {
+	pods []backendmetrics.PodMetrics
+}
+
+func (ds *mockDatastore) PoolGet() (*v1.InferencePool, error)                { return nil, nil }
+func (ds *mockDatastore) ObjectiveGet(_ string) *v1alpha2.InferenceObjective { return nil }
+func (ds *mockDatastore) PodList(predicate func(backendmetrics.PodMetrics) bool) []backendmetrics.PodMetrics {
+	res := []backendmetrics.PodMetrics{}
+	for _, pod := range ds.pods {
+		if predicate(pod) {
+			res = append(res, pod)
+		}
+	}
+
+	return res
+}
+
 func TestDirector_HandleRequest(t *testing.T) {
 	ctx := logutil.NewTestLoggerIntoContext(context.Background())
 
@@ -425,125 +443,78 @@ func TestDirector_HandleRequest(t *testing.T) {
 func TestGetCandidatePodsForScheduling(t *testing.T) {
 	var makeFilterMetadata = func(data []any) map[string]any {
 		return map[string]any{
-			"envoy.lb.subset_hint": map[string]any{
-				"x-gateway-destination-endpoint-subset": data,
+			subsetHintNamespace: map[string]any{
+				subsetHintKey: data,
 			},
 		}
 	}
 
-	testInput := []*corev1.Pod{
-		{
-			ObjectMeta: metav1.ObjectMeta{
-				Name: "pod1",
-			},
-			Status: corev1.PodStatus{
-				PodIP: "10.0.0.1",
-			},
-		},
-		{
-			ObjectMeta: metav1.ObjectMeta{
-				Name: "pod2",
-			},
-			Status: corev1.PodStatus{
-				PodIP: "10.0.0.2",
-			},
-		},
-	}
-
-	outputPod1 := &backend.Pod{
+	pod1 := &backend.Pod{
 		NamespacedName: types.NamespacedName{Name: "pod1"},
 		Address:        "10.0.0.1",
 		Labels:         map[string]string{},
 	}
 
-	outputPod2 := &backend.Pod{
+	pod2 := &backend.Pod{
 		NamespacedName: types.NamespacedName{Name: "pod2"},
 		Address:        "10.0.0.2",
 		Labels:         map[string]string{},
 	}
 
+	testInput := []backendmetrics.PodMetrics{
+		&backendmetrics.FakePodMetrics{Pod: pod1},
+		&backendmetrics.FakePodMetrics{Pod: pod2},
+	}
+
 	tests := []struct {
 		name     string
 		metadata map[string]any
-		output   []schedulingtypes.Pod
+		output   []backendmetrics.PodMetrics
 	}{
 		{
 			name:     "SubsetFilter, filter not present — return all pods",
 			metadata: map[string]any{},
-			output: []schedulingtypes.Pod{
-				&schedulingtypes.PodMetrics{
-					Pod:          outputPod1,
-					MetricsState: backendmetrics.NewMetricsState(),
-				},
-				&schedulingtypes.PodMetrics{
-					Pod:          outputPod2,
-					MetricsState: backendmetrics.NewMetricsState(),
-				},
-			},
+			output:   testInput,
 		},
 		{
 			name:     "SubsetFilter, namespace present filter not present — return all pods",
-			metadata: map[string]any{"envoy.lb.subset_hint": map[string]any{}},
-			output: []schedulingtypes.Pod{
-				&schedulingtypes.PodMetrics{
-					Pod:          outputPod1,
-					MetricsState: backendmetrics.NewMetricsState(),
-				},
-				&schedulingtypes.PodMetrics{
-					Pod:          outputPod2,
-					MetricsState: backendmetrics.NewMetricsState(),
-				},
-			},
+			metadata: map[string]any{subsetHintNamespace: map[string]any{}},
+			output:   testInput,
 		},
 		{
 			name:     "SubsetFilter, filter present with empty list — return error",
 			metadata: makeFilterMetadata([]any{}),
-			output:   []schedulingtypes.Pod{},
+			output:   []backendmetrics.PodMetrics{},
 		},
 		{
 			name:     "SubsetFilter, subset with one matching pod",
 			metadata: makeFilterMetadata([]any{"10.0.0.1"}),
-			output: []schedulingtypes.Pod{
-				&schedulingtypes.PodMetrics{
-					Pod:          outputPod1,
-					MetricsState: backendmetrics.NewMetricsState(),
+			output: []backendmetrics.PodMetrics{
+				&backendmetrics.FakePodMetrics{
+					Pod: pod1,
 				},
 			},
 		},
 		{
 			name:     "SubsetFilter, subset with multiple matching pods",
 			metadata: makeFilterMetadata([]any{"10.0.0.1", "10.0.0.2", "10.0.0.3"}),
-			output: []schedulingtypes.Pod{
-				&schedulingtypes.PodMetrics{
-					Pod:          outputPod1,
-					MetricsState: backendmetrics.NewMetricsState(),
-				},
-				&schedulingtypes.PodMetrics{
-					Pod:          outputPod2,
-					MetricsState: backendmetrics.NewMetricsState(),
-				},
-			},
+			output:   testInput,
 		},
 		{
 			name:     "SubsetFilter, subset with no matching pods",
 			metadata: makeFilterMetadata([]any{"10.0.0.3"}),
-			output:   []schedulingtypes.Pod{},
+			output:   []backendmetrics.PodMetrics{},
 		},
 	}
 
-	pmf := backendmetrics.NewPodMetricsFactory(&backendmetrics.FakePodMetricsClient{}, time.Second, time.Second*2)
-	ds := datastore.NewDatastore(t.Context(), pmf)
-	for _, testPod := range testInput {
-		ds.PodUpdateOrAddIfNotExist(testPod)
-	}
-
+	ds := &mockDatastore{pods: testInput}
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
 			director := NewDirectorWithConfig(ds, &mockScheduler{}, &mockSaturationDetector{}, NewConfig())
 
 			got := director.getCandidatePodsForScheduling(context.Background(), test.metadata)
 
-			diff := cmp.Diff(test.output, got, cmpopts.SortSlices(func(a, b schedulingtypes.Pod) bool {
+			diff := cmp.Diff(test.output, got, cmpopts.SortSlices(func(a, b backendmetrics.PodMetrics) bool {
 				return a.GetPod().NamespacedName.String() < b.GetPod().NamespacedName.String()
 			}))
 			if diff != "" {
@@ -567,8 +538,8 @@ func TestRandomWeightedDraw(t *testing.T) {
 			model: &v1alpha2.InferenceObjective{
 				Spec: v1alpha2.InferenceObjectiveSpec{
 					TargetModels: []v1alpha2.TargetModel{
-						{Name: "canary", Weight: pointer(50)},
-						{Name: "v1", Weight: pointer(50)},
+						{Name: "canary", Weight: ptr.To(int32(50))},
+						{Name: "v1", Weight: ptr.To(int32(50))},
 					},
 				},
 			},
@@ -579,9 +550,9 @@ func TestRandomWeightedDraw(t *testing.T) {
 			model: &v1alpha2.InferenceObjective{
 				Spec: v1alpha2.InferenceObjectiveSpec{
 					TargetModels: []v1alpha2.TargetModel{
-						{Name: "canary", Weight: pointer(25)},
-						{Name: "v1.1", Weight: pointer(55)},
-						{Name: "v1", Weight: pointer(50)},
+						{Name: "canary", Weight: ptr.To(int32(25))},
+						{Name: "v1.1", Weight: ptr.To(int32(55))},
+						{Name: "v1", Weight: ptr.To(int32(50))},
 					},
 				},
 			},
@@ -592,9 +563,9 @@ func TestRandomWeightedDraw(t *testing.T) {
 			model: &v1alpha2.InferenceObjective{
 				Spec: v1alpha2.InferenceObjectiveSpec{
 					TargetModels: []v1alpha2.TargetModel{
-						{Name: "canary", Weight: pointer(20)},
-						{Name: "v1.1", Weight: pointer(20)},
-						{Name: "v1", Weight: pointer(10)},
+						{Name: "canary", Weight: ptr.To(int32(20))},
+						{Name: "v1.1", Weight: ptr.To(int32(20))},
+						{Name: "v1", Weight: ptr.To(int32(10))},
 					},
 				},
 			},
@@ -672,10 +643,6 @@ func TestGetRandomPod(t *testing.T) {
 	}
 }
 
-func pointer(v int32) *int32 {
-	return &v
-}
-
 func TestDirector_HandleResponse(t *testing.T) {
 	pr1 := newTestPostResponse("pr1")