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

@@ -65,6 +65,12 @@ type Datastore interface {
 	PodUpdateOrAddIfNotExist(pod *corev1.Pod, pool *v1alpha2.InferencePool) bool
 	PodDelete(namespacedName types.NamespacedName)
 	PodResyncAll(ctx context.Context, ctrlClient client.Client, pool *v1alpha2.InferencePool)
+	// PodGetRandom will grab a random pod, used for selecting a passthrough endpoint.
+	PodGetRandom() *backendmetrics.Pod
+
+	RandomWeightedDraw(logger logr.Logger, model *v1alpha2.InferenceModel) string
+	// This func should only be called for tests, to set a deterministic seed.
+	SetRand(seed int64)
 
 	// Clears the store state, happens when the pool gets deleted.
 	Clear()
@@ -77,6 +83,7 @@ func NewDatastore(parentCtx context.Context, pmf *backendmetrics.PodMetricsFacto
 		models:          make(map[string]*v1alpha2.InferenceModel),
 		pods:            &sync.Map{},
 		pmf:             pmf,
+		rand:            rand.New(rand.NewSource(rand.Int63())),
 	}
 	return store
 }
@@ -92,6 +99,7 @@ type datastore struct {
 	// key: types.NamespacedName, value: backendmetrics.PodMetrics
 	pods *sync.Map
 	pmf  *backendmetrics.PodMetricsFactory
+	rand *rand.Rand
 }
 
 func (ds *datastore) Clear() {
@@ -292,6 +300,17 @@ 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()
+}
+
+// This func should only be called for tests to set a deterministic seed
+func (ds *datastore) SetRand(seed int64) {
+	ds.rand = rand.New(rand.NewSource(seed))
+}
+
 func selectorFromInferencePoolSelector(selector map[v1alpha2.LabelKey]v1alpha2.LabelValue) labels.Selector {
 	return labels.SelectorFromSet(stripLabelKeyAliasFromLabelMap(selector))
 }
@@ -304,16 +323,10 @@ func stripLabelKeyAliasFromLabelMap(labels map[v1alpha2.LabelKey]v1alpha2.LabelV
 	return outMap
 }
 
-func RandomWeightedDraw(logger logr.Logger, model *v1alpha2.InferenceModel, seed int64) string {
-	source := rand.NewSource(rand.Int63())
-	if seed > 0 {
-		source = rand.NewSource(seed)
-	}
-	r := rand.New(source)
-
+func (ds *datastore) RandomWeightedDraw(logger logr.Logger, model *v1alpha2.InferenceModel) string {
 	// all the weight values are nil, then we should return random model name
 	if model.Spec.TargetModels[0].Weight == nil {
-		index := r.Int31n(int32(len(model.Spec.TargetModels)))
+		index := ds.rand.Int31n(int32(len(model.Spec.TargetModels)))
 		return model.Spec.TargetModels[index].Name
 	}
 
@@ -322,7 +335,7 @@ func RandomWeightedDraw(logger logr.Logger, model *v1alpha2.InferenceModel, seed
 		weights += *model.Weight
 	}
 	logger.V(logutil.TRACE).Info("Weights for model computed", "model", model.Name, "weights", weights)
-	randomVal := r.Int31n(weights)
+	randomVal := ds.rand.Int31n(weights)
 	// TODO: optimize this without using loop
 	for _, model := range model.Spec.TargetModels {
 		if randomVal < *model.Weight {

pkg/epp/datastore/datastore_test.go

@@ -313,10 +313,12 @@ func TestRandomWeightedDraw(t *testing.T) {
 		},
 	}
 	var seedVal int64 = 420
+	var ds *datastore = &datastore{}
+	ds.SetRand(seedVal)
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
 			for range 10000 {
-				model := RandomWeightedDraw(logger, test.model, seedVal)
+				model := ds.RandomWeightedDraw(logger, test.model)
 				if model != test.want {
 					t.Errorf("Model returned: %v != %v", model, test.want)
 					break

pkg/epp/handlers/request.go

@@ -69,7 +69,7 @@ func (s *Server) HandleRequestBody(
 		return nil, errutil.Error{Code: errutil.BadConfiguration, Msg: fmt.Sprintf("error finding a model object in InferenceModel for input %v", model)}
 	}
 	if len(modelObj.Spec.TargetModels) > 0 {
-		modelName = datastore.RandomWeightedDraw(logger, modelObj, 0)
+		modelName = s.datastore.RandomWeightedDraw(logger, modelObj)
 		if modelName == "" {
 			return nil, errutil.Error{Code: errutil.BadConfiguration, Msg: fmt.Sprintf("error getting target model name for model %v", modelObj.Name)}
 		}

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
+)