Clean up

Author: liu-cong

cmd/epp/main.go

@@ -121,7 +121,6 @@ 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),
@@ -192,18 +191,15 @@ func run() error {
 
 	scheduler := scheduling.NewScheduler(datastore)
 	if schedulerV2 == "true" {
-		queueConfig := scorer.QueueScorerConfig{
-			Weight: envutil.GetEnvInt("QUEUE_SCORE_WEIGHT", scorer.DefaultQueueScorerWeight, setupLog),
-		}
-		kvCacheConfig := scorer.KVCacheScorerConfig{
-			Weight: envutil.GetEnvInt("KV_CACHE_SCORE_WEIGHT", scorer.DefaultKVCacheScorerWeight, setupLog),
-		}
+		queueScorerWeight := envutil.GetEnvInt("QUEUE_SCORE_WEIGHT", scorer.DefaultQueueScorerWeight, setupLog)
+		kvCacheScorerWeight := envutil.GetEnvInt("KV_CACHE_SCORE_WEIGHT", scorer.DefaultKVCacheScorerWeight, setupLog)
 		schedConfigOpts := []scheduling.ConfigOption{
-			scheduling.WithQueuePlugin(queueConfig),
-			scheduling.WithKVCachePlugin(kvCacheConfig),
+			scheduling.AddScorer(&scorer.QueueScorer{}, queueScorerWeight),
+			scheduling.AddScorer(&scorer.KVCacheScorer{}, kvCacheScorerWeight),
 		}
 		if prefixCacheScheduling == "true" {
-			schedConfigOpts = append(schedConfigOpts, scheduling.WithPrefixPlugin(loadPrefixCacheConfig()))
+			prefixScorerWeight := envutil.GetEnvInt("PREFIX_CACHE_WEIGHT", prefix.DefaultScorerWeight, setupLog)
+			schedConfigOpts = append(schedConfigOpts, scheduling.AddPrefixPlugin(loadPrefixCacheConfig(), prefixScorerWeight))
 		}
 		schedulerConfig := scheduling.CreateConfig(schedConfigOpts...)
 		scheduler = scheduling.NewSchedulerWithConfig(datastore, schedulerConfig)

pkg/epp/scheduling/config.go

@@ -20,7 +20,6 @@ import (
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/plugins"
 	"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"
 )
 
 // NewSchedulerConfig creates a new SchedulerConfig object with the given plugins.
@@ -74,25 +73,17 @@ func CreateConfig(opts ...ConfigOption) *SchedulerConfig {
 
 type ConfigOption func(*SchedulerConfig)
 
-func WithPrefixPlugin(prefixConfig prefix.Config) ConfigOption {
+func AddScorer(scorer plugins.Scorer, weight int) 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
+		cfg.scorers[scorer] = weight
 	}
 }
 
-func WithKVCachePlugin(kvCacheConfig scorer.KVCacheScorerConfig) ConfigOption {
+func AddPrefixPlugin(prefixConfig prefix.Config, weight int) ConfigOption {
 	return func(cfg *SchedulerConfig) {
-		kvCachePlugin := &scorer.KVCacheScorer{}
-		cfg.scorers[kvCachePlugin] = kvCacheConfig.Weight
+		prefixPlugin := prefix.New(prefixConfig)
+		cfg.preSchedulePlugins = append(cfg.preSchedulePlugins, prefixPlugin)
+		cfg.postSchedulePlugins = append(cfg.postSchedulePlugins, prefixPlugin)
+		cfg.scorers[prefixPlugin] = weight
 	}
 }

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

@@ -58,7 +58,6 @@ const (
 )
 
 type Config struct {
-	Weight int
 	// The input prompt is broken into sizes of HashBlockSize to calculate block hashes . Requests
 	// with length shorter than the block size will be ignored.
 	HashBlockSize int

pkg/epp/scheduling/plugins/scorer/kvcache.go

@@ -24,10 +24,6 @@ const (
 	DefaultKVCacheScorerWeight = 2
 )
 
-type KVCacheScorerConfig struct {
-	Weight int
-}
-
 type KVCacheScorer struct{}
 
 func (ss *KVCacheScorer) Name() string {

pkg/epp/scheduling/plugins/scorer/queue.go

@@ -26,10 +26,6 @@ const (
 	DefaultQueueScorerWeight = 1
 )
 
-type QueueScorerConfig struct {
-	Weight int
-}
-
 type QueueScorer struct{}
 
 func (q *QueueScorer) Name() string {