Fix streamed request being called one final time after request complete, add predictor check to the beginning of each requestcontrol hook

Author: BenjaminBraunDev

pkg/epp/requestcontrol/director.go

@@ -269,8 +269,9 @@ func (d *Director) HandleResponseBodyStreaming(ctx context.Context, reqCtx *hand
 	logger := log.FromContext(ctx).WithValues("stage", "bodyChunk")
 	logger.V(logutil.TRACE).Info("Entering HandleResponseBodyChunk")
 	response := &Response{
-		RequestId: reqCtx.Request.Headers[requtil.RequestIdHeaderKey],
-		Headers:   reqCtx.Response.Headers,
+		RequestId:   reqCtx.Request.Headers[requtil.RequestIdHeaderKey],
+		Headers:     reqCtx.Response.Headers,
+		EndOfStream: reqCtx.ResponseComplete,
 	}
 
 	d.runResponseStreamingPlugins(ctx, reqCtx.SchedulingRequest, response, reqCtx.TargetPod)

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

@@ -113,6 +113,9 @@ func (t *SLOAwareRouter) PreRequest(ctx context.Context, request *schedulingtype
 	}
 
 	targetPod := schedulingResult.ProfileResults[schedulingResult.PrimaryProfileName].TargetPods[0].GetPod()
+	if !t.CheckPredictor(logger, targetPod) {
+		return
+	}
 
 	podName := types.NamespacedName{
 		Name:      targetPod.NamespacedName.Name,
@@ -153,6 +156,10 @@ 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)
+	if !t.CheckPredictor(logger, targetPod) {
+		return
+	}
+
 	id := request.Headers[requtil.RequestIdHeaderKey]
 
 	sloCtx, err := t.getSLOContextForRequest(request)
@@ -161,10 +168,6 @@ func (t *SLOAwareRouter) ResponseReceived(ctx context.Context, request *scheduli
 		return
 	}
 
-	if !t.CheckPredictor(logger, targetPod) {
-		return
-	}
-
 	if err := ProcessHeaderForLatencyPrediction(ctx, t.latencypredictor, sloCtx); err != nil {
 		logger.V(logutil.DEBUG).Error(err, "ProcessHeader in latencypredictor failed")
 	}
@@ -173,7 +176,7 @@ func (t *SLOAwareRouter) ResponseReceived(ctx context.Context, request *scheduli
 
 func (t *SLOAwareRouter) ResponseStreaming(ctx context.Context, request *schedulingtypes.LLMRequest, response *requestcontrol.Response, pod *backend.Pod) {
 	logger := log.FromContext(ctx)
-	if !t.CheckPredictor(logger, pod) {
+	if !t.CheckPredictor(logger, pod) || response.EndOfStream {
 		return
 	}
 
@@ -248,11 +251,11 @@ func (t *SLOAwareRouter) ResponseComplete(ctx context.Context, request *scheduli
 
 func (t *SLOAwareRouter) CheckPredictor(logger logr.Logger, targetPod *backend.Pod) bool {
 	if targetPod == nil {
-		logger.V(logutil.TRACE).Info("SLOAwareRouter: Skipping PostResponse because no target pod was provided.")
+		logger.V(logutil.TRACE).Info("SLOAwareRouter: Skipping hook because no target pod was provided.")
 		return false
 	}
 	if t.latencypredictor == nil {
-		logger.V(logutil.TRACE).Info("SLOAwareRouter: Skipping PostResponse because predictor missing")
+		logger.V(logutil.TRACE).Info("SLOAwareRouter: Skipping hook because predictor missing")
 		return false
 	}
 	return true

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

@@ -150,6 +150,9 @@ func TestSLOAwareRouter_PreRequest_EmptySchedulingResult(t *testing.T) {
 
 func TestSLOAwareRouter_PreRequest_Success(t *testing.T) {
 	router := createTestRouter()
+	mockPredictor := new(mockPredictor)
+	router.latencypredictor = mockPredictor
+
 	ctx := context.Background()
 	pod := createTestPod("test-pod", 1, 1, 1)
 	request := createTestLLMRequest("test", 100, 50, true)
@@ -180,6 +183,9 @@ func TestSLOAwareRouter_PreRequest_Success(t *testing.T) {
 
 func TestSLOAwareRouter_PreRequest_AddsToQueue(t *testing.T) {
 	router := createTestRouter()
+	mockPredictor := new(mockPredictor)
+	router.latencypredictor = mockPredictor
+
 	ctx := context.Background()
 	pod := createTestPod("test-pod", 1, 1, 1)
 	request := createTestLLMRequest("test", 100, 50, true)
@@ -201,6 +207,9 @@ func TestSLOAwareRouter_PreRequest_AddsToQueue(t *testing.T) {
 
 func TestSLOAwareRouter_PreRequest_QueueAlreadyExists(t *testing.T) {
 	router := createTestRouter()
+	mockPredictor := new(mockPredictor)
+	router.latencypredictor = mockPredictor
+
 	ctx := context.Background()
 	pod := createTestPod("test-pod", 1, 1, 1)
 	request1 := createTestLLMRequest("test-id-1", 100, 50, true)
@@ -729,6 +738,9 @@ func TestSLOAwareRouter_ConcurrentContextAccess(t *testing.T) {
 
 func TestSLOAwareRouter_MultipleRequests_SamePod(t *testing.T) {
 	router := createTestRouter()
+	mockPredictor := new(mockPredictor)
+	router.latencypredictor = mockPredictor
+
 	ctx := context.Background()
 	pod := createTestPod("test-pod", 1, 1, 1)
 
@@ -807,6 +819,9 @@ func TestSLOAwareRouter_RequestLifecycle_Complete(t *testing.T) {
 
 func TestSLOAwareRouter_MultipleRequests_DifferentPods(t *testing.T) {
 	router := createTestRouter()
+	mockPredictor := new(mockPredictor)
+	router.latencypredictor = mockPredictor
+
 	ctx := context.Background()
 
 	pod1 := createTestPod("test-pod-1", 1, 1, 1)