Add flow controller.

Author: LukeAVanDrie

cmd/epp/main.go

@@ -17,6 +17,7 @@ limitations under the License.
 package main
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"net"
@@ -41,8 +42,10 @@ import (
 	"sigs.k8s.io/gateway-api-inference-extension/internal/runnable"
 	backendmetrics "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend/metrics"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/datastore"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/flowcontroller"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/metrics"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/metrics/collectors"
+	"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/plugins"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/plugins/filter"
@@ -157,13 +160,20 @@ func run() error {
 	})
 	setupLog.Info("Flags processed", "flags", flags)
 
-	// Init runtime.
+	// --- Load Configurations from Environment Variables ---
+	// Note: Scheduler config is loaded via its package init currently. We may
+	// want to load it here explicitly:
+	fcConfig, flowControllerEnabled := flowcontroller.LoadConfigFromEnv()
+	sdConfig := saturationdetector.LoadConfigFromEnv()
+
+	// --- Get Kubernetes Config ---
 	cfg, err := ctrl.GetConfig()
 	if err != nil {
-		setupLog.Error(err, "Failed to get rest config")
+		setupLog.Error(err, "Failed to get Kubernetes rest config")
 		return err
 	}
 
+	// --- Setup Manager ---
 	poolNamespacedName := types.NamespacedName{
 		Name:      *poolName,
 		Namespace: *poolNamespace,
@@ -174,7 +184,7 @@ func run() error {
 		return err
 	}
 
-	// Set up mapper for metric scraping.
+	// --- Setup Datastore ---
 	mapping, err := backendmetrics.NewMetricMapping(
 		*totalQueuedRequestsMetric,
 		*kvCacheUsagePercentageMetric,
@@ -185,14 +195,12 @@ func run() error {
 		return err
 	}
 	verifyMetricMapping(*mapping, setupLog)
-
 	pmf := backendmetrics.NewPodMetricsFactory(&backendmetrics.PodMetricsClientImpl{MetricMapping: mapping}, *refreshMetricsInterval)
-	// Setup runner.
 	ctx := ctrl.SetupSignalHandler()
+	appDatastore := datastore.NewDatastore(ctx, pmf)
 
-	datastore := datastore.NewDatastore(ctx, pmf)
-
-	scheduler := scheduling.NewScheduler(datastore)
+	// --- Initialize EPP Components ---
+	appScheduler := scheduling.NewScheduler(appDatastore)
 	if schedulerV2 == "true" {
 		queueScorerWeight := envutil.GetEnvInt("QUEUE_SCORE_WEIGHT", scorer.DefaultQueueScorerWeight, setupLog)
 		kvCacheScorerWeight := envutil.GetEnvInt("KV_CACHE_SCORE_WEIGHT", scorer.DefaultKVCacheScorerWeight, setupLog)
@@ -213,41 +221,78 @@ func run() error {
 			[]plugins.PostSchedule{},
 			[]plugins.PostResponse{},
 			schedConfigOpts...)
-		scheduler = scheduling.NewSchedulerWithConfig(datastore, schedulerConfig)
+			appScheduler = scheduling.NewSchedulerWithConfig(appDatastore, schedulerConfig)
+	}
+
+	appSaturationDetector, err := saturationdetector.NewDetector(
+		*sdConfig,
+		appDatastore,
+		ctrl.Log.WithName("saturation-detector"),
+	)
+	if err != nil {
+		setupLog.Error(err, "Failed to create SaturationDetector")
+		return err
+	}
+
+	var appFlowController *flowcontroller.FlowController
+	if flowControllerEnabled {
+		appFlowController, err = flowcontroller.NewFlowController(
+			appSaturationDetector,
+			fcConfig,
+		)
+		if err != nil {
+			setupLog.Error(err, "Failed to create FlowController")
+			return err
+		}
+		setupLog.Info("FlowController enabled and initialized.")
+	} else {
+		setupLog.Info("FlowController is disabled via configuration.")
 	}
+
+	// --- Setup ExtProc Server Runner ---
 	serverRunner := &runserver.ExtProcServerRunner{
 		GrpcPort:                                 *grpcPort,
 		DestinationEndpointHintMetadataNamespace: *destinationEndpointHintMetadataNamespace,
 		DestinationEndpointHintKey:               *destinationEndpointHintKey,
 		PoolNamespacedName:                       poolNamespacedName,
-		Datastore:                                datastore,
+		Datastore:                                appDatastore,
 		SecureServing:                            *secureServing,
 		CertPath:                                 *certPath,
 		RefreshPrometheusMetricsInterval:         *refreshPrometheusMetricsInterval,
-		Scheduler:                                scheduler,
+		Scheduler:                                appScheduler,
+		FlowController:                           appFlowController, // Pass instance (can be nil)
+		SaturationDetector:                       appSaturationDetector,
+		FlowControllerEnabled:                    flowControllerEnabled,
 	}
 	if err := serverRunner.SetupWithManager(ctx, mgr); err != nil {
-		setupLog.Error(err, "Failed to setup ext-proc controllers")
+		setupLog.Error(err, "Failed to setup EPP controllers")
+		return err
+	}
+
+	// --- Add Runnables to Manager ---
+
+	// Register FlowController Run loop.
+	if err := registerFlowController(mgr, appFlowController); err != nil {
 		return err
 	}
 
 	// Register health server.
-	if err := registerHealthServer(mgr, ctrl.Log.WithName("health"), datastore, *grpcHealthPort); err != nil {
+	if err := registerHealthServer(mgr, ctrl.Log.WithName("health"), appDatastore, *grpcHealthPort); err != nil {
 		return err
 	}
 
 	// Register ext-proc server.
-	if err := mgr.Add(serverRunner.AsRunnable(ctrl.Log.WithName("ext-proc"))); err != nil {
-		setupLog.Error(err, "Failed to register ext-proc gRPC server")
+	if err := registerExtProcServer(mgr, serverRunner, ctrl.Log.WithName("ext-proc")); err != nil {
 		return err
 	}
 
 	// Register metrics handler.
-	if err := registerMetricsHandler(mgr, *metricsPort, cfg, datastore); err != nil {
+	if err := registerMetricsHandler(mgr, *metricsPort, cfg, appDatastore); err != nil {
 		return err
 	}
 
-	// Start the manager. This blocks until a signal is received.
+	// --- Start Manager ---
+	// This blocks until a signal is received.
 	setupLog.Info("Controller manager starting")
 	if err := mgr.Start(ctx); err != nil {
 		setupLog.Error(err, "Error starting controller manager")
@@ -275,6 +320,39 @@ func initLogging(opts *zap.Options) {
 	ctrl.SetLogger(logger)
 }
 
+// registerFlowController adds the FlowController Run loop as a Runnable to the
+// manager.
+func registerFlowController(mgr manager.Manager, fc *flowcontroller.FlowController) error {
+	if fc == nil {
+		setupLog.Info("FlowController is nil (disabled), skipping registration.")
+		return nil // Not an error if it's intentionally disabled
+	}
+	if err := mgr.Add(manager.RunnableFunc(func(runCtx context.Context) error {
+		fcLog := ctrl.Log.WithName("flowcontroller-runnable")
+		fcLog.Info("Starting FlowController Run loop")
+		// Run the FlowController; it handles context cancellation internally.
+		fc.Run(runCtx)
+		fcLog.Info("FlowController Run loop stopped")
+		return nil
+	})); err != nil {
+		setupLog.Error(err, "Failed to add FlowController runnable to manager")
+		return err
+	}
+	setupLog.Info("FlowController Run loop added to manager.")
+	return nil
+}
+
+// registerExtProcServer adds the ExtProcServerRunner as a Runnable to the
+// manager.
+func registerExtProcServer(mgr manager.Manager, runner *runserver.ExtProcServerRunner, logger logr.Logger) error {
+	if err := mgr.Add(runner.AsRunnable(logger)); err != nil {
+		setupLog.Error(err, "Failed to register ext-proc gRPC server runnable")
+		return err
+	}
+	setupLog.Info("ExtProc server runner added to manager.")
+	return nil
+}
+
 // registerHealthServer adds the Health gRPC server as a Runnable to the given manager.
 func registerHealthServer(mgr manager.Manager, logger logr.Logger, ds datastore.Datastore, port int) error {
 	srv := grpc.NewServer()
@@ -364,5 +442,4 @@ func verifyMetricMapping(mapping backendmetrics.MetricMapping, logger logr.Logge
 	if mapping.LoraRequestInfo == nil {
 		logger.Info("Not scraping metric: LoraRequestInfo")
 	}
-
 }

pkg/epp/flowcontroller/config.go

@@ -0,0 +1,181 @@
+/*
+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 flowcontroller
+
+import (
+	"fmt"
+	"time"
+
+	"sigs.k8s.io/controller-runtime/pkg/log"
+	v1alpha2 "sigs.k8s.io/gateway-api-inference-extension/api/v1alpha2"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/flowcontroller/types"
+	envutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/env"
+)
+
+// PolicyType strings
+const (
+	PolicyFCFS        = "FCFS"
+	PolicyRoundRobin  = "RoundRobin"
+	PolicyEvictOldest = "EvictOldest"
+	PolicyNone        = "None"
+)
+
+// Default byte capacities per criticality and other settings
+const (
+	DefaultCriticalCapacityBytes  uint64 = 1000 * 1024 * 1024 // 1000MB
+	DefaultStandardCapacityBytes  uint64 = 500 * 1024 * 1024  // 500MB
+	DefaultSheddableCapacityBytes uint64 = 500 * 1024 * 1024  // 500MB
+	DefaultQueueTTL                      = 10 * time.Second
+	DefaultExpiryCleanupInterval         = time.Second
+	DefaultFcEnabled                     = false // Default to disabled
+)
+
+// Environment variable names for FlowController configuration
+const (
+	EnvFcMaxBytesCritical      = "FC_MAX_BYTES_CRITICAL"
+	EnvFcMaxBytesStandard      = "FC_MAX_BYTES_STANDARD"
+	EnvFcMaxBytesSheddable     = "FC_MAX_BYTES_SHEDDABLE"
+	EnvFcMaxGlobalBytes        = "FC_MAX_GLOBAL_BYTES"
+	EnvFcDefaultQueueTTL       = "FC_DEFAULT_QUEUE_TTL"
+	EnvFcExpiryCleanupInterval = "FC_EXPIRY_CLEANUP_INTERVAL"
+	EnvFcFairnessPolicy        = "FC_FAIRNESS_POLICY"
+	EnvFcPreemptionStrategy    = "FC_PREEMPTION_STRATEGY"
+	EnvFcEnabled               = "FC_ENABLED"
+)
+
+// Supported policy values
+var (
+	SupportedInterModelFairnessPolicies = map[string]struct{}{
+		PolicyFCFS:       {},
+		PolicyRoundRobin: {},
+	}
+	SupportedPreemptionStrategies = map[string]struct{}{
+		PolicyEvictOldest: {},
+		PolicyNone:        {},
+	}
+)
+
+// FlowControllerConfig holds the configuration for the FlowController.
+type FlowControllerConfig struct {
+	// MaxBytesPerCriticality maps criticality levels to their maximum allowed
+	// byte size in the queue.
+	MaxBytesPerCriticality map[v1alpha2.Criticality]uint64
+	// MaxGlobalBytes defines the overall limit on the total byte size across
+	// all criticality bands. 0 means no global limit (beyond the sum of
+	// per-criticality limits).
+	MaxGlobalBytes uint64
+	// DefaultQueueTTL is the TTL applied to requests if not specified elsewhere
+	// (eventually this comes from InferenceModel). Defaults to 10s.
+	DefaultQueueTTL time.Duration
+	// ExpiryCleanupInterval is the interval for cleaning up expired requests.
+	// Defaults to 1s.
+	ExpiryCleanupInterval time.Duration
+	// InterModelFairnessPolicy defines the fairness algorithm ("FCFS",
+	// "RoundRobin") by string. This is used if CustomFairnessPolicies are not
+	// provided for a criticality.
+	// Defaults to "FCFS".
+	InterModelFairnessPolicy string
+	// PreemptionStrategy defines the preemption algorithm ("None", ...).
+	// Defaults to "None".
+	PreemptionStrategy string
+	// CustomFairnessPolicies allows injecting specific fairness policy instances
+	// for each criticality band. If a policy is provided here for a given
+	// criticality, it will be used instead of the one specified by
+	// InterModelFairnessPolicy string.
+	CustomFairnessPolicies map[v1alpha2.Criticality]types.FairnessPolicy
+	// TODO: Add CustomPreemptionStrategies map[v1alpha2.Criticality]types.PreemptionStrategy
+	// if similar custom injection is desired for preemption strategies.
+}
+
+// validateAndApplyDefaults validates the provided config and assigns default
+// values. Called by NewFlowController.
+func (cfg *FlowControllerConfig) validateAndApplyDefaults() error {
+	if cfg.MaxBytesPerCriticality == nil {
+		cfg.MaxBytesPerCriticality = make(map[v1alpha2.Criticality]uint64)
+	}
+
+	// Apply default capacities only if not specified.
+	if _, ok := cfg.MaxBytesPerCriticality[v1alpha2.Critical]; !ok {
+		cfg.MaxBytesPerCriticality[v1alpha2.Critical] = DefaultCriticalCapacityBytes
+	}
+	if _, ok := cfg.MaxBytesPerCriticality[v1alpha2.Standard]; !ok {
+		cfg.MaxBytesPerCriticality[v1alpha2.Standard] = DefaultStandardCapacityBytes
+	}
+	if _, ok := cfg.MaxBytesPerCriticality[v1alpha2.Sheddable]; !ok {
+		cfg.MaxBytesPerCriticality[v1alpha2.Sheddable] = DefaultSheddableCapacityBytes
+	}
+
+	// Apply default TTL and intervals if zero/negative.
+	if cfg.DefaultQueueTTL <= 0 {
+		cfg.DefaultQueueTTL = DefaultQueueTTL
+	}
+	if cfg.ExpiryCleanupInterval <= 0 {
+		cfg.ExpiryCleanupInterval = DefaultExpiryCleanupInterval
+	}
+
+	// Apply default policies if empty and validate provided ones.
+	if cfg.InterModelFairnessPolicy == "" {
+		cfg.InterModelFairnessPolicy = PolicyFCFS
+	} else {
+		if _, ok := SupportedInterModelFairnessPolicies[cfg.InterModelFairnessPolicy]; !ok {
+			return fmt.Errorf("unsupported or misconfigured InterModelFairnessPolicy: %s", cfg.InterModelFairnessPolicy)
+		}
+	}
+
+	if cfg.PreemptionStrategy == "" {
+		cfg.PreemptionStrategy = PolicyNone
+	} else {
+		if _, ok := SupportedPreemptionStrategies[cfg.PreemptionStrategy]; !ok {
+			return fmt.Errorf("unsupported PreemptionStrategy: %s", cfg.PreemptionStrategy)
+		}
+	}
+
+	return nil
+}
+
+// LoadConfigFromEnv loads FlowControllerConfig from environment variables.
+// It also returns a boolean indicating if the FlowController is enabled.
+func LoadConfigFromEnv() (*FlowControllerConfig, bool) {
+	// Use a default logger for initial configuration loading.
+	logger := log.Log.WithName("flow-controller-config")
+
+	cfg := &FlowControllerConfig{
+		MaxBytesPerCriticality: make(map[v1alpha2.Criticality]uint64),
+		// Custom policies are not loaded from env vars
+	}
+
+	cfg.MaxBytesPerCriticality[v1alpha2.Critical] = envutil.GetEnvUint64(EnvFcMaxBytesCritical, DefaultCriticalCapacityBytes, logger)
+	cfg.MaxBytesPerCriticality[v1alpha2.Standard] = envutil.GetEnvUint64(EnvFcMaxBytesStandard, DefaultStandardCapacityBytes, logger)
+	cfg.MaxBytesPerCriticality[v1alpha2.Sheddable] = envutil.GetEnvUint64(EnvFcMaxBytesSheddable, DefaultSheddableCapacityBytes, logger)
+
+	cfg.MaxGlobalBytes = envutil.GetEnvUint64(EnvFcMaxGlobalBytes, 0 /*Default is 0*/, logger)
+	cfg.DefaultQueueTTL = envutil.GetEnvDuration(EnvFcDefaultQueueTTL, DefaultQueueTTL, logger)
+	cfg.ExpiryCleanupInterval = envutil.GetEnvDuration(EnvFcExpiryCleanupInterval, DefaultExpiryCleanupInterval, logger)
+
+	cfg.InterModelFairnessPolicy = envutil.GetEnvString(EnvFcFairnessPolicy, PolicyFCFS, logger)
+	cfg.PreemptionStrategy = envutil.GetEnvString(EnvFcPreemptionStrategy, PolicyNone, logger)
+
+	fcEnabled := envutil.GetEnvBool(EnvFcEnabled, DefaultFcEnabled, logger)
+
+	// Note: We don't call cfg.validateAndApplyDefaults() here.
+	// It is called within NewFlowController to ensure defaults are applied just
+	// before the controller is created, using the final config object.
+	logger.Info("FlowController configuration loaded from env",
+		"config", fmt.Sprintf("%+v", cfg),
+		"fcEnabled", fcEnabled)
+	return cfg, fcEnabled
+}