Clean up errors from rebase, add running request metric to datasource, add predictor to new 2 phase configuration parser

Author: BenjaminBraunDev

cmd/epp/runner/runner.go

@@ -227,7 +227,20 @@ func (r *Runner) Run(ctx context.Context) error {
 		return err
 	}
 
-	rawConfig, err := r.parseConfigurationPhaseOne(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"))
+	} else {
+		setupLog.Info("Latency predictor is disabled.")
+		predictor = nil // This will be a true nil interface
+	}
+	// ===================================================================
+
+	rawConfig, err := r.parseConfigurationPhaseOne(ctx, predictor)
 	if err != nil {
 		setupLog.Error(err, "Failed to parse configuration")
 		return err
@@ -322,32 +335,6 @@ func (r *Runner) Run(ctx context.Context) error {
 		runtime.SetBlockProfileRate(1)
 	}
 
-	// ===================================================================
-	// == 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 the configuration")
-		return err
-	}
-
 	// --- Initialize Core EPP Components ---
 	if r.schedulerConfig == nil {
 		err := errors.New("scheduler config must be set either by config api or through code")
@@ -420,6 +407,12 @@ func (r *Runner) Run(ctx context.Context) error {
 		return err
 	}
 
+	if *enableLatencyPredictor && predictor != nil {
+		if err := registerLatencyPredictorServer(mgr, predictor); err != nil {
+			return err
+		}
+	}
+
 	// --- Start Manager ---
 	// This blocks until a signal is received.
 	setupLog.Info("Controller manager starting")
@@ -454,7 +447,7 @@ func (r *Runner) registerLatencyPredictorPlugins(predictor latencypredictor.Pred
 	plugins.Register(profile.SLOAwareProfileHandlerType, profile.SLOAwareProfileHandlerFactory)
 }
 
-func (r *Runner) parseConfigurationPhaseOne(ctx context.Context) (*configapi.EndpointPickerConfig, error) {
+func (r *Runner) parseConfigurationPhaseOne(ctx context.Context, predictor latencypredictor.PredictorInterface) (*configapi.EndpointPickerConfig, error) {
 	if *configText == "" && *configFile == "" {
 		return nil, nil // configuring through code, not through file
 	}
@@ -683,6 +676,18 @@ func registerHealthServer(mgr manager.Manager, logger logr.Logger, ds datastore.
 	return nil
 }
 
+// registerLatencyPredictorServer adds the Latency Predictor server as a Runnable to the given manager.
+func registerLatencyPredictorServer(mgr manager.Manager, predictor latencypredictor.PredictorInterface) error {
+	// 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
+	}
+	setupLog.Info("Latency predictor runnable added to manager.")
+	return nil
+}
+
 func validateFlags() error {
 	if *poolName == "" {
 		return fmt.Errorf("required %q flag not set", "poolName")

pkg/epp/datalayer/metrics/datasource_test.go

@@ -29,7 +29,7 @@ import (
 
 func TestDatasource(t *testing.T) {
 	source := NewDataSource("https", "/metrics", true, nil)
-	extractor, err := NewExtractor(defaultTotalQueuedRequestsMetric, "", "", "")
+	extractor, err := NewExtractor(defaultTotalQueuedRequestsMetric, "", "", "", "")
 	assert.Nil(t, err, "failed to create extractor")
 
 	name := source.Name()

pkg/epp/datalayer/metrics/extractor_test.go

@@ -31,6 +31,7 @@ import (
 const (
 	// use hardcoded values - importing causes cycle
 	defaultTotalQueuedRequestsMetric    = "vllm:num_requests_waiting"
+	defaultTotalRunningRequestsMetric   = "vllm:num_requests_running"
 	defaultKvCacheUsagePercentageMetric = "vllm:gpu_cache_usage_perc"
 	defaultLoraInfoMetric               = "vllm:lora_requests_info"
 	defaultCacheInfoMetric              = "vllm:cache_config_info"
@@ -39,11 +40,11 @@ const (
 func TestExtractorExtract(t *testing.T) {
 	ctx := context.Background()
 
-	if _, err := NewExtractor("vllm: dummy", "", "", ""); err == nil {
+	if _, err := NewExtractor("vllm: dummy", "", "", "", ""); err == nil {
 		t.Error("expected to fail to create extractor with invalid specification")
 	}
 
-	extractor, err := NewExtractor(defaultTotalQueuedRequestsMetric,
+	extractor, err := NewExtractor(defaultTotalQueuedRequestsMetric, defaultTotalRunningRequestsMetric,
 		defaultKvCacheUsagePercentageMetric, defaultLoraInfoMetric, defaultCacheInfoMetric)
 	if err != nil {
 		t.Fatalf("failed to create extractor: %v", err)

pkg/epp/scheduling/framework/plugins/multi/slo_aware_router/requestcontrol_hooks.go

@@ -114,11 +114,7 @@ func (t *SLOAwareRouter) PreRequest(ctx context.Context, request *schedulingtype
 	}
 
 	targetPod := schedulingResult.ProfileResults[schedulingResult.PrimaryProfileName].TargetPods[0].GetPod()
-<<<<<<< HEAD
 	if !t.checkPredictor(logger, targetPod) {
-=======
-	if !t.CheckPredictor(logger, targetPod) {
->>>>>>> b2a7d45 (Fix streamed request being called one final time after request complete, add predictor check to the beginning of each requestcontrol hook)
 		return
 	}
 
@@ -161,11 +157,7 @@ func (t *SLOAwareRouter) PreRequest(ctx context.Context, request *schedulingtype
 
 func (t *SLOAwareRouter) ResponseReceived(ctx context.Context, request *schedulingtypes.LLMRequest, response *requestcontrol.Response, targetPod *backend.Pod) {
 	logger := log.FromContext(ctx)
-<<<<<<< HEAD
 	if !t.checkPredictor(logger, targetPod) {
-=======
-	if !t.CheckPredictor(logger, targetPod) {
->>>>>>> b2a7d45 (Fix streamed request being called one final time after request complete, add predictor check to the beginning of each requestcontrol hook)
 		return
 	}
 
@@ -177,23 +169,15 @@ func (t *SLOAwareRouter) ResponseReceived(ctx context.Context, request *scheduli
 		return
 	}
 
-<<<<<<< HEAD
 	if err := processHeaderForLatencyPrediction(ctx, t.latencypredictor, sloCtx); err != nil {
-=======
-	if err := ProcessHeaderForLatencyPrediction(ctx, t.latencypredictor, sloCtx); err != nil {
->>>>>>> b2a7d45 (Fix streamed request being called one final time after request complete, add predictor check to the beginning of each requestcontrol hook)
 		logger.V(logutil.DEBUG).Error(err, "ProcessHeader in latencypredictor failed")
 	}
 
 }
 
 func (t *SLOAwareRouter) ResponseStreaming(ctx context.Context, request *schedulingtypes.LLMRequest, response *requestcontrol.Response, pod *backend.Pod) {
 	logger := log.FromContext(ctx)
-<<<<<<< HEAD
 	if !t.checkPredictor(logger, pod) || response.EndOfStream {
-=======
-	if !t.CheckPredictor(logger, pod) || response.EndOfStream {
->>>>>>> b2a7d45 (Fix streamed request being called one final time after request complete, add predictor check to the beginning of each requestcontrol hook)
 		return
 	}