diff --git a/pkg/coordinator/gateway/client.go b/pkg/coordinator/gateway/client.go index a07b28dca5..d4b7835854 100644 --- a/pkg/coordinator/gateway/client.go +++ b/pkg/coordinator/gateway/client.go @@ -92,7 +92,7 @@ func (c *Client) Request(ctx context.Context, method, path string, body []byte, logger := log.FromContext(ctx).WithName("gateway") if body != nil { if v := logger.V(logutil.TRACE); v.Enabled() { - v.Info("request body", "method", method, "path", path, "headers", httplog.RedactedHeaders(req.Header), "body", redactBody(body)) + v.Info("request body", "method", method, "path", path, "headers", httplog.RedactedHeaders(req.Header), "body", RedactBody(body)) } } @@ -109,7 +109,7 @@ func (c *Client) Request(ctx context.Context, method, path string, body []byte, if err != nil { return nil, fmt.Errorf("reading response from gateway: %w", err) } - v.Info("response body", "status", resp.StatusCode, "body", redactBody(respBody)) + v.Info("response body", "status", resp.StatusCode, "body", RedactBody(respBody)) resp.Body = io.NopCloser(bytes.NewReader(respBody)) } @@ -146,10 +146,11 @@ const ( maxRedactDepth = 32 ) -// redactBody parses JSON and redacts oversized string values (see +// RedactBody parses JSON and redacts oversized string values (see // maxRedactStringLen); a body that is not valid JSON is returned verbatim, or -// truncated at maxRedactRawBodyLen. -func redactBody(data []byte) any { +// truncated at maxRedactRawBodyLen. It is exported for the decode reverse-proxy +// path, which bypasses Client.Request and so redacts its own body before logging. +func RedactBody(data []byte) any { var v any if err := json.Unmarshal(data, &v); err != nil { if len(data) > maxRedactRawBodyLen { diff --git a/pkg/coordinator/gateway/client_test.go b/pkg/coordinator/gateway/client_test.go index cf4a418f38..1bf557d929 100644 --- a/pkg/coordinator/gateway/client_test.go +++ b/pkg/coordinator/gateway/client_test.go @@ -147,9 +147,9 @@ func TestRedactBody(t *testing.T) { t.Fatal(err) } - got, ok := redactBody(body).(map[string]any) + got, ok := RedactBody(body).(map[string]any) if !ok { - t.Fatalf("expected map[string]any, got %T", redactBody(body)) + t.Fatalf("expected map[string]any, got %T", RedactBody(body)) } if got["model"] != "m" { t.Errorf("short value changed: %v", got["model"]) @@ -160,23 +160,23 @@ func TestRedactBody(t *testing.T) { }) t.Run("non-JSON under max length returned verbatim", func(t *testing.T) { - if got := redactBody([]byte("not json")); got != "not json" { + if got := RedactBody([]byte("not json")); got != "not json" { t.Errorf("redactBody = %v, want verbatim string", got) } }) t.Run("non-JSON exactly max length returned verbatim", func(t *testing.T) { raw := strings.Repeat("q", maxRedactRawBodyLen) - if got := redactBody([]byte(raw)); got != raw { + if got := RedactBody([]byte(raw)); got != raw { t.Errorf("redactBody = %v, want verbatim (boundary is > not >=)", got) } }) t.Run("non-JSON over max length truncated", func(t *testing.T) { raw := strings.Repeat("q", maxRedactRawBodyLen+50) - got, ok := redactBody([]byte(raw)).(string) + got, ok := RedactBody([]byte(raw)).(string) if !ok { - t.Fatalf("expected string, got %T", redactBody([]byte(raw))) + t.Fatalf("expected string, got %T", RedactBody([]byte(raw))) } if want := raw[:maxRedactRawBodyLen] + "..."; got != want { t.Errorf("redactBody truncation = %q, want %q", got, want) diff --git a/pkg/coordinator/steps/decode_proxy.go b/pkg/coordinator/steps/decode_proxy.go index e6783c5b4e..e7c42866f1 100644 --- a/pkg/coordinator/steps/decode_proxy.go +++ b/pkg/coordinator/steps/decode_proxy.go @@ -78,6 +78,14 @@ func newDecodeProxyRequest(ctx context.Context, logger logr.Logger, step string, v.Info("request body", "method", "POST", "path", reqCtx.OriginalPath, "bodyLen", len(bodyBytes), "headers", httplog.RedactedHeaders(proxyReq.Header)) } + // The decode leg is proxied straight to the gateway and never passes through + // gateway.Client.Request, which is what logs full request bodies at TRACE for + // the prefill and encode legs. Log the redacted body here so the decode leg's + // fields (min_tokens, max_tokens, ...) are observable at the same verbosity. + if v := logger.V(logutil.TRACE); v.Enabled() { + v.Info("request body", "method", "POST", "path", reqCtx.OriginalPath, "headers", httplog.RedactedHeaders(proxyReq.Header), "body", gateway.RedactBody(bodyBytes)) + } + return proxyReq, nil } diff --git a/pkg/coordinator/steps/encode.go b/pkg/coordinator/steps/encode.go index 2556003e13..ccb99ad349 100644 --- a/pkg/coordinator/steps/encode.go +++ b/pkg/coordinator/steps/encode.go @@ -215,10 +215,10 @@ func (s *EncodeStep) buildEncodeBody(reqCtx *pipeline.RequestContext, tokenIDs [ }, }, } - reqcommon.PrimeSingleTokenRequest(body, reqCtx.Body) + capSingleTokenOutput(body, format) return body default: - return map[string]any{ + body := map[string]any{ "model": reqCtx.Model, "token_ids": tokenIDs, "features": map[string]any{ @@ -226,8 +226,9 @@ func (s *EncodeStep) buildEncodeBody(reqCtx *pipeline.RequestContext, tokenIDs [ "mm_placeholders": map[string][]any{ModalityImage: {map[string]any{"offset": 1, "length": entry.Placeholder.Length}}}, "kwargs_data": mmKwargsField([]string{entry.KwargsData}), }, - reqcommon.FieldSamplingParams: map[string]any{reqcommon.FieldMaxTokens: 1}, } + capSingleTokenOutput(body, format) + return body } } diff --git a/pkg/coordinator/steps/encode_test.go b/pkg/coordinator/steps/encode_test.go index 48a5a7827f..be9732bc9f 100644 --- a/pkg/coordinator/steps/encode_test.go +++ b/pkg/coordinator/steps/encode_test.go @@ -352,10 +352,10 @@ func TestEncodeStep_ChatCompletionsFormat(t *testing.T) { } // TestEncodeStep_ChatCompletionsFormat_CapsMaxCompletionTokens is a -// regression test: buildEncodeBody used to always emit max_tokens=1 without -// regard to which token-limit field the client used, so a reasoning-model -// client's max_completion_tokens was never capped in the encoder sub-request. -func TestEncodeStep_ChatCompletionsFormat_CapsMaxCompletionTokens(t *testing.T) { +// The encode chat sub-request is built fresh from the request context and does +// not carry the client's sampling fields, so max_completion_tokens is not +// propagated and is never injected: max_tokens=1 alone caps output. +func TestEncodeStep_ChatCompletionsFormat_OmitsMaxCompletionTokens(t *testing.T) { var receivedBody map[string]any server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -401,8 +401,11 @@ func TestEncodeStep_ChatCompletionsFormat_CapsMaxCompletionTokens(t *testing.T) t.Fatalf("unexpected error: %v", err) } - if receivedBody["max_completion_tokens"] != float64(1) { - t.Fatalf("expected encode sub-request max_completion_tokens capped to 1, got %v", receivedBody["max_completion_tokens"]) + if receivedBody["max_tokens"] != float64(1) { + t.Fatalf("expected encode sub-request max_tokens capped to 1, got %v", receivedBody["max_tokens"]) + } + if _, ok := receivedBody["max_completion_tokens"]; ok { + t.Fatalf("expected encode sub-request to omit max_completion_tokens, got %v", receivedBody["max_completion_tokens"]) } } @@ -575,3 +578,44 @@ func TestEncodeStep_BuildsCorrectTokenIDs(t *testing.T) { } } } + +// TestEncodeStep_GenerateFormat_CapsSingleToken verifies the generate-format +// encoder sub-request caps output to a single token: sampling_params carries +// max_tokens=1 and strips min_tokens (it defaults to 0, keeping min_tokens <= +// max_tokens). +func TestEncodeStep_GenerateFormat_CapsSingleToken(t *testing.T) { + var samplingParams map[string]any + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + var parsed map[string]any + _ = json.Unmarshal(body, &parsed) + samplingParams, _ = parsed["sampling_params"].(map[string]any) + _ = json.NewEncoder(w).Encode(map[string]any{ + "ec_transfer_params": map[string]any{"h1": map[string]any{"peer_port": 5501}}, + }) + })) + defer server.Close() + + gwClient := gateway.New(config.GatewayConfig{Address: server.URL}) + step, _ := NewEncodeStep(gwClient, map[string]any{"use_openai_format": false}) + + reqCtx := &pipeline.RequestContext{ + RequestID: "req-gen-cap", + Model: "test", + TokenIDs: []int{1, 32000, 32000, 32000, 2345}, + MultimodalEntries: []pipeline.MultimodalEntry{ + {Index: 0, Hash: "h1", KwargsData: "dGVzdA==", Placeholder: pipeline.PlaceholderRange{Offset: 1, Length: 3}}, + }, + } + + if err := step.Execute(context.Background(), reqCtx); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if samplingParams["max_tokens"] != float64(1) { + t.Fatalf("expected sampling_params.max_tokens=1, got %v", samplingParams["max_tokens"]) + } + if _, ok := samplingParams["min_tokens"]; ok { + t.Fatalf("expected sampling_params.min_tokens to be stripped, got %v", samplingParams["min_tokens"]) + } +} diff --git a/pkg/coordinator/steps/prefill.go b/pkg/coordinator/steps/prefill.go index 167ad2caf4..e003ee6e5d 100644 --- a/pkg/coordinator/steps/prefill.go +++ b/pkg/coordinator/steps/prefill.go @@ -137,7 +137,7 @@ func (s *PrefillStep) buildPrefillBody(ctx context.Context, reqCtx *pipeline.Req switch format { case gateway.FormatChatCompletions: body := maps.Clone(reqCtx.Body) - reqcommon.PrimeSingleTokenRequest(body, reqCtx.Body) + capSingleTokenOutput(body, format) tokens := map[string]any{ "token_ids": reqCtx.TokenIDs, } @@ -164,9 +164,9 @@ func (s *PrefillStep) buildPrefillBody(ctx context.Context, reqCtx *pipeline.Req "request_id": reqCtx.RequestID, "model": reqCtx.Model, "prompt": prompt, - reqcommon.FieldMaxTokens: 1, reqcommon.FieldKVTransferParams: kvParams, } + capSingleTokenOutput(body, format) if features != nil { body["features"] = features } @@ -186,6 +186,7 @@ func (s *PrefillStep) buildPrefillBody(ctx context.Context, reqCtx *pipeline.Req "model": reqCtx.Model, reqcommon.FieldSamplingParams: sampling, } + capSingleTokenOutput(body, format) if features != nil { body["features"] = features } diff --git a/pkg/coordinator/steps/prefill_test.go b/pkg/coordinator/steps/prefill_test.go index f352a7b5b6..9ab2e46d96 100644 --- a/pkg/coordinator/steps/prefill_test.go +++ b/pkg/coordinator/steps/prefill_test.go @@ -133,6 +133,9 @@ func TestPrefillStep_SendsCorrectGenerateRequest(t *testing.T) { if samplingParams["max_tokens"] != float64(1) { t.Fatalf("expected sampling_params.max_tokens=1, got %v", samplingParams["max_tokens"]) } + if _, ok := samplingParams["min_tokens"]; ok { + t.Fatalf("expected sampling_params.min_tokens to be stripped, got %v", samplingParams["min_tokens"]) + } extraArgs, ok := samplingParams["extra_args"].(map[string]any) if !ok { t.Fatal("expected sampling_params.extra_args in generate format") @@ -222,6 +225,14 @@ func TestPrefillStep_CompletionsFormat(t *testing.T) { if prefillBody["request_id"] != "req-compl" { t.Fatalf("expected request_id, got %v", prefillBody["request_id"]) } + // Prefill leg caps output to a single token: max_tokens is pinned to 1 and + // min_tokens is stripped (it defaults to 0, keeping min_tokens <= max_tokens). + if prefillBody["max_tokens"] != float64(1) { + t.Fatalf("expected max_tokens=1, got %v", prefillBody["max_tokens"]) + } + if _, ok := prefillBody["min_tokens"]; ok { + t.Fatalf("expected min_tokens to be stripped, got %v", prefillBody["min_tokens"]) + } // Completions format has top-level kv_transfer_params kvParams, ok := prefillBody["kv_transfer_params"].(map[string]any) if !ok { @@ -470,6 +481,56 @@ func TestPrefillStep_ChatCompletionsFormat_CapsMaxCompletionTokens(t *testing.T) } } +// TestPrefillStep_ChatCompletionsFormat_StripsClientMinTokens is a regression +// test for the one coordinator path where a client-supplied min_tokens survives +// into the capped body: chat-completions clones reqCtx.Body, so a client +// min_tokens > 1 would leave min_tokens > max_tokens=1 and vLLM rejects the leg. +// The generate and completions legs build fresh bodies that never carry it. +func TestPrefillStep_ChatCompletionsFormat_StripsClientMinTokens(t *testing.T) { + var prefillBody map[string]any + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + _ = json.Unmarshal(body, &prefillBody) + _ = json.NewEncoder(w).Encode(map[string]any{ + "kv_transfer_params": map[string]any{"block_id": "block-5"}, + }) + })) + defer server.Close() + + gwClient := gateway.New(config.GatewayConfig{Address: server.URL}) + step, err := NewPrefillStep(gwClient, map[string]any{}) + if err != nil { + t.Fatal(err) + } + + reqCtx := &pipeline.RequestContext{ + RequestID: "req-chat-min-tokens", + OriginalPath: gateway.PathChatCompletions, + Model: "test-model", + Body: map[string]any{ + "model": "test-model", + "min_tokens": 50, + "max_tokens": 200, + "messages": []any{ + map[string]any{"role": "user", "content": "hello"}, + }, + }, + KVTransferParams: make(map[string]any), + } + + if err := step.Execute(context.Background(), reqCtx); err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if prefillBody["max_tokens"] != float64(1) { + t.Fatalf("expected prefill request max_tokens=1, got %v", prefillBody["max_tokens"]) + } + if _, ok := prefillBody["min_tokens"]; ok { + t.Fatalf("expected client min_tokens to be stripped, got %v", prefillBody["min_tokens"]) + } +} + // TestSharedStorage_OmitsECTransferParams_InPrefillBody verifies that the // ec-shared-storage EC connector never emits an ec_transfer_params field on // the prefill body, in every prefill wire format (chat-completions, diff --git a/pkg/coordinator/steps/utils.go b/pkg/coordinator/steps/utils.go index f632088696..3db8cf9aca 100644 --- a/pkg/coordinator/steps/utils.go +++ b/pkg/coordinator/steps/utils.go @@ -74,6 +74,40 @@ func resolveFormat(useOpenAIFormat bool, path string) gateway.RequestFormat { return detected } +// capSingleTokenOutput rewrites body into a single-output-token, non-streaming +// request for the synthetic prefill and encode legs. +// +// Intentionally distinct from the sidecar's reqcommon.PrimeSingleTokenRequest +// for now. +// TODO: unify the two into one shared single-token helper in a future refactor. +func capSingleTokenOutput(body map[string]any, format gateway.RequestFormat) { + target := body + if format == gateway.FormatGenerate { + sp, ok := body[reqcommon.FieldSamplingParams].(map[string]any) + if !ok { + sp = map[string]any{} + body[reqcommon.FieldSamplingParams] = sp + } + target = sp + } + + target[reqcommon.FieldMaxTokens] = 1 + // Strip rather than clamp min_tokens: it defaults to 0 in vLLM, so removing it + // keeps min_tokens <= max_tokens=1 without raising the floor above the cap. + delete(target, reqcommon.FieldMinTokens) + + if _, ok := body[reqcommon.FieldMaxCompletionTokens]; ok { + body[reqcommon.FieldMaxCompletionTokens] = 1 + } + + // TODO: max_output_tokens is another client-supplied output cap (Responses + // API) that a client can send instead of max_tokens/max_completion_tokens; it + // should be capped to 1 here as well so the synthetic legs stay single-token. + + body[reqcommon.FieldStream] = false + delete(body, reqcommon.FieldStreamOptions) +} + // buildMMFeatures builds the multimodal features map (mm_hashes, mm_placeholders, // and optionally kwargs_data) from the request's multimodal entries. It returns // nil when there are no entries. diff --git a/pkg/coordinator/steps/utils_test.go b/pkg/coordinator/steps/utils_test.go index 0629882144..c482d6604c 100644 --- a/pkg/coordinator/steps/utils_test.go +++ b/pkg/coordinator/steps/utils_test.go @@ -19,9 +19,11 @@ package steps import ( "encoding/json" "errors" + "reflect" "strings" "testing" + "github.com/llm-d/llm-d-router/pkg/coordinator/gateway" "github.com/llm-d/llm-d-router/pkg/coordinator/pipeline" ) @@ -45,6 +47,104 @@ func TestReadErrorBody_ReturnsSmallBodyVerbatim(t *testing.T) { } } +func TestCapSingleTokenOutput(t *testing.T) { + tests := []struct { + name string + format gateway.RequestFormat + body map[string]any + want map[string]any + }{ + { + name: "chat completions caps output fields and forces non-streaming", + format: gateway.FormatChatCompletions, + body: map[string]any{ + "model": "m", + "max_tokens": 100, + "min_tokens": 5, + "max_completion_tokens": 100, + "stream": true, + "stream_options": map[string]any{"include_usage": true}, + }, + want: map[string]any{ + "model": "m", + "max_tokens": 1, + "max_completion_tokens": 1, + "stream": false, + }, + }, + { + name: "max_completion_tokens is not added when the client omitted it", + format: gateway.FormatChatCompletions, + body: map[string]any{"model": "m"}, + want: map[string]any{ + "model": "m", + "max_tokens": 1, + "stream": false, + }, + }, + { + name: "completions caps max_tokens, strips min_tokens, forces non-streaming", + format: gateway.FormatCompletions, + body: map[string]any{"model": "m", "max_tokens": 100, "min_tokens": 5}, + want: map[string]any{"model": "m", "max_tokens": 1, "stream": false}, + }, + { + name: "streaming is forced false and stream_options stripped", + format: gateway.FormatCompletions, + body: map[string]any{"stream": true, "stream_options": map[string]any{"include_usage": true}}, + want: map[string]any{"stream": false, "max_tokens": 1}, + }, + { + name: "generate caps max_tokens and strips min_tokens inside sampling_params", + format: gateway.FormatGenerate, + body: map[string]any{ + "model": "m", + "sampling_params": map[string]any{"max_tokens": 100, "min_tokens": 5}, + }, + want: map[string]any{ + "model": "m", + "sampling_params": map[string]any{"max_tokens": 1}, + "stream": false, + }, + }, + { + name: "generate synthesizes sampling_params when absent", + format: gateway.FormatGenerate, + body: map[string]any{"model": "m"}, + want: map[string]any{ + "model": "m", + "sampling_params": map[string]any{"max_tokens": 1}, + "stream": false, + }, + }, + { + name: "generate preserves other sampling_params entries", + format: gateway.FormatGenerate, + body: map[string]any{ + "sampling_params": map[string]any{ + "extra_args": map[string]any{"kv_transfer_params": "x"}, + }, + }, + want: map[string]any{ + "sampling_params": map[string]any{ + "max_tokens": 1, + "extra_args": map[string]any{"kv_transfer_params": "x"}, + }, + "stream": false, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + capSingleTokenOutput(tt.body, tt.format) + if !reflect.DeepEqual(tt.body, tt.want) { + t.Fatalf("got %v, want %v", tt.body, tt.want) + } + }) + } +} + func TestExtractTokenIDs(t *testing.T) { tests := []struct { name string diff --git a/test/coordinator/e2e/coordinator/configs_test.go b/test/coordinator/e2e/coordinator/configs_test.go index f030b0f290..9c71b22ecd 100644 --- a/test/coordinator/e2e/coordinator/configs_test.go +++ b/test/coordinator/e2e/coordinator/configs_test.go @@ -16,6 +16,8 @@ limitations under the License. package coordinate2e +import "strings" + // coordinatorConfigNIXL is the coordinator pipeline config for the e-p-d-pools topology. // ${NAMESPACE} and ${VLLM_RENDER_PORT} are substituted by createCoordinator before the ConfigMap is built. const coordinatorConfigNIXL = `log_level: 5 @@ -55,6 +57,14 @@ pipeline: - type: decode ` +// coordinatorConfigNIXLGenerate is coordinatorConfigNIXL with OpenAI passthrough +// disabled. With use_openai_format: false a chat/completions request collapses to +// the native generate wire format on the encode and prefill legs, which build a +// fresh sampling_params carrying only max_tokens: 1. This exercises a different +// min_tokens-stripping path than the OpenAI clone-and-cap path. +var coordinatorConfigNIXLGenerate = strings.Replace( + coordinatorConfigNIXL, "use_openai_format: true", "use_openai_format: false", 1) + // eppConfig is the scheduling config for the single EPP that serves all three // phases. Each request runs exactly one scheduling profile, named by its // EPP-Profile header value (see header-profile-handler); the role filters diff --git a/test/coordinator/e2e/coordinator/coordinator_test.go b/test/coordinator/e2e/coordinator/coordinator_test.go index 329866b681..8ff1edc374 100644 --- a/test/coordinator/e2e/coordinator/coordinator_test.go +++ b/test/coordinator/e2e/coordinator/coordinator_test.go @@ -51,35 +51,49 @@ var _ = ginkgo.Describe("Coordinator pipeline", func() { runCoordinatorPipeline(gateway.PathChatCompletions, []byte(fmt.Sprintf( `{"model":%q,"messages":[{"role":"user","content":"hello"}]}`, modelName, - )), textOnlySteps, 0) + )), textOnlySteps, 0, 0, 0) + }) + + ginkgo.It("forwards the client min_tokens/max_tokens to decode and caps them on prefill and encode", func() { + runCoordinatorPipeline(gateway.PathChatCompletions, []byte(fmt.Sprintf( + `{"model":%q,"messages":[{"role":"user","content":[{"type":"image_url","image_url":{"url":%q},"uuid":"image-0"},{"type":"text","text":"Describe what you see."}]}],"min_tokens":3,"max_tokens":5}`, + modelName, inlineImageDataURI, + )), allSteps, 1, 3, 5) + }) + + ginkgo.It("forwards min_tokens/max_tokens to decode and caps them on prefill and encode with OpenAI passthrough disabled", func() { + runCoordinatorPipeline(gateway.PathChatCompletions, []byte(fmt.Sprintf( + `{"model":%q,"messages":[{"role":"user","content":[{"type":"image_url","image_url":{"url":%q},"uuid":"image-0"},{"type":"text","text":"Describe what you see."}]}],"min_tokens":3,"max_tokens":5}`, + modelName, inlineImageDataURI, + )), allSteps, 1, 3, 5, coordinatorConfigNIXLGenerate) }) ginkgo.It("routes a multimodal image chat completion end-to-end", func() { runCoordinatorPipeline(gateway.PathChatCompletions, []byte(fmt.Sprintf( `{"model":%q,"messages":[{"role":"user","content":[{"type":"image_url","image_url":{"url":%q},"uuid":"image-0"},{"type":"text","text":"Describe what you see."}]}],"max_tokens":150}`, modelName, testImageURL, - )), allSteps, 1) + )), allSteps, 1, 0, 0) }) ginkgo.It("routes a multimodal chat completion with two images end-to-end", func() { runCoordinatorPipeline(gateway.PathChatCompletions, []byte(fmt.Sprintf( `{"model":%q,"messages":[{"role":"user","content":[{"type":"image_url","image_url":{"url":%q},"uuid":"image-0"},{"type":"image_url","image_url":{"url":%q},"uuid":"image-1"},{"type":"text","text":"What is in these two images?"}]}],"max_tokens":150}`, modelName, testImageURL, testImageURL2, - )), allSteps, 2) + )), allSteps, 2, 0, 0) }) ginkgo.It("routes a multimodal chat completion with an inline base64 image end-to-end", func() { runCoordinatorPipeline(gateway.PathChatCompletions, []byte(fmt.Sprintf( `{"model":%q,"messages":[{"role":"user","content":[{"type":"image_url","image_url":{"url":%q},"uuid":"image-0"},{"type":"text","text":"Describe what you see."}]}],"max_tokens":150}`, modelName, inlineImageDataURI, - )), allSteps, 1) + )), allSteps, 1, 0, 0) }) ginkgo.It("routes a multimodal chat completion with one inline and one remote image end-to-end", func() { runCoordinatorPipeline(gateway.PathChatCompletions, []byte(fmt.Sprintf( `{"model":%q,"messages":[{"role":"user","content":[{"type":"image_url","image_url":{"url":%q},"uuid":"image-0"},{"type":"image_url","image_url":{"url":%q},"uuid":"image-1"},{"type":"text","text":"Describe what you see in both images."}]}],"max_tokens":150}`, modelName, inlineImageDataURI, testImageURL, - )), allSteps, 2) + )), allSteps, 2, 0, 0) }) }) @@ -94,8 +108,13 @@ const inlineImageDataURI = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAEAAAA // asserts a 200 with a non-empty body, verifies that the coordinator logs show // all expected pipeline steps completed, then tears the workload down. // expectedImages is the number of images in the request; when > 0 the encoder -// log assertions are also verified. -func runCoordinatorPipeline(path string, body []byte, expectedSteps []string, expectedImages int) { +// log assertions are also verified. When wantMaxTokens > 0 the token-limit +// contract is additionally asserted: decode forwards the client's +// min_tokens/max_tokens unchanged, while prefill (and encode, for multimodal +// requests) cap max_tokens to 1 and strip min_tokens. +// coordinatorConfig, when supplied, overrides the default coordinatorConfigNIXL +// pipeline config for the coordinator deployment. +func runCoordinatorPipeline(path string, body []byte, expectedSteps []string, expectedImages, wantMinTokens, wantMaxTokens int, coordinatorConfig ...string) { nsName := getNamespace() var ( coordinator []string @@ -141,7 +160,11 @@ func runCoordinatorPipeline(path string, body []byte, expectedSteps []string, ex gomega.Expect(prefillPods).Should(gomega.HaveLen(prefillReplicas)) gomega.Expect(decodePods).Should(gomega.HaveLen(decodeReplicas)) - coordinator = createCoordinator(coordinatorConfigNIXL) + cfg := coordinatorConfigNIXL + if len(coordinatorConfig) > 0 { + cfg = coordinatorConfig[0] + } + coordinator = createCoordinator(cfg) req, err := http.NewRequest(http.MethodPost, gatewayBaseURL()+path, @@ -160,20 +183,30 @@ func runCoordinatorPipeline(path string, body []byte, expectedSteps []string, ex gomega.Expect(resp.StatusCode).To(gomega.Equal(http.StatusOK), "coordinator returned non-200: body=%s", string(raw)) gomega.Expect(raw).NotTo(gomega.BeEmpty(), "coordinator returned empty body") - verifyCoordinatorSteps(nsName, expectedSteps, expectedImages, true, true) + + logs := fetchCoordinatorLogs(nsName) + verifyCoordinatorSteps(logs, expectedSteps, expectedImages, true, true) + if wantMaxTokens > 0 { + // Prefill is always capped; encode is capped only when the request carries + // images, since it fires one sub-request per image. + capLegs := []string{"prefill"} + if expectedImages > 0 { + capLegs = append(capLegs, "encode") + } + verifyTokenLimits(logs, wantMinTokens, wantMaxTokens, capLegs) + } } -// verifyCoordinatorSteps fetches the coordinator pod logs and asserts that -// every expected step has a "step complete" log entry. When kvNIXL is set it +// verifyCoordinatorSteps asserts that the coordinator logs show every expected +// step has a "step complete" log entry. When kvNIXL is set it // also asserts kv_transfer_params on the prefill and decode legs, and when // expectedImages > 0 it asserts the encode step completed all image // sub-requests, plus, when ecNIXL is set, the ec_transfer_params total via the // "merged encode response" marker. The kv/ec params surface in the logs only // for the NIXL connectors, so those checks are gated accordingly. -func verifyCoordinatorSteps(nsName string, expectedSteps []string, expectedImages int, kvNIXL, ecNIXL bool) { +func verifyCoordinatorSteps(logs string, expectedSteps []string, expectedImages int, kvNIXL, ecNIXL bool) { ginkgo.By("Verifying coordinator logs contain all pipeline steps") - logs := fetchCoordinatorLogs(nsName) for _, step := range expectedSteps { stepField := `"step":"` + step + `"` gomega.Expect(logHasLine(logs, `"msg":"step complete"`, stepField)).To(gomega.BeTrue(), @@ -239,6 +272,30 @@ func fetchCoordinatorLogs(nsName string) string { return string(out) } +// verifyTokenLimits asserts the pipeline's token-limit contract: the decode leg +// forwards the client's min_tokens/max_tokens unchanged, while the synthetic +// prefill and encode legs (capLegs) cap output to a single token (max_tokens=1) +// and strip min_tokens. Leg request bodies surface only at TRACE, so this relies +// on the coordinator running at log_level 5. +func verifyTokenLimits(logs string, wantMin, wantMax int, capLegs []string) { + ginkgo.By("Verifying decode leg forwards the client min_tokens/max_tokens") + gomega.Expect(logHasLine(logs, `"msg":"request body"`, `"epp-profile":"decode"`, + fmt.Sprintf(`"min_tokens":%d`, wantMin), fmt.Sprintf(`"max_tokens":%d`, wantMax))).To(gomega.BeTrue(), + "coordinator logs have no decode request body carrying min_tokens=%d and max_tokens=%d", wantMin, wantMax) + + for _, phase := range capLegs { + phaseField := `"epp-profile":"` + phase + `"` + + ginkgo.By("Verifying " + phase + " leg caps max_tokens to 1") + gomega.Expect(logHasLine(logs, `"msg":"request body"`, phaseField, `"max_tokens":1`)).To(gomega.BeTrue(), + "coordinator logs have no %s request body carrying max_tokens=1", phase) + + ginkgo.By("Verifying " + phase + " leg strips min_tokens") + gomega.Expect(logHasLine(logs, `"msg":"request body"`, phaseField, `"min_tokens"`)).To(gomega.BeFalse(), + "%s request body must not carry min_tokens", phase) + } +} + // verifyEncodeSkipped asserts the coordinator skipped the encode fan-out on the // generate path. The skip marker is logged immediately before the step returns, // so its presence is dispositive: the prefill worker encodes inline from diff --git a/test/coordinator/e2e/coordinator/generate_endpoint_test.go b/test/coordinator/e2e/coordinator/generate_endpoint_test.go index 632f5940f5..9689a90721 100644 --- a/test/coordinator/e2e/coordinator/generate_endpoint_test.go +++ b/test/coordinator/e2e/coordinator/generate_endpoint_test.go @@ -47,7 +47,7 @@ var generateSteps = []string{"render", "prefill", "decode"} var _ = ginkgo.Describe("Coordinator pipeline - generate endpoint", func() { ginkgo.It("routes a text-only generate end-to-end", func() { runCoordinatorPipeline(gateway.DefaultGeneratePath, - generateBody(modelName, nil), generateSteps, 0) + generateBody(modelName, nil), generateSteps, 0, 0, 0) }) ginkgo.It("routes a single-image generate end-to-end", func() { @@ -55,7 +55,7 @@ var _ = ginkgo.Describe("Coordinator pipeline - generate endpoint", func() { {Hash: "e2e-gen-hash-0", Offset: 1, Length: 3}, } runCoordinatorPipeline(gateway.DefaultGeneratePath, - generateBody(modelName, images), generateSteps, 0) + generateBody(modelName, images), generateSteps, 0, 0, 0) verifyEncodeSkipped(getNamespace()) }) })