Merge branch 'main' into 2169-registry-capacity-counters

Author: LukeAVanDrie

pkg/epp/flowcontrol/benchmark/benchmark_test.go

@@ -346,7 +346,7 @@ func BenchmarkFlowController_FullPath(b *testing.B) {
 			//    detector reads to compute saturation.
 			infReq := &scheduling.InferenceRequest{
 				RequestID: reqID,
-				Body:      &requesthandling.InferenceRequestBody{TokenizedPrompt: &requesthandling.TokenizedPrompt{PerPromptTokens: [][]uint32{benchTokenIDs}}},
+				Body:      &requesthandling.InferenceRequestBody{TokenizedRequest: &requesthandling.TokenizedRequest{Prompts: []requesthandling.PromptTokens{{TokenIDs: benchTokenIDs}}}},
 			}
 			schedResult := &scheduling.SchedulingResult{ProfileResults: profileResults}
 			_ = h.producer.PreRequest(ctx, infReq, schedResult)

pkg/epp/flowcontrol/integration_test.go

@@ -87,7 +87,7 @@ func TestConcurrentSaturationReads(t *testing.T) {
 			req := &fwksched.InferenceRequest{
 				RequestID: fmt.Sprintf("req-%d", i),
 				Body: &fwkrh.InferenceRequestBody{
-					TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, 10)}},
+					TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, 10)}}},
 				},
 			}
 			result := &fwksched.SchedulingResult{
@@ -164,7 +164,7 @@ func TestSaturationFullLoop(t *testing.T) {
 		req := &fwksched.InferenceRequest{
 			RequestID: fmt.Sprintf("prefill-%d", i),
 			Body: &fwkrh.InferenceRequestBody{
-				TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, 50)}},
+				TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, 50)}}},
 			},
 		}
 		result := &fwksched.SchedulingResult{
@@ -453,7 +453,7 @@ func TestUsageLimitThresholdGatesDispatch(t *testing.T) {
 		req := &fwksched.InferenceRequest{
 			RequestID: fmt.Sprintf("inflight-%d", i),
 			Body: &fwkrh.InferenceRequestBody{
-				TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, 10)}},
+				TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, 10)}}},
 			},
 		}
 		result := &fwksched.SchedulingResult{
@@ -1003,7 +1003,7 @@ func TestEndpointReregistrationSaturationAccuracy(t *testing.T) {
 	oldReq := &fwksched.InferenceRequest{
 		RequestID: "old-req",
 		Body: &fwkrh.InferenceRequestBody{
-			TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, 50)}},
+			TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, 50)}}},
 		},
 	}
 	oldResult := &fwksched.SchedulingResult{
@@ -1059,7 +1059,7 @@ func TestEndpointReregistrationSaturationAccuracy(t *testing.T) {
 	newReq := &fwksched.InferenceRequest{
 		RequestID: "new-req",
 		Body: &fwkrh.InferenceRequestBody{
-			TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, 50)}},
+			TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, 50)}}},
 		},
 	}
 	newResult := &fwksched.SchedulingResult{
@@ -1120,7 +1120,7 @@ func TestEndpointIdentityCollisionDuringPodReplacement(t *testing.T) {
 	req := &fwksched.InferenceRequest{
 		RequestID: "new-pod-req",
 		Body: &fwkrh.InferenceRequestBody{
-			TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, 50)}},
+			TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, 50)}}},
 		},
 	}
 	result := &fwksched.SchedulingResult{

pkg/epp/framework/interface/requesthandling/types.go

@@ -113,9 +113,9 @@ type InferenceRequestBody struct {
 	// If the payload is unmarshaled, we can perform advanced processing (like prefix cache aware routing).
 	// If it remains as raw bytes, such processing may not be supported.
 	Payload RequestPayload `json:"-"`
-	// TokenizedPrompt contains parser-derived tokenization results when available.
+	// TokenizedRequest contains parser-derived tokenization results when available.
 	// It is nil when the request was not already tokenized.
-	TokenizedPrompt *TokenizedPrompt `json:"-"`
+	TokenizedRequest *TokenizedRequest `json:"-"`
 
 	// Stream indicates whether the request specifies a streaming response (e.g., via a stream field).
 	// This typically implies the model server's response will be streamed.
@@ -189,32 +189,36 @@ func MaxOutputTokensFromPayload(m PayloadMap, keys ...string) *int64 {
 	return nil
 }
 
-// TokenizedPrompt contains the result of tokenizing the request prompt.
+// TokenizedRequest contains the result of tokenizing the request prompt.
 // It is consumed by scheduling and request-control plugins that benefit from
 // actual token data such as prefix-cache awareness.
-type TokenizedPrompt struct {
-	// PerPromptTokens holds the token IDs for each prompt in the request.
-	// Single-prompt requests (chat, generate, single-string completions) use a
-	// length-1 outer slice. Multi-string completions use one inner slice per
-	// prompt string.
-	PerPromptTokens [][]uint32
-	// MultiModalFeatures holds one entry per multimodal item in prompt order.
-	// Nil if the prompt contains no multimodal content. Offsets are relative
-	// to PerPromptTokens[0] (always single-prompt when multimodal content is
-	// present).
-	MultiModalFeatures []MultiModalFeature
+type TokenizedRequest struct {
+	// Prompts holds the per-prompt token data. Single-prompt requests (chat,
+	// generate, single-string completions) use a length-1 slice. Multi-string
+	// completions use one entry per prompt string.
+	Prompts []PromptTokens
 	// CacheSalt isolates prefix caches across requests. Populated by the token-producer.
 	CacheSalt string
 }
 
+// PromptTokens bundles the token IDs and multimodal features for a single
+// prompt in the request.
+type PromptTokens struct {
+	// TokenIDs holds the token IDs for this prompt.
+	TokenIDs []uint32
+	// MultiModalFeatures holds multimodal items for this prompt, ordered by
+	// token position. Nil if the prompt contains no multimodal content.
+	MultiModalFeatures []MultiModalFeature
+}
+
 // TokenCount returns the total number of tokens across all prompts.
-func (tp *TokenizedPrompt) TokenCount() int {
+func (tp *TokenizedRequest) TokenCount() int {
 	if tp == nil {
 		return 0
 	}
 	n := 0
-	for _, pp := range tp.PerPromptTokens {
-		n += len(pp)
+	for _, p := range tp.Prompts {
+		n += len(p.TokenIDs)
 	}
 	return n
 }
@@ -227,9 +231,10 @@ type MultiModalFeature struct {
 	Modality Modality
 	// Hash is the content hash of the item, used for KV-cache reuse across requests.
 	Hash string
-	// Offset is the index of the first placeholder token for this item in TokenIDs.
+	// Offset is the index of the first placeholder token for this item
+	// in the owning PromptTokens.TokenIDs slice.
 	Offset int
-	// Length is the number of placeholder tokens this item occupies in TokenIDs.
+	// Length is the number of placeholder tokens this item occupies.
 	Length int
 }
 

pkg/epp/framework/interface/scheduling/types.go

@@ -33,7 +33,7 @@ type Modality = fwkrh.Modality
 
 const ModalityImage = fwkrh.ModalityImage
 
-type TokenizedPrompt = fwkrh.TokenizedPrompt
+type TokenizedRequest = fwkrh.TokenizedRequest
 
 type MultiModalFeature = fwkrh.MultiModalFeature
 

pkg/epp/framework/observability/multimodal/multimodal.go

@@ -67,8 +67,12 @@ func Summary(req *scheduling.InferenceRequest) (modality string, hashCount int)
 }
 
 func requestMMFeatures(req *scheduling.InferenceRequest) []fwkrh.MultiModalFeature {
-	if req == nil || req.Body == nil || req.Body.TokenizedPrompt == nil {
+	if req == nil || req.Body == nil || req.Body.TokenizedRequest == nil {
 		return nil
 	}
-	return req.Body.TokenizedPrompt.MultiModalFeatures
+	var features []fwkrh.MultiModalFeature
+	for _, p := range req.Body.TokenizedRequest.Prompts {
+		features = append(features, p.MultiModalFeatures...)
+	}
+	return features
 }

pkg/epp/framework/observability/multimodal/multimodal_test.go

@@ -109,8 +109,10 @@ func TestSpanAttributes(t *testing.T) {
 func requestWith(features ...fwkrh.MultiModalFeature) *scheduling.InferenceRequest {
 	return &scheduling.InferenceRequest{
 		Body: &fwkrh.InferenceRequestBody{
-			TokenizedPrompt: &fwkrh.TokenizedPrompt{
-				MultiModalFeatures: features,
+			TokenizedRequest: &fwkrh.TokenizedRequest{
+				Prompts: []fwkrh.PromptTokens{{
+					MultiModalFeatures: features,
+				}},
 			},
 		},
 	}

pkg/epp/framework/plugins/flowcontrol/saturationdetector/concurrency/detector_test.go

@@ -809,7 +809,7 @@ func makeTokenRequest(requestID string, inputTokens int) *fwksched.InferenceRequ
 	return &fwksched.InferenceRequest{
 		RequestID: requestID,
 		Body: &fwkrh.InferenceRequestBody{
-			TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{make([]uint32, inputTokens)}},
+			TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: make([]uint32, inputTokens)}}},
 		},
 	}
 }

pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin.go

@@ -126,12 +126,12 @@ func (p *dataProducer) Produces() map[plugin.DataKey]any {
 	return map[plugin.DataKey]any{p.dk: attrprefix.PrefixCacheMatchInfo{}}
 }
 
-// Consumes declares the TokenizedPrompt dependency so the data-layer DAG orders
+// Consumes declares the TokenizedRequest dependency so the data-layer DAG orders
 // the token-producer before this producer runs and auto-creates one when none
 // is configured.
 func (p *dataProducer) Consumes() plugin.DataDependencies {
 	return plugin.DataDependencies{
-		Required: map[plugin.DataKey]any{tokenproducer.TokenizedPromptDataKey: fwksched.TokenizedPrompt{}},
+		Required: map[plugin.DataKey]any{tokenproducer.TokenizedPromptDataKey: fwksched.TokenizedRequest{}},
 	}
 }
 

pkg/epp/framework/plugins/requestcontrol/dataproducer/approximateprefix/plugin_test.go

@@ -53,7 +53,7 @@ func disableMinBlockSizeClamp(t *testing.T) {
 // tokenizedBody returns a request body carrying only a tokenized prompt.
 func tokenizedBody(tokenIDs []uint32) *fwkrh.InferenceRequestBody {
 	return &fwkrh.InferenceRequestBody{
-		TokenizedPrompt: &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{tokenIDs}},
+		TokenizedRequest: &fwkrh.TokenizedRequest{Prompts: []fwkrh.PromptTokens{{TokenIDs: tokenIDs}}},
 	}
 }
 
@@ -663,8 +663,8 @@ func TestProduce_MultiPrompt(t *testing.T) {
 		RequestID:   uuid.NewString(),
 		TargetModel: "test-model",
 		Body: &fwkrh.InferenceRequestBody{
-			TokenizedPrompt: &fwkrh.TokenizedPrompt{
-				PerPromptTokens: [][]uint32{{1, 2, 3}, {4, 5}},
+			TokenizedRequest: &fwkrh.TokenizedRequest{
+				Prompts: []fwkrh.PromptTokens{{TokenIDs: []uint32{1, 2, 3}}, {TokenIDs: []uint32{4, 5}}},
 			},
 		},
 	}
@@ -706,8 +706,8 @@ func TestMultiPromptMatchAggregation(t *testing.T) {
 		RequestID:   uuid.NewString(),
 		TargetModel: "test-model",
 		Body: &fwkrh.InferenceRequestBody{
-			TokenizedPrompt: &fwkrh.TokenizedPrompt{
-				PerPromptTokens: [][]uint32{{1, 2, 3}, {4, 5}},
+			TokenizedRequest: &fwkrh.TokenizedRequest{
+				Prompts: []fwkrh.PromptTokens{{TokenIDs: []uint32{1, 2, 3}}, {TokenIDs: []uint32{4, 5}}},
 			},
 		},
 	}
@@ -725,8 +725,8 @@ func TestMultiPromptMatchAggregation(t *testing.T) {
 		RequestID:   uuid.NewString(),
 		TargetModel: "test-model",
 		Body: &fwkrh.InferenceRequestBody{
-			TokenizedPrompt: &fwkrh.TokenizedPrompt{
-				PerPromptTokens: [][]uint32{{1, 2, 3}, {4, 5}},
+			TokenizedRequest: &fwkrh.TokenizedRequest{
+				Prompts: []fwkrh.PromptTokens{{TokenIDs: []uint32{1, 2, 3}}, {TokenIDs: []uint32{4, 5}}},
 			},
 		},
 	}
@@ -759,8 +759,8 @@ func TestMultiPromptPartialMatch(t *testing.T) {
 		RequestID:   uuid.NewString(),
 		TargetModel: "test-model",
 		Body: &fwkrh.InferenceRequestBody{
-			TokenizedPrompt: &fwkrh.TokenizedPrompt{
-				PerPromptTokens: [][]uint32{{1, 2}, {3, 4}},
+			TokenizedRequest: &fwkrh.TokenizedRequest{
+				Prompts: []fwkrh.PromptTokens{{TokenIDs: []uint32{1, 2}}, {TokenIDs: []uint32{3, 4}}},
 			},
 		},
 	}
@@ -778,8 +778,8 @@ func TestMultiPromptPartialMatch(t *testing.T) {
 		RequestID:   uuid.NewString(),
 		TargetModel: "test-model",
 		Body: &fwkrh.InferenceRequestBody{
-			TokenizedPrompt: &fwkrh.TokenizedPrompt{
-				PerPromptTokens: [][]uint32{{1, 2}, {5, 6}},
+			TokenizedRequest: &fwkrh.TokenizedRequest{
+				Prompts: []fwkrh.PromptTokens{{TokenIDs: []uint32{1, 2}}, {TokenIDs: []uint32{5, 6}}},
 			},
 		},
 	}

pkg/epp/framework/plugins/requestcontrol/dataproducer/burstprefix/plugin.go

@@ -73,11 +73,11 @@ func (p *dataProducer) Produces() map[plugin.DataKey]any {
 	return map[plugin.DataKey]any{p.dk: attrprefix.PrefixCacheMatchInfo{}}
 }
 
-// Consumes declares the TokenizedPrompt dependency so the token-producer runs
+// Consumes declares the TokenizedRequest dependency so the token-producer runs
 // before this producer and one is auto-created when none is configured.
 func (p *dataProducer) Consumes() plugin.DataDependencies {
 	return plugin.DataDependencies{
-		Required: map[plugin.DataKey]any{tokenproducer.TokenizedPromptDataKey: fwksched.TokenizedPrompt{}},
+		Required: map[plugin.DataKey]any{tokenproducer.TokenizedPromptDataKey: fwksched.TokenizedRequest{}},
 	}
 }