refactor: unify API type handling across sidecar and coordinator - #2743
refactor: unify API type handling across sidecar and coordinator#2743revit13 wants to merge 24 commits into
Conversation
Signed-off-by: Revital Sur <eres@il.ibm.com>
Let CapSingleToken own sampling_params in the generate prefill body, derive the sidecar route table from DetectAPIType, inline the single-caller capMaxTokensField, and merge the duplicated APIType test tables. Signed-off-by: Revital Sur <eres@il.ibm.com>
CapSingleToken named each API's cap fields in an if-chain while APIType.TokenLimitFields already listed them as data, and the two disagreed: a /v1/responses leg was capped on max_tokens, which that API ignores, and left max_output_tokens as the client sent it, so the leg ran a full generation. Drive the caps from TokenLimitFields. min_tokens is a floor rather than a cap and is still stripped for every API, so chat completions, completions and generate are unchanged. Signed-off-by: Revital Sur <eres@il.ibm.com>
NIXLv2 capped the client's own body for the prefill leg and restored every field afterwards for the decode leg, the only connector doing so. Copy the body one level and cap the copy, as the shared-storage, mooncake and p2p connectors already do, so the shared CapSingleToken helper covers all five. The concurrent WRITE-mode path gains the apiType it previously ignored, which had it capping chat-completions fields on generate and Responses bodies. min_tokens is now stripped from the prefill leg rather than set to 1, matching every other connector. Signed-off-by: Revital Sur <eres@il.ibm.com>
Five sites hand-rolled the same one-level map copy with make plus a range loop. Use maps.Clone, which the same package already uses elsewhere. buildEncoderRequest documented its copy as deep when it is one level, which matters because nested values stay shared with the client's body. Signed-off-by: Revital Sur <eres@il.ibm.com>
A JSON "null" body unmarshals without error into a nil map. Connectors clone the parsed body and write the transfer params into the clone, and a write to a nil map panics. Reject a non-object body in the parser, where every connector sees it, and route the remaining connectors through readJSONBody so the check cannot be skipped. Signed-off-by: Revital Sur <eres@il.ibm.com>
The sidecar caps a leg straight off the parsed client body, with no equivalent of the coordinator's validateSamplingParams ahead of it, so a non-object or null sampling_params reaches TokenLimitMap. Both land on the synthesize branch and the field is replaced; pin that. State the one-level-copy guarantee in the CapSingleToken doc comment, where callers that clone before capping read it, instead of inside the function body. Signed-off-by: Revital Sur <eres@il.ibm.com>
Signed-off-by: Revital Sur <eres@il.ibm.com>
main landed its own raw-JSON body parser (decodeRequestBody plus inspectedRequestFields), which overlaps the body-reading refactor on this branch. The resolution keeps the branch helpers readJSONBody/bodyAsJSON as the single entry point and main's decodeRequestBody as the parser beneath them, so uninspected fields still reach the backends byte-for-byte. Follow-on fixes the resolution required: - bodyAsJSON dropped its parsed == nil guard; decodeRequestBody already refuses a body that is not a JSON object. - proxy.go keeps requestFieldMaxOutputTokens, which inspectedRequestFields needs to decode the Responses API output cap. - coordinator e2e assertions use the "body" log key that the OpenTelemetry JSON logger emits. Signed-off-by: Revital Sur <eres@il.ibm.com>
Signed-off-by: Revital Sur <eres@il.ibm.com>
Resolve the sidecar proxy conflicts against the completionRequest -> body rename (llm-d#2775) and the shared concurrent P/D dispatch helper (llm-d#2589): - connector_nixlv2.go: keep the parser3 clone-and-cap prefill body and the apiType plumbing, renamed to body. - connector_shared_storage.go: keep CapSingleToken(body, apiType) over the removed PrimeSingleTokenRequest. - connector_sglang.go: drop the fmt import orphaned by both sides removing their only fmt.Errorf callers. Signed-off-by: Revital Sur <eres@il.ibm.com>
roytman
left a comment
There was a problem hiding this comment.
Several small comments, plus you need to update the PR description.
|
|
||
| // TokenLimitFields returns the output token cap field names the API uses. | ||
| // The returned slices are shared package-level vars; callers must not mutate them. | ||
| func (a APIType) TokenLimitFields() []string { |
There was a problem hiding this comment.
TokenLimitFields and TokenLimitMap (L:126) are exported but only called from CapSingleToken in this package (tokens.go). Can these be lowercase? Path, String, and DetectAPIType have real external callers and should stay exported, but these two don't need to be exported.
| // Chat completions lists both max_tokens and max_completion_tokens: vLLM and | ||
| // SGLang accept the two together and prefer max_completion_tokens, so capping | ||
| // both bounds the leg regardless of which field the engine consults. | ||
| // | ||
| // body is rewritten in place, so the caller passes its own copy. A one-level | ||
| // copy is enough: the generate API caps inside sampling_params, and | ||
| // TokenLimitMap replaces that nested map rather than writing through it, so a | ||
| // body that still shares it with the decode leg keeps the client's limits. |
There was a problem hiding this comment.
This restates the chat-completions dual-field rationale from apitype.go:100-102 and the TokenLimitMap replacement contract from apitype.go:122-125. Could this live in one place? The min_tokens-as-floor paragraph above (lines 25-28) is the only part here that isn't said elsewhere.
| if minTokensOk { | ||
| body[requestFieldMinTokens] = minTokensValue | ||
| } | ||
| // body still carries the client's streaming flags and token limits: only the |
There was a problem hiding this comment.
"the copy above" - body vs prefillRequest is already stated at line 430. This line doesn't add anything past that.
| @@ -478,6 +418,7 @@ retryLoop: | |||
| func (s *Server) runNIXLProtocolV2WriteParallel( | |||
There was a problem hiding this comment.
All three MoRIIOParallelDispatch test cases in connector_nixlv2_test.go (lines 1028, 1071, 1105) send chat-completions requests. This function is the one that used to hardcode chat-completions field names regardless of apiType - for chat-completions, old and new behavior coincide except on min_tokens, so nothing here would catch a regression back to the hardcoded fields for a /v1/responses or generate request. Worth one case posting to GeneratePath or ResponsesPath with MoRIIOWriteMode on, asserting the prefill leg caps sampling_params.max_tokens (or max_output_tokens) rather than a stray top-level field.
Signed-off-by: Revital Sur <eres@il.ibm.com>
Reconcile the unified reqcommon.APIType with the EC apiType threading that landed on main (llm-d#2749): - connector_ec_common.go, connector_ec_nixl.go, connector_ec_shared_storage.go, proxy.go: keep main's apiType parameter, typed reqcommon.APIType in place of the local APIType this branch removes. The EC legs carry the client's API instead of the hardcoded chat type. - connector_ec_common_test.go (new on main) and the EC connector tests: port the local APIType constants to reqcommon, and expect the prefill leg to drop min_tokens rather than cap it to 1, which is CapSingleToken's contract. Signed-off-by: Revital Sur <eres@il.ibm.com>
Signed-off-by: Revital Sur <eres@il.ibm.com>
Signed-off-by: Revital Sur <eres@il.ibm.com>
The sidecar treated both the zero value and an unrecognized APIType as chat completions, while the shared type fell back to generate. Use the sidecar's default everywhere. LookupAPIType reports whether a path is known, so the coordinator render and format steps still skip or collapse to generate for paths the router does not register. Signed-off-by: Revital Sur <eres@il.ibm.com>
Flatten resolveFormat and state the chat completions fallback once in apitype.go. Document why resolveFormat falls back to generate. Fix a dupword lint error in render_test.go. CapSingleToken always strips min_tokens, so assert that the prefill leg has no min_tokens instead of guarding a check that can never run. Signed-off-by: Revital Sur <eres@il.ibm.com>
The coordinator router sends only the chat completions, completions, and generate routes to the steps, so the known flag guarded a path that production never takes. An unrecognized path now resolves as chat completions in the steps too. TestPrefillStep_GatewayError relied on the empty path falling back to generate, so it sets the generate path. Signed-off-by: Revital Sur <eres@il.ibm.com>
The sampling_params sharing rule and the min_tokens stripping rule were each explained in several helper and test comments. Keep the reasoning in tokenLimitMap and CapSingleToken, and make the other comments name what they test and point to those docs. Signed-off-by: Revital Sur <eres@il.ibm.com>
|
|
||
| const ( | ||
| // APITypeChatCompletions is the Chat Completions API (/v1/chat/completions) | ||
| // and the Anthropic Messages API (/v1/messages), which share its field names. |
There was a problem hiding this comment.
in the original code it was
// APITypeChatCompletions is the Chat Completions API (/v1/chat/completions, /v1/completions)"Chat completions and completion processes are the same, but I'w define separate entries, and of course not for messages
we can define a separate entry for /v1/messages
| PathGenerate = "/inference/v1/generate" | ||
| ) | ||
|
|
||
| // APIType is the inference API a request speaks. It selects the JSON field |
There was a problem hiding this comment.
I'd write just "APIType is the inference API a request was sent. "
| func (s *RenderStep) Execute(ctx context.Context, reqCtx *pipeline.RequestContext) error { | ||
| if reqCtx.OriginalPath == gateway.DefaultGeneratePath { | ||
| // The Responses API carries no token_ids to normalize, so it skips the step. | ||
| switch reqcommon.DetectAPIType(reqCtx.OriginalPath) { |
There was a problem hiding this comment.
The switch has no default:, so the "skip" fallthrough below (line 136-137) is now reached only by APITypeResponses. On main, an unrecognized path fell through to "skipping render step"; here it falls through to nothing, i.e. executeChatCompletions never runs but the switch also never explicitly names what's being skipped for.
Concretely: the moment a new route is registered on the coordinator, RenderStep will silently do the wrong thing instead of skipping, since there's no default: arm to catch it. Suggest:
switch reqcommon.DetectAPIType(reqCtx.OriginalPath) {
case reqcommon.APITypeGenerate:
return s.executeGenerate(ctx, reqCtx)
case reqcommon.APITypeCompletions:
return s.executeCompletions(ctx, reqCtx)
case reqcommon.APITypeChatCompletions:
return s.executeChatCompletions(ctx, reqCtx)
default:
logger := log.FromContext(ctx).WithName(RenderStepName)
logger.V(logutil.DEFAULT).Info("skipping render step", "path", reqCtx.OriginalPath)
return nil
}Dormant today since only 3 exact paths are registered in pkg/coordinator/server/server.go, but worth closing before it's a live bug.
There was a problem hiding this comment.
This is a leftover exact-equality check (reqCtx.OriginalPath == gateway.DefaultGeneratePath) that wasn't updated when resolveFormat a few lines below switched to DetectAPIType's substring match. On a prefixed generate path (e.g. /prefix/inference/v1/generate), the request is classified as generate everywhere else in the pipeline, but this check misses it — so the encode fan-out runs and re-ships the oversized preprocessed pixel tensor this very check exists to avoid (vllm#46722).
Suggest using the same DetectAPIType(reqCtx.OriginalPath) == reqcommon.APITypeGenerate check here for consistency with the rest of the unification.
| @@ -253,10 +253,7 @@ func decodeRequestBody(raw []byte) (map[string]any, error) { | |||
| func (s *Server) readJSONBody(r *http.Request, w http.ResponseWriter) ([]byte, map[string]any, bool) { | |||
There was a problem hiding this comment.
Now that mooncake/p2p/SGLang were consolidated onto this one helper, it's the natural place to add the one log line that was missing from all of them: none of the 9 call sites through readJSONBody currently log the underlying read/decode error server-side — only the client-visible 400 exists. An operator grepping logs for a request with a malformed body finds nothing.
Suggest logging inside the if err != nil branch here, e.g.:
if err != nil {
s.logger.V(logging.DEBUG).Info("invalid request body", "error", err)
if writeErr := errorJSONInvalid(err, w); writeErr != nil {
s.logger.Error(writeErr, "failed to send error response to client")
}
return nil, nil, false
}Not a regression from this PR (the gap predates it in every connector), but the consolidation makes it a one-line fix instead of a fix-in-9-places.
| // APIType is the inference API a request speaks. It selects the JSON field | ||
| // names a request carries and the path a synthesized request is sent to. A | ||
| // value outside the constants below degrades to APITypeChatCompletions. | ||
| type APIType int |
There was a problem hiding this comment.
Two things worth a second look here, not blocking:
-
APITypeis a bareintwith no validating constructor, and the "degrades to chat completions" rule is repeated independently inPath,tokenLimitFields,tokenLimitMap, andDetectAPITyperather than centralized. There's also noexhaustivelint rule enabled for this package, so a 5thAPITypeconstant added later would compile cleanly and silently degrade in every one of these switches exceptprefill.go, which treats an unhandled value as a hard error. Worth enablingexhaustivefor this package if the closed-set invariant is meant to be enforced rather than conventional.
(We can return to it later) -
The PR description's test plan says unknown/out-of-range values "fall back to generate." They actually fall back to chat completions (confirmed by this file and by
apitype_test.go:56,74). Worth fixing the PR description so reviewers aren't misled about which API type is the actual default.
| } | ||
| case KVConnectorSGLang: | ||
| s.handlePDConnector = func(w http.ResponseWriter, r *http.Request, host string, _ string, _ APIType) { | ||
| s.handlePDConnector = func(w http.ResponseWriter, r *http.Request, host string, _ string, _ reqcommon.APIType) { |
There was a problem hiding this comment.
_ reqcommon.APIType here reads like an oversight from the mechanical signature update, but it's actually correct: handleSGLang sends the identical marshaled body to both legs by design, so there's nothing to cap. Worth a one-line comment saying so, otherwise this looks like a spot that got missed when apiType was threaded through the other four connectors.
| // sampling_params for the generate API, body itself otherwise. The generate map | ||
| // is always replaced with one body owns, so a caller that writes into the result | ||
| // never reaches a nested map the body was cloned from. | ||
| func (a APIType) tokenLimitMap(body map[string]any) map[string]any { |
There was a problem hiding this comment.
do we need this separate method only to handle a specific API type? It is calledpkg/common/request/tokens.go:33 from only; I think it will be better to copy its functionality directly into CapSingleToken, it will improve readability. wdyt?
| const ( | ||
| // ChatCompletionsPath is the OpenAI chat completions path | ||
| ChatCompletionsPath = "/v1/chat/completions" | ||
| ChatCompletionsPath = reqcommon.PathChatCompletions |
There was a problem hiding this comment.
Do we need to redefine all these constants?
The file is called chat_completions, but it includes all possible paths ;-) (we can handle it separately)
There was a problem hiding this comment.
Removed them in the coordinator too. pkg/coordinator/gateway/paths.go now keeps only the header and phase constants, and the coordinator code and tests use
reqcommon.PathChatCompletions, reqcommon.PathCompletions, and reqcommon.PathGenerate directly. Thanks
|
|
||
| func openAIAPIAttr(apiType APIType) attribute.KeyValue { | ||
| func openAIAPIAttr(apiType reqcommon.APIType) attribute.KeyValue { | ||
| return attribute.String("llm_d.openai.api", apiType.String()) |
There was a problem hiding this comment.
This is wrong for none of OpenAI's calls, we can handle it separately
Give /v1/messages its own APIType so its prefill leg caps only max_tokens, the one output cap the Anthropic Messages API defines. Use the path constants in pkg/common/request directly in the sidecar and the coordinator instead of keeping aliases, and inline tokenLimitMap into CapSingleToken, its only caller. Make the render step skip explicit with a default case, classify the encode generate skip with DetectAPIType so a prefixed path matches, log the reason readJSONBody refuses a body at DEBUG, and document why the SGLang connector ignores the API type. Signed-off-by: Revital Sur <eres@il.ibm.com>
Signed-off-by: Revital Sur <eres@il.ibm.com>
|
|
||
| encoderRequest["messages"] = messages | ||
| reqcommon.PrimeSingleTokenRequest(encoderRequest) | ||
| reqcommon.CapSingleToken(encoderRequest, reqcommon.APITypeChatCompletions) |
There was a problem hiding this comment.
maybe we need a comment why we always pass reqcommon.APITypeChatCompletion
| path: ChatCompletionsPath, | ||
| body: `{"model":"m","messages":[{"role":"user","content":"hello"}],"max_tokens":80,"max_completion_tokens":90,"min_tokens":5}`, | ||
| tokenFields: []string{"max_tokens", "max_completion_tokens", "min_tokens"}, | ||
| tokenFields: []string{"max_tokens", "max_completion_tokens"}, |
There was a problem hiding this comment.
Please replace the literals with the constants e.g eqcommon.FieldMaxTokens, ....
There was a problem hiding this comment.
Done. The test uses reqcommon.FieldMaxTokens, reqcommon.FieldMaxCompletionTokens, reqcommon.FieldMaxOutputTokens, and the other reqcommon.Field* constants for every field name in
Go code. The JSON request bodies stay as literals, since they are the raw bytes the test sends.
| // TestBuildEncoderRequest_MinTokens is a regression test: a client-supplied | ||
| // min_tokens above the encoder leg's max_tokens=1 cap trips vLLM's | ||
| // min_tokens<=max_tokens validation. | ||
| // TestBuildEncoderRequest_MinTokens is a regression test for stripping a |
There was a problem hiding this comment.
Move all three TestBuildEncoderRequest* functions from connector_ec_shared_storage_test.go to connector_ec_common_test.go (which already imports assert/testing, so it's a straight cut-paste, no new imports). Keep the content as-is — the coverage itself is sound, just misplaced.
| savedTokenValues[i] = savedField{field: field} | ||
| } | ||
| } | ||
| // Keeps the client's body intact for the decode leg below. |
There was a problem hiding this comment.
I think, there was a request to replace legs with requests
There was a problem hiding this comment.
I couldn't find an earlier comment on this PR asking to rename "leg". The only comments that mention it are yours, and two of them use "leg" too. main already uses "leg" widely in the
sidecar and coordinator (for example 14 lines in connector_nixlv2.go, 9 in proxy.go, 9 in options.go). If I renamed only the lines this PR adds, those files would use both terms,
so I kept "leg" here. Thanks
| } | ||
| func (s *Server) handleP2P(w http.ResponseWriter, r *http.Request, prefillPodHostPort, kvCacheSource string, apiType reqcommon.APIType) { | ||
| _, requestData, ok := s.readJSONBody(r, w) | ||
| if !ok { |
Document why buildEncoderRequest caps the encoder request as chat completions: it carries the item in messages and is always sent to /v1/chat/completions, whatever API the client used. Keep the buildEncoderRequest tests next to the code they cover in connector_ec_common_test.go, and name request fields there with the reqcommon.Field constants. Signed-off-by: Revital Sur <eres@il.ibm.com>
| @@ -94,10 +79,7 @@ func (s *Server) handleP2P(w http.ResponseWriter, r *http.Request, prefillPodHos | |||
|
|
|||
| // Decode leg: pull KV from the prefiller's OffloadingConnector P2P tier. Original body | |||
| } | ||
| // 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). | ||
| // The leg body is built from RequestContext, so this guards against the branch |
| } | ||
| // The legacy Completions API does not define max_completion_tokens. | ||
| if _, ok := prefillBody["max_completion_tokens"]; ok { | ||
| t.Fatalf("completions leg carries max_completion_tokens=%v", prefillBody["max_completion_tokens"]) |
There was a problem hiding this comment.
pls replace leg with request
|
|
||
| encoderRequest["messages"] = messages | ||
| reqcommon.PrimeSingleTokenRequest(encoderRequest) | ||
| // The encoder leg carries the item in messages and is sent to |
There was a problem hiding this comment.
Replaced "leg" with "request" in the comments and test names this PR adds, including the ones in handleP2P. The older uses on main (flag help, MORIIO_README, the rest of the sidecar) are unchanged; I can open a follow-up issue to rename them package-wide.
Address review comments: the prefill, decode, and encoder requests are described as requests, not legs, in the code this change adds. Signed-off-by: Revital Sur <eres@il.ibm.com>
/kind bug
/kind cleanup
What this PR does / why we need it:
The sidecar and the coordinator each carried their own copy of "which inference API
is this request speaking". The sidecar had
proxy.APITypewithtokenLimitFieldsForAPITypeandtokenLimitMap. The coordinator hadgateway.RequestFormatwithDetectFormatandPathForFormat. Each side also hadits own path constants and its own single-token capping helper
(
PrimeSingleTokenRequestin the sidecar,capSingleTokenOutputin thecoordinator). The two enums did not cover the same APIs: the sidecar had no
completions member and the coordinator had no responses member.
This PR moves the concept into
pkg/common/request: oneAPITypewithDetectAPITypeandPath, and oneCapSingleToken(body, apiType)that replacesboth capping helpers. The per-API list of cap fields, and the map that holds them,
stay internal to the package. An unknown path, and a value outside the defined API
types, fall back to chat completions. The sidecar and the coordinator now agree on
the path-to-API mapping and on which fields a synthetic leg caps.
Bugs fixed:
/v1/responsesprefill and encode legs were capped onmax_tokens, which that APIignores.
max_output_tokenskept the client's value, so the leg ran a fullgeneration on the prefiller.
--moriio-parallel-dispatch) ignored theapiTypeit was given and always capped the chat-completions fields. Generate andResponses bodies were capped on fields they do not define and left uncapped on the
ones they do.
/v1/completionslegs carriedmax_completion_tokens=1, a field the legacyCompletions API does not define, because the sidecar mapped that path to chat
completions.
APITypeCompletionsnow has its own field list./v1/messageslegs carriedmax_completion_tokens=1, a field the AnthropicMessages API does not define, for the same reason.
APITypeMessagesnow caps onlymax_tokens.Supporting refactors:
leg and then restored every field to rebuild the decode leg. It now caps a
one-level copy, like the other connectors, so all five share
CapSingleToken.min_tokensis removed from the NIXLv2 prefill leg instead of set to 1, whichmatches the other connectors and vLLM's rule that
min_tokens <= max_tokens.buildEncoderRequest, two in mooncake,two in p2p, one in SGLang) use
maps.Clone. ThebuildEncoderRequestcommentcalled its copy deep; it is one level, so nested values stay shared with the
client's body, and the comment now says so.
readJSONBody, like the other connectors, instead of repeating the read, decodeand error write. Their responses do not change.
readJSONBodynow answers anunreadable body with the same
BadRequestErrorJSON envelope it already used forinvalid JSON, instead of a plain-text
400. The status is unchanged. It also logsthe refusal reason at DEBUG.
DetectAPIType,like the other steps.
From the merge with main:
follow the encoder. This branch keeps that behavior, typed with the shared
APIType. ItsTestECPipelineTokenLimitsexpectedmin_tokens=1on the prefillleg; the expectation now follows
CapSingleToken, which removes it.Not in scope:
(sidecar: encoder disaggregation treats every request as chat completions #2742).
/v1/responsesbodies.resolveFormatstill collapsesResponses to generate. feat(coordinator) Initial implementation for handling API /v1/responses #2360 adds that support and edits
gateway/paths.go, whichthis PR reduces to aliases of the shared constants, so whichever merges second
needs a rebase.
Which issue(s) this PR fixes:
Fixes #TODO
Release note:
Test plan:
APITypename, path, detection and cap fields for every API, including anunknown path and an out-of-range value, which both fall back to chat
completions
sampling_paramsthat isabsent, null or not an object
CapSingleTokenper API: the right fields capped to 1,min_tokensremoved,streaming off,
stream_optionsremoved, and a nestedsampling_paramssharedwith another leg left intact
sampling_paramsonthe prefill leg, carries no chat cap field at the top level, and keeps the
client's limits on the decode leg
/v1/messagesprefill leg capsmax_tokensand carries nomax_completion_tokenspoints, including the EC connectors and the decoder-only paths
leg is sent
max_completion_tokens