Replace scheduler v2 with config v2

Author: liu-cong

cmd/epp/main.go

@@ -111,19 +111,26 @@ var (
 	setupLog = ctrl.Log.WithName("setup")
 
 	// Environment variables
-	schedulerV2       = envutil.GetEnvString("EXPERIMENTAL_USE_SCHEDULE_V2", "false", setupLog)
-	prefixCacheConfig = loadPrefixCacheConfig()
+	schedulerV2 = envutil.GetEnvString("EXPERIMENTAL_USE_SCHEDULER_V2", "false", setupLog)
 )
 
 func loadPrefixCacheConfig() prefix.Config {
-	// logger := zap.New(zap.RawZapOpts(uberzap.AddCaller()))
-	// log.SetLogger(logger)
 	baseLogger := log.Log.WithName("env-config")
 
 	return prefix.Config{
-		HashBlockSize:          envutil.GetEnvInt("PREFIX_CACHE_HASH_BLOCK_SIZE", prefix.DefaultCacheBlockSize, 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_MAX_CACHE_SIZE_MB", prefix.DefaultLRUIndexerCapacity, baseLogger),
+		LRUIndexerCapacity:     envutil.GetEnvInt("PREFIX_CACHE_LRU_CAPACITY", prefix.DefaultLRUIndexerCapacity, baseLogger),
+	}
+}
+
+func loadSchedulingScorerWeights() scheduling.ScorerWeights {
+	baseLogger := log.Log.WithName("env-config")
+
+	return scheduling.ScorerWeights{
+		Prefix:  envutil.GetEnvInt("PREFIX_CACHE_SCORE_WEIGHT", 3, baseLogger),
+		Queue:   envutil.GetEnvInt("QUEUE_SCORE_WEIGHT", 2, baseLogger),
+		KVCache: envutil.GetEnvInt("KV_CACHE_SCORE_WEIGHT", 1, baseLogger),
 	}
 }
 
@@ -191,8 +198,9 @@ func run() error {
 
 	scheduler := scheduling.NewScheduler(datastore)
 	if schedulerV2 == "true" {
-		setupLog.Info("Creating scheduler with prefixCache plugin", "prefix cache config", prefixCacheConfig)
-		scheduler = scheduling.NewSchedulerV2(datastore, prefixCacheConfig)
+		schedConfig := scheduling.CreateConfig(loadSchedulingScorerWeights(), loadPrefixCacheConfig())
+		setupLog.Info("Creating scheduler", "config", *schedConfig)
+		scheduler = scheduling.NewSchedulerWithConfig(datastore, schedConfig)
 	}
 	serverRunner := &runserver.ExtProcServerRunner{
 		GrpcPort:                                 *grpcPort,

pkg/epp/scheduling/scheduler_v2.go renamed to pkg/epp/scheduling/config_v2.go

@@ -26,27 +26,32 @@ import (
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
 )
 
-func NewSchedulerV2(datastore Datastore, prefixConfig prefix.Config) *Scheduler {
+func CreateConfig(weights ScorerWeights, prefixConfig prefix.Config) *SchedulerConfig {
 	prefixPlugin := prefix.New(prefixConfig)
 	queuePlugin := &scorer.QueueScorer{}
 	kvCachePlugin := &scorer.KVCacheScorer{}
-	configV2 := &SchedulerConfig{
+	config := &SchedulerConfig{
 		PreSchedulePlugins:  []plugins.PreSchedule{prefixPlugin},
 		PostSchedulePlugins: []plugins.PostSchedule{prefixPlugin},
 		Scorers: map[plugins.Scorer]int{
-			prefixPlugin:  3,
-			queuePlugin:   1,
-			kvCachePlugin: 1,
+			prefixPlugin:  weights.Prefix,
+			queuePlugin:   weights.Queue,
+			kvCachePlugin: weights.KVCache,
 		},
 		Filters: []plugins.Filter{&sheddableRequestFilterV2{}},
 		Picker:  &picker.MaxScorePicker{},
 	}
-	return NewSchedulerWithConfig(datastore, configV2)
+	return config
 }
 
-type sheddableRequestFilterV2 struct {
+type ScorerWeights struct {
+	Prefix  int
+	Queue   int
+	KVCache int
 }
 
+type sheddableRequestFilterV2 struct{}
+
 func (p *sheddableRequestFilterV2) Name() string {
 	return "sheddableRequestFilterV2"
 }

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

@@ -33,7 +33,7 @@ const (
 	// We optimistically search more than one to give more candidates for the scheduler to choose.
 	DefaultNumServersToMatch = 2
 	// vLLM default token block size is 16, and a good guess of average characters per token is 4.
-	DefaultCacheBlockSize  = 64
+	DefaultHashBlockSize   = 64
 	DefaultMaxPrefixBlocks = 128
 	// Assume each request reaches DefaultMaxPrefixBlocks = 128, and each BlockHash is cached onto 2
 	// servers due to load balancing, then it requires 256 entries per request.
@@ -60,7 +60,7 @@ type Config struct {
 }
 
 var DefaultConfig = Config{
-	HashBlockSize:          DefaultCacheBlockSize,
+	HashBlockSize:          DefaultHashBlockSize,
 	MaxPrefixBlocksToMatch: DefaultMaxPrefixBlocks,
 	LRUIndexerCapacity:     DefaultLRUIndexerCapacity,
 }