Skip to content

feat: dedicated disaggregated serving router front (B3b2b)#193

Merged
inureyes merged 1 commit into
mainfrom
feature/issue-126-b3b2b-router-front
Jun 9, 2026
Merged

feat: dedicated disaggregated serving router front (B3b2b)#193
inureyes merged 1 commit into
mainfrom
feature/issue-126-b3b2b-router-front

Conversation

@inureyes

@inureyes inureyes commented Jun 9, 2026

Copy link
Copy Markdown
Member

Summary

B3b2b of the disaggregated serving productionization (epic #116, capstone #126). Adds a dedicated --node-role router process 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): a RouterState that loads only a tokenizer + chat template (never model weights), a background demux task routing each ResultFrame to its per-request channel by request_id, RequestRouter-based prefill 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 model worker uses, so the prompt tokens and sampling match a single-node run. Serves POST /v1/chat/completions (streaming SSE and non-streaming JSON) + /health.
  • startup.rs: the "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.

Notes

  • The router registers only its prefill/decode peers in the routing registry, never itself, so route_to_prefill cannot 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.)
  • Reuses RequestRouter (node selection) and StreamBridge (phase/ordering validation); the prefill/decode nodes and TCP wire protocol are unchanged from B3a/B3b2a.
  • Known follow-ups (documented): the router does not yet apply the chat stream filter (thinking/tool markers pass through, so the E2E uses 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/completions is implemented.

Verification (orchestrator-local, the real metal gate)

  • Cold clippy -D warnings --all-targets both feature sets (metal,accelerate and +test-utils): clean
  • cargo fmt --check: clean
  • 3-process router E2E (tests/disaggregated_router_e2e.rs): router + prefill + decode reproduce "2 + 2 = 4." over SSE, byte-identical to a hybrid single-node run
  • Full lib suite: 3086 passed, 0 failed

Part of #126

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
@inureyes
inureyes merged commit 0a2b7b9 into main Jun 9, 2026
5 checks passed
@inureyes
inureyes deleted the feature/issue-126-b3b2b-router-front branch June 9, 2026 15:29
@inureyes inureyes self-assigned this Jun 10, 2026
@inureyes inureyes added status:done Completed type:enhancement New features, capabilities, or significant additions priority:medium Medium priority labels Jun 10, 2026
@inureyes inureyes added this to the 0.2 milestone Jun 21, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

priority:medium Medium priority status:done Completed type:enhancement New features, capabilities, or significant additions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant