Address comments

Author: liu-cong

cmd/epp/main.go

@@ -121,6 +121,7 @@ func loadPrefixCacheConfig() prefix.Config {
 	baseLogger := log.Log.WithName("env-config")
 
 	return prefix.Config{
+		Weight:                 envutil.GetEnvInt("PREFIX_CACHE_WEIGHT", prefix.DefaultScorerWeight, baseLogger),
 		HashBlockSize:          envutil.GetEnvInt("PREFIX_CACHE_HASH_BLOCK_SIZE", prefix.DefaultHashBlockSize, baseLogger),
 		MaxPrefixBlocksToMatch: envutil.GetEnvInt("PREFIX_CACHE_MAX_PREFIX_BLOCKS", prefix.DefaultMaxPrefixBlocks, baseLogger),
 		LRUIndexerCapacity:     envutil.GetEnvInt("PREFIX_CACHE_LRU_CAPACITY", prefix.DefaultLRUIndexerCapacity, baseLogger),

pkg/epp/metrics/metrics.go

@@ -18,7 +18,6 @@ package metrics
 
 import (
 	"context"
-	"runtime/debug"
 	"sync"
 	"time"
 
@@ -210,17 +209,6 @@ var (
 		[]string{"plugin_type", "plugin_name"},
 	)
 
-	// Info Metrics
-	InferenceExtensionInfo = compbasemetrics.NewGaugeVec(
-		&compbasemetrics.GaugeOpts{
-			Subsystem:      InferenceExtension,
-			Name:           "info",
-			Help:           "General information of the current build of Inference Extension.",
-			StabilityLevel: compbasemetrics.ALPHA,
-		},
-		[]string{"commit"},
-	)
-
 	// Prefix indexer Metrics
 	PrefixCacheSize = compbasemetrics.NewGaugeVec(
 		&compbasemetrics.GaugeOpts{
@@ -254,6 +242,17 @@ var (
 		},
 		[]string{},
 	)
+
+	// Info Metrics
+	InferenceExtensionInfo = compbasemetrics.NewGaugeVec(
+		&compbasemetrics.GaugeOpts{
+			Subsystem:      InferenceExtension,
+			Name:           "info",
+			Help:           "General information of the current build of Inference Extension.",
+			StabilityLevel: compbasemetrics.ALPHA,
+		},
+		[]string{"commit"},
+	)
 )
 
 var registerMetrics sync.Once
@@ -414,21 +413,3 @@ func RecordInferenceExtensionInfo() {
 		InferenceExtensionInfo.WithLabelValues(CommitSHA).Set(1)
 	}
 }
-
-func init() {
-	info, ok := debug.ReadBuildInfo()
-	if !ok {
-		return
-	}
-
-	var Commit = func(i *debug.BuildInfo) string {
-		for _, setting := range i.Settings {
-			if setting.Key == "vcs.revision" {
-				return setting.Value
-			}
-		}
-		return ""
-	}(info)
-
-	CommitSHA = Commit
-}

pkg/epp/scheduling/config.go

@@ -16,7 +16,14 @@ limitations under the License.
 
 package scheduling
 
-import "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/plugins"
+import (
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/plugins"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/plugins/filter"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/plugins/picker"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/plugins/prefix"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/plugins/scorer"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
+)
 
 // NewSchedulerConfig creates a new SchedulerConfig object with the given plugins.
 func NewSchedulerConfig(preSchedulePlugins []plugins.PreSchedule, filters []plugins.Filter, scorers map[plugins.Scorer]int,
@@ -52,3 +59,58 @@ var defaultConfig = &SchedulerConfig{
 	picker:              defPlugin,
 	postSchedulePlugins: []plugins.PostSchedule{},
 }
+
+func CreateConfig(opts ...ConfigOption) *SchedulerConfig {
+	config := &SchedulerConfig{
+		preSchedulePlugins:  []plugins.PreSchedule{},
+		postSchedulePlugins: []plugins.PostSchedule{},
+		scorers:             map[plugins.Scorer]int{},
+		filters:             []plugins.Filter{&sheddableRequestFilterV2{}},
+		picker:              &picker.MaxScorePicker{},
+	}
+	for _, opt := range opts {
+		opt(config)
+	}
+	return config
+}
+
+type ConfigOption func(*SchedulerConfig)
+
+func WithPrefixPlugin(prefixConfig prefix.Config) ConfigOption {
+	return func(cfg *SchedulerConfig) {
+		prefixPlugin := prefix.New(prefixConfig)
+		cfg.preSchedulePlugins = append(cfg.preSchedulePlugins, prefixPlugin)
+		cfg.postSchedulePlugins = append(cfg.postSchedulePlugins, prefixPlugin)
+		cfg.scorers[prefixPlugin] = prefixConfig.Weight
+	}
+}
+
+func WithQueuePlugin(queueConfig scorer.QueueScorerConfig) ConfigOption {
+	return func(cfg *SchedulerConfig) {
+		queuePlugin := &scorer.QueueScorer{}
+		cfg.scorers[queuePlugin] = queueConfig.Weight
+	}
+}
+
+func WithKVCachePlugin(kvCacheConfig scorer.KVCacheScorerConfig) ConfigOption {
+	return func(cfg *SchedulerConfig) {
+		kvCachePlugin := &scorer.KVCacheScorer{}
+		cfg.scorers[kvCachePlugin] = kvCacheConfig.Weight
+	}
+}
+
+type sheddableRequestFilterV2 struct{}
+
+func (p *sheddableRequestFilterV2) Name() string {
+	return "sheddableRequestFilterV2"
+}
+
+func (p *sheddableRequestFilterV2) Filter(ctx *types.SchedulingContext, pods []types.Pod) []types.Pod {
+	if ctx.Req.Critical {
+		// Allow all pods to pass through if the request is critical, even if all pods reach their capacity.
+		return pods
+	}
+
+	// Only allow pods that have enough capacity to handle the request.
+	return filter.HasCapacityFilter.Filter(ctx, pods)
+}

pkg/epp/scheduling/config_v2.go

@@ -256,14 +256,6 @@ var HasCapacityFilter = &baseFilter{
 	filter: toFilterFunc(queueThresholdPredicate(config.Conf.QueueThresholdCritical).and(kvCacheThresholdPredicate(config.Conf.KVCacheThreshold))),
 }
 
-// NoopFilter is a filter that does not filter out any pods.
-var NoopFilter = &baseFilter{
-	name: "noop",
-	filter: toFilterFunc(func(req *types.LLMRequest, pod types.Pod) bool {
-		return true
-	}),
-}
-
 // podPredicate is a filter function to check whether a pod is desired.
 type podPredicate func(req *types.LLMRequest, pod types.Pod) bool
 

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

@@ -24,14 +24,13 @@ import (
 
 	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/metrics"
-	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
 )
 
 func newIndexer(maxCacheSize int) *indexer {
 	t := &indexer{
 		maxCacheSize: maxCacheSize,
-		table:        make(map[types.BlockHash]map[types.ServerID]*node),
+		table:        make(map[BlockHash]map[ServerID]*node),
 		list:         newLinkedList(),
 	}
 	go t.ReportCacheSize(time.Second)
@@ -43,15 +42,15 @@ func newIndexer(maxCacheSize int) *indexer {
 type indexer struct {
 	mu           sync.RWMutex
 	maxCacheSize int
-	table        map[types.BlockHash]map[types.ServerID]*node // from any prefix cache to the cache entry to find the server
-	list         *linkedList                                  // LRU list to keep track of the order of entries
+	table        map[BlockHash]map[ServerID]*node // from any prefix cache to the cache entry to find the server
+	list         *linkedList                      // LRU list to keep track of the order of entries
 }
 
 // Get returns the set of servers that have the given prefix hash cached.
-func (i *indexer) Get(hash types.BlockHash) map[types.ServerID]bool {
+func (i *indexer) Get(hash BlockHash) map[ServerID]bool {
 	i.mu.RLock()
 	defer i.mu.RUnlock()
-	res := map[types.ServerID]bool{}
+	res := map[ServerID]bool{}
 	for server := range i.table[hash] {
 		res[server] = true
 	}
@@ -61,15 +60,15 @@ func (i *indexer) Get(hash types.BlockHash) map[types.ServerID]bool {
 // Add adds a list of prefix hashes of a single request to the server the request was sent to.
 // The intuition is that this server is likely to have the prefix cached, so next time a request
 // sharing the longest prefix should be sent to the same server to take advantage of the cache hit.
-func (i *indexer) Add(hashes []types.BlockHash, server types.ServerID) {
+func (i *indexer) Add(hashes []BlockHash, server ServerID) {
 	i.mu.Lock()
 	defer i.mu.Unlock()
 	for _, hash := range hashes {
 		i.add(hash, server)
 	}
 }
 
-func (i *indexer) check(hash types.BlockHash, server types.ServerID) (*node, bool) {
+func (i *indexer) check(hash BlockHash, server ServerID) (*node, bool) {
 	servers, ok := i.table[hash]
 	if !ok {
 		return nil, false
@@ -78,7 +77,7 @@ func (i *indexer) check(hash types.BlockHash, server types.ServerID) (*node, boo
 	return n, ok
 }
 
-func (i *indexer) add(hash types.BlockHash, server types.ServerID) {
+func (i *indexer) add(hash BlockHash, server ServerID) {
 	node, exists := i.check(hash, server)
 	if exists {
 		i.list.moveToTail(node)
@@ -87,7 +86,7 @@ func (i *indexer) add(hash types.BlockHash, server types.ServerID) {
 	}
 }
 
-func (i *indexer) create(hash types.BlockHash, server types.ServerID) {
+func (i *indexer) create(hash BlockHash, server ServerID) {
 	n := &node{
 		hash:   hash,
 		server: server,
@@ -99,7 +98,7 @@ func (i *indexer) create(hash types.BlockHash, server types.ServerID) {
 	}
 
 	if _, ok := i.table[hash]; !ok {
-		i.table[hash] = make(map[types.ServerID]*node)
+		i.table[hash] = make(map[ServerID]*node)
 	}
 	i.table[hash][server] = n
 	i.list.add(n)

pkg/epp/scheduling/plugins/prefix/indexer.go

@@ -19,28 +19,27 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
 )
 
 func TestIndexer_AddAndGet(t *testing.T) {
 	cache := newIndexer(2)
 
-	hash1 := types.BlockHash(1)
-	server := types.ServerID{Namespace: "default", Name: "server1"}
+	hash1 := BlockHash(1)
+	server := ServerID{Namespace: "default", Name: "server1"}
 
 	// Add an entry to the cache
-	cache.Add([]types.BlockHash{hash1}, server)
+	cache.Add([]BlockHash{hash1}, server)
 
 	// Retrieve the entry
 	assert.Equal(t, 1, cache.list.size, "Cache size should be 1 after adding an entry")
 	servers := cache.Get(hash1)
 	assert.Contains(t, servers, server, "Cache should contain the added server")
 
 	// Add another entry to the cache, the cache size should be incremented to 2.
-	cache.Add([]types.BlockHash{types.BlockHash(2)}, server)
+	cache.Add([]BlockHash{BlockHash(2)}, server)
 	assert.Equal(t, 2, cache.list.size, "Cache size should  be 2 after adding an entry")
 
 	// Add another entry to the cache, which should evict the first one due to max size.
-	cache.Add([]types.BlockHash{types.BlockHash(3)}, server)
+	cache.Add([]BlockHash{BlockHash(3)}, server)
 	assert.Equal(t, 2, cache.list.size, "Cache size should still be 2 after adding an entry")
 }

pkg/epp/scheduling/plugins/prefix/indexer_test.go

@@ -16,10 +16,6 @@ limitations under the License.
 
 package prefix
 
-import (
-	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
-)
-
 type linkedList struct {
 	dummyHead *node // The head of the linked list (dummy node).
 	tail      *node // The tail of the linked list.
@@ -40,8 +36,8 @@ func newLinkedList() *linkedList {
 type node struct {
 	prev   *node
 	next   *node
-	server types.ServerID
-	hash   types.BlockHash
+	server ServerID
+	hash   BlockHash
 }
 
 // add adds a node to the end of the linked list.