move AllPodsPredicate to datastore package (#1939)

Author: elevran

cmd/epp/runner/runner.go

@@ -479,7 +479,7 @@ func (r *Runner) parseConfigurationPhaseOne(ctx context.Context) (*configapi.End
 // Return a function that can be used in the EPP Handle to list pod names.
 func makePodListFunc(ds datastore.Datastore) func() []types.NamespacedName {
 	return func() []types.NamespacedName {
-		pods := ds.PodList(backendmetrics.AllPodsPredicate)
+		pods := ds.PodList(datastore.AllPodsPredicate)
 		names := make([]types.NamespacedName, 0, len(pods))
 
 		for _, p := range pods {

pkg/epp/backend/metrics/types.go

@@ -27,10 +27,6 @@ import (
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/datalayer"
 )
 
-var (
-	AllPodsPredicate = func(PodMetrics) bool { return true }
-)
-
 func PodsWithFreshMetrics(stalenessThreshold time.Duration) func(PodMetrics) bool {
 	return func(pm PodMetrics) bool {
 		if pm == nil {

pkg/epp/controller/inferencepool_reconciler_test.go

@@ -182,8 +182,8 @@ type diffStoreParams struct {
 	wantObjectives []*v1alpha2.InferenceObjective
 }
 
-func diffStore(datastore datastore.Datastore, params diffStoreParams) string {
-	gotPool, _ := datastore.PoolGet()
+func diffStore(store datastore.Datastore, params diffStoreParams) string {
+	gotPool, _ := store.PoolGet()
 	if diff := cmp.Diff(params.wantPool, gotPool); diff != "" {
 		return "inferencePool:" + diff
 	}
@@ -193,7 +193,7 @@ func diffStore(datastore datastore.Datastore, params diffStoreParams) string {
 		params.wantPods = []string{}
 	}
 	gotPods := []string{}
-	for _, pm := range datastore.PodList(backendmetrics.AllPodsPredicate) {
+	for _, pm := range store.PodList(datastore.AllPodsPredicate) {
 		gotPods = append(gotPods, pm.GetPod().NamespacedName.Name)
 	}
 	if diff := cmp.Diff(params.wantPods, gotPods, cmpopts.SortSlices(func(a, b string) bool { return a < b })); diff != "" {
@@ -205,7 +205,7 @@ func diffStore(datastore datastore.Datastore, params diffStoreParams) string {
 		params.wantObjectives = []*v1alpha2.InferenceObjective{}
 	}
 
-	if diff := cmp.Diff(params.wantObjectives, datastore.ObjectiveGetAll(), cmpopts.SortSlices(func(a, b *v1alpha2.InferenceObjective) bool {
+	if diff := cmp.Diff(params.wantObjectives, store.ObjectiveGetAll(), cmpopts.SortSlices(func(a, b *v1alpha2.InferenceObjective) bool {
 		return a.Name < b.Name
 	})); diff != "" {
 		return "models:" + diff
@@ -328,8 +328,8 @@ type xDiffStoreParams struct {
 	wantObjectives []*v1alpha2.InferenceObjective
 }
 
-func xDiffStore(datastore datastore.Datastore, params xDiffStoreParams) string {
-	gotPool, _ := datastore.PoolGet()
+func xDiffStore(store datastore.Datastore, params xDiffStoreParams) string {
+	gotPool, _ := store.PoolGet()
 	if gotPool == nil && params.wantPool == nil {
 		return ""
 	}
@@ -343,7 +343,7 @@ func xDiffStore(datastore datastore.Datastore, params xDiffStoreParams) string {
 		params.wantPods = []string{}
 	}
 	gotPods := []string{}
-	for _, pm := range datastore.PodList(backendmetrics.AllPodsPredicate) {
+	for _, pm := range store.PodList(datastore.AllPodsPredicate) {
 		gotPods = append(gotPods, pm.GetPod().NamespacedName.Name)
 	}
 	if diff := cmp.Diff(params.wantPods, gotPods, cmpopts.SortSlices(func(a, b string) bool { return a < b })); diff != "" {
@@ -355,7 +355,7 @@ func xDiffStore(datastore datastore.Datastore, params xDiffStoreParams) string {
 		params.wantObjectives = []*v1alpha2.InferenceObjective{}
 	}
 
-	if diff := cmp.Diff(params.wantObjectives, datastore.ObjectiveGetAll(), cmpopts.SortSlices(func(a, b *v1alpha2.InferenceObjective) bool {
+	if diff := cmp.Diff(params.wantObjectives, store.ObjectiveGetAll(), cmpopts.SortSlices(func(a, b *v1alpha2.InferenceObjective) bool {
 		return a.Name < b.Name
 	})); diff != "" {
 		return "models:" + diff

pkg/epp/controller/pod_reconciler_test.go

@@ -213,7 +213,7 @@ func TestPodReconciler(t *testing.T) {
 			}
 
 			var gotPods []*corev1.Pod
-			for _, pm := range store.PodList(backendmetrics.AllPodsPredicate) {
+			for _, pm := range store.PodList(datastore.AllPodsPredicate) {
 				pod := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{Name: pm.GetPod().PodName, Namespace: pm.GetPod().NamespacedName.Namespace}, Status: corev1.PodStatus{PodIP: pm.GetPod().GetIPAddress()}}
 				gotPods = append(gotPods, pod)
 			}

pkg/epp/datastore/datastore.go

@@ -40,6 +40,7 @@ import (
 
 var (
 	errPoolNotSynced = errors.New("InferencePool is not initialized in data store")
+	AllPodsPredicate = func(_ datalayer.Endpoint) bool { return true }
 )
 
 // The datastore is a local cache of relevant data for the given InferencePool (currently all pulled from k8s-api)

pkg/epp/datastore/datastore_test.go

@@ -335,7 +335,7 @@ func TestMetrics(t *testing.T) {
 			}
 			time.Sleep(1 * time.Second) // Give some time for the metrics to be fetched.
 			if test.predict == nil {
-				test.predict = backendmetrics.AllPodsPredicate
+				test.predict = AllPodsPredicate
 			}
 			assert.EventuallyWithT(t, func(t *assert.CollectT) {
 				got := ds.PodList(test.predict)
@@ -407,7 +407,7 @@ func TestPods(t *testing.T) {
 
 			test.op(ctx, ds)
 			var gotPods []*corev1.Pod
-			for _, pm := range ds.PodList(backendmetrics.AllPodsPredicate) {
+			for _, pm := range ds.PodList(AllPodsPredicate) {
 				pod := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{Name: pm.GetPod().PodName, Namespace: pm.GetPod().NamespacedName.Namespace}, Status: corev1.PodStatus{PodIP: pm.GetPod().GetIPAddress()}}
 				gotPods = append(gotPods, pod)
 			}
@@ -591,7 +591,7 @@ func TestPodInfo(t *testing.T) {
 
 			test.op(ctx, ds)
 			var gotPodInfos []*datalayer.PodInfo
-			for _, pm := range ds.PodList(backendmetrics.AllPodsPredicate) {
+			for _, pm := range ds.PodList(AllPodsPredicate) {
 				gotPodInfos = append(gotPodInfos, pm.GetPod())
 			}
 			if diff := cmp.Diff(test.wantPodInfos, gotPodInfos, cmpopts.SortSlices(func(a, b *datalayer.PodInfo) bool { return a.NamespacedName.Name < b.NamespacedName.Name })); diff != "" {

pkg/epp/metrics/collectors/inference_pool.go

@@ -20,7 +20,6 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 	compbasemetrics "k8s.io/component-base/metrics"
 
-	backendmetrics "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend/metrics"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/datastore"
 	metricsutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/metrics"
 )
@@ -63,7 +62,7 @@ func (c *inferencePoolMetricsCollector) Collect(ch chan<- prometheus.Metric) {
 		return
 	}
 
-	podMetrics := c.ds.PodList(backendmetrics.AllPodsPredicate)
+	podMetrics := c.ds.PodList(datastore.AllPodsPredicate)
 	if len(podMetrics) == 0 {
 		return
 	}

pkg/epp/requestcontrol/director.go

@@ -31,6 +31,7 @@ import (
 	"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/datalayer"
+	"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/metadata"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/metrics"
@@ -240,13 +241,13 @@ func (d *Director) getCandidatePodsForScheduling(ctx context.Context, requestMet
 
 	subsetMap, found := requestMetadata[metadata.SubsetFilterNamespace].(map[string]any)
 	if !found {
-		return d.datastore.PodList(backendmetrics.AllPodsPredicate)
+		return d.datastore.PodList(datastore.AllPodsPredicate)
 	}
 
 	// Check if endpoint key is present in the subset map and ensure there is at least one value
 	endpointSubsetList, found := subsetMap[metadata.SubsetFilterKey].([]any)
 	if !found {
-		return d.datastore.PodList(backendmetrics.AllPodsPredicate)
+		return d.datastore.PodList(datastore.AllPodsPredicate)
 	} else if len(endpointSubsetList) == 0 {
 		loggerTrace.Info("found empty subset filter in request metadata, filtering all pods")
 		return []backendmetrics.PodMetrics{}
@@ -362,7 +363,7 @@ func (d *Director) HandleResponseBodyComplete(ctx context.Context, reqCtx *handl
 }
 
 func (d *Director) GetRandomPod() *backend.Pod {
-	pods := d.datastore.PodList(backendmetrics.AllPodsPredicate)
+	pods := d.datastore.PodList(datastore.AllPodsPredicate)
 	if len(pods) == 0 {
 		return nil
 	}

pkg/epp/requestcontrol/director_test.go

@@ -661,7 +661,7 @@ func TestDirector_HandleRequest(t *testing.T) {
 			config = config.WithAdmissionPlugins(newMockAdmissionPlugin("test-admit-plugin", test.admitRequestDenialError))
 			director := NewDirectorWithConfig(ds, mockSched, test.mockAdmissionController, config)
 			if test.name == "successful request with model rewrite" {
-				director.datastore = &mockDatastore{pods: ds.PodList(backendmetrics.AllPodsPredicate), rewrites: []*v1alpha2.InferenceModelRewrite{rewrite}}
+				director.datastore = &mockDatastore{pods: ds.PodList(datastore.AllPodsPredicate), rewrites: []*v1alpha2.InferenceModelRewrite{rewrite}}
 			}
 
 			reqCtx := &handlers.RequestContext{

test/integration/epp/hermetic_test.go

@@ -1096,7 +1096,7 @@ func setUpHermeticServer(t *testing.T, podAndMetrics map[*backend.Pod]*backendme
 
 	// check if all pods are synced to datastore
 	assert.EventuallyWithT(t, func(t *assert.CollectT) {
-		assert.Len(t, serverRunner.Datastore.PodList(backendmetrics.AllPodsPredicate), len(podAndMetrics), "Datastore not synced")
+		assert.Len(t, serverRunner.Datastore.PodList(datastore.AllPodsPredicate), len(podAndMetrics), "Datastore not synced")
 	}, 10*time.Second, time.Second)
 
 	// Create a grpc connection