feat: dedicated disaggregated serving router front (B3b2b)#193
Merged
Conversation
Add a --node-role router process that completes topology A: a client's chat-completions request hits the router's HTTP front, is routed to a prefill node over TCP, handed off to a decode node, and the merged stream (prefill first token plus decode continuation) is returned to the client over SSE. - router_front.rs: a model-free RouterState (loads only a tokenizer and chat template, never model weights), a background demux task that routes each ResultFrame to its per-request channel by request_id, RequestRouter-based prefill node selection over a registry of the configured peers, and StreamBridge ordering validation. The router tokenizes with the same add_special rule and build_generate_options the worker uses, so the prompt tokens and sampling match a single-node run. - startup.rs: a "router" node role maps to ServingMode::Router (a router is not a cluster inference role, so no NodeRole variant is added); the serve path branches before the model load to bind the role transport, build the router app, and serve, skipping the model worker and warmup. - tests/disaggregated_router_e2e.rs: spawns three real mlxcel-server processes (router + prefill + decode) and asserts the streamed SSE content is byte-identical to the single-node greedy reference. The router registers only its prefill and decode peers in the routing registry, never itself, so route_to_prefill cannot select the router's own address. Part of #126
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
B3b2b of the disaggregated serving productionization (epic #116, capstone #126). Adds a dedicated
--node-role routerprocess that completes topology A: client HTTP → router → prefill node (TCP) → decode node (TCP) → router merges (prefill first token + decode continuation) → client SSE. A three-process E2E proves the streamed output is byte-identical to a single-node run.What this does
router_front.rs(model-free): aRouterStatethat loads only a tokenizer + chat template (never model weights), a background demux task routing eachResultFrameto its per-request channel byrequest_id,RequestRouter-based prefill selection over a registry of the configured peers, andStreamBridgeordering validation. The router tokenizes with the sameadd_specialrule andbuild_generate_optionsthe model worker uses, so the prompt tokens and sampling match a single-node run. ServesPOST /v1/chat/completions(streaming SSE and non-streaming JSON) +/health.startup.rs: the"router"node role maps toServingMode::Router(a router is not a cluster inference role, so noNodeRolevariant is added); the serve path branches before the model load to bind the role transport, build the router app, and serve, skipping the model worker and warmup.tests/disaggregated_router_e2e.rs: spawns three realmlxcel-serverprocesses (router + prefill + decode) and asserts the streamed SSE content is byte-identical to the single-node greedy reference.Notes
route_to_prefillcannot select the router's own address. (An earlier draft registered the router as a Hybrid node, which let routing nondeterministically pick the router itself; the real-model SSE E2E caught it.)RequestRouter(node selection) andStreamBridge(phase/ordering validation); the prefill/decode nodes and TCP wire protocol are unchanged from B3a/B3b2a.enable_thinking=false); the decode continuation arrives as one frame and is re-emitted per token (true per-token decode streaming is a later enhancement); only/v1/chat/completionsis implemented.Verification (orchestrator-local, the real metal gate)
-D warnings --all-targetsboth feature sets (metal,accelerateand+test-utils): cleancargo fmt --check: cleantests/disaggregated_router_e2e.rs): router + prefill + decode reproduce"2 + 2 = 4."over SSE, byte-identical to a hybrid single-node runPart of #126