Experimental SLO-Aware Routing and Latency Prediction (#1568)

* add latency predictor

* add cv in model and update epp deployment

* bug fix

* track mape for predictions

* add running queue size to metrics

* add xgboost regressor and update tpot sampling logic

* emit predicted and actual ttft tpot in body

* seperate servers for training and prediction

* add latency predictor

put the predictor functions in director in a helper function

add scores to reqcxt

record prediction duration metrics

add prefix cache score to model input

slo based routing changes

retreive  request priority queue from the datastore

update scoring logic

* better inital implemenation

Add scheduling profile, working state

remove latencypredictor from director

Move all latency prediction logic out of director and into scheduling profile. Make all Request/Response plugins take in RequestContext

* progress towards fixing up merge conflicts from latency predictor merge

* More refactor progress, fixing and adding tests

* working state, latency prediction

* Clean up changes, remove unneeded files, working functionality without latency flag and scheduling plugins

* Rebase cleanup, remove duplicate lines

* Integrate new alpha-beta slo scoring into scoring plugin

* Fix prefix cache scoring for slo-aware routing

* Add pycache or latency predictor to gitignore

* Rebase with main

* Fix prefix cache scoring being piped to latencyprediction_helper

* add dependancies in scorer

* chage to single profile

* chage to single profile

* restore two profiles

* restore two profiles

* restore two profiles

* update admit request to shed based on predictions

* add TODOs for future changes

* Change artifact registry references to personal compiled images

* Fix existing non-slo aware routing unit tests

* update latency predictor with better eval metrics

* Fix saturation detector unit test

* Change naming of SLO headers and prediction based routing header

* Remove port 9002 service on InferencePool causing make test to fail

* Fix epp hermetic integration test to expect ProcessingMode Send in response header

---------

Co-authored-by: kaushikmitr <[email protected]>

Author: BenjaminBraunDev

.gitignore

@@ -7,6 +7,7 @@
 bin/*
 Dockerfile.cross
 artifacts
+latencypredictor-v1/__pycache__
 
 # Test binary, built with `go test -c`
 *.test

cmd/epp/runner/runner.go

@@ -19,6 +19,7 @@ package runner
 import (
 	"context"
 	"crypto/tls"
+	"encoding/json"
 	"errors"
 	"flag"
 	"fmt"
@@ -50,10 +51,12 @@ import (
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/datalayer"
 	dlmetrics "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/datalayer/metrics"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/datastore"
+	latencypredictor "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/latencypredictorasync"
 	"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/plugins"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/requestcontrol"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/requestcontrol/plugins/slorequest"
 	"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/plugins/multi/prefix"
@@ -89,6 +92,7 @@ var (
 		"then a self-signed certificate is used.")
 	// metric flags
 	totalQueuedRequestsMetric    = flag.String("total-queued-requests-metric", runserver.DefaultTotalQueuedRequestsMetric, "Prometheus metric for the number of queued requests.")
+	totalRunningRequestsMetric   = flag.String("total-running-requests-metric", runserver.DefaultTotalRunningRequestsMetric, "Prometheus metric for the number of running requests.")
 	kvCacheUsagePercentageMetric = flag.String("kv-cache-usage-percentage-metric", runserver.DefaultKvCacheUsagePercentageMetric, "Prometheus metric for the fraction of KV-cache blocks currently in use (from 0 to 1).")
 	// LoRA metrics
 	loraInfoMetric = flag.String("lora-info-metric", runserver.DefaultLoraInfoMetric, "Prometheus metric for the LoRA info metrics (must be in vLLM label format).")
@@ -107,6 +111,9 @@ var (
 	modelServerMetricsHttpsInsecureSkipVerify = flag.Bool("model-server-metrics-https-insecure-skip-verify", true, "When using 'https' scheme for 'model-server-metrics-scheme', configure 'InsecureSkipVerify' (default to true)")
 	haEnableLeaderElection                    = flag.Bool("ha-enable-leader-election", false, "Enables leader election for high availability. When enabled, readiness probes will only pass on the leader.")
 
+	// Latency Predictor Flag
+	enableLatencyPredictor = flag.Bool("enable-latency-predictor", false, "Enable the regression-based latency predictor and scheduler scorer.")
+
 	setupLog = ctrl.Log.WithName("setup")
 )
 
@@ -233,9 +240,29 @@ func (r *Runner) Run(ctx context.Context) error {
 		runtime.SetBlockProfileRate(1)
 	}
 
-	err = r.parsePluginsConfiguration(ctx)
+	// ===================================================================
+	// == Latency Predictor Integration
+	// ===================================================================
+	var predictor latencypredictor.PredictorInterface // Use the interface type
+	if *enableLatencyPredictor {
+		setupLog.Info("Latency predictor is enabled. Initializing...")
+		predictor = latencypredictor.New(latencypredictor.ConfigFromEnv(), ctrl.Log.WithName("latency-predictor"))
+
+		// For the runnable, you'll need to type assert back to the concrete type
+		concretePredictor := predictor.(*latencypredictor.Predictor)
+		if err := mgr.Add(runnable.NoLeaderElection(&predictorRunnable{predictor: concretePredictor})); err != nil {
+			setupLog.Error(err, "Failed to register latency predictor runnable")
+			return err
+		}
+	} else {
+		setupLog.Info("Latency predictor is disabled.")
+		predictor = nil // This will be a true nil interface
+	}
+	// ===================================================================
+
+	err = r.parsePluginsConfiguration(ctx, predictor, datastore)
 	if err != nil {
-		setupLog.Error(err, "Failed to parse plugins configuration")
+		setupLog.Error(err, "Failed to parse the configuration")
 		return err
 	}
 
@@ -268,6 +295,7 @@ func (r *Runner) Run(ctx context.Context) error {
 		Director:                         director,
 		SaturationDetector:               saturationDetector,
 		UseExperimentalDatalayerV2:       useDatalayerV2, // pluggable data layer feature flag
+		LatencyPredictor:                 predictor,
 	}
 	if err := serverRunner.SetupWithManager(ctx, mgr); err != nil {
 		setupLog.Error(err, "Failed to setup EPP controllers")
@@ -310,7 +338,20 @@ func (r *Runner) registerInTreePlugins() {
 	plugins.Register(testfilter.HeaderBasedTestingFilterType, testfilter.HeaderBasedTestingFilterFactory)
 }
 
-func (r *Runner) parsePluginsConfiguration(ctx context.Context) error {
+func (r *Runner) registerLatencyPredictorPlugins(predictor latencypredictor.PredictorInterface, datastore datastore.Datastore) {
+	// Register the SLO request tracker and scorer plugin, these plugins need access to the predictor and datastore.
+	// We have to specify a custom factory function to create the plugins with the correct dependencies.
+	plugins.Register(slorequest.SLORequestTrackerPluginType, func(name string, _ json.RawMessage, _ plugins.Handle) (plugins.Plugin, error) {
+		return slorequest.New(predictor, datastore).WithName(name), nil
+	})
+	plugins.Register(scorer.SLOScorerPluginType, func(name string, _ json.RawMessage, _ plugins.Handle) (plugins.Plugin, error) {
+		return scorer.NewSLOScorer(predictor, datastore, scorer.HeadroomSelectionStrategy).WithName(name), nil
+	})
+	plugins.Register(profile.SLOAwareProfileHandlerType, profile.SLOAwareProfileHandlerFactory)
+	plugins.Register(picker.WeightedRandomPickerType, picker.WeightedRandomPickerFactory)
+}
+
+func (r *Runner) parsePluginsConfiguration(ctx context.Context, predictor latencypredictor.PredictorInterface, datastore datastore.Datastore) error {
 	if *configText == "" && *configFile == "" {
 		return nil // configuring through code, not through file
 	}
@@ -329,6 +370,12 @@ func (r *Runner) parsePluginsConfiguration(ctx context.Context) error {
 	}
 
 	r.registerInTreePlugins()
+	// If we have a latency predictor enabled and predictor and datastore are not nil,
+	// register the latency predictor plugins (currently just the SLO scorer).
+	if *enableLatencyPredictor && predictor != nil && datastore != nil {
+		setupLog.Info("Registering latency predictor plugins")
+		r.registerLatencyPredictorPlugins(predictor, datastore)
+	}
 	handle := plugins.NewEppHandle(ctx)
 	config, err := loader.LoadConfig(configBytes, handle, logger)
 	if err != nil {
@@ -358,6 +405,7 @@ func (r *Runner) setupMetricsCollection(setupLog logr.Logger, useExperimentalDat
 func setupMetricsV1(setupLog logr.Logger) (datalayer.EndpointFactory, error) {
 	mapping, err := backendmetrics.NewMetricMapping(
 		*totalQueuedRequestsMetric,
+		*totalRunningRequestsMetric,
 		*kvCacheUsagePercentageMetric,
 		*loraInfoMetric,
 	)
@@ -402,6 +450,7 @@ func setupDatalayer() (datalayer.EndpointFactory, error) {
 		*modelServerMetricsHttpsInsecureSkipVerify,
 		nil)
 	extractor, err := dlmetrics.NewExtractor(*totalQueuedRequestsMetric,
+		*totalRunningRequestsMetric,
 		*kvCacheUsagePercentageMetric,
 		*loraInfoMetric)
 
@@ -510,3 +559,21 @@ func setupPprofHandlers(mgr ctrl.Manager) error {
 	}
 	return nil
 }
+
+// ===================================================================
+// == Latency Predictor Plugin and Helpers
+// ===================================================================
+
+// predictorRunnable implements controller-runtime's Runnable interface to manage the predictor's lifecycle.
+type predictorRunnable struct {
+	predictor *latencypredictor.Predictor
+}
+
+func (p *predictorRunnable) Start(ctx context.Context) error {
+	setupLog.Info("Starting latency predictor...")
+	p.predictor.Start(ctx)
+	<-ctx.Done()
+	setupLog.Info("Stopping latency predictor...")
+	p.predictor.Stop()
+	return nil
+}

config/manifests/gateway/gke/gcp-backend-policy.yaml

@@ -4,7 +4,7 @@ metadata:
   name: inferencepool-backend-policy
 spec:
   targetRef:
-    group: "inference.networking.k8s.io"
+    group: "inference.networking.x-k8s.io"
     kind: InferencePool
     name: vllm-llama3-8b-instruct
   default:

config/manifests/gateway/gke/healthcheck.yaml

@@ -5,7 +5,7 @@ metadata:
   namespace: default
 spec:
   targetRef:
-    group: "inference.networking.k8s.io"
+    group: "inference.networking.x-k8s.io"
     kind: InferencePool
     name: vllm-llama3-8b-instruct
   default:

config/manifests/gateway/gke/httproute.yaml

@@ -9,7 +9,7 @@ spec:
     name: inference-gateway
   rules:
   - backendRefs:
-    - group: inference.networking.k8s.io
+    - group: inference.networking.x-k8s.io
       kind: InferencePool
       name: vllm-llama3-8b-instruct
     matches: