feat(zmq): teach the direct-backend path to speak TokenSpeed

[slin1237]

Description

Problem

The ZMQ direct-backend path (#2000, built on #2015) speaks only vLLM's EngineCore protocol. TokenSpeed (and later sglang) engines need the same same-host ipc:// fast path, but the crate's transport, connector, and gateway adapter all hard-assumed the vLLM wire format — and several workflow steps hard-assumed "ZMQ worker ⇒ vLLM".

Solution

Make the ZMQ stack engine-neutral and add TokenSpeed as the second protocol:

• engine-zmq-client: a new EngineProtocol trait seams the shared transport/connector (handshake, ROUTER/DEALER identity framing, output loop, abort-on-drop) away from the per-engine wire structs. VllmProtocol keeps the existing behavior; TokenSpeedProtocol adds the sglang-family msgpack tuples (WireTokenizedGenerateReq 5-tuple, WireSamplingParams 13-tuple, WireBatchTokenIDOut 8-tuple with sampled-token logprob columns). Handshake structs move to a neutral protocol/handshake.rs (re-exported for compatibility). A neutral EngineLoad replaces the vLLM-specific stats type in the shared seam.
• Gateway adapter (zmq_client.rs): ZmqEngineClient selects the protocol from the worker's explicit runtime_type (tokenspeed vs vllm; anything else is rejected before the handshake and at registration). Streams map both protocols to the existing vLLM-proto pipeline — chunks carry incremental tokens/logprobs, the terminal Complete carries the cumulative set, and a finish-tick's tokens are emitted as a chunk first so streaming never loses the last token.
• No silent narrowing: requests the wire cannot honor fail loudly with invalid_argument instead of degrading — structured-output constraints, n>1, top-k/prompt logprobs, stop strings (TokenSpeed), logit_bias (TokenSpeed), nonzero data_parallel_rank. Sampled-token logprobs are wired end-to-end on both ZMQ protocols.
• Runtime plumbing: detect_backend/discover_metadata no longer force ZMQ workers to vLLM — an explicitly configured runtime_type survives to the built worker (unspecified still defaults to vLLM with a warning). BackendClient::runtime_type() reports the actual ZMQ backend runtime.
• wfaas fix: the DAG scheduler could fail a workflow with a spurious "Workflow deadlocked" when a step completed between the completion-drain and the tracker read (instant-completing ZMQ detection steps hit this routinely). The deadlock branch now re-drains the completion channel before failing, and the run_if paths send their completion inside the tracker lock scope. Regression tests included.
• smg serve: per-user ZMQ socket dir (SMG_ZMQ_SOCKET_DIR override), FNV handshake-port derivation pinned by conformance vectors against the Python mirror.

Changes

• crates/engine_zmq_client: protocol/mod.rs (EngineProtocol/EngineOutput/EngineBatch/EngineLoad), protocol/tokenspeed/{mod,request,sampling,output}.rs, protocol/handshake.rs, generic connector.rs/transport.rs, EngineCoreReadyResponse.max_num_batched_tokens widened to i64 (TokenSpeed sends -1 = disabled).
• model_gateway: routers/grpc/zmq_client.rs (ZmqBackend enum, per-protocol translate/map with loud rejection boundary, logprob accumulation), routers/grpc/backend_client.rs (runtime passthrough), worker/worker.rs (runtime param on connect), workflow/steps/local/{detect_backend,discover_metadata,create_worker}.rs (runtime preservation + ZMQ runtime validation).
• crates/workflow: deadlock-detector re-drain + run_if in-lock completion send + regression tests.
• bindings/python/src/smg/serve.py: per-user socket dir.

Test Plan

• cargo clippy --workspace --all-targets --all-features -- -D warnings, cargo +nightly fmt --all -- --check clean; engine-zmq-client (incl. mock-engine e2e for both protocols), wfaas, and smg zmq_client suites green.
• Cross-language contract: Python msgspec-encoded fixtures decode in the Rust codec and vice versa (field order, arity, float64 logprob columns).
• Live e2e on GB300: real TokenSpeed engine (Qwen3-0.6B, branch lightseekorg/tokenspeed#feat/zmq-msgpack) registered via POST /workers {"url":"ipc://...","connection_mode":"zmq","runtime_type":"tokenspeed"} — chat (content + reasoning_content), completions, sampled-token logprobs (--enable-output-logprobs), top-k/prompt-logprob rejections, and repeated worker registration with zero workflow-deadlock failures. The existing vLLM ZMQ path re-validated unchanged.

Follow-ups: engine-side handshake retry, DP>1 (coordinator/wave, task tracked), legacy /v1/completions logprobs rendering (pre-existing gap for all backends).

Checklist

• cargo +nightly fmt passes
• cargo clippy --all-targets --all-features -- -D warnings passes
• (Optional) Documentation updated
• (Optional) Please join us on Slack #sig-smg to discuss, review, and merge PRs

[coderabbitai]

Caution

Review failed

The pull request is closed.

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 05451134-7ab7-4562-9893-479480a5d656

📥 Commits

Reviewing files that changed from the base of the PR and between 244fd62 and 4443132.

📒 Files selected for processing (31)

• bindings/python/src/lib.rs
• bindings/python/src/smg/router.py
• bindings/python/src/smg/router_args.py
• bindings/python/src/smg/serve.py
• bindings/python/tests/test_serve.py
• crates/engine_zmq_client/src/connector.rs
• crates/engine_zmq_client/src/lib.rs
• crates/engine_zmq_client/src/mock_engine.rs
• crates/engine_zmq_client/src/protocol/handshake.rs
• crates/engine_zmq_client/src/protocol/mod.rs
• crates/engine_zmq_client/src/protocol/tokenspeed/mod.rs
• crates/engine_zmq_client/src/protocol/tokenspeed/output.rs
• crates/engine_zmq_client/src/protocol/tokenspeed/request.rs
• crates/engine_zmq_client/src/protocol/tokenspeed/sampling.rs
• crates/engine_zmq_client/src/protocol/vllm/mod.rs
• crates/engine_zmq_client/src/transport.rs
• crates/protocols/src/worker.rs
• crates/workflow/src/engine.rs
• crates/workflow/tests/workflow_test.rs
• model_gateway/src/config/builder.rs
• model_gateway/src/config/types.rs
• model_gateway/src/main.rs
• model_gateway/src/routers/grpc/backend_client.rs
• model_gateway/src/routers/grpc/common/stages/encode.rs
• model_gateway/src/routers/grpc/zmq_client.rs
• model_gateway/src/worker/builder.rs
• model_gateway/src/worker/worker.rs
• model_gateway/src/workflow/job_queue.rs
• model_gateway/src/workflow/steps/local/create_worker.rs
• model_gateway/src/workflow/steps/local/detect_backend.rs
• model_gateway/src/workflow/steps/local/discover_metadata.rs

---

📝 Walkthrough

Summary by CodeRabbit

• New Features

  • Added TokenSpeed as a supported serving backend alongside vLLM.
  • Added TokenSpeed ZMQ worker launching, request handling, streaming output, sampling, and logprob support.
  • Added backend-aware runtime selection and optional ZMQ handshake address configuration.
  • Added support for multiple generated choices and backend-specific output handling.

• Bug Fixes

  • Improved workflow completion reliability for skipped and conditionally executed steps.

• Tests

  • Added comprehensive coverage for TokenSpeed integration, ZMQ routing, handshake configuration, and workflow regressions.

Walkthrough

Changes

TokenSpeed ZMQ integration

Layer / File(s)	Summary
Generic engine protocol and transport crates/engine_zmq_client/src/protocol/*, crates/engine_zmq_client/src/connector.rs, crates/engine_zmq_client/src/transport.rs	The ZMQ client now supports generic engine protocols, batches, outputs, and load data. The vLLM adapter uses the shared protocol contracts.
TokenSpeed wire protocol crates/engine_zmq_client/src/protocol/tokenspeed/*	TokenSpeed request, sampling, output, framing, positional MessagePack encoding, decoding, normalization, and compatibility tests were added.
Gateway runtime and worker integration model_gateway/src/main.rs, model_gateway/src/routers/grpc/*, model_gateway/src/config/*, model_gateway/src/worker/*, model_gateway/src/workflow/*	The gateway selects vLLM or TokenSpeed runtimes, maps streams, propagates handshake settings, and preserves runtime metadata.
Python serving and router wiring bindings/python/src/*, bindings/python/tests/test_serve.py	Python serving adds the TokenSpeed launcher, CLI parsing, ZMQ validation, backend passthrough, runtime configuration, and tests.

Workflow completion signaling

Layer / File(s)	Summary
Completion signaling and regression coverage crates/workflow/src/engine.rs, crates/workflow/tests/workflow_test.rs	The scheduler drains queued completions before deadlock detection. Skip and failure signals are emitted with tracker updates. Regression tests cover skipped dependency chains.

Estimated code review effort: 5 (Critical) | ~120 minutes

Possibly related issues

• ZMQ direct backend connections: bypass the Python servicer/AsyncLLM for same-host deployments #2000 — The PR implements the issue’s ZMQ direct-backend work for vLLM and TokenSpeed connectors and runtime wiring.

Possibly related PRs

• smg-project/smg#2015 — Extends the earlier ZMQ backend integration with generalized client, worker launcher, runtime, and protocol support for TokenSpeed.
• smg-project/smg#2019 — Shares the BackendClient and router backend-runtime handling used by this integration.
• smg-project/smg#2032 — Covers the same TokenSpeed ZMQ support components and functions.

Suggested reviewers: catherinesue, key4ng

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch feat/zmq-tokenspeed

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}