simplify config

Signed-off-by: Nir Rozenbaum <nirro@il.ibm.com>

Author: nirrozenbaum

cmd/epp/main.go

@@ -178,11 +178,10 @@ func run() error {
 	// Setup runner.
 	ctx := ctrl.SetupSignalHandler()
 
-	schedCfg := config.NewConfig(setupLog)
-	schedCfg.LoadConfig()
+	schedulerConfig := config.LoadConfig(setupLog)
 
 	datastore := datastore.NewDatastore(ctx, pmf)
-	scheduler, err := pd.NewScheduler(ctx, schedCfg, datastore)
+	scheduler, err := pd.NewScheduler(ctx, schedulerConfig, datastore)
 	if err != nil {
 		setupLog.Error(err, "Failed to create PD scheduler")
 		return err

pkg/config/config.go

@@ -3,8 +3,6 @@
 package config
 
 import (
-	"math"
-
 	"github.com/go-logr/logr"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/env"
 )
@@ -16,6 +14,10 @@ const (
 	//  - "PREFILL_ENABLE_" + pluginName  Enables the named plugin for prefill processing
 	//  - "PREFILL_" + pluginName + "_WEIGHT"  The weight for a scorer in prefill processing
 
+	prefillPrefix = "PREFILL_"
+	enablePrefix  = "ENABLE_"
+	weightSuffix  = "_WEIGHT"
+
 	// KVCacheScorerName name of the kv-cache scorer in configuration
 	KVCacheScorerName = "KVCACHE_AWARE_SCORER"
 	// LoadAwareScorerName name of the load aware scorer in configuration
@@ -25,10 +27,6 @@ const (
 	// SessionAwareScorerName name of the session aware scorer in configuration
 	SessionAwareScorerName = "SESSION_AWARE_SCORER"
 
-	prefillPrefix = "PREFILL_"
-	enablePrefix  = "ENABLE_"
-	weightSuffix  = "_WEIGHT"
-
 	// Plugins from Upstream
 
 	// GIELeastKVCacheFilterName name of the GIE least kv-cache filter in configuration
@@ -58,47 +56,34 @@ const (
 
 // Config contains scheduler configuration, currently configuration is loaded from environment variables
 type Config struct {
-	logger                  logr.Logger
 	DecodeSchedulerPlugins  map[string]int
 	PrefillSchedulerPlugins map[string]int
-
-	PDEnabled       bool
-	PDThreshold     int
-	PrefixBlockSize int
+	PDEnabled               bool
+	PDThreshold             int
+	PrefixBlockSize         int
 }
 
-// NewConfig creates a new instance if Config
-func NewConfig(logger logr.Logger) *Config {
-	return &Config{
-		logger:                  logger,
-		DecodeSchedulerPlugins:  map[string]int{},
-		PrefillSchedulerPlugins: map[string]int{},
-		PDEnabled:               false,
-		PDThreshold:             math.MaxInt,
-		PrefixBlockSize:         prefixScorerBlockSizeDefault,
-	}
-}
-
-// LoadConfig loads configuration from environment variables
-func (c *Config) LoadConfig() {
-	c.loadPluginInfo(c.DecodeSchedulerPlugins, false,
-		KVCacheScorerName, LoadAwareScorerName, PrefixScorerName, SessionAwareScorerName,
-		GIELeastKVCacheFilterName, GIELeastQueueFilterName, GIELoraAffinityFilterName,
-		GIELowQueueFilterName, GIESheddableCapacityFilterName,
-		GIEKVCacheUtilizationScorerName, GIEQueueScorerName, GIEPrefixScorerName)
-
-	c.loadPluginInfo(c.PrefillSchedulerPlugins, true,
+// LoadConfig loads configuration from environment variables and returns a new instance of Config
+func LoadConfig(logger logr.Logger) *Config {
+	pluginNames := []string{
 		KVCacheScorerName, LoadAwareScorerName, PrefixScorerName, SessionAwareScorerName,
 		GIELeastKVCacheFilterName, GIELeastQueueFilterName, GIELoraAffinityFilterName,
 		GIELowQueueFilterName, GIESheddableCapacityFilterName,
-		GIEKVCacheUtilizationScorerName, GIEQueueScorerName, GIEPrefixScorerName)
+		GIEKVCacheUtilizationScorerName, GIEQueueScorerName, GIEPrefixScorerName,
+	}
 
-	c.PDEnabled = env.GetEnvString(pdEnabledEnvKey, "false", c.logger) == "true"
-	c.PDThreshold = env.GetEnvInt(pdPromptLenThresholdEnvKey, pdPromptLenThresholdDefault, c.logger)
-	c.PrefixBlockSize = env.GetEnvInt(prefixScorerBlockSizeEnvKey, prefixScorerBlockSizeDefault, c.logger)
+	return &Config{
+		DecodeSchedulerPlugins:  loadPluginInfo(logger, false, pluginNames...),
+		PrefillSchedulerPlugins: loadPluginInfo(logger, true, pluginNames...),
+		PDEnabled:               env.GetEnvString(pdEnabledEnvKey, "false", logger) == "true",
+		PDThreshold:             env.GetEnvInt(pdPromptLenThresholdEnvKey, pdPromptLenThresholdDefault, logger),
+		PrefixBlockSize:         env.GetEnvInt(prefixScorerBlockSizeEnvKey, prefixScorerBlockSizeDefault, logger),
+	}
 }
 
-func (c *Config) loadPluginInfo(plugins map[string]int, prefill bool, pluginNames ...string) {
+func loadPluginInfo(logger logr.Logger, prefill bool, pluginNames ...string) map[string]int {
+	result := map[string]int{}
+
 	for _, pluginName := range pluginNames {
 		var enablementKey string
 		var weightKey string
@@ -110,13 +95,15 @@ func (c *Config) loadPluginInfo(plugins map[string]int, prefill bool, pluginName
 			weightKey = pluginName + weightSuffix
 		}
 
-		if env.GetEnvString(enablementKey, "false", c.logger) != "true" {
-			c.logger.Info("Skipping plugin creation as it is not enabled", "name", pluginName)
+		if env.GetEnvString(enablementKey, "false", logger) != "true" {
+			logger.Info("Skipping plugin creation as it is not enabled", "name", pluginName)
 		} else {
-			weight := env.GetEnvInt(weightKey, 1, c.logger)
+			weight := env.GetEnvInt(weightKey, 1, logger)
 
-			plugins[pluginName] = weight
-			c.logger.Info("Initialized plugin", "plugin", pluginName, "weight", weight)
+			result[pluginName] = weight
+			logger.Info("Initialized plugin", "plugin", pluginName, "weight", weight)
 		}
 	}
+
+	return result
 }

pkg/scheduling/pd/scheduler.go

@@ -57,26 +57,26 @@ type Datastore interface {
 
 // NewScheduler returns a new disaggregated Prefill/Decode filter, using the
 // provided configuration.
-func NewScheduler(ctx context.Context, schedCfg *config.Config, ds Datastore) (*Scheduler, error) {
+func NewScheduler(ctx context.Context, schedulerConfig *config.Config, ds Datastore) (*Scheduler, error) {
 	prefixConfig := scorer.DefaultPrefixStoreConfig()
-	prefixConfig.BlockSize = schedCfg.PrefixBlockSize
+	prefixConfig.BlockSize = schedulerConfig.PrefixBlockSize
 
 	scheduler := &Scheduler{
-		threshold:    schedCfg.PDThreshold,
-		pdEnabled:    schedCfg.PDEnabled,
+		threshold:    schedulerConfig.PDThreshold,
+		pdEnabled:    schedulerConfig.PDEnabled,
 		store:        ds,
 		prefixScorer: scorer.NewPrefixAwareScorer(ctx, prefixConfig),
 	}
 
 	scheduler.prefill = scheduling.NewSchedulerWithConfig(
 		ds,
-		scheduler.generateSchedulerConfig(ctx, schedCfg.PrefillSchedulerPlugins,
+		scheduler.generateSchedulerConfig(ctx, schedulerConfig.PrefillSchedulerPlugins,
 			&filter.PrefillFilter{}),
 	)
 
 	scheduler.decode = scheduling.NewSchedulerWithConfig(
 		ds,
-		scheduler.generateSchedulerConfig(ctx, schedCfg.DecodeSchedulerPlugins,
+		scheduler.generateSchedulerConfig(ctx, schedulerConfig.DecodeSchedulerPlugins,
 			&filter.DecodeFilter{}),
 	)
 

pkg/scheduling/pd/scheduler_test.go

@@ -136,13 +136,17 @@ func TestPDSchedule(t *testing.T) {
 	logger := testr.New(t)
 	ctx = log.IntoContext(ctx, logger)
 
-	schedCfg := config.NewConfig(logger)
-	schedCfg.PDEnabled = true
-	schedCfg.PDThreshold = 5
+	schedulderConfig := &config.Config{
+		DecodeSchedulerPlugins:  map[string]int{},
+		PrefillSchedulerPlugins: map[string]int{},
+		PDEnabled:               true,
+		PDThreshold:             5,
+		PrefixBlockSize:         256,
+	}
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			scheduler, _ := pd.NewScheduler(ctx, schedCfg, &fakeDataStore{pods: test.input})
+			scheduler, _ := pd.NewScheduler(ctx, schedulderConfig, &fakeDataStore{pods: test.input})
 			got, err := scheduler.Schedule(ctx, test.req)
 
 			if test.err != (err != nil) {