Enable EPP to support endpoint discovery using pod selector

cmd/epp/runner/runner.go

@@ -26,6 +26,8 @@ import (
 	"net/http/pprof"
 	"os"
 	"runtime"
+	"strconv"
+	"strings"
 	"sync/atomic"
 
 	"github.com/go-logr/logr"
@@ -34,16 +36,18 @@ import (
 	"go.uber.org/zap/zapcore"
 	"google.golang.org/grpc"
 	healthPb "google.golang.org/grpc/health/grpc_health_v1"
+	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/apimachinery/pkg/util/sets"
 	"k8s.io/client-go/rest"
+
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/controller-runtime/pkg/log/zap"
 	"sigs.k8s.io/controller-runtime/pkg/manager"
 	"sigs.k8s.io/controller-runtime/pkg/metrics/filters"
 	metricsserver "sigs.k8s.io/controller-runtime/pkg/metrics/server"
-
 	"sigs.k8s.io/gateway-api-inference-extension/internal/runnable"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/common"
 	backendmetrics "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend/metrics"
@@ -100,6 +104,8 @@ var (
 	poolName            = flag.String("pool-name", runserver.DefaultPoolName, "Name of the InferencePool this Endpoint Picker is associated with.")
 	poolGroup           = flag.String("pool-group", runserver.DefaultPoolGroup, "group of the InferencePool this Endpoint Picker is associated with.")
 	poolNamespace       = flag.String("pool-namespace", "", "Namespace of the InferencePool this Endpoint Picker is associated with.")
+	endpointSelector    = flag.String("endpoint-selector", "", "selector to filter model server pods on, only key=value paris is supported. Format: a comma-separated list of key value paris,  e.g., 'app=vllm-llama3-8b-instruct,env=prod'.")
+	endpointTargetPorts = flag.String("endpoint-target-ports", "", "target ports of model server pods. Format: a comma-separated list of numbers, e.g., '3000,3001,3002'")
 	logVerbosity        = flag.Int("v", logging.DEFAULT, "number for the log level verbosity")
 	secureServing       = flag.Bool("secure-serving", runserver.DefaultSecureServing, "Enables secure serving. Defaults to true.")
 	healthChecking      = flag.Bool("health-checking", runserver.DefaultHealthChecking, "Enables health checking")
@@ -209,14 +215,20 @@ func (r *Runner) Run(ctx context.Context) error {
 		setupLog.Error(err, "Failed to get Kubernetes rest config")
 		return err
 	}
+	// Setup EndpointPool
+	endpointPool, err := setupEndpointPool(setupLog)
+	if err != nil {
+		setupLog.Error(err, "Failed to set up Endpoints Pool")
+		return err
+	}
 
 	// --- Setup Datastore ---
 	useDatalayerV2 := env.GetEnvBool(enableExperimentalDatalayerV2, false, setupLog)
 	epf, err := r.setupMetricsCollection(setupLog, useDatalayerV2)
 	if err != nil {
 		return err
 	}
-	datastore := datastore.NewDatastore(ctx, epf, int32(*modelServerMetricsPort))
+	datastore := datastore.NewDatastore(ctx, epf, int32(*modelServerMetricsPort), endpointPool)
 
 	// --- Setup Metrics Server ---
 	customCollectors := []prometheus.Collector{collectors.NewInferencePoolMetricsCollector(datastore)}
@@ -241,34 +253,10 @@ func (r *Runner) Run(ctx context.Context) error {
 		}(),
 	}
 
-	// Determine pool namespace: if --pool-namespace is non-empty, use it; else NAMESPACE env var; else default
-	resolvePoolNamespace := func() string {
-		if *poolNamespace != "" {
-			return *poolNamespace
-		}
-		if nsEnv := os.Getenv("NAMESPACE"); nsEnv != "" {
-			return nsEnv
-		}
-		return runserver.DefaultPoolNamespace
-	}
-	resolvedPoolNamespace := resolvePoolNamespace()
-	poolNamespacedName := types.NamespacedName{
-		Name:      *poolName,
-		Namespace: resolvedPoolNamespace,
-	}
-	poolGroupKind := schema.GroupKind{
-		Group: *poolGroup,
-		Kind:  "InferencePool",
-	}
-	poolGKNN := common.GKNN{
-		NamespacedName: poolNamespacedName,
-		GroupKind:      poolGroupKind,
-	}
-
 	isLeader := &atomic.Bool{}
 	isLeader.Store(false)
 
-	mgr, err := runserver.NewDefaultManager(poolGKNN, cfg, metricsServerOptions, *haEnableLeaderElection)
+	mgr, err := runserver.NewDefaultManager(endpointPool, cfg, metricsServerOptions, *haEnableLeaderElection)
 	if err != nil {
 		setupLog.Error(err, "Failed to create controller manager")
 		return err
@@ -357,8 +345,7 @@ func (r *Runner) Run(ctx context.Context) error {
 	// --- Setup ExtProc Server Runner ---
 	serverRunner := &runserver.ExtProcServerRunner{
 		GrpcPort:                         *grpcPort,
-		PoolNamespacedName:               poolNamespacedName,
-		PoolGKNN:                         poolGKNN,
+		EndpointPool:                     endpointPool,
 		Datastore:                        datastore,
 		SecureServing:                    *secureServing,
 		HealthChecking:                   *healthChecking,
@@ -396,6 +383,68 @@ func (r *Runner) Run(ctx context.Context) error {
 	return nil
 }
 
+func setupEndpointPool(setupLog logr.Logger) (*datalayer.EndpointPool, error) {
+	endpointPool := datalayer.NewEndpointPool(false, common.GKNN{})
+	if *poolName != "" {
+		// Determine pool namespace: if --pool-namespace is non-empty, use it; else NAMESPACE env var; else default
+		resolvePoolNamespace := func() string {
+			if *poolNamespace != "" {
+				return *poolNamespace
+			}
+			if nsEnv := os.Getenv("NAMESPACE"); nsEnv != "" {
+				return nsEnv
+			}
+			return runserver.DefaultPoolNamespace
+		}
+		resolvedPoolNamespace := resolvePoolNamespace()
+		poolNamespacedName := types.NamespacedName{
+			Name:      *poolName,
+			Namespace: resolvedPoolNamespace,
+		}
+		poolGroupKind := schema.GroupKind{
+			Group: *poolGroup,
+			Kind:  "InferencePool",
+		}
+		poolGKNN := common.GKNN{
+			NamespacedName: poolNamespacedName,
+			GroupKind:      poolGroupKind,
+		}
+		endpointPool.GKNN = poolGKNN
+	}
+
+	if *endpointSelector != "" {
+		labelsMap, err := labels.ConvertSelectorToLabelsMap(*endpointSelector)
+		if err != nil {
+			setupLog.Error(err, "Failed to parse flag %q with error: %w", "endpoint-selector", err)
+			return nil, err
+		}
+		endpointPool.EndPoints.Selector = labelsMap
+		endpointPool.EndPoints.TargetPorts, err = strToUniqueIntSlice(*endpointTargetPorts)
+		if err != nil {
+			setupLog.Error(err, "Failed to parse flag %q with error: %w", "endpoint-target-ports", err)
+		}
+		endpointPool.DisableK8sCrdReconcile = true
+
+		// Determine EPP namespace: NAMESPACE env var; else default
+		eppNsEnv := os.Getenv("NAMESPACE")
+		if eppNsEnv == "" {
+			setupLog.Error(err, "Failed to get environment variable EPP_NAMESPACE")
+		}
+		// Determine EPP name: EPP_NAME env var
+		eppNameEnv := os.Getenv("EPP_NAME")
+		if eppNameEnv == "" {
+			setupLog.Error(err, "Failed to get environment variable EPP_NAME")
+
+		}
+		endpointPool.GKNN = common.GKNN{
+			NamespacedName: types.NamespacedName{Namespace: eppNsEnv, Name: eppNameEnv},
+			GroupKind:      schema.GroupKind{Kind: "apps", Group: "Deployment"},
+		}
+
+	}
+	return endpointPool, nil
+}
+
 // registerInTreePlugins registers the factory functions of all known plugins
 func (r *Runner) registerInTreePlugins() {
 	plugins.Register(prefix.PrefixCachePluginType, prefix.PrefixCachePluginFactory)
@@ -575,9 +624,19 @@ func registerHealthServer(mgr manager.Manager, logger logr.Logger, ds datastore.
 }
 
 func validateFlags() error {
-	if *poolName == "" {
-		return fmt.Errorf("required %q flag not set", "poolName")
+	if (*poolName != "" && *endpointSelector != "") || (*poolName == "" && *endpointSelector == "") {
+		return errors.New("either poolName or endpoint-selector must be set")
+	}
+	if *endpointSelector != "" {
+		targetPortsList, err := strToUniqueIntSlice(*endpointTargetPorts)
+		if err != nil {
+			return fmt.Errorf("unexpected value for %q flag with error %w", "endpoint-target-ports", err)
+		}
+		if len(targetPortsList) == 0 || len(targetPortsList) > 8 {
+			return fmt.Errorf("flag %q should have length from 1 to 8", "endpoint-target-ports")
+		}
 	}
+
 	if *configText != "" && *configFile != "" {
 		return fmt.Errorf("both the %q and %q flags can not be set at the same time", "configText", "configFile")
 	}
@@ -588,6 +647,34 @@ func validateFlags() error {
 	return nil
 }
 
+func strToUniqueIntSlice(s string) ([]int, error) {
+	seen := sets.NewInt()
+	var intList []int
+
+	if s == "" {
+		return intList, nil
+	}
+
+	strList := strings.Split(s, ",")
+
+	for _, str := range strList {
+		trimmedStr := strings.TrimSpace(str)
+		if trimmedStr == "" {
+			continue
+		}
+		portInt, err := strconv.Atoi(trimmedStr)
+		if err != nil {
+			return nil, fmt.Errorf("invalid number: '%s' is not an integer", trimmedStr)
+		}
+
+		if _, ok := seen[portInt]; !ok {
+			seen[portInt] = struct{}{}
+			intList = append(intList, portInt)
+		}
+	}
+	return intList, nil
+}
+
 func verifyMetricMapping(mapping backendmetrics.MetricMapping, logger logr.Logger) {
 	if mapping.TotalQueuedRequests == nil {
 		logger.Info("Not scraping metric: TotalQueuedRequests")

config/charts/inferencepool/templates/epp-deployment.yaml

@@ -147,6 +147,10 @@ spec:
           valueFrom:
             fieldRef:
               fieldPath: metadata.namespace
+        - name: EPP_NAME
+          valueFrom:
+            fieldRef:
+              fieldPath: metadata.name
         {{- if .Values.inferenceExtension.tracing.enabled }}
         - name: OTEL_SERVICE_NAME
           value: "gateway-api-inference-extension"

pkg/epp/backend/metrics/logger.go

@@ -97,7 +97,7 @@ func refreshPrometheusMetrics(logger logr.Logger, datastore datalayer.PoolInfo,
 	}
 
 	podTotalCount := len(podMetrics)
-	metrics.RecordInferencePoolAvgKVCache(pool.Name, kvCacheTotal/float64(podTotalCount))
-	metrics.RecordInferencePoolAvgQueueSize(pool.Name, float64(queueTotal/podTotalCount))
-	metrics.RecordInferencePoolReadyPods(pool.Name, float64(podTotalCount))
+	metrics.RecordInferencePoolAvgKVCache(pool.GKNN.Name, kvCacheTotal/float64(podTotalCount))
+	metrics.RecordInferencePoolAvgQueueSize(pool.GKNN.Name, float64(queueTotal/podTotalCount))
+	metrics.RecordInferencePoolReadyPods(pool.GKNN.Name, float64(podTotalCount))
 }

pkg/epp/backend/metrics/pod_metrics_test.go

@@ -25,7 +25,6 @@ import (
 	"github.com/stretchr/testify/assert"
 	"k8s.io/apimachinery/pkg/types"
 
-	v1 "sigs.k8s.io/gateway-api-inference-extension/api/v1"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/datalayer"
 )
 
@@ -86,8 +85,8 @@ func TestMetricsRefresh(t *testing.T) {
 
 type fakeDataStore struct{}
 
-func (f *fakeDataStore) PoolGet() (*v1.InferencePool, error) {
-	return &v1.InferencePool{Spec: v1.InferencePoolSpec{TargetPorts: []v1.Port{{Number: 8000}}}}, nil
+func (f *fakeDataStore) PoolGet() (*datalayer.EndpointPool, error) {
+	return &datalayer.EndpointPool{}, nil
 }
 
 func (f *fakeDataStore) PodList(func(PodMetrics) bool) []PodMetrics {

pkg/epp/controller/inferenceobjective_reconciler.go

@@ -21,12 +21,12 @@ import (
 	"fmt"
 
 	"k8s.io/apimachinery/pkg/api/errors"
+
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/event"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/controller-runtime/pkg/predicate"
-
 	"sigs.k8s.io/gateway-api-inference-extension/apix/v1alpha2"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/common"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/datastore"

pkg/epp/controller/inferenceobjective_reconciler_test.go

@@ -27,25 +27,27 @@ import (
 	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/apimachinery/pkg/types"
 	clientgoscheme "k8s.io/client-go/kubernetes/scheme"
+
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/client/fake"
-
 	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/common"
 	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/util/pool"
 	utiltest "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/testing"
 )
 
 var (
-	pool          = utiltest.MakeInferencePool("test-pool1").Namespace("ns1").ObjRef()
+	inferencePool = utiltest.MakeInferencePool("test-pool1").Namespace("ns1").ObjRef()
 	infObjective1 = utiltest.MakeInferenceObjective("model1").
-			Namespace(pool.Namespace).
+			Namespace(inferencePool.Namespace).
 			Priority(1).
 			CreationTimestamp(metav1.Unix(1000, 0)).
-			PoolName(pool.Name).
+			PoolName(inferencePool.Name).
 			PoolGroup("inference.networking.k8s.io").ObjRef()
 	infObjective1Pool2 = utiltest.MakeInferenceObjective(infObjective1.Name).
 				Namespace(infObjective1.Namespace).
@@ -57,24 +59,24 @@ var (
 				Namespace(infObjective1.Namespace).
 				Priority(2).
 				CreationTimestamp(metav1.Unix(1003, 0)).
-				PoolName(pool.Name).
+				PoolName(inferencePool.Name).
 				PoolGroup("inference.networking.k8s.io").ObjRef()
 	infObjective1Deleted = utiltest.MakeInferenceObjective(infObjective1.Name).
 				Namespace(infObjective1.Namespace).
 				CreationTimestamp(metav1.Unix(1004, 0)).
 				DeletionTimestamp().
-				PoolName(pool.Name).
+				PoolName(inferencePool.Name).
 				PoolGroup("inference.networking.k8s.io").ObjRef()
 	infObjective1DiffGroup = utiltest.MakeInferenceObjective(infObjective1.Name).
-				Namespace(pool.Namespace).
+				Namespace(inferencePool.Namespace).
 				Priority(1).
 				CreationTimestamp(metav1.Unix(1005, 0)).
-				PoolName(pool.Name).
+				PoolName(inferencePool.Name).
 				PoolGroup("inference.networking.x-k8s.io").ObjRef()
 	infObjective2 = utiltest.MakeInferenceObjective("model2").
-			Namespace(pool.Namespace).
+			Namespace(inferencePool.Namespace).
 			CreationTimestamp(metav1.Unix(1000, 0)).
-			PoolName(pool.Name).
+			PoolName(inferencePool.Name).
 			PoolGroup("inference.networking.k8s.io").ObjRef()
 )
 
@@ -120,7 +122,7 @@ func TestInferenceObjectiveReconciler(t *testing.T) {
 		{
 			name:               "Objective not found, no matching existing objective to delete",
 			objectivessInStore: []*v1alpha2.InferenceObjective{infObjective1},
-			incomingReq:        &types.NamespacedName{Name: "non-existent-objective", Namespace: pool.Namespace},
+			incomingReq:        &types.NamespacedName{Name: "non-existent-objective", Namespace: inferencePool.Namespace},
 			wantObjectives:     []*v1alpha2.InferenceObjective{infObjective1},
 		},
 		{
@@ -160,17 +162,18 @@ func TestInferenceObjectiveReconciler(t *testing.T) {
 				WithObjects(initObjs...).
 				Build()
 			pmf := backendmetrics.NewPodMetricsFactory(&backendmetrics.FakePodMetricsClient{}, time.Second)
-			ds := datastore.NewDatastore(t.Context(), pmf, 0)
+			ds := datastore.NewDatastore(t.Context(), pmf, 0, datalayer.NewEndpointPool(false, pool.ToGKNN(inferencePool)))
 			for _, m := range test.objectivessInStore {
 				ds.ObjectiveSet(m)
 			}
-			_ = ds.PoolSet(context.Background(), fakeClient, pool)
+			endpointPool := pool.InferencePoolToEndpointPool(inferencePool)
+			_ = ds.PoolSet(context.Background(), fakeClient, endpointPool)
 			reconciler := &InferenceObjectiveReconciler{
 				Reader:    fakeClient,
 				Datastore: ds,
 				PoolGKNN: common.GKNN{
-					NamespacedName: types.NamespacedName{Name: pool.Name, Namespace: pool.Namespace},
-					GroupKind:      schema.GroupKind{Group: pool.GroupVersionKind().Group, Kind: pool.GroupVersionKind().Kind},
+					NamespacedName: types.NamespacedName{Name: inferencePool.Name, Namespace: inferencePool.Namespace},
+					GroupKind:      schema.GroupKind{Group: inferencePool.GroupVersionKind().Group, Kind: inferencePool.GroupVersionKind().Kind},
 				},
 			}
 			if test.incomingReq == nil {
@@ -190,8 +193,7 @@ func TestInferenceObjectiveReconciler(t *testing.T) {
 			if len(test.wantObjectives) != len(ds.ObjectiveGetAll()) {
 				t.Errorf("Unexpected; want: %d, got:%d", len(test.wantObjectives), len(ds.ObjectiveGetAll()))
 			}
-
-			if diff := diffStore(ds, diffStoreParams{wantPool: pool, wantObjectives: test.wantObjectives}); diff != "" {
+			if diff := diffStore(ds, diffStoreParams{wantPool: endpointPool, wantObjectives: test.wantObjectives}); diff != "" {
 				t.Errorf("Unexpected diff (+got/-want): %s", diff)
 			}