[Metrics] Add streaming support for metrics

Address #178

Author: JeffLuoo

Diff for: pkg/epp/handlers/request.go

@@ -61,6 +61,38 @@ func (s *Server) HandleRequestBody(
 	loggerVerbose.Info("Model requested", "model", model)
 	modelName := model
 
+	// Resolve streaming options
+
+	streaming, ok := rb["stream"].(bool)
+	if !ok {
+		// streaming not set, no-op
+	} else {
+		reqCtx.Streaming = streaming
+	}
+
+	if reqCtx.Streaming {
+		type Usage struct {
+			IncludeUsage string `json:"include_usage,omitempty"`
+		}
+		if streamOption, ok := rb["stream_options"]; ok {
+			includeUsage := Usage{}
+
+			// Parsing `stream_options` won't reject the request.
+			optionJson, err := json.Marshal(streamOption)
+			if err != nil {
+				logger.V(logutil.DEFAULT).Error(err, "Error unmarshaling stream_options")
+			}
+			if err := json.Unmarshal(optionJson, &includeUsage); err != nil {
+				logger.V(logutil.DEFAULT).Error(err, "Error unmarshaling stream_options")
+			}
+			if usageEnabled, err := strconv.ParseBool(includeUsage.IncludeUsage); err != nil {
+				logger.V(logutil.DEFAULT).Error(err, "Error fetching include_usage")
+			} else {
+				reqCtx.StreamingIncludeUsage = usageEnabled
+			}
+		}
+	}
+
 	// NOTE: The nil checking for the modelObject means that we DO allow passthrough currently.
 	// This might be a security risk in the future where adapters not registered in the InferenceModel
 	// are able to be requested by using their distinct name.

Diff for: pkg/epp/handlers/response.go

@@ -20,9 +20,11 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"strings"
 
 	configPb "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
 	extProcPb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
+	"github.com/go-logr/logr"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 	errutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/error"
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
@@ -132,22 +134,19 @@ func (s *Server) HandleResponseBody(
 ) (*extProcPb.ProcessingResponse, error) {
 	logger := log.FromContext(ctx)
 	loggerVerbose := logger.V(logutil.VERBOSE)
-	loggerVerbose.Info("Processing HandleResponseBody")
 	body := req.Request.(*extProcPb.ProcessingRequest_ResponseBody)
 
-	res := Response{}
-	if err := json.Unmarshal(body.ResponseBody.Body, &res); err != nil {
-		return nil, errutil.Error{Code: errutil.Internal, Msg: fmt.Sprintf("unmarshaling response body: %v", err)}
+	if reqCtx.Streaming {
+		logger.V(logutil.DEBUG).Info("Processing HandleResponseBody")
+		if err := s.HandleStreaming(ctx, reqCtx, body, loggerVerbose); err != nil {
+			return nil, err
+		}
+	} else {
+		loggerVerbose.Info("Processing HandleResponseBody")
+		if err := s.HandleNonStreaming(ctx, reqCtx, body, loggerVerbose); err != nil {
+			return nil, err
+		}
 	}
-	reqCtx.Response = res
-	reqCtx.ResponseSize = len(body.ResponseBody.Body)
-	// ResponseComplete is to indicate the response is complete. In non-streaming
-	// case, it will be set to be true once the response is processed; in
-	// streaming case, it will be set to be true once the last chunk is processed.
-	// TODO(https://github.com/kubernetes-sigs/gateway-api-inference-extension/issues/178)
-	// will add the processing for streaming case.
-	reqCtx.ResponseComplete = true
-	loggerVerbose.Info("Response generated", "response", res)
 
 	resp := &extProcPb.ProcessingResponse{
 		Response: &extProcPb.ProcessingResponse_ResponseBody{
@@ -159,6 +158,74 @@ func (s *Server) HandleResponseBody(
 	return resp, nil
 }
 
+func (s *Server) HandleNonStreaming(
+	ctx context.Context,
+	reqCtx *RequestContext,
+	body *extProcPb.ProcessingRequest_ResponseBody,
+	loggerVerbose logr.Logger,
+) error {
+	loggerVerbose.Info("Processing HandleResponseBody")
+
+	res := Response{}
+	if err := json.Unmarshal(body.ResponseBody.Body, &res); err != nil {
+		return errutil.Error{Code: errutil.Internal, Msg: fmt.Sprintf("unmarshaling response body: %v", err)}
+	}
+	reqCtx.Response = res
+	reqCtx.ResponseSize = len(body.ResponseBody.Body)
+	reqCtx.ResponseComplete = true
+	loggerVerbose.Info("Response generated", "response", res)
+	return nil
+}
+
+func (s *Server) HandleStreaming(
+	ctx context.Context,
+	reqCtx *RequestContext,
+	body *extProcPb.ProcessingRequest_ResponseBody,
+	loggerVerbose logr.Logger,
+) error {
+	responseText := string(body.ResponseBody.Body)
+	// Example message if "stream_options": {"include_usage": "true"} is included in the request:
+	// data: {"id":"...","object":"text_completion","created":1739400043,"model":"tweet-summary-0","choices":[],
+	// "usage":{"prompt_tokens":7,"total_tokens":17,"completion_tokens":10}}
+	//
+	// data: [DONE]
+	//
+	// Noticed that vLLM returns two entries in one response.
+	// We need to strip the `data:` prefix and next Data: [DONE] from the message to fetch response data.
+	//
+	// If include_usage is not included in the request, `data: [DONE]` is returned separately, which
+	// indicates end of streaming.
+	if strings.Contains(responseText, "data: [DONE]") {
+		response := Response{}
+
+		if reqCtx.StreamingIncludeUsage {
+			lines := strings.Split(responseText, "\n")
+			if len(lines) < 2 {
+				loggerVerbose.Info("model server returns invalid response", "message", responseText)
+				return nil
+			}
+
+			content := strings.TrimPrefix(lines[0], "data: ")
+			byteSlice := []byte(content)
+			if err := json.Unmarshal(byteSlice, &response); err != nil {
+				loggerVerbose.Error(err, "unmarshaling response body")
+				return nil
+			}
+		}
+
+		reqCtx.Response = response
+	}
+
+	if body.ResponseBody.EndOfStream {
+		loggerVerbose.Info("Streaming is completed")
+		reqCtx.ResponseComplete = true
+	} else {
+		reqCtx.ResponseSize += len(body.ResponseBody.Body)
+	}
+
+	return nil
+}
+
 type Response struct {
 	Usage Usage `json:"usage"`
 }

Diff for: pkg/epp/handlers/response_test.go

@@ -49,6 +49,13 @@ const (
 		}
 	}
 	`
+
+	streamingBodyWithoutUsage = `data: {"id":"cmpl-41764c93-f9d2-4f31-be08-3ba04fa25394","object":"text_completion","created":1740002445,"model":"tweet-summary-0","choices":[],"usage":null}
+	`
+
+	streamingBodyWithUsage = `data: {"id":"cmpl-41764c93-f9d2-4f31-be08-3ba04fa25394","object":"text_completion","created":1740002445,"model":"tweet-summary-0","choices":[],"usage":{"prompt_tokens":7,"total_tokens":17,"completion_tokens":10}}
+data: [DONE]
+	`
 )
 
 func TestHandleResponseBody(t *testing.T) {
@@ -57,6 +64,7 @@ func TestHandleResponseBody(t *testing.T) {
 	tests := []struct {
 		name    string
 		req     *extProcPb.ProcessingRequest_ResponseBody
+		reqCtx  *RequestContext
 		want    Response
 		wantErr bool
 	}{
@@ -84,12 +92,49 @@ func TestHandleResponseBody(t *testing.T) {
 			},
 			wantErr: true,
 		},
+		{
+			name: "streaming request without usage",
+			req: &extProcPb.ProcessingRequest_ResponseBody{
+				ResponseBody: &extProcPb.HttpBody{
+					Body: []byte(streamingBodyWithoutUsage),
+				},
+			},
+			reqCtx: &RequestContext{
+				Streaming:             true,
+				StreamingIncludeUsage: false,
+			},
+			wantErr: false,
+			// In the middle of streaming response, so request context response is not set yet.
+		},
+		{
+			name: "streaming request with usage",
+			req: &extProcPb.ProcessingRequest_ResponseBody{
+				ResponseBody: &extProcPb.HttpBody{
+					Body: []byte(streamingBodyWithUsage),
+				},
+			},
+			reqCtx: &RequestContext{
+				Streaming:             true,
+				StreamingIncludeUsage: true,
+			},
+			wantErr: false,
+			want: Response{
+				Usage: Usage{
+					PromptTokens:     7,
+					TotalTokens:      17,
+					CompletionTokens: 10,
+				},
+			},
+		},
 	}
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
 			server := &Server{}
-			reqCtx := &RequestContext{}
+			reqCtx := test.reqCtx
+			if reqCtx == nil {
+				reqCtx = &RequestContext{}
+			}
 			_, err := server.HandleResponseBody(ctx, reqCtx, &extProcPb.ProcessingRequest{Request: test.req})
 			if err != nil {
 				if !test.wantErr {

Diff for: pkg/epp/handlers/server.go

@@ -124,7 +124,11 @@ func (s *Server) Process(srv extProcPb.ExternalProcessor_ProcessServer) error {
 				metrics.RecordInputTokens(reqCtx.Model, reqCtx.ResolvedTargetModel, reqCtx.Response.Usage.PromptTokens)
 				metrics.RecordOutputTokens(reqCtx.Model, reqCtx.ResolvedTargetModel, reqCtx.Response.Usage.CompletionTokens)
 			}
-			loggerVerbose.Info("Request context after HandleResponseBody", "context", reqCtx)
+			if reqCtx.Streaming {
+				logger.V(logutil.DEBUG).Info("Request context after HandleResponseBody", "context", reqCtx)
+			} else {
+				loggerVerbose.Info("Request context after HandleResponseBody", "context", reqCtx)
+			}
 		default:
 			logger.V(logutil.DEFAULT).Error(nil, "Unknown Request type", "request", v)
 			return status.Error(codes.Unknown, "unknown request type")
@@ -182,7 +186,11 @@ func (s *Server) Process(srv extProcPb.ExternalProcessor_ProcessServer) error {
 			}
 		}
 
-		loggerVerbose.Info("Response generated", "response", resp)
+		if !reqCtx.Streaming {
+			loggerVerbose.Info("Response generated", "response", resp)
+		} else {
+			logger.V(logutil.DEBUG).Info("Response generated", "response", resp)
+		}
 		if err := srv.Send(resp); err != nil {
 			logger.V(logutil.DEFAULT).Error(err, "Send failed")
 			return status.Errorf(codes.Unknown, "failed to send response back to Envoy: %v", err)
@@ -203,4 +211,6 @@ type RequestContext struct {
 	ResponseSize              int
 	ResponseComplete          bool
 	ResponseStatusCode        string
+	Streaming                 bool
+	StreamingIncludeUsage     bool
 }

Diff for: site-src/guides/metrics.md

@@ -4,13 +4,29 @@ This guide describes the current state of exposed metrics and how to scrape them
 
 ## Requirements
 
-Response metrics are only supported in non-streaming mode, with the follow up [issue](https://github.com/kubernetes-sigs/gateway-api-inference-extension/issues/178) to address streaming mode.
+For non-streaming request, enable `Buffered` for response in `EnvoyExtensionPolicy`:
 
-Currently there are two options:
-- If requests don't use response streaming, then you can enable `Buffered` mode for response in `EnvoyExtensionPolicy`, this will buffer the response body at the proxy and forward it to the endpoint picker, which allows the endpoint picker to report response metrics.
-
-- If requests use response streaming, then it is not recommended to enable `Buffered` mode, the response body processing mode should be left empty in the `EnvoyExtensionPolicy` (default). In this case response bodies will not be forwarded to the endpoint picker, and therefore response metrics will not be reported.
+```
+apiVersion: gateway.envoyproxy.io/v1alpha1
+kind: EnvoyExtensionPolicy
+metadata:
+  name: ext-proc-policy
+  namespace: default
+spec:
+  extProc:
+    - backendRefs:
+      - group: ""
+        kind: Service
+        name: inference-gateway-ext-proc
+        port: 9002
+      processingMode:
+        request:
+          body: Buffered
+        response:
+          body: Buffered
+```
 
+For streaming request, enable `Streamed` for response in `EnvoyExtensionPolicy`:
 
 ```
 apiVersion: gateway.envoyproxy.io/v1alpha1
@@ -29,7 +45,20 @@ spec:
         request:
           body: Buffered
         response:
-          body: Buffered
+          body: Streamed
+```
+
+If you want to include usage metrics for vLLM model server, send the request with `include_usage`:
+
+```
+curl -i ${IP}:${PORT}/v1/completions -H 'Content-Type: application/json' -d '{
+"model": "tweet-summary",
+"prompt": "whats your fav movie?",
+"max_tokens": 10,
+"temperature": 0,
+"stream": true,
+"stream_options": {"include_usage": "true"}
+}'
 ```
 
 ## Exposed metrics