Externalize Scheduler's saturation logic

This commit refactors the request processing pipeline, externalizing
saturation detection and criticality-based service differentiation
from the Scheduler. These responsibilities are now primarily managed by
the RequestControl.Director.

This change is a preparatory step for the introduction of a new
Flow Controller component, which will eventually absorb these admission
control duties.

Key changes include:

- Introduced `PreDispatch` method to `RequestControl.Director` It
  utilizes the `SaturationDetector` for admission control of
  non-critical requests and handles request criticality to determine if
  saturation checks are bypassed.
- The saturation detection logic for dropping non-critical requests
  is intentionally preserved within the `Director` at this stage.
  This allows the option to bypass the future Flow Controller
  component during its maturation, ensuring the existing saturation
  and sheddable request behavior can be maintained as a fallback.
- Updated `main.go` to instantiate the `SaturationDetector`, wiring it
  into the request handling flow.
- Updated `director_test.go` to align with the new component
  responsibilities, adding additional coverage where necessary.

This refactoring leads to a cleaner architecture, making the `Scheduler`
a more focused component and centralizing initial admission control
logic while paving the way for the future Flow Controller.

This is aligned with the direction in `0683-epp-architecture-proposal`
and should be nearly no-op in terms of EPP behavior.

Author: LukeAVanDrie

cmd/epp/main.go

@@ -40,6 +40,7 @@ import (
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/datastore"
 	"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/framework"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/framework/plugins/filter"
@@ -151,14 +152,17 @@ func run() error {
 	})
 	setupLog.Info("Flags processed", "flags", flags)
 
-	// Init runtime.
+	// --- Load Configurations from Environment Variables ---
+	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
 	}
 
-	// Set up mapper for metric scraping.
+	// --- Setup Datastore ---
 	mapping, err := backendmetrics.NewMetricMapping(
 		*totalQueuedRequestsMetric,
 		*kvCacheUsagePercentageMetric,
@@ -169,13 +173,11 @@ func run() error {
 		return err
 	}
 	verifyMetricMapping(*mapping, setupLog)
-
 	pmf := backendmetrics.NewPodMetricsFactory(&backendmetrics.PodMetricsClientImpl{MetricMapping: mapping}, *refreshMetricsInterval)
-	// Setup runner.
 	ctx := ctrl.SetupSignalHandler()
-
 	datastore := datastore.NewDatastore(ctx, pmf)
 
+	// --- Setup Metrics Server ---
 	customCollectors := []prometheus.Collector{collectors.NewInferencePoolMetricsCollector(datastore)}
 	metrics.Register(customCollectors...)
 	metrics.RecordInferenceExtensionInfo()
@@ -199,6 +201,7 @@ func run() error {
 		return err
 	}
 
+	// --- Initialize Core EPP Components ---
 	scheduler := scheduling.NewScheduler(datastore)
 	if schedulerV2 == "true" {
 		queueScorerWeight := envutil.GetEnvInt("QUEUE_SCORE_WEIGHT", scorer.DefaultQueueScorerWeight, setupLog)
@@ -221,6 +224,10 @@ func run() error {
 		schedulerConfig := scheduling.NewSchedulerConfig(profilepicker.NewAllProfilesPicker(), map[string]*framework.SchedulerProfile{"schedulerv2": schedulerProfile})
 		scheduler = scheduling.NewSchedulerWithConfig(datastore, schedulerConfig)
 	}
+
+	saturationDetector := saturationdetector.NewDetector(sdConfig, datastore, ctrl.Log)
+
+	// --- Setup ExtProc Server Runner ---
 	serverRunner := &runserver.ExtProcServerRunner{
 		GrpcPort:                                 *grpcPort,
 		DestinationEndpointHintMetadataNamespace: *destinationEndpointHintMetadataNamespace,
@@ -231,24 +238,26 @@ func run() error {
 		CertPath:                                 *certPath,
 		RefreshPrometheusMetricsInterval:         *refreshPrometheusMetricsInterval,
 		Scheduler:                                scheduler,
+		SaturationDetector:                       saturationDetector,
 	}
 	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 health server.
 	if err := registerHealthServer(mgr, ctrl.Log.WithName("health"), datastore, *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
 	}
 
-	// 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")
@@ -276,6 +285,16 @@ func initLogging(opts *zap.Options) {
 	ctrl.SetLogger(logger)
 }
 
+// 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()
@@ -309,5 +328,4 @@ func verifyMetricMapping(mapping backendmetrics.MetricMapping, logger logr.Logge
 	if mapping.LoraRequestInfo == nil {
 		logger.Info("Not scraping metric: LoraRequestInfo")
 	}
-
 }

pkg/epp/requestcontrol/director.go

@@ -14,6 +14,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
+// Package requestcontrol defines the Director component responsible for orchestrating request processing after initial
+// parsing.
 package requestcontrol
 
 import (
@@ -34,33 +36,45 @@ import (
 	requtil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/request"
 )
 
+// Scheduler defines the interface required by the Director for scheduling.
 type Scheduler interface {
 	Schedule(ctx context.Context, b *schedulingtypes.LLMRequest) (result map[string]*schedulingtypes.Result, err error)
 	OnResponse(ctx context.Context, resp *schedulingtypes.LLMResponse, targetPodName string)
 }
 
+// SaturationDetector provides a signal indicating whether the backends are considered saturated.
+type SaturationDetector interface {
+	IsSaturated(ctx context.Context) bool
+}
+
+// Director orchestrates the request handling flow, including scheduling.
 type Director struct {
-	datastore datastore.Datastore
-	scheduler Scheduler
+	datastore          datastore.Datastore
+	scheduler          Scheduler
+	saturationDetector SaturationDetector
 }
 
-func NewDirector(datastore datastore.Datastore, scheduler Scheduler) *Director {
-	return &Director{
-		datastore: datastore,
-		scheduler: scheduler,
-	}
+// NewDirector creates a new Director instance with all dependencies.
+func NewDirector(datastore datastore.Datastore, scheduler Scheduler, saturationDetector SaturationDetector) *Director {
+	return &Director{datastore, scheduler, saturationDetector}
 }
 
-// HandleRequest always returns the requestContext even in the error case, as the request context is used in error handling.
+// HandleRequest orchestrates the request lifecycle:
+//  1. Parses request details.
+//  2. Calls PreDispatch for admission control.
+//  3. Calls Dispatch (which calls Scheduler) if request is approved.
+//  4. Calls PostDispatch to populate RequestContext with results.
+//
+// It always returns the requestContext even in the error case, as the request context is used in error handling.
 func (d *Director) HandleRequest(ctx context.Context, reqCtx *handlers.RequestContext) (*handlers.RequestContext, error) {
 	logger := log.FromContext(ctx)
 
-	// Resolve target models.
+	// --- 1. Parse Request, Resolve Target Models, and Determine Parameters ---
 	var ok bool
 	requestBodyMap := reqCtx.Request.Body
 	reqCtx.Model, ok = requestBodyMap["model"].(string)
 	if !ok {
-		return reqCtx, errutil.Error{Code: errutil.BadRequest, Msg: "model not found in request"}
+		return reqCtx, errutil.Error{Code: errutil.BadRequest, Msg: "model not found in request body"}
 	}
 	prompt, err := requtil.ExtractPromptFromRequestBody(requestBodyMap)
 	if err != nil {
@@ -84,29 +98,69 @@ func (d *Director) HandleRequest(ctx context.Context, reqCtx *handlers.RequestCo
 		reqCtx.Request.Body["model"] = reqCtx.ResolvedTargetModel // Update target model in the body.
 	}
 
+	requestCriticality := v1alpha2.Standard
+	if modelObj.Spec.Criticality != nil {
+		requestCriticality = *modelObj.Spec.Criticality
+	}
+
+	// Prepare LLMRequest (needed for both saturation detection and Scheduler)
 	llmReq := &schedulingtypes.LLMRequest{
 		TargetModel: reqCtx.ResolvedTargetModel,
 		RequestId:   reqCtx.Request.Headers[requtil.RequestIdHeaderKey],
-		Critical:    modelObj.Spec.Criticality != nil && *modelObj.Spec.Criticality == v1alpha2.Critical,
+		Critical:    requestCriticality == v1alpha2.Critical,
 		Prompt:      prompt,
 		Headers:     reqCtx.Request.Headers,
 	}
-	logger.V(logutil.DEBUG).Info("LLM request assembled", "request", llmReq)
-	results, err := d.Dispatch(ctx, llmReq)
-	if err != nil {
-		return reqCtx, err
+	logger = logger.WithValues(
+		"model", reqCtx.Model,
+		"resolvedTargetModel", llmReq.TargetModel,
+		"criticality", requestCriticality,
+	)
+	ctx = log.IntoContext(ctx, logger)
+	logger.V(logutil.DEBUG).Info("LLM request assembled")
+
+	// --- 2. Saturation Check ---
+	preDispatchErr := d.PreDispatch(ctx, reqCtx, requestCriticality)
+	if preDispatchErr != nil {
+		return reqCtx, preDispatchErr
 	}
 
+	// --- 3. Dispatch (Calls Scheduler) ---
+	results, dispatchErr := d.Dispatch(ctx, llmReq)
+	if dispatchErr != nil {
+		return reqCtx, dispatchErr
+	}
+
+	// --- 4. PostDispatch (Populates RequestContext) ---
 	// Insert target endpoint to instruct Envoy to route requests to the specified target pod.
-	// Attach the port number
-	reqCtx, err = d.PostDispatch(ctx, reqCtx, results)
-	if err != nil {
-		return reqCtx, err
+	// Attach the port number.
+	reqCtx, postDispatchErr := d.PostDispatch(ctx, reqCtx, results)
+	if postDispatchErr != nil {
+		return reqCtx, postDispatchErr
 	}
 
 	return reqCtx, nil
 }
 
+// PreDispatch handles admission control before dispatch.
+func (d *Director) PreDispatch(ctx context.Context, reqCtx *handlers.RequestContext, reqCriticality v1alpha2.Criticality) error {
+	logger := log.FromContext(ctx)
+
+	if reqCriticality == v1alpha2.Critical {
+		logger.V(logutil.DEBUG).Info("Critical request bypassing saturation check.")
+		return nil
+	}
+
+	logger.V(logutil.DEBUG).Info("Performing saturation check for non-critical request.")
+	if d.saturationDetector.IsSaturated(ctx) { // Assuming non-nil Saturation Detector
+		return errutil.Error{
+			Code: errutil.InferencePoolResourceExhausted,
+			Msg:  "system saturated, non-critical request dropped",
+		}
+	}
+	return nil
+}
+
 // Dispatch runs one or many scheduling cycles.
 func (d *Director) Dispatch(ctx context.Context, llmReq *schedulingtypes.LLMRequest) (map[string]*schedulingtypes.Result, error) {
 	var err error
@@ -118,6 +172,7 @@ func (d *Director) Dispatch(ctx context.Context, llmReq *schedulingtypes.LLMRequ
 	return res, nil // TODO handle multi cycle result after defining the PostDispatch extension point
 }
 
+// PostDispatch populates the RequestContext based on scheduling results.
 func (d *Director) PostDispatch(ctx context.Context, reqCtx *handlers.RequestContext, results map[string]*schedulingtypes.Result) (*handlers.RequestContext, error) {
 	logger := log.FromContext(ctx)
 	// currently only get a single result. Will refactor to pluggably implement the PostSchedule