Conversation
slin1237
requested review from
CatherineSue,
gongwei-130 and
key4ng
as code owners
August 2, 2026 21:31
📝 WalkthroughSummary by CodeRabbit
WalkthroughChangesTokenSpeed CLI and runtime configuration
Generic ZMQ protocol and gateway
Worker and workflow changes
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related issues
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
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. Comment |
github-actions
Bot
added
python-bindings
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Thorough review of all 22 changed files (~2,260 lines added). No issues found — this is a high-quality PR.
Key observations:
EngineProtocoltrait seam is well-designed: the connector and transport are cleanly generic over the protocol, with zero-sized type implementations for each engine family. No runtime dispatch overhead in the hot path.- Loud rejection boundary in both
translate_requestandtranslate_request_tokenspeedcorrectly fails requests the wire cannot honor (structured output,n>1, prompt logprobs, stop strings, logit_bias, nonzero DP rank for TokenSpeed) instead of silently narrowing. - Finish-tick chunk/complete split correctly handles the case where the terminal output carries new tokens: emits a
Chunkfirst (so streaming frontends decode the last token) and holds back the cumulativeCompletefor the next poll. - Deadlock fix in the workflow engine is correct: the re-drain catches completions that arrive between Phase 0's drain and Phase 1's tracker read, and the
run_ifpaths now atomically send their completion inside the tracker write-lock scope to close the race window. The 50-iteration stress test provides confidence. discover_metadatareturningNonefor ZMQ runtime preserves an explicitly configured runtime (previously it always overwrote with "vllm"), which is the key plumbing change that makes TokenSpeed registration work.- Cross-language FNV pinning (same vectors in Rust
derive_handshake_portand Python_zmq_handshake_port) is a nice contract test that prevents silent port-agreement breakage. - f64→f32 logprob downcast for TokenSpeed is an acceptable precision tradeoff given the proto column type constraint, and is documented.
0 🔴 Important · 0 🟡 Nit · 0 🟣 Pre-existing
slin1237
force-pushed
the
feat/zmq-tokenspeed
branch
from
August 3, 2026 03:08
d4d092a to
5baa27d
Compare
CatherineSue
approved these changes
Aug 3, 2026
| sub.request_id = format!("{}-{i}", req.request_id); | ||
| if let Some(sp) = sub.sampling_params.as_mut() { | ||
| sp.n = 1; | ||
| sp.seed = sp.seed.map(|seed| seed.wrapping_add(i as i32)); |
There was a problem hiding this comment.
🟡 Nit: wrapping_add can produce a negative i32 when seed is close to i32::MAX. On the vLLM path this is fine (negative i64 seeds are still deterministic), but on the TokenSpeed path translate_sampling_tokenspeed silently drops negative seeds via u64::try_from(seed).ok(), making those sub-requests non-deterministic when the user expected determinism. Practically impossible to hit (requires seed ≈ i32::MAX with n > 1), but saturating_add would be a safer choice — it would clamp to i32::MAX and give identical samples for the overflow subs (undesirable but at least predictable).
Review Summary0 🔴 Important · 1 🟡 Nit · 0 🟣 Pre-existing Thoroughly reviewed all 31 files in this PR. The engine-neutral The one nit posted is on seed derivation in No blocking issues found. Not approving per synchronize-event policy. |
slin1237
force-pushed
the
feat/zmq-tokenspeed
branch
from
August 3, 2026 05:44
5baa27d to
dcedacc
Compare
Make the ZMQ direct-backend stack engine-neutral and add TokenSpeed as the second wire protocol, speaking its msgpack-native tagged msgspec structs directly — a same-host TokenSpeed scheduler is driven over ipc:// with no Python servicer hop. engine-zmq-client: - EngineProtocol trait seams the shared transport/connector (handshake, ROUTER/DEALER identity framing, output loop, abort) away from the per-engine wire structs; the vLLM protocol keeps its behavior. - TokenSpeed protocol speaks the engine's native tagged structs: the tokenized generate request is emitted as the tagged positional prefix through `stream` (nested native SamplingParams, normalized frontend- side), and the per-step output decodes the tagged slim batch struct (token ids, finish reasons, token counts, sampled-token logprob columns). Tag-validated decode, trailing-field tolerance, and pinned cross-language byte vectors from the Python encoder. - Handshake structs move to a neutral protocol/handshake.rs; a neutral EngineLoad replaces engine-specific stats in the shared seam; max_num_batched_tokens widened to i64 (-1 = disabled). Gateway: - ZmqEngineClient selects the protocol from the worker's explicit runtime_type; unsupported runtimes are rejected before the handshake and at registration. BackendClient reports the actual runtime. - n>1 is fanned out frontend-side: n single-sample wire requests with per-sub rids and deterministic seed derivation, streams interleaved with per-choice proto indexes, cumulative usage counted once, drop aborts every sub. - Streams emit a finish-tick's tokens as a chunk before the cumulative Complete, so streaming never drops the last token. Sampled-token logprobs are wired end-to-end on both ZMQ protocols. - No silent narrowing: structured-output constraints, top-k/prompt logprobs, stop strings, logit_bias, and nonzero data_parallel_rank fail loudly with invalid_argument. - Workers honor an optional WorkerSpec.zmq_handshake_address bind override; the FNV port derivation from the ipc path remains the no-config default (doc + conformance vectors pinned against the Python launcher mirror). - detect_backend/discover_metadata preserve an explicitly configured ZMQ runtime; unspecified still defaults to vLLM with a warning. smg serve / config: - `smg serve --backend tokenspeed --connection-mode zmq` launches the engine headless (`python -m tokenspeed.cli serve --headless` with the derived --data-parallel-rpc-port), mirroring the vLLM zmq launcher; dense data parallelism = N independent workers. - `--backend` pins the startup ZMQ worker runtime through both config conversion paths and the Python bindings; per-user ZMQ socket dir (SMG_ZMQ_SOCKET_DIR override). workflow engine: - Fix a spurious "Workflow deadlocked" failure: completions landing between the drain and the tracker read are now re-drained in the deadlock branch, and the run_if paths send their completion inside the tracker lock scope. Regression tests included. Validated live on GB300 against a real TokenSpeed engine (Qwen3-0.6B): chat (content + reasoning_content), completions, streaming, n=2, sampled-token logprobs, loud rejections, engine-side invalid-request aborts, ENGINE_CORE_DEAD death detection, and the handshake-address override with a bare-default engine. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
slin1237
force-pushed
the
feat/zmq-tokenspeed
branch
from
August 3, 2026 15:40
dcedacc to
4443132
Compare
|
👋 The PR description doesn't fully follow
Please update the PR description so reviewers have the context they need. |
Member
Author
|
Superseded by a fresh PR against main. This PR was created stacked on #2015 (feat/zmq-direct-backend); after #2015 merged, GitHub keeps it flagged as a stacked PR and refuses both the sync GraphQL and REST merge endpoints. Reopening the identical rebased commit as a non-stacked PR to unblock the merge. |
2 of 4 tasks
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
crates/workflow/src/engine.rs (1)
696-716: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win🔴 Important: Persist the terminal step state before publishing its completion.
The scheduler can process the queued completion and finalize the workflow before
state_store.update(...).awaitfinishes.wait_for_completioncan then callcleanup_if_terminal, and the ignored update can fail after the workflow reports completion. This can leave a skipped or failedrun_ifstep without its terminal state.
crates/workflow/src/engine.rs#L696-L716: PersistStepStatus::Skippedbefore removing the step fromrunningand sendingStepResult::Skipinside the tracker lock.crates/workflow/src/engine.rs#L728-L749: PersistStepStatus::Failedandlast_errorbefore removing the step fromrunningand sendingStepResult::Failureinside the tracker lock.Handle state-store and completion-send errors explicitly. As per coding guidelines, do not silently fall back to
Noneor a default when configuration validation should fail loudly.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/workflow/src/engine.rs` around lines 696 - 716, The skip and failure completion paths in crates/workflow/src/engine.rs#L696-L716 and crates/workflow/src/engine.rs#L728-L749 must persist terminal state before publishing completion. In both paths, await and explicitly handle the state_store.update result for StepStatus::Skipped or StepStatus::Failed with last_error, then remove the step from running and send the completion while handling send errors explicitly; do not ignore failures or substitute None/default values when validation should fail loudly.Source: Coding guidelines
🧹 Nitpick comments (3)
model_gateway/src/routers/grpc/zmq_client.rs (1)
90-109: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win🟡 Nit: Add tests for the two connect-time rejection branches.
Lines 90-96 reject a TokenSpeed backend with
engine_count > 1. Lines 100-109 reject any runtime other thanVllm,TokenSpeed, orUnspecified. The test module covers the request-level rejections but not these two. Both branches return before the handshake, so a unit test needs no mock engine and stays fast.💚 Proposed tests
#[tokio::test] async fn tokenspeed_rejects_multi_engine_before_handshake() { let err = ZmqEngineClient::connect( "tcp://127.0.0.1:1", "ipc:///tmp/unused-in", "ipc:///tmp/unused-out", 2, "m".to_string(), RuntimeType::TokenSpeed, Duration::from_millis(50), ) .await .expect_err("DP>1 must be rejected"); assert!(err.to_string().contains("single engine"), "{err}"); } #[tokio::test] async fn unsupported_runtime_is_rejected_before_handshake() { let err = ZmqEngineClient::connect( "tcp://127.0.0.1:1", "ipc:///tmp/unused-in", "ipc:///tmp/unused-out", 1, "m".to_string(), RuntimeType::Sglang, Duration::from_millis(50), ) .await .expect_err("sglang has no ZMQ engine adapter"); assert!(err.to_string().contains("no engine implementation"), "{err}"); }As per coding guidelines: "Run the pr-test-analyzer agent to verify that tests adequately cover new or changed functionality."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@model_gateway/src/routers/grpc/zmq_client.rs` around lines 90 - 109, Add unit tests in the existing test module for ZmqEngineClient::connect covering both connect-time rejection branches: TokenSpeed with engine_count greater than one and an unsupported runtime such as RuntimeType::Sglang. Use unreachable endpoints and a short timeout to verify each returns before the handshake, and assert the error messages identify the single-engine and missing-engine-implementation conditions; run the pr-test-analyzer agent to confirm coverage.Source: Coding guidelines
crates/engine_zmq_client/src/lib.rs (1)
38-42: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value🟡 Nit — Re-export
EngineLoadnext to the other protocol types.
Client::engine_loadreturnsOption<EngineLoad>, butEngineLoadis not in the crate-root re-export list. Callers must reach it throughprotocol::EngineLoadwhileEngineBatchandEngineOutputare available at the root. Add it for a consistent public surface.♻️ Proposed change
-pub use protocol::{EngineBatch, EngineOutput, EngineProtocol}; +pub use protocol::{EngineBatch, EngineLoad, EngineOutput, EngineProtocol};🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/engine_zmq_client/src/lib.rs` around lines 38 - 42, Update the crate-root protocol re-exports to include EngineLoad alongside EngineBatch, EngineOutput, and EngineProtocol, so the type returned by Client::engine_load is publicly available consistently.crates/engine_zmq_client/src/protocol/tokenspeed/mod.rs (1)
109-112: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win🟡 Nit —
validatedoes not enforce the documentedn == 1invariant.
crates/engine_zmq_client/src/protocol/tokenspeed/sampling.rsline 90 documents that "the transport validatesn == 1". Thisvalidateimplementation accepts every request unconditionally. A caller that setssampling_params.n > 1then reaches the engine, which storesnand never fans out, so the caller silently receives one completion instead ofn. Either enforce the check here or correct the doc comment insampling.rs.♻️ Proposed fix to enforce the invariant
- fn validate(_request: &Self::Request) -> Result<()> { - // The tokenized text path has no fields this client cannot represent. - Ok(()) - } + fn validate(request: &Self::Request) -> Result<()> { + // n > 1 is fanned out by the gateway; the engine stores n without + // acting on it, so a value above 1 would silently drop completions. + if request.sampling_params.n != 1 { + return Err(Error::UnsupportedField { + context: "TokenizedGenerateReqInput", + field: "sampling_params.n", + }); + } + Ok(()) + }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/engine_zmq_client/src/protocol/tokenspeed/mod.rs` around lines 109 - 112, Update tokenspeed’s `validate` method to inspect the request’s sampling parameters and reject any request whose `n` is not 1, preserving success for valid single-completion requests. Keep the documented transport invariant in `sampling.rs` aligned with this enforcement.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/engine_zmq_client/src/connector.rs`:
- Around line 242-253: Update the submit flow around register and P::encode_add
so encoding occurs before the registry entry is created, or otherwise remove the
request_id entry when encoding fails. Preserve the existing rollback for
send_to_engine failures and ensure successful requests still register before
sending.
- Around line 628-634: Replace the std::mem::forget(ns) call in the IpcNamespace
setup with a scoped binding that keeps ns alive for the full test while allowing
it to be dropped normally for cleanup. Preserve the endpoint values and ensure
ns remains in scope through all connection and assertion logic.
In `@crates/workflow/tests/workflow_test.rs`:
- Around line 993-1094: Add a deterministic regression test near
test_run_if_false_mid_chain_completes using a custom StateStore whose
get_context() returns an error; configure a run_if step to exercise that failure
path, then assert the workflow reaches its expected failure status and completes
without a deadlock or timeout.
In `@model_gateway/src/main.rs`:
- Line 1465: Add Go binding support for startup_worker_runtime_type alongside
the existing builder call, matching the Python binding’s ZMQ-only
Vllm/Tokenspeed mapping. Expose the field in the public Go configuration API and
propagate it through the relevant Go-to-Rust conversion or builder path,
preserving existing behavior for other runtime types.
- Around line 1401-1414: Update the ZMQ branch that initializes
startup_worker_runtime_type to reject explicitly selected backends other than
Backend::Vllm and Backend::Tokenspeed before leaving the runtime unpinned;
return the existing startup error type with a clear unsupported-backend message.
Preserve None for non-ZMQ connections and unspecified ZMQ backends, and add
coverage for direct --backend trtllm with --worker-urls ipc://....
In `@model_gateway/src/workflow/steps/local/create_worker.rs`:
- Around line 380-391: Update validate_zmq_handshake_override to validate the
configured zmq_handshake_address scheme for ZMQ workers, accepting only tcp://
values and returning an error for ipc:// or any other non-TCP value. Preserve
the existing error for overrides used with non-ZMQ connection modes and ensure
invalid configurations fail before registration.
---
Outside diff comments:
In `@crates/workflow/src/engine.rs`:
- Around line 696-716: The skip and failure completion paths in
crates/workflow/src/engine.rs#L696-L716 and
crates/workflow/src/engine.rs#L728-L749 must persist terminal state before
publishing completion. In both paths, await and explicitly handle the
state_store.update result for StepStatus::Skipped or StepStatus::Failed with
last_error, then remove the step from running and send the completion while
handling send errors explicitly; do not ignore failures or substitute
None/default values when validation should fail loudly.
---
Nitpick comments:
In `@crates/engine_zmq_client/src/lib.rs`:
- Around line 38-42: Update the crate-root protocol re-exports to include
EngineLoad alongside EngineBatch, EngineOutput, and EngineProtocol, so the type
returned by Client::engine_load is publicly available consistently.
In `@crates/engine_zmq_client/src/protocol/tokenspeed/mod.rs`:
- Around line 109-112: Update tokenspeed’s `validate` method to inspect the
request’s sampling parameters and reject any request whose `n` is not 1,
preserving success for valid single-completion requests. Keep the documented
transport invariant in `sampling.rs` aligned with this enforcement.
In `@model_gateway/src/routers/grpc/zmq_client.rs`:
- Around line 90-109: Add unit tests in the existing test module for
ZmqEngineClient::connect covering both connect-time rejection branches:
TokenSpeed with engine_count greater than one and an unsupported runtime such as
RuntimeType::Sglang. Use unreachable endpoints and a short timeout to verify
each returns before the handshake, and assert the error messages identify the
single-engine and missing-engine-implementation conditions; run the
pr-test-analyzer agent to confirm coverage.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 7ff49401-9a52-482f-ad3f-17c662df0e33
📒 Files selected for processing (31)
bindings/python/src/lib.rsbindings/python/src/smg/router.pybindings/python/src/smg/router_args.pybindings/python/src/smg/serve.pybindings/python/tests/test_serve.pycrates/engine_zmq_client/src/connector.rscrates/engine_zmq_client/src/lib.rscrates/engine_zmq_client/src/mock_engine.rscrates/engine_zmq_client/src/protocol/handshake.rscrates/engine_zmq_client/src/protocol/mod.rscrates/engine_zmq_client/src/protocol/tokenspeed/mod.rscrates/engine_zmq_client/src/protocol/tokenspeed/output.rscrates/engine_zmq_client/src/protocol/tokenspeed/request.rscrates/engine_zmq_client/src/protocol/tokenspeed/sampling.rscrates/engine_zmq_client/src/protocol/vllm/mod.rscrates/engine_zmq_client/src/transport.rscrates/protocols/src/worker.rscrates/workflow/src/engine.rscrates/workflow/tests/workflow_test.rsmodel_gateway/src/config/builder.rsmodel_gateway/src/config/types.rsmodel_gateway/src/main.rsmodel_gateway/src/routers/grpc/backend_client.rsmodel_gateway/src/routers/grpc/common/stages/encode.rsmodel_gateway/src/routers/grpc/zmq_client.rsmodel_gateway/src/worker/builder.rsmodel_gateway/src/worker/worker.rsmodel_gateway/src/workflow/job_queue.rsmodel_gateway/src/workflow/steps/local/create_worker.rsmodel_gateway/src/workflow/steps/local/detect_backend.rsmodel_gateway/src/workflow/steps/local/discover_metadata.rs
Comment on lines
242
to
253
| let receiver = self.inner.registry.lock().register(request_id.clone())?; | ||
|
|
||
| let payload = encode_msgpack(&request)?; | ||
| // Text path carries no aux tensor frames. | ||
| let (payload, aux_frames) = P::encode_add(&request)?; | ||
| if let Err(error) = self | ||
| .inner | ||
| .send_to_engine(&engine_id, EngineCoreRequestType::Add, payload, Vec::new()) | ||
| .send_to_engine(&engine_id, P::add_frame(), payload, aux_frames) | ||
| .await | ||
| { | ||
| // Roll back the registry entry so a failed send doesn't leak it. | ||
| self.inner.registry.lock().remove_all([&request_id]); | ||
| return Err(error); | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🔴 Important — Registry entry leaks when encode_add fails.
register runs at line 242, before P::encode_add at line 244. The rollback at line 251 only covers a send failure. If encode_add returns an error, submit returns early and leaves the request_id in the registry. The receiver is dropped, so no output can clear the entry, and every later submit with the same request_id fails with DuplicateRequestId until the client is dropped. Encode the payload before you register, or roll back on both error paths.
🐛 Proposed fix: encode before registering
let request_id = P::request_id(&request).to_string();
+ let (payload, aux_frames) = P::encode_add(&request)?;
let receiver = self.inner.registry.lock().register(request_id.clone())?;
- let (payload, aux_frames) = P::encode_add(&request)?;
if let Err(error) = self📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| let receiver = self.inner.registry.lock().register(request_id.clone())?; | |
| let payload = encode_msgpack(&request)?; | |
| // Text path carries no aux tensor frames. | |
| let (payload, aux_frames) = P::encode_add(&request)?; | |
| if let Err(error) = self | |
| .inner | |
| .send_to_engine(&engine_id, EngineCoreRequestType::Add, payload, Vec::new()) | |
| .send_to_engine(&engine_id, P::add_frame(), payload, aux_frames) | |
| .await | |
| { | |
| // Roll back the registry entry so a failed send doesn't leak it. | |
| self.inner.registry.lock().remove_all([&request_id]); | |
| return Err(error); | |
| } | |
| let (payload, aux_frames) = P::encode_add(&request)?; | |
| let receiver = self.inner.registry.lock().register(request_id.clone())?; | |
| if let Err(error) = self | |
| .inner | |
| .send_to_engine(&engine_id, P::add_frame(), payload, aux_frames) | |
| .await | |
| { | |
| // Roll back the registry entry so a failed send doesn't leak it. | |
| self.inner.registry.lock().remove_all([&request_id]); | |
| return Err(error); | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@crates/engine_zmq_client/src/connector.rs` around lines 242 - 253, Update the
submit flow around register and P::encode_add so encoding occurs before the
registry entry is created, or otherwise remove the request_id entry when
encoding fails. Preserve the existing rollback for send_to_engine failures and
ensure successful requests still register before sending.
Comment on lines
+628
to
+634
| let ns = IpcNamespace::new().unwrap(); | ||
| let (handshake, input, output) = ( | ||
| ns.handshake_endpoint(), | ||
| ns.input_endpoint(), | ||
| ns.output_endpoint(), | ||
| ); | ||
| std::mem::forget(ns); |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🟡 Nit — Replace std::mem::forget(ns) with a scoped binding.
IpcNamespace owns the socket tempdir, as documented at lines 437-438. std::mem::forget leaks the temp directory and the ipc socket files on every run of this test. A plain binding keeps the namespace alive for the whole test and still cleans up on drop, which matches the connect() helper.
🧹 Proposed fix
- let ns = IpcNamespace::new().unwrap();
- let (handshake, input, output) = (
- ns.handshake_endpoint(),
- ns.input_endpoint(),
- ns.output_endpoint(),
- );
- std::mem::forget(ns);
+ let _ns = IpcNamespace::new().unwrap();
+ let (handshake, input, output) = (
+ _ns.handshake_endpoint(),
+ _ns.input_endpoint(),
+ _ns.output_endpoint(),
+ );📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| let ns = IpcNamespace::new().unwrap(); | |
| let (handshake, input, output) = ( | |
| ns.handshake_endpoint(), | |
| ns.input_endpoint(), | |
| ns.output_endpoint(), | |
| ); | |
| std::mem::forget(ns); | |
| let _ns = IpcNamespace::new().unwrap(); | |
| let (handshake, input, output) = ( | |
| _ns.handshake_endpoint(), | |
| _ns.input_endpoint(), | |
| _ns.output_endpoint(), | |
| ); |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@crates/engine_zmq_client/src/connector.rs` around lines 628 - 634, Replace
the std::mem::forget(ns) call in the IpcNamespace setup with a scoped binding
that keeps ns alive for the full test while allowing it to be dropped normally
for cleanup. Preserve the endpoint values and ensure ns remains in scope through
all connection and assertion logic.
Comment on lines
+993
to
+1094
| /// Regression test: instantly-skipped run_if steps feeding a dependent chain | ||
| /// must not trigger a spurious "Workflow deadlocked" failure. The skip signal | ||
| /// must be sent in the same tracker lock scope as the removal from `running`, | ||
| /// otherwise the scheduler can observe running == 0 with the completion unsent. | ||
| #[tokio::test] | ||
| async fn test_run_if_skip_chain_no_spurious_deadlock() { | ||
| for i in 0..50 { | ||
| let engine: WorkflowEngine<TestWorkflowData> = WorkflowEngine::new(); | ||
|
|
||
| let workflow = WorkflowDefinition::new( | ||
| "run_if_skip_chain_workflow", | ||
| "Run If Skip Chain Deadlock Regression", | ||
| ) | ||
| .add_step( | ||
| StepDefinition::new("root", "Root", Arc::new(AlwaysSucceedStep)).run_if(|_ctx| false), | ||
| ) | ||
| .add_step( | ||
| StepDefinition::new("mid", "Mid", Arc::new(AlwaysSucceedStep)) | ||
| .depends_on(&["root"]) | ||
| .run_if(|_ctx| false), | ||
| ) | ||
| .add_step( | ||
| StepDefinition::new("leaf", "Leaf", Arc::new(AlwaysSucceedStep)).depends_on(&["mid"]), | ||
| ); | ||
|
|
||
| let workflow_id = workflow.id.clone(); | ||
| engine.register_workflow(workflow).unwrap(); | ||
|
|
||
| let instance_id = engine | ||
| .start_workflow(workflow_id, TestWorkflowData::default()) | ||
| .await | ||
| .unwrap(); | ||
|
|
||
| let result = engine | ||
| .wait_for_completion(instance_id, "skip-chain", Duration::from_secs(5)) | ||
| .await; | ||
| assert!(result.is_ok(), "iteration {i} failed: {result:?}"); | ||
| } | ||
| } | ||
|
|
||
| /// A -> B(run_if=false) -> C: the chain completes with B skipped and C executed. | ||
| #[tokio::test] | ||
| async fn test_run_if_false_mid_chain_completes() { | ||
| use tokio::time::sleep; | ||
|
|
||
| let engine: WorkflowEngine<TestWorkflowData> = WorkflowEngine::new(); | ||
|
|
||
| let executed = Arc::new(AtomicU32::new(0)); | ||
| let executed_clone = Arc::clone(&executed); | ||
|
|
||
| struct TrackingStep { | ||
| counter: Arc<AtomicU32>, | ||
| } | ||
|
|
||
| #[async_trait::async_trait] | ||
| impl StepExecutor<TestWorkflowData> for TrackingStep { | ||
| async fn execute( | ||
| &self, | ||
| _context: &mut WorkflowContext<TestWorkflowData>, | ||
| ) -> WorkflowResult<StepResult> { | ||
| self.counter.fetch_add(1, Ordering::SeqCst); | ||
| Ok(StepResult::Success) | ||
| } | ||
| } | ||
|
|
||
| let workflow = WorkflowDefinition::new("run_if_mid_chain_workflow", "Run If Mid Chain Test") | ||
| .add_step(StepDefinition::new("a", "A", Arc::new(AlwaysSucceedStep))) | ||
| .add_step( | ||
| StepDefinition::new("b", "B", Arc::new(AlwaysSucceedStep)) | ||
| .depends_on(&["a"]) | ||
| .run_if(|_ctx| false), | ||
| ) | ||
| .add_step( | ||
| StepDefinition::new("c", "C", Arc::new(TrackingStep { counter: executed })) | ||
| .depends_on(&["b"]), | ||
| ); | ||
|
|
||
| let workflow_id = workflow.id.clone(); | ||
| engine.register_workflow(workflow).unwrap(); | ||
|
|
||
| let instance_id = engine | ||
| .start_workflow(workflow_id, TestWorkflowData::default()) | ||
| .await | ||
| .unwrap(); | ||
|
|
||
| // Poll instead of wait_for_completion so step states survive for inspection. | ||
| let mut state = engine.get_status(instance_id).await.unwrap(); | ||
| for _ in 0..100 { | ||
| if state.status != WorkflowStatus::Running && state.status != WorkflowStatus::Pending { | ||
| break; | ||
| } | ||
| sleep(Duration::from_millis(50)).await; | ||
| state = engine.get_status(instance_id).await.unwrap(); | ||
| } | ||
|
|
||
| assert_eq!(state.status, WorkflowStatus::Completed); | ||
| assert_eq!( | ||
| state.step_states.get(&StepId::new("b")).unwrap().status, | ||
| StepStatus::Skipped | ||
| ); | ||
| assert_eq!(executed_clone.load(Ordering::SeqCst), 1); | ||
| } |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
ast-grep outline crates/workflow --items all --match 'StateStore|WorkflowEngine'
rg -n --type rust -C 3 \
'get_context.*Err|Err\(.*get_context|impl.*StateStore|Mock.*StateStore' \
crates/workflow
rg -n --type rust -C 3 \
'run_if.*context|context.*run_if|run_if' \
crates/workflow/testsRepository: smg-project/smg
Length of output: 12976
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== state.rs relevant section =="
sed -n '1,140p' crates/workflow/src/state.rs
echo
echo "== engine run_if/context references =="
rg -n --type rust -C 4 \
'get_context|RunIf|run_if|skip|Skipped|deadlock|Workflow deadlocked' \
crates/workflow/src crates/workflow/tests
echo
echo "== WorkflowContext/get_workflow_context definitions =="
rg -n --type rust -C 4 \
'struct WorkflowContext|impl WorkflowContext|get_workflow_context|WorkflowContext<' \
crates/workflow/srcRepository: smg-project/smg
Length of output: 50371
🟡 Nit: Add coverage for run_if context retrieval failure.
run_if fails the step when state_store.get_context() returns an error, but the tests only cover true/false/context-data conditions. Add a deterministic test using a custom StateStore that returns an error for get_context() and assert the workflow/failure status without deadlock.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@crates/workflow/tests/workflow_test.rs` around lines 993 - 1094, Add a
deterministic regression test near test_run_if_false_mid_chain_completes using a
custom StateStore whose get_context() returns an error; configure a run_if step
to exercise that failure path, then assert the workflow reaches its expected
failure status and completes without a deadlock or timeout.
Source: Coding guidelines
Comment on lines
+1401
to
+1414
| // `--backend` normally only steers the routing mode. Over ZMQ it | ||
| // additionally pins the startup workers' runtime: the shared EngineCore | ||
| // handshake carries no engine identity, so the wire protocol cannot be | ||
| // probed and must be declared. HTTP/gRPC keep auto-detection (None). | ||
| let startup_worker_runtime_type = if connection_mode == ConnectionMode::Zmq { | ||
| match self.backend { | ||
| Some(Backend::Vllm) => Some(RuntimeType::Vllm), | ||
| Some(Backend::Tokenspeed) => Some(RuntimeType::TokenSpeed), | ||
| _ => None, | ||
| } | ||
| } else { | ||
| None | ||
| }; | ||
|
|
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== repo files around model_gateway =="
git ls-files | rg '(^|/)model_gateway/src/main\.rs|detect_backend\.rs|parse_serve_args|types|bindings/python/src/lib\.rs' || true
echo
echo "== main.rs outline around relevant symbols =="
if [ -f model_gateway/src/main.rs ]; then
wc -l model_gateway/src/main.rs
ast-grep outline model_gateway/src/main.rs --match backend --view expanded || true
ast-grep outline model_gateway/src/main.rs --match startup_worker_runtime_type --view expanded || true
echo
sed -n '60,115p' model_gateway/src/main.rs
echo
sed -n '1380,1425p' model_gateway/src/main.rs
echo
sed -n '1450,1475p' model_gateway/src/main.rs
echo
sed -n '1938,2015p' model_gateway/src/main.rs
fi
echo
echo "== backend/zmq validation symbols =="
rg -n "startup_worker_runtime_type|detect_backend|runtime_type|ConfigError::InvalidValue|Trtllm|Tokenspeed|Vllm|Backend::" -S .Repository: smg-project/smg
Length of output: 50371
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== job_queue startup pin/default handling =="
sed -n '540,585p' model_gateway/src/workflow/job_queue.rs
sed -n '780,820p' model_gateway/src/workflow/job_queue.rs
echo
echo "== detect_backend runtime default =="
sed -n '1,120p' model_gateway/pkg/modules/local/detect_backend.rs 2>/dev/null || sed -n '1,140p' model_gateway/src/workflow/steps/local/detect_backend.rs
echo
echo "== Python launcher direct-backend/ZMQ restriction search =="
rg -n "tokenspeed|direct-backend|default_worker|worker_url|startup_worker_runtime_type|BackendType::|backend ==" bindings/python/src/smg/serve.py bindings/python/src/lib.rs clients/python/smg_client -S
echo
echo "== Rust binding startup_worker_runtime_type construction =="
sed -n '510,565p' bindings/python/src/lib.rs
sed -n '735,765p' bindings/python/src/lib.rs
echo
echo "== Go API types/search candidate files =="
git ls-files | rg 'bindings/golang|golang|go-sdk|sdk|grpc|model_gateway' | sed -n '1,200p'
rg -n "runtimeType|startupWorkerRuntimeType|runtime_type|BackendType|Tokenspeed|Trtllm|Sglang" bindings/golang crates clients examples -S 2>/dev/null | sed -n '1,220p' || true
echo
echo "== deterministic config conversion behavior for unsupported ZMQ backends =="
python3 - <<'PY'
from pathlib import Path
import re
main = Path("model_gateway/src/main.rs").read_text()
m = re.search(r'let startup_worker_runtime_type = if connection_mode == ConnectionMode::Zmq \{(?P<body>.*?)\n \} else \{(?P<else>.*?)\n \};', main, re.S)
print("found_startup_mapping:", bool(m))
if m:
body = re.sub(r'\s+', ' ', m.group('body'))
print("has_vllm_branch:", "Some(Backend::Vllm) => Some(RuntimeType::Vllm)" in body)
print("has_tokenspeed_branch:", "Some(Backend::Tokenspeed) => Some(RuntimeType::TokenSpeed)" in body)
print("has_else_default:", "_ => None" in body and "RuntimeType::External" not in body)
print("has_immediate_error:", "return Err(ConfigError::InvalidValue" in body)
if Path("bindings/python/tests/test_serve.py").exists():
txt = Path("bindings/python/tests/test_serve.py").read_text()
print("python_test_has_trtllm_unsupported_zmq_error:", ("TrtllmWorkerLauncher" in txt) and ("ZMQ" in txt or "ipc" in txt) and ("ValueError" in txt or "TypeError" in txt or "only support" in txt.lower()))
PYRepository: smg-project/smg
Length of output: 50371
🔴 Important — Reject unsupported ZMQ backends before leaving startup runtime unpinned.
startup_worker_runtime_type only maps Backend::Vllm and Backend::Tokenspeed for ZMQ; every other backend falls to _ => None. Direct CLI invocation can accept --backend trtllm --worker-urls ipc://..., and the unpinned ZMQ path later defaults to vLLM during backend detection. That turns an explicit unsupported backend request into a wire-protocol mismatch or hang instead of a startup error. Reject unsupported ZMQ backends immediately, and add a test covering --backend trtllm --worker-urls ipc://....
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@model_gateway/src/main.rs` around lines 1401 - 1414, Update the ZMQ branch
that initializes startup_worker_runtime_type to reject explicitly selected
backends other than Backend::Vllm and Backend::Tokenspeed before leaving the
runtime unpinned; return the existing startup error type with a clear
unsupported-backend message. Preserve None for non-ZMQ connections and
unspecified ZMQ backends, and add coverage for direct --backend trtllm with
--worker-urls ipc://....
Source: Coding guidelines
| .mode(mode) | ||
| .policy(policy) | ||
| .connection_mode(connection_mode) | ||
| .startup_worker_runtime_type(startup_worker_runtime_type) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Confirm startup_worker_runtime_type parity across bindings.
rg -n -C5 'startup_worker_runtime_type' bindings/python/src/lib.rs
fd -e go . bindings/golang 2>/dev/null | xargs -r rg -n -i 'runtime_type|RuntimeType' 2>/dev/nullRepository: smg-project/smg
Length of output: 1557
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== model_gateway/src/main.rs relevant section =="
rg -n -C8 'startup_worker_runtime_type|connection_mode|BackendType|RuntimeType|Tokenspeed' model_gateway/src/main.rs | sed -n '1,260p'
echo
echo "== Python bindings occurrences =="
rg -n -C8 'backend|BackendType|connection_mode|RuntimeType|Tokenspeed|startup_worker_runtime_type' bindings/python/src/lib.rs | sed -n '1,220p'
echo
echo "== Go SDK files =="
git ls-files 'bindings/golang/*' 'bindings/golang/**/*' 2>/dev/null | sed -n '1,120p'
echo
echo "== Go SDK runtime/config occurrences =="
fd -e go . bindings/golang 2>/dev/null | xargs -r rg -n -C6 'runtime_type|RuntimeType|StartupWorker|backend|Backend|ConnectionMode|connection_mode|Tokenspeed|tokenspeed|VLLM|vllm' 2>/dev/null || trueRepository: smg-project/smg
Length of output: 23287
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== go-specific backend/connection config symbols =="
rg -n -C5 -i 'backend|backend_type|connection_mode|worker_url|worker_urls|vllm|tokenspeed|runtimetype|startup_worker' bindings/golang -g '*.go' -g '*.rs' 2>/dev/null || true
echo
echo "== go binding constructor / config method definitions =="
ast-grep outline bindings/golang/src/lib.rs 2>/dev/null | sed -n '1,220p' || true
rg -n -C8 '#\[pyo3|class Config|to_router_config|RouterConfig|backend|connection_mode|worker_urls' bindings/golang/src/lib.rs bindings/golang/src/client.rs bindings/golang/client.go 2>/dev/null || true
echo
echo "== Go example config backend usage =="
rg -n -C5 -i 'Backend|backend|backend_type|WorkerUrls|worker_urls' bindings/golang/examples/oai_server/config config.go bindings/golang/examples/oai_server/main.go 2>/dev/null || trueRepository: smg-project/smg
Length of output: 10084
Add Golang binding support for startup_worker_runtime_type.
The Python binding applies the same ZMQ-only Vllm/Tokenspeed mapping, but the Go binding does not expose this config field. Add equivalent Go support so the public config parity applies consistently across bindings.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@model_gateway/src/main.rs` at line 1465, Add Go binding support for
startup_worker_runtime_type alongside the existing builder call, matching the
Python binding’s ZMQ-only Vllm/Tokenspeed mapping. Expose the field in the
public Go configuration API and propagate it through the relevant Go-to-Rust
conversion or builder path, preserving existing behavior for other runtime
types.
Source: Coding guidelines
Comment on lines
+380
to
+391
| fn validate_zmq_handshake_override( | ||
| config: &WorkerSpec, | ||
| connection_mode: ConnectionMode, | ||
| ) -> Result<(), String> { | ||
| if config.zmq_handshake_address.is_some() && connection_mode != ConnectionMode::Zmq { | ||
| return Err(format!( | ||
| "worker {} sets zmq_handshake_address but its connection mode is \ | ||
| {connection_mode:?}: the field is only meaningful for ZMQ workers", | ||
| config.url | ||
| )); | ||
| } | ||
| Ok(()) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🔴 Important: Reject non-TCP handshake overrides during registration.
validate_zmq_handshake_override accepts ipc:// and other non-TCP values for ZMQ workers. The field contract requires tcp://. Reject these values here so invalid worker configuration fails before the worker enters the registration and health-check flow.
Proposed fix
fn validate_zmq_handshake_override(
config: &WorkerSpec,
connection_mode: ConnectionMode,
) -> Result<(), String> {
- if config.zmq_handshake_address.is_some() && connection_mode != ConnectionMode::Zmq {
- return Err(format!(
- "worker {} sets zmq_handshake_address but its connection mode is \
- {connection_mode:?}: the field is only meaningful for ZMQ workers",
- config.url
- ));
+ if let Some(address) = &config.zmq_handshake_address {
+ if connection_mode != ConnectionMode::Zmq {
+ return Err(format!(
+ "worker {} sets zmq_handshake_address but its connection mode is \
+ {connection_mode:?}: the field is only meaningful for ZMQ workers",
+ config.url
+ ));
+ }
+ if !address.starts_with("tcp://") {
+ return Err(format!(
+ "worker {} sets invalid zmq_handshake_address {address:?}: \
+ expected a tcp:// address",
+ config.url
+ ));
+ }
}
Ok(())
}As per coding guidelines, "Prioritize logic errors, production-breaking bugs, security vulnerabilities, missing error handling, broken cross-references, and incorrect defaults or configuration values."
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| fn validate_zmq_handshake_override( | |
| config: &WorkerSpec, | |
| connection_mode: ConnectionMode, | |
| ) -> Result<(), String> { | |
| if config.zmq_handshake_address.is_some() && connection_mode != ConnectionMode::Zmq { | |
| return Err(format!( | |
| "worker {} sets zmq_handshake_address but its connection mode is \ | |
| {connection_mode:?}: the field is only meaningful for ZMQ workers", | |
| config.url | |
| )); | |
| } | |
| Ok(()) | |
| fn validate_zmq_handshake_override( | |
| config: &WorkerSpec, | |
| connection_mode: ConnectionMode, | |
| ) -> Result<(), String> { | |
| if let Some(address) = &config.zmq_handshake_address { | |
| if connection_mode != ConnectionMode::Zmq { | |
| return Err(format!( | |
| "worker {} sets zmq_handshake_address but its connection mode is \ | |
| {connection_mode:?}: the field is only meaningful for ZMQ workers", | |
| config.url | |
| )); | |
| } | |
| if !address.starts_with("tcp://") { | |
| return Err(format!( | |
| "worker {} sets invalid zmq_handshake_address {address:?}: \ | |
| expected a tcp:// address", | |
| config.url | |
| )); | |
| } | |
| } | |
| Ok(()) | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@model_gateway/src/workflow/steps/local/create_worker.rs` around lines 380 -
391, Update validate_zmq_handshake_override to validate the configured
zmq_handshake_address scheme for ZMQ workers, accepting only tcp:// values and
returning an error for ipc:// or any other non-TCP value. Preserve the existing
error for overrides used with non-ZMQ connection modes and ensure invalid
configurations fail before registration.
Source: Coding guidelines
This was referenced Aug 3, 2026
This was referenced Aug 5, 2026
This was referenced Aug 13, 2026
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
2 participants
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.
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 newEngineProtocoltrait seams the shared transport/connector (handshake, ROUTER/DEALER identity framing, output loop, abort-on-drop) away from the per-engine wire structs.VllmProtocolkeeps the existing behavior;TokenSpeedProtocoladds the sglang-family msgpack tuples (WireTokenizedGenerateReq5-tuple,WireSamplingParams13-tuple,WireBatchTokenIDOut8-tuple with sampled-token logprob columns). Handshake structs move to a neutralprotocol/handshake.rs(re-exported for compatibility). A neutralEngineLoadreplaces the vLLM-specific stats type in the shared seam.zmq_client.rs):ZmqEngineClientselects the protocol from the worker's explicitruntime_type(tokenspeedvsvllm; 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 terminalCompletecarries the cumulative set, and a finish-tick's tokens are emitted as a chunk first so streaming never loses the last token.invalid_argumentinstead of degrading — structured-output constraints,n>1, top-k/prompt logprobs, stop strings (TokenSpeed),logit_bias(TokenSpeed), nonzerodata_parallel_rank. Sampled-token logprobs are wired end-to-end on both ZMQ protocols.detect_backend/discover_metadatano longer force ZMQ workers to vLLM — an explicitly configuredruntime_typesurvives to the built worker (unspecified still defaults to vLLM with a warning).BackendClient::runtime_type()reports the actual ZMQ backend runtime.wfaasfix: 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 therun_ifpaths send their completion inside the tracker lock scope. Regression tests included.smg serve: per-user ZMQ socket dir (SMG_ZMQ_SOCKET_DIRoverride), 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, genericconnector.rs/transport.rs,EngineCoreReadyResponse.max_num_batched_tokenswidened toi64(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_ifin-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 -- --checkclean;engine-zmq-client(incl. mock-engine e2e for both protocols),wfaas, andsmgzmq_client suites green.msgspec-encoded fixtures decode in the Rust codec and vice versa (field order, arity, float64 logprob columns).lightseekorg/tokenspeed#feat/zmq-msgpack) registered viaPOST /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/completionslogprobs rendering (pre-existing gap for all backends).Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspasses