Allow requests with no body to passthrough EPP

Author: kfswain

config/manifests/inferencepool.yaml

@@ -43,7 +43,7 @@ spec:
     spec:
       containers:
       - name: epp
-        image: us-east1-docker.pkg.dev/kfswain-gke-dev/test-repo/ext-proc:test-mar-51
+        image: us-east1-docker.pkg.dev/kfswain-gke-dev/test-repo/ext-proc:test-mar-53
         imagePullPolicy: Always
         args:
         - -poolName

pkg/epp/datastore/datastore.go

@@ -62,6 +62,8 @@ type Datastore interface {
 	PodGetAll() []backendmetrics.PodMetrics
 	// PodList lists pods matching the given predicate.
 	PodList(func(backendmetrics.PodMetrics) bool) []backendmetrics.PodMetrics
+	// PodGetRandom will grab a random pod, used for selecting a passthrough endpoint.
+	PodGetRandom() *backendmetrics.Pod
 	PodUpdateOrAddIfNotExist(pod *corev1.Pod, pool *v1alpha2.InferencePool) bool
 	PodDelete(namespacedName types.NamespacedName)
 	PodResyncAll(ctx context.Context, ctrlClient client.Client, pool *v1alpha2.InferencePool)
@@ -292,6 +294,12 @@ func (ds *datastore) PodDelete(namespacedName types.NamespacedName) {
 	}
 }
 
+func (ds *datastore) PodGetRandom() *backendmetrics.Pod {
+	pods := ds.PodGetAll()
+	pod := pods[rand.Intn(len(pods))]
+	return pod.GetPod()
+}
+
 func selectorFromInferencePoolSelector(selector map[v1alpha2.LabelKey]v1alpha2.LabelValue) labels.Selector {
 	return labels.SelectorFromSet(stripLabelKeyAliasFromLabelMap(selector))
 }

pkg/epp/handlers/response.go

@@ -85,9 +85,7 @@ func (s *Server) HandleResponseHeaders(
 		if header.Key == "content-type" {
 			contentType := header.RawValue
 			if strings.Contains(string(contentType), "text/event-stream") {
-				reqCtx.Streaming = true
-			} else {
-				reqCtx.Streaming = false
+				reqCtx.modelServerStreaming = true
 			}
 			typeFound = true
 		}
@@ -155,7 +153,7 @@ func (s *Server) HandleResponseBody(
 	loggerVerbose := logger.V(logutil.VERBOSE)
 	body := req.Request.(*extProcPb.ProcessingRequest_ResponseBody)
 
-	if reqCtx.Streaming {
+	if reqCtx.modelServerStreaming {
 		logger.V(logutil.DEBUG).Info("Processing HandleResponseBody")
 		if err := s.HandleStreaming(ctx, reqCtx, body, loggerVerbose); err != nil {
 			return nil, err
@@ -189,7 +187,7 @@ func (s *Server) HandleNonStreaming(
 	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.Usage = res.Usage
 	reqCtx.ResponseSize = len(body.ResponseBody.Body)
 	reqCtx.ResponseComplete = true
 	loggerVerbose.Info("Response generated", "response", res)
@@ -205,7 +203,7 @@ func (s *Server) HandleStreaming(
 	responseText := string(body.ResponseBody.Body)
 	if strings.Contains(responseText, streamingEndMsg) {
 		parsedResp := ParseRespForUsage(ctx, responseText, loggerVerbose)
-		reqCtx.Response = parsedResp
+		reqCtx.Usage = parsedResp.Usage
 	}
 
 	if body.ResponseBody.EndOfStream {

pkg/epp/handlers/response_test.go

@@ -65,7 +65,7 @@ func TestHandleResponseBody(t *testing.T) {
 		name    string
 		req     *extProcPb.ProcessingRequest_ResponseBody
 		reqCtx  *RequestContext
-		want    Response
+		want    Usage
 		wantErr bool
 	}{
 		{
@@ -75,12 +75,10 @@ func TestHandleResponseBody(t *testing.T) {
 					Body: []byte(body),
 				},
 			},
-			want: Response{
-				Usage: Usage{
-					PromptTokens:     11,
-					TotalTokens:      111,
-					CompletionTokens: 100,
-				},
+			want: Usage{
+				PromptTokens:     11,
+				TotalTokens:      111,
+				CompletionTokens: 100,
 			},
 		},
 		{
@@ -100,7 +98,7 @@ func TestHandleResponseBody(t *testing.T) {
 				},
 			},
 			reqCtx: &RequestContext{
-				Streaming: true,
+				modelServerStreaming: true,
 			},
 			wantErr: false,
 			// In the middle of streaming response, so request context response is not set yet.
@@ -113,15 +111,13 @@ func TestHandleResponseBody(t *testing.T) {
 				},
 			},
 			reqCtx: &RequestContext{
-				Streaming: true,
+				modelServerStreaming: true,
 			},
 			wantErr: false,
-			want: Response{
-				Usage: Usage{
-					PromptTokens:     7,
-					TotalTokens:      17,
-					CompletionTokens: 10,
-				},
+			want: Usage{
+				PromptTokens:     7,
+				TotalTokens:      17,
+				CompletionTokens: 10,
 			},
 		},
 	}
@@ -141,7 +137,7 @@ func TestHandleResponseBody(t *testing.T) {
 				return
 			}
 
-			if diff := cmp.Diff(test.want, reqCtx.Response); diff != "" {
+			if diff := cmp.Diff(test.want, reqCtx.Usage); diff != "" {
 				t.Errorf("HandleResponseBody returned unexpected response, diff(-want, +got): %v", diff)
 			}
 		})

pkg/epp/handlers/server.go

@@ -128,10 +128,10 @@ func (s *Server) Process(srv extProcPb.ExternalProcessor_ProcessServer) error {
 				reqCtx.ResponseCompleteTimestamp = time.Now()
 				metrics.RecordRequestLatencies(ctx, reqCtx.Model, reqCtx.ResolvedTargetModel, reqCtx.RequestReceivedTimestamp, reqCtx.ResponseCompleteTimestamp)
 				metrics.RecordResponseSizes(reqCtx.Model, reqCtx.ResolvedTargetModel, reqCtx.ResponseSize)
-				metrics.RecordInputTokens(reqCtx.Model, reqCtx.ResolvedTargetModel, reqCtx.Response.Usage.PromptTokens)
-				metrics.RecordOutputTokens(reqCtx.Model, reqCtx.ResolvedTargetModel, reqCtx.Response.Usage.CompletionTokens)
+				metrics.RecordInputTokens(reqCtx.Model, reqCtx.ResolvedTargetModel, reqCtx.Usage.PromptTokens)
+				metrics.RecordOutputTokens(reqCtx.Model, reqCtx.ResolvedTargetModel, reqCtx.Usage.CompletionTokens)
 			}
-			if reqCtx.Streaming {
+			if reqCtx.modelServerStreaming {
 				logger.V(logutil.DEBUG).Info("Request context after HandleResponseBody", "context", reqCtx)
 			} else {
 				loggerVerbose.Info("Request context after HandleResponseBody", "context", reqCtx)
@@ -149,7 +149,7 @@ func (s *Server) Process(srv extProcPb.ExternalProcessor_ProcessServer) error {
 			}
 		}
 
-		if !reqCtx.Streaming {
+		if !reqCtx.modelServerStreaming {
 			loggerVerbose.Info("Response generated", "response", resp)
 		} else {
 			logger.V(logutil.DEBUG).Info("Response generated", "response", resp)
@@ -224,9 +224,32 @@ type RequestContext struct {
 	RequestReceivedTimestamp  time.Time
 	ResponseCompleteTimestamp time.Time
 	RequestSize               int
-	Response                  Response
+	Usage                     Usage
 	ResponseSize              int
 	ResponseComplete          bool
 	ResponseStatusCode        string
-	Streaming                 bool
+
+	RequestState         StreamRequestState
+	modelServerStreaming bool
+
+	reqHeaderResp  *extProcPb.ProcessingResponse
+	reqBodyResp    *extProcPb.ProcessingResponse
+	reqTrailerResp *extProcPb.ProcessingResponse
+
+	respHeaderResp  *extProcPb.ProcessingResponse
+	respBodyResp    *extProcPb.ProcessingResponse
+	respTrailerResp *extProcPb.ProcessingResponse
 }
+
+type StreamRequestState int
+
+const (
+	RequestReceived                  StreamRequestState = 0
+	HeaderRequestResponseComplete    StreamRequestState = 1
+	BodyRequestResponsesComplete     StreamRequestState = 2
+	TrailerRequestResponsesComplete  StreamRequestState = 3
+	ResponseRecieved                 StreamRequestState = 4
+	HeaderResponseResponseComplete   StreamRequestState = 5
+	BodyResponseResponsesComplete    StreamRequestState = 6
+	TrailerResponseResponsesComplete StreamRequestState = 7
+)