Add subsetting logic for epp

Author: rlakhtakia

cmd/epp/runner/runner.go

@@ -44,6 +44,7 @@ import (
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/saturationdetector"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/framework"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/framework/plugins/filter"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/framework/plugins/multi/prefix"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/framework/plugins/picker"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/framework/plugins/profile"
@@ -270,11 +271,13 @@ func (r *Runner) initializeScheduler(datastore datastore.Datastore) (*scheduling
 	if schedulerV2 {
 		queueScorerWeight := envutil.GetEnvInt("QUEUE_SCORE_WEIGHT", scorer.DefaultQueueScorerWeight, setupLog)
 		kvCacheScorerWeight := envutil.GetEnvInt("KV_CACHE_SCORE_WEIGHT", scorer.DefaultKVCacheScorerWeight, setupLog)
+		endpointSubsetFilter := filter.NewSubsetFilter()
 
 		schedulerProfile := framework.NewSchedulerProfile().
 			WithScorers(framework.NewWeightedScorer(scorer.NewQueueScorer(), queueScorerWeight),
 				framework.NewWeightedScorer(scorer.NewKVCacheScorer(), kvCacheScorerWeight)).
-			WithPicker(picker.NewMaxScorePicker())
+			WithPicker(picker.NewMaxScorePicker()).
+			WithFilters(endpointSubsetFilter)
 
 		if prefixCacheScheduling {
 			prefixScorerWeight := envutil.GetEnvInt("PREFIX_CACHE_SCORE_WEIGHT", prefix.DefaultScorerWeight, setupLog)

pkg/epp/handlers/server.go

@@ -111,8 +111,9 @@ type RequestContext struct {
 }
 
 type Request struct {
-	Headers map[string]string
-	Body    map[string]interface{}
+	Headers  map[string]string
+	Body     map[string]interface{}
+	Metadata map[string]any
 }
 type Response struct {
 	Headers map[string]string
@@ -141,8 +142,9 @@ func (s *StreamingServer) Process(srv extProcPb.ExternalProcessor_ProcessServer)
 	reqCtx := &RequestContext{
 		RequestState: RequestReceived,
 		Request: &Request{
-			Headers: make(map[string]string),
-			Body:    make(map[string]interface{}),
+			Headers:  make(map[string]string),
+			Body:     make(map[string]interface{}),
+			Metadata: make(map[string]any),
 		},
 		Response: &Response{
 			Headers: make(map[string]string),
@@ -185,6 +187,8 @@ func (s *StreamingServer) Process(srv extProcPb.ExternalProcessor_ProcessServer)
 			return status.Errorf(codes.Unknown, "cannot receive stream request: %v", err)
 		}
 
+		reqCtx.Request.Metadata = requtil.ExtractMetadataValues(req, requtil.MetadataSubsetKey)
+
 		switch v := req.Request.(type) {
 		case *extProcPb.ProcessingRequest_RequestHeaders:
 			if requestId := requtil.ExtractHeaderValue(v, requtil.RequestIdHeaderKey); len(requestId) > 0 {

pkg/epp/requestcontrol/director.go

@@ -117,6 +117,8 @@ func (d *Director) HandleRequest(ctx context.Context, reqCtx *handlers.RequestCo
 		Prompt:      prompt,
 		Headers:     reqCtx.Request.Headers,
 	}
+	reqCtx.SchedulingRequest.SetMetadata(reqCtx.Request.Metadata)
+
 	logger = logger.WithValues(
 		"model", reqCtx.Model,
 		"resolvedTargetModel", reqCtx.ResolvedTargetModel,

pkg/epp/scheduling/framework/plugins/filter/filter_test.go

@@ -247,3 +247,108 @@ func TestLoRASoftAffinityDistribution(t *testing.T) {
 			actualAvailablePercent, availableLowerBound, availableUpperBound)
 	}
 }
+
+func TestSubsettingFilter(t *testing.T) {
+	var makeFilterMetadata = func(data []interface{}) map[string]any {
+		return map[string]any{
+			"envoy.lb.subset_hint": map[string]any{
+				"x-gateway-destination-endpoint-subset": data,
+			},
+		}
+	}
+
+	tests := []struct {
+		name     string
+		metadata map[string]any
+		filter   framework.Filter
+		input    []types.Pod
+		output   []types.Pod
+	}{
+		{
+			name:     "SubsetFilter, filter not present — return all pods",
+			filter:   &SubsetFilter{},
+			metadata: map[string]any{},
+			input: []types.Pod{
+				&types.PodMetrics{
+					Pod: &backend.Pod{Address: "10.0.0.1"},
+				},
+				&types.PodMetrics{
+					Pod: &backend.Pod{Address: "10.0.0.2"},
+				},
+			},
+			output: []types.Pod{
+				&types.PodMetrics{
+					Pod: &backend.Pod{Address: "10.0.0.1"},
+				},
+				&types.PodMetrics{
+					Pod: &backend.Pod{Address: "10.0.0.2"},
+				},
+			},
+		},
+		{
+			name:     "SubsetFilter, subset with one matching pod",
+			metadata: makeFilterMetadata([]interface{}{"10.0.0.1"}),
+			filter:   &SubsetFilter{},
+			input: []types.Pod{
+				&types.PodMetrics{
+					Pod: &backend.Pod{Address: "10.0.0.1"},
+				},
+				&types.PodMetrics{
+					Pod: &backend.Pod{Address: "10.0.0.2"},
+				},
+			},
+			output: []types.Pod{
+				&types.PodMetrics{
+					Pod: &backend.Pod{Address: "10.0.0.1"},
+				},
+			},
+		},
+		{
+			name:     "SubsetFilter, subset with multiple matching pods",
+			metadata: makeFilterMetadata([]interface{}{"10.0.0.1", "10.0.0.2", "10.0.0.3"}),
+			filter:   &SubsetFilter{},
+			input: []types.Pod{
+				&types.PodMetrics{
+					Pod: &backend.Pod{Address: "10.0.0.1"},
+				},
+				&types.PodMetrics{
+					Pod: &backend.Pod{Address: "10.0.0.2"},
+				},
+			},
+			output: []types.Pod{
+				&types.PodMetrics{
+					Pod: &backend.Pod{Address: "10.0.0.1"},
+				},
+				&types.PodMetrics{
+					Pod: &backend.Pod{Address: "10.0.0.2"},
+				},
+			},
+		},
+		{
+			name:     "SubsetFilter, subset with no matching pods",
+			metadata: makeFilterMetadata([]interface{}{"10.0.0.3"}),
+			filter:   &SubsetFilter{},
+			input: []types.Pod{
+				&types.PodMetrics{
+					Pod: &backend.Pod{Address: "10.0.0.1"},
+				},
+				&types.PodMetrics{
+					Pod: &backend.Pod{Address: "10.0.0.2"},
+				},
+			},
+			output: []types.Pod{},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			req := &types.LLMRequest{}
+			req.SetMetadata(test.metadata)
+			got := test.filter.Filter(context.Background(), req, types.NewCycleState(), test.input)
+
+			if diff := cmp.Diff(test.output, got); diff != "" {
+				t.Errorf("Unexpected output (-want +got): %v", diff)
+			}
+		})
+	}
+}

pkg/epp/scheduling/framework/plugins/filter/subsetting_filter.go

@@ -0,0 +1,78 @@
+/*
+Copyright 2025 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package filter
+
+import (
+	"context"
+	"strings"
+
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/framework"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
+)
+
+const (
+	subsetHintNamespace = "envoy.lb.subset_hint"
+	subsetHintKey       = "x-gateway-destination-endpoint-subset"
+)
+
+// compile-time type assertion
+var _ framework.Filter = &SubsetFilter{}
+
+// NewSubsetFilter initializes a new SubsetFilter.
+func NewSubsetFilter() *SubsetFilter {
+	return &SubsetFilter{}
+}
+
+// SubsetFilter filters Pods based on the subset hint provided by the proxy via filterMetadata.
+type SubsetFilter struct{}
+
+// Name returns the name of the filter.
+func (f *SubsetFilter) Name() string {
+	return "subset-hint"
+}
+
+// Filter filters out pods that are not in the subset provided in filterMetadata.
+func (f *SubsetFilter) Filter(_ context.Context, request *types.LLMRequest, _ *types.CycleState, pods []types.Pod) []types.Pod {
+	// Check if envoy.lb.subset_hint is present in the metadata map
+	subsetMap, found := request.GetMetadata()[subsetHintNamespace].(map[string]any)
+	if !found {
+		return pods
+	}
+
+	// Check if endpoint key is present in the subset map and ensure there is at least one value
+	endpointSubsetList, found := subsetMap[subsetHintKey].([]interface{})
+	if !found || len(endpointSubsetList) == 0 {
+		return pods
+	}
+
+	// Create a map of endpoint addys for easy lookup
+	endpoints := make(map[string]bool)
+	for _, endpoint := range endpointSubsetList {
+		epStr := strings.Split(endpoint.(string), ":")[0]
+		endpoints[epStr] = true
+	}
+
+	// Filter based on address
+	filteredPods := []types.Pod{}
+	for _, pod := range pods {
+		if _, found := endpoints[pod.GetPod().Address]; found {
+			filteredPods = append(filteredPods, pod)
+		}
+	}
+
+	return filteredPods
+}

pkg/epp/scheduling/scheduler.go

@@ -38,6 +38,7 @@ func NewScheduler(datastore Datastore) *Scheduler {
 	// When the scheduler is initialized with NewScheduler function, thw below config will be used as default.
 	// it's possible to call NewSchedulerWithConfig to pass a different scheduler config.
 	// For build time plugins changes, it's recommended to call in main.go to NewSchedulerWithConfig.
+	endpointSubsetFilter := filter.NewSubsetFilter()
 	loraAffinityFilter := filter.NewLoraAffinityFilter()
 	leastQueueFilter := filter.NewLeastQueueFilter()
 	leastKvCacheFilter := filter.NewLeastKVCacheFilter()
@@ -65,7 +66,7 @@ func NewScheduler(datastore Datastore) *Scheduler {
 	}
 
 	defaultProfile := framework.NewSchedulerProfile().
-		WithFilters(lowLatencyFilter).
+		WithFilters(endpointSubsetFilter, lowLatencyFilter).
 		WithPicker(&picker.RandomPicker{})
 
 	profilePicker := profilepicker.NewSingleProfileHandler()

pkg/epp/scheduling/types/types.go

@@ -33,12 +33,23 @@ type LLMRequest struct {
 	Prompt string
 	// Headers is a map of the request headers.
 	Headers map[string]string
+
+	// filterMetadata is a map of metadata in the request
+	metadata map[string]any
 }
 
 func (r *LLMRequest) String() string {
 	return fmt.Sprintf("RequestID: %s, TargetModel: %s, PromptLength: %d, Headers: %v", r.RequestId, r.TargetModel, len(r.Prompt), r.Headers)
 }
 
+func (r *LLMRequest) GetMetadata() map[string]any {
+	return r.metadata
+}
+
+func (r *LLMRequest) SetMetadata(metadata map[string]any) {
+	r.metadata = metadata
+}
+
 type Pod interface {
 	GetPod() *backend.Pod
 	GetMetrics() *backendmetrics.MetricsState

pkg/epp/util/request/metadata.go

@@ -0,0 +1,38 @@
+/*
+Copyright 2025 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package request
+
+import (
+	extProcPb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
+)
+
+const (
+	MetadataSubsetKey = "envoy.lb.subset_hint"
+)
+
+func ExtractMetadataValues(req *extProcPb.ProcessingRequest, metadataKey string) map[string]any {
+	metadata := make(map[string]any)
+	if req != nil && req.MetadataContext != nil && req.MetadataContext.FilterMetadata != nil {
+		for key, val := range req.MetadataContext.FilterMetadata {
+			if key == metadataKey {
+				// envoy.lb.subset_hint key is always a map[string]any
+				metadata[key] = val.AsMap()
+			}
+		}
+	}
+	return metadata
+}