feat(zmq): teach the direct-backend path to speak SGLang - #2042
Conversation
📝 WalkthroughSummary by CodeRabbit
WalkthroughSGLang direct-ZMQ support now spans Python scheduler launch, MessagePack request and output protocols, PUSH/PULL transport, gateway request streaming, runtime selection, and worker IPC socket cleanup. Existing vLLM and TokenSpeed paths now support optional protocol frames and ready responses. ChangesSGLang direct ZMQ integration
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant ZmqGateway
participant SglangScheduler
Client->>ZmqGateway: submit generation request
ZmqGateway->>SglangScheduler: send translated request over PUSH/PULL
SglangScheduler-->>ZmqGateway: stream MessagePack outputs
ZmqGateway-->>Client: return streaming response chunks
Possibly related PRs
Suggested labels: 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}
| .. | ||
| } = self; | ||
|
|
||
| Ok(rids | ||
| .into_iter() | ||
| .zip(finished_reasons) | ||
| .zip(output_ids) | ||
| .zip(prompt_tokens) | ||
| .zip(reasoning_tokens) |
There was a problem hiding this comment.
🟡 Nit: per_request_logprobs independently filter_maps None entries out of the val and idx columns. If SGLang ever produces misaligned None patterns between the two columns (e.g. val has a None at position 1 but idx doesn't), the resulting vecs would have different lengths, producing silently misaligned logprob data downstream.
In practice the columns should always have the same None structure (both describe the same token positions), so this is a theoretical concern — but a debug_assert_eq!(output.output_logprobs_val.len(), output.output_logprobs_idx.len()) in into_outputs would catch it early if the assumption ever breaks.
There was a problem hiding this comment.
Thorough review of the SGLang direct-ZMQ backend. This is a well-crafted PR — the protocol implementation is clean, the SlimTick/SlimOutput refactor nicely shares the stream-mapping code between TokenSpeed and SGLang, and the test coverage is strong (pinned wire vectors, round-trip codecs, e2e adapter tests, stale-socket cleanup).
Summary: 0 🔴 Important · 1 🟡 Nit · 0 🟣 Pre-existing
The one nit is a theoretical logprob alignment concern in the output decoder's filter_map — not a real bug given SGLang's wire guarantees, but a debug_assert would catch any future divergence cheaply.
54172f5 to
91f234d
Compare
Extend the same-host direct-ZMQ backend to SGLang engines. On this path SMG plays the SGLang TokenizerManager role: it binds the PUSH input and PULL output ipc sockets, sends already-tokenized generate requests, and receives batched token outputs straight off the wire, bypassing the Python gRPC servicer. - protocol/sglang: msgpack codec for the tokenizer<->scheduler wire — request, output, sampling, and token-id encoders that emit already-normalized SamplingParams (SMG owns normalization) - transport: PUSH/PULL sockets wait for the engine to attach via a socket monitor before admitting requests, since PUSH returns to sender rather than buffering when no peer is connected - connector: tag-dispatched output handling generalized across engines - gateway adapter: proto <-> sglang request/response translation - serve.py + _sglang_zmq_launcher.py: launch a headless SGLang scheduler (skip-tokenizer-init) wired to SMG's two sockets, forcing msgpack over the ZMQ wire Validated end-to-end against a live SGLang scheduler on GB300. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
91f234d to
8c5000b
Compare
There was a problem hiding this comment.
Incremental review (force-push → fell back to full-diff review)
I reviewed all 21 changed files across the SGLang ZMQ backend implementation. The code follows the established vLLM/TokenSpeed patterns well: the new protocol::sglang codec, the connect_push_pull no-handshake transport, the SlimTick/SlimOutput generalization, and the launcher/serve plumbing are all clean and internally consistent.
Findings: 0 🔴 · 0 🟡 · 0 🟣
The previous nit about independent per_request_logprobs filtering (output.rs:210) remains valid and is unchanged by this push — not re-posting.
Highlights:
- The
SlimTick/SlimOutputabstraction is a good refactor — it unifies the TokenSpeed and SGLang stream mappers with no duplication. - The
remove_stale_ipc_sockethelper correctly handles the crashed-predecessor case (socket-type check, no-op on missing, refuses non-sockets). translate_sampling_sglangcorrectly resolves greedy decoding (temperature < ε → argmax) since theis_normalizedflag skips the scheduler's own__post_init__, and the stop lists are pre-resolved as empty.- The
wait_for_input_peersocket-monitor barrier ensures no first-request drops on the bound PUSH socket. - Comprehensive test coverage: pinned Python wire vectors, round-trip codecs, E2E
generate_e2e_translates_and_streams_sglang, and validation rejection tests.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (10)
model_gateway/src/worker/worker.rs (1)
2848-2886: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win🟡 Nit: Exercise the production reconnect path.
These tests validate
remove_stale_ipc_socketin isolation. They do not verify cleanup ordering beforeZmqEngineClient::connect, cleanup of both data-plane paths, or reconnect after health-check eviction.Add an integration test that creates stale input and output sockets, runs the production connection path with a mock engine, and verifies that both endpoints are usable.
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/worker/worker.rs` around lines 2848 - 2886, The current tests only exercise remove_stale_ipc_socket directly; add an integration test through the production ZmqEngineClient::connect/reconnect path. Create stale input and output socket files, use a mock engine, verify cleanup occurs before connection and both data-plane endpoints are usable, then cover reconnect after health-check eviction. Run the pr-test-analyzer agent to confirm adequate coverage.Source: Coding guidelines
crates/engine_zmq_client/src/connector.rs (2)
762-764: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value🟣 Pre-existing — The test leaks the socket tempdir.
std::mem::forget(ns)drops theIpcNamespaceguard without running its destructor, so the temporary directory and itsipcsocket files stay on disk after the test. The same pattern already exists intokenspeed_client_submits_and_streamsat Line 639. Bind the guard to a local instead; that keeps the endpoints alive for the whole test and still cleans up.♻️ Proposed fix
let ns = IpcNamespace::new().unwrap(); let (input, output) = (ns.input_endpoint(), ns.output_endpoint()); - std::mem::forget(ns); + // Hold the namespace so the ipc files outlive the client and are removed. + let _ns = ns;🤖 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 762 - 764, Remove std::mem::forget(ns) in the test setup and retain the IpcNamespace guard in a local variable for the test’s lifetime, as already done in tokenspeed_client_submits_and_streams. Ensure the guard remains alive while input and output endpoints are used so its destructor cleans up the temporary directory and socket files.
738-760: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win🟡 Nit — The pinned SGLang wire fixtures are duplicated across two crates. Both tests copy the same
OUTPUT_STILL_GENERATINGandOUTPUT_FINISHEDhex strings and the samefrom_hexhelper. A third copy offrom_hexalready exists incrates/engine_zmq_client/src/protocol/sglang/request.rs. When the SGLang wire format changes, all copies must be edited together, and a missed copy produces a test that passes against a stale format. Publish the fixtures once from the crate that owns the protocol.
crates/engine_zmq_client/src/connector.rs#L738-L760: move the two constants andfrom_hexinto a shared test-fixture module besidemock_engine, export them, and import them here.model_gateway/src/routers/grpc/zmq_client.rs#L1511-L1533: delete the local constants andfrom_hex, and import the exported fixtures fromengine_zmq_client.🤖 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 738 - 760, Centralize the duplicated SGLang fixtures: in crates/engine_zmq_client/src/connector.rs#L738-L760, move OUTPUT_STILL_GENERATING, OUTPUT_FINISHED, and from_hex into the shared test-fixture module beside mock_engine and export them; in model_gateway/src/routers/grpc/zmq_client.rs#L1511-L1533, remove the local copies and import the exported fixtures from engine_zmq_client.model_gateway/src/routers/grpc/zmq_client.rs (1)
896-975: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win🟡 Nit —
translate_request_sglangduplicatestranslate_request_tokenspeed.Lines 900-960 repeat the tokenized-input extraction, the DP-rank guard, and the five rejection checks from
translate_request_tokenspeedat Lines 772-836, with only the backend name changed in the messages. The two copies will drift as either backend gains a capability. Extract the shared validation into one helper that takes the backend label, and keep only the wire-struct construction per backend./// Shared pre-translation checks for the slim ZMQ backends. fn validate_slim_request(req: &vllm::GenerateRequest, backend: &str) -> Result<Vec<u32>, String> { // tokenized-input extraction, DP-rank guard, logprobs/constraint/stop/logit_bias rejects }🤖 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 896 - 975, Extract the duplicated tokenized-input extraction and validation from translate_request_tokenspeed and translate_request_sglang into a shared validate_slim_request helper accepting the request and backend label. Preserve all existing DP-rank, logprobs, prompt-logprobs, constraint, stop, and logit_bias checks while using the label in backend-specific errors, then have both translation functions call the helper and retain only their backend-specific wire-struct construction.bindings/python/src/smg/serve.py (1)
40-48: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win🟡 Nit — Keep the ZMQ socket-name derivation in one place.
_zmq_data_addressesduplicates the Rustzmq_socket_addresses-in.sock/-out.sockrule. Move the suffix constants or endpoint derivation to shared config/docs, or add a Python test that assertsmodel_gateway/src/worker/worker.rsoutputs.🤖 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 `@bindings/python/src/smg/serve.py` around lines 40 - 48, The ZMQ endpoint suffix rule is duplicated between `_zmq_data_addresses` and Rust’s `zmq_socket_addresses`; centralize the `-in.sock`/`-out.sock` derivation in shared configuration or documentation, or add a Python test that verifies `_zmq_data_addresses` matches `zmq_socket_addresses` output, while preserving the existing input/output endpoint order.crates/engine_zmq_client/src/protocol/sglang/output.rs (1)
253-368: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win🟡 Nit: Add a test for partially-
Nonelogprob columns.The existing tests use fully-populated inner lists. Add a case where one column has a
Noneat a position and the other does not. That test pins the behavior chosen for the issue raised at Lines 239-251.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 `@crates/engine_zmq_client/src/protocol/sglang/output.rs` around lines 253 - 368, Add a test alongside into_outputs_flattens_optional_logprobs that supplies logprob value and index columns with mismatched inner None entries at the same position, then assert the behavior selected by the implementation around BatchTokenIDOutput::into_outputs. Ensure the test specifically covers partial-None columns rather than only fully populated lists, and run the pr-test-analyzer agent to verify coverage.Source: Coding guidelines
crates/engine_zmq_client/src/protocol/mod.rs (1)
26-65: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value🟡 Nit:
drain_trailingis unused by the new SGLang decoders.The SGLang decoders inline the same loop instead of calling this helper:
sglang/request.rsLine 140 and Line 200,sglang/output.rsLine 101, andsglang/sampling.rsLine 240. Calldrain_trailing(&mut seq)?at those sites to keep one implementation.🤖 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/mod.rs` around lines 26 - 65, Replace the duplicated trailing-element drain loops in the SGLang decoder paths of request, output, and sampling with calls to the shared drain_trailing helper. Update both request decoding sites and the corresponding output and sampling sites, passing each sequence access mutably and propagating errors with ?. Remove only the redundant inline loops.crates/engine_zmq_client/src/protocol/sglang/token_ids.rs (1)
50-96: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low value🟡 Nit: Decoding through
rmpv::Valuecopies the payload on the per-step output path.
Value::deserializematerializes the whole 2-element array and copies the binary payload into an ownedVec<u8>before the token ids are parsed.BatchTokenIDOutputcontains oneTokenIdArrayper request per scheduler step, so this runs on the streaming decode path. AVisitorwithvisit_seqplusserde_bytes::ByteBuf(or a borrowed&[u8]element) avoids the intermediate dynamic value and keeps the same validation.Payloads are small (8 bytes per token), so the current form is acceptable if you prefer the simpler code. The comment on Line 52 states that serde cannot map a msgpack
binwithoutserde_bytes; adding that dependency is the direct fix.As per coding guidelines: "Avoid unnecessary clone() calls in gRPC streaming hot paths, especially during per-token response processing."
🤖 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/sglang/token_ids.rs` around lines 50 - 96, The custom TokenIdArray::deserialize implementation materializes the payload through rmpv::Value, causing an avoidable allocation on the streaming path. Replace the Value-based decoding with a serde Visitor implementing visit_seq, deserialize the binary element through serde_bytes::ByteBuf (or an equivalent borrowed byte representation), and preserve the existing two-element, typecode, length, and u32-range validation.Source: Coding guidelines
crates/engine_zmq_client/src/protocol/sglang/sampling.rs (1)
196-242: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value🟡 Nit: Unify the positional read style, and confirm the silent default fallback is intended.
Two points in this visitor:
- The
field!macro is applied only to the non-Optionfields. TheOptionfields use the expanded form. Both forms behave the same, becausenext_element::<Option<T>>returnsSome(None)for a wirenilandNoneonly when the array ends. Use the macro for every field to keep one style.- A short array falls back to defaults for every missing position. The shared helper
next_fieldinprotocol/mod.rsfails loudly for missing modeled fields. This decoder is the opposite policy for the same wire family. SMG is the producer on this path, so the risk is low today. If a truncated array indicates version skew, prefer a decode error over silent defaults.As per coding guidelines: "Do not silently fall back to None or a default when configuration validation should fail loudly."
♻️ Uniform macro use
- max_new_tokens: seq.next_element()?.unwrap_or(default.max_new_tokens), - stop: seq.next_element()?.unwrap_or(default.stop), - stop_token_ids: seq.next_element()?.unwrap_or(default.stop_token_ids), - stop_regex: seq.next_element()?.unwrap_or(default.stop_regex), + max_new_tokens: field!(max_new_tokens), + stop: field!(stop), + stop_token_ids: field!(stop_token_ids), + stop_regex: field!(stop_regex),🤖 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/sglang/sampling.rs` around lines 196 - 242, Update the visit_seq deserializer to use the field! macro for every modeled SamplingParams field, including Option-valued fields, so positional reads share one style. Replace silent defaulting for missing modeled positions with a decode error consistent with next_field in protocol/mod.rs, while preserving explicit wire nil handling and ignoring only fields beyond the modeled prefix.Source: Coding guidelines
model_gateway/src/main.rs (1)
1409-1418: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win🟡 Nit: Add SGLang coverage to the ZMQ startup runtime pin test.
Backend::Sglangmaps toRuntimeType::Sglanginmodel_gateway, and the Python binding mirrors that mapping.zmq_backend_pins_startup_worker_runtime_in_both_configsonly covers--backend tokenspeed; add a--backend sglangSGLang case so the ZMQ pin is covered for all current runtime-specific backends.🤖 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 1409 - 1418, Extend the test zmq_backend_pins_startup_worker_runtime_in_both_configs to include a ZMQ configuration using --backend sglang, and assert that it pins the startup worker to RuntimeType::Sglang. Keep the existing tokenspeed coverage and both configuration cases unchanged.Source: Coding guidelines
🤖 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 `@bindings/python/src/smg/_sglang_zmq_launcher.py`:
- Around line 111-117: Update the post-block_until_scheduler_exits() handling to
always propagate a non-zero exit status because scheduler termination means the
engine is unavailable, regardless of individual exit codes. Remove the
any(proc.exitcode) condition and raise SystemExit(1) unconditionally after
block_until_scheduler_exits() returns; do not use exitcode values, which may be
None for live processes.
In `@crates/engine_zmq_client/src/protocol/sglang/output.rs`:
- Around line 239-251: Update per_request_logprobs and its callers to process
output_token_logprobs_val and output_token_logprobs_idx together, retaining only
positions where both entries are Some so values and token IDs remain aligned. Do
not flatten the two columns independently; alternatively reject ragged columns
using the same contract as into_outputs.
In `@model_gateway/src/routers/grpc/zmq_client.rs`:
- Around line 299-307: Update get_model_info so max_context_length falls back to
SGLang’s worker/server configuration context_length when ready_response.as_ref()
is unavailable, rather than defaulting to 0. Preserve the existing
ready_response max_model_len path for engines that provide a handshake response,
and use the established configuration source already exposed by the engine or
worker.
In `@model_gateway/src/worker/worker.rs`:
- Around line 221-224: Update the stale IPC socket cleanup around remove_file to
treat a NotFound error as successful completion, while preserving fail(...) for
all other removal errors.
- Around line 190-224: Update remove_stale_ipc_socket and its
connect_zmq_backend call path to verify that the IPC endpoint is not owned by a
live binder before removing it. Replace the FileTypeExt::is_socket() check alone
with an ownership/liveness check that distinguishes a crashed predecessor from
an active endpoint, and preserve errors for non-socket collisions.
---
Nitpick comments:
In `@bindings/python/src/smg/serve.py`:
- Around line 40-48: The ZMQ endpoint suffix rule is duplicated between
`_zmq_data_addresses` and Rust’s `zmq_socket_addresses`; centralize the
`-in.sock`/`-out.sock` derivation in shared configuration or documentation, or
add a Python test that verifies `_zmq_data_addresses` matches
`zmq_socket_addresses` output, while preserving the existing input/output
endpoint order.
In `@crates/engine_zmq_client/src/connector.rs`:
- Around line 762-764: Remove std::mem::forget(ns) in the test setup and retain
the IpcNamespace guard in a local variable for the test’s lifetime, as already
done in tokenspeed_client_submits_and_streams. Ensure the guard remains alive
while input and output endpoints are used so its destructor cleans up the
temporary directory and socket files.
- Around line 738-760: Centralize the duplicated SGLang fixtures: in
crates/engine_zmq_client/src/connector.rs#L738-L760, move
OUTPUT_STILL_GENERATING, OUTPUT_FINISHED, and from_hex into the shared
test-fixture module beside mock_engine and export them; in
model_gateway/src/routers/grpc/zmq_client.rs#L1511-L1533, remove the local
copies and import the exported fixtures from engine_zmq_client.
In `@crates/engine_zmq_client/src/protocol/mod.rs`:
- Around line 26-65: Replace the duplicated trailing-element drain loops in the
SGLang decoder paths of request, output, and sampling with calls to the shared
drain_trailing helper. Update both request decoding sites and the corresponding
output and sampling sites, passing each sequence access mutably and propagating
errors with ?. Remove only the redundant inline loops.
In `@crates/engine_zmq_client/src/protocol/sglang/output.rs`:
- Around line 253-368: Add a test alongside
into_outputs_flattens_optional_logprobs that supplies logprob value and index
columns with mismatched inner None entries at the same position, then assert the
behavior selected by the implementation around BatchTokenIDOutput::into_outputs.
Ensure the test specifically covers partial-None columns rather than only fully
populated lists, and run the pr-test-analyzer agent to verify coverage.
In `@crates/engine_zmq_client/src/protocol/sglang/sampling.rs`:
- Around line 196-242: Update the visit_seq deserializer to use the field! macro
for every modeled SamplingParams field, including Option-valued fields, so
positional reads share one style. Replace silent defaulting for missing modeled
positions with a decode error consistent with next_field in protocol/mod.rs,
while preserving explicit wire nil handling and ignoring only fields beyond the
modeled prefix.
In `@crates/engine_zmq_client/src/protocol/sglang/token_ids.rs`:
- Around line 50-96: The custom TokenIdArray::deserialize implementation
materializes the payload through rmpv::Value, causing an avoidable allocation on
the streaming path. Replace the Value-based decoding with a serde Visitor
implementing visit_seq, deserialize the binary element through
serde_bytes::ByteBuf (or an equivalent borrowed byte representation), and
preserve the existing two-element, typecode, length, and u32-range validation.
In `@model_gateway/src/main.rs`:
- Around line 1409-1418: Extend the test
zmq_backend_pins_startup_worker_runtime_in_both_configs to include a ZMQ
configuration using --backend sglang, and assert that it pins the startup worker
to RuntimeType::Sglang. Keep the existing tokenspeed coverage and both
configuration cases unchanged.
In `@model_gateway/src/routers/grpc/zmq_client.rs`:
- Around line 896-975: Extract the duplicated tokenized-input extraction and
validation from translate_request_tokenspeed and translate_request_sglang into a
shared validate_slim_request helper accepting the request and backend label.
Preserve all existing DP-rank, logprobs, prompt-logprobs, constraint, stop, and
logit_bias checks while using the label in backend-specific errors, then have
both translation functions call the helper and retain only their
backend-specific wire-struct construction.
In `@model_gateway/src/worker/worker.rs`:
- Around line 2848-2886: The current tests only exercise remove_stale_ipc_socket
directly; add an integration test through the production
ZmqEngineClient::connect/reconnect path. Create stale input and output socket
files, use a mock engine, verify cleanup occurs before connection and both
data-plane endpoints are usable, then cover reconnect after health-check
eviction. Run the pr-test-analyzer agent to confirm adequate 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: 0640fe66-29c1-4888-84cb-18e4596c1950
📒 Files selected for processing (21)
bindings/python/src/lib.rsbindings/python/src/smg/_sglang_zmq_launcher.pybindings/python/src/smg/serve.pycrates/engine_zmq_client/examples/live_probe.rscrates/engine_zmq_client/src/connector.rscrates/engine_zmq_client/src/lib.rscrates/engine_zmq_client/src/mock_engine.rscrates/engine_zmq_client/src/protocol/mod.rscrates/engine_zmq_client/src/protocol/sglang/mod.rscrates/engine_zmq_client/src/protocol/sglang/output.rscrates/engine_zmq_client/src/protocol/sglang/request.rscrates/engine_zmq_client/src/protocol/sglang/sampling.rscrates/engine_zmq_client/src/protocol/sglang/token_ids.rscrates/engine_zmq_client/src/protocol/tokenspeed/mod.rscrates/engine_zmq_client/src/protocol/vllm/mod.rscrates/engine_zmq_client/src/transport.rscrates/mock_worker/src/zmq.rsmodel_gateway/src/main.rsmodel_gateway/src/routers/grpc/zmq_client.rsmodel_gateway/src/worker/worker.rsmodel_gateway/src/workflow/steps/local/create_worker.rs
| result.block_until_scheduler_exits() | ||
|
|
||
| # A scheduler exiting means the engine is gone; propagate a non-zero status | ||
| # (rather than a silent exit 0) so the parent launcher sees the failure | ||
| # instead of leaving the router pushing to a dead socket. | ||
| if any(proc.exitcode for proc in procs): | ||
| raise SystemExit(1) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🔴 Important — A clean scheduler exit still returns status 0.
block_until_scheduler_exits() returns only after a scheduler process ends. At that point the engine is gone in every case. The guard at Line 116 raises SystemExit(1) only when some exitcode is non-zero. If every scheduler exits with 0, this process exits 0, and the parent launcher treats the engine as a normal shutdown while the router still pushes to a dead socket. The comment above the guard states the opposite intent.
Also note exitcode is None for a process that is still alive, which is falsy and therefore indistinguishable from a clean exit in this check.
🛠️ Proposed fix
result.block_until_scheduler_exits()
- # A scheduler exiting means the engine is gone; propagate a non-zero status
- # (rather than a silent exit 0) so the parent launcher sees the failure
- # instead of leaving the router pushing to a dead socket.
- if any(proc.exitcode for proc in procs):
- raise SystemExit(1)
+ # A scheduler exiting means the engine is gone, whatever its status; always
+ # propagate a non-zero status so the parent launcher sees the failure
+ # instead of leaving the router pushing to a dead socket.
+ exitcodes = [proc.exitcode for proc in procs]
+ logger.error("SGLang scheduler(s) exited with %s; shutting down", exitcodes)
+ raise SystemExit(1)As per coding guidelines: "Run the silent-failure-hunter agent on changed files to detect swallowed errors, inappropriate fallbacks, and missing error propagation."
📝 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.
| result.block_until_scheduler_exits() | |
| # A scheduler exiting means the engine is gone; propagate a non-zero status | |
| # (rather than a silent exit 0) so the parent launcher sees the failure | |
| # instead of leaving the router pushing to a dead socket. | |
| if any(proc.exitcode for proc in procs): | |
| raise SystemExit(1) | |
| result.block_until_scheduler_exits() | |
| # A scheduler exiting means the engine is gone, whatever its status; always | |
| # propagate a non-zero status so the parent launcher sees the failure | |
| # instead of leaving the router pushing to a dead socket. | |
| exitcodes = [proc.exitcode for proc in procs] | |
| logger.error("SGLang scheduler(s) exited with %s; shutting down", exitcodes) | |
| raise SystemExit(1) |
🤖 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 `@bindings/python/src/smg/_sglang_zmq_launcher.py` around lines 111 - 117,
Update the post-block_until_scheduler_exits() handling to always propagate a
non-zero exit status because scheduler termination means the engine is
unavailable, regardless of individual exit codes. Remove the any(proc.exitcode)
condition and raise SystemExit(1) unconditionally after
block_until_scheduler_exits() returns; do not use exitcode values, which may be
None for live processes.
Source: Coding guidelines
| /// Flatten request `index`'s logprob column into a dense vec, dropping the | ||
| /// per-token `Option` nesting (`None` inner values are omitted). Returns an | ||
| /// empty vec when the column is absent or the request had no logprobs. | ||
| fn per_request_logprobs<T: Copy>( | ||
| column: Option<&Vec<Option<Vec<Option<T>>>>>, | ||
| index: usize, | ||
| ) -> Vec<T> { | ||
| column | ||
| .and_then(|rows| rows.get(index)) | ||
| .and_then(Option::as_ref) | ||
| .map(|values| values.iter().filter_map(|value| *value).collect()) | ||
| .unwrap_or_default() | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Inspect SGLang logprob consumption in the gateway to confirm the pairing assumption.
set -euo pipefail
rg -n -C 8 'output_logprobs_val|output_logprobs_idx' --type=rustRepository: smg-project/smg
Length of output: 153
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== candidate file =="
git ls-files | rg 'crates/engine_zmq_client/src/protocol/sglang/output\.rs|sglang' || true
echo "== outline =="
ast-grep outline crates/engine_zmq_client/src/protocol/sglang/output.rs --view expanded 2>/dev/null || true
echo "== relevant source =="
wc -l crates/engine_zmq_client/src/protocol/sglang/output.rs
sed -n '1,330p' crates/engine_zmq_client/src/protocol/sglang/output.rs
echo "== references in repository =="
rg -n -C 4 'per_request_logprobs|output_token_logprobs_val|output_token_logprobs_idx|output_token_logprobs' . --glob '*.rs' || trueRepository: smg-project/smg
Length of output: 50371
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== generator logprob references =="
rg -n -C 6 'output_logprobs|output_token_logprobs|convert_generate_output_logprobs|current_output_logprobs|final_output_logprobs' model_gateway model_gateway/src | head -n 260
echo "== deterministic flattening probe =="
python3 - <<'PY'
def per_request_logprobs(values, index):
rows = values.get("rows")
if index >= len(rows):
return []
if rows[index] is None:
return []
result = []
for value in rows[index]:
if value is not None:
result.append(value)
return result
val_case = {"rows": [[None, "a"], ["b"]] for _ in ["_"]}
val_case = {"rows": [[None, "a"]]}
idx_case = {"rows": [[1, None], [2]]}
idx_case = {"rows": [[1, None]]}
print("value", per_request_logprobs({"rows": [[None, "a"]]}, 0))
print("index", per_request_logprobs({"rows": [[1, None]]}, 0))
print("aligned?", per_request_logprobs({"rows": [[None, "a"]]}, 0) == per_request_logprobs({"rows": [[1, None]]}, 0))
PY
echo "== tests around partial nested optional values =="
sed -n '330,370p' crates/engine_zmq_client/src/protocol/sglang/output.rsRepository: smg-project/smg
Length of output: 26116
Zip and reject incomplete logprob pairs before flattening.
output_token_logprobs_val and output_token_logprobs_idx are flattened separately, so position k = Some(value)/None produces a logprob without a token id or a token id without a logprob. Use a single pass that keeps only positions where both columns are Some, or reject the batch for the same column-ragged contract already used by into_outputs.
🤖 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/sglang/output.rs` around lines 239 -
251, Update per_request_logprobs and its callers to process
output_token_logprobs_val and output_token_logprobs_idx together, retaining only
positions where both entries are Some so values and token IDs remain aligned. Do
not flatten the two columns independently; alternatively reject ragged columns
using the same contract as into_outputs.
| /// so those come from worker config). SGLang's no-handshake wire has no ready | ||
| /// response, so the context length falls back to `0` (unknown). | ||
| pub fn get_model_info(&self) -> vllm::GetModelInfoResponse { | ||
| let max_context_length = self | ||
| .engines() | ||
| .first() | ||
| .map(|e| e.ready_response.max_model_len) | ||
| .and_then(|e| e.ready_response.as_ref()) | ||
| .map(|ready| ready.max_model_len) | ||
| .unwrap_or(0); |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Find readers of max_context_length to assess the impact of the 0 fallback.
rg -n -C 4 'max_context_length' --type=rust -g '!target'Repository: smg-project/smg
Length of output: 153
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== git status/stat =="
git diff --stat || true
echo "== locate zmq_client.rs =="
fd -a 'zmq_client.rs' . || true
echo "== find GetModelInfoResponse and related symbols =="
rg -n -C 3 'GetModelInfoResponse|get_model_info|max_model_len|max_context_length' . \
--glob '!target/**' --glob '!node_modules/**' || true
echo "== candidate file outline/section =="
if [ -f model_gateway/src/routers/grpc/zmq_client.rs ]; then
wc -l model_gateway/src/routers/grpc/zmq_client.rs
sed -n '1,360p' model_gateway/src/routers/grpc/zmq_client.rs | cat -n
fiRepository: smg-project/smg
Length of output: 50372
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== zmq_client relevant section =="
sed -n '260,330p' model_gateway/src/routers/grpc/zmq_client.rs | cat -n
echo "== deterministic search over Rust files for MaxContextLength/readers named context =="
python3 - <<'PY'
from pathlib import Path
import re
for path in Path('.').rglob('*.rs'):
p = str(path)
if 'target' in p.split('/'):
continue
try:
text = path.read_text()
except Exception:
continue
hits = []
for i,line in enumerate(text.splitlines(), 1):
if 'MaxContextLength' in line or 'max_context_length' in line:
hits.append((i,line))
if hits:
print(f"\n-- {p}")
for i,line in hits:
print(f"{i}: {line.strip()}")
print("\n== deterministic model_gateway text search for model info context fields ==")
root = Path('model_gateway')
if not root.exists():
raise SystemExit("missing model_gateway")
for path in root.rglob('*'):
if not path.is_file() or path.suffix not in {'.rs','.proto','.toml'} or 'target' in str(path):
continue
text = path.read_text(errors='ignore')
if 'context_length' in text or 'MaxContextLength' in text or 'max_model_len' in text or 'get_model_info' in text or 'GetModelInfo' in text:
print(f"\n-- {path}")
lines=text.splitlines()
for idx,line in enumerate(lines,1):
if any(s in line for s in ['context_length','MaxContextLength','max_model_len','get_model_info','GetModelInfo']):
print(f"{idx}: {line.strip()}")
PY
echo "== inspect model_gateway grpc routers for zmq_client usages =="
rg -n -C 3 'get_model_info|GetModelInfo|ModelInfo|Zmq|zmq_client' model_gateway/src --glob '*.rs' || trueRepository: smg-project/smg
Length of output: 50371
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== zmq_client fields =="
rg -n -C 4 'ready_response|data_parallel_size|context_length|max_model_len|engine_info|WorkerCard|WorkerSpec' model_gateway/src/routers/grpc/zmq_client.rs | cat -n
echo "== ModelInfo labels flat_labels context fields =="
sed -n '780,905p' model_gateway/src/routers/grpc/client.rs | cat -n
echo "== ZmqEngineClient connect model_id/context_length labels context =="
rg -n -C 5 'impl ZmqEngineClient|pub struct ZmqEngineClient|model_id|context_length|max_context_length|max_model_len|routable|WorkerCard|EngineCoreReadyResponse' model_gateway/src/routers/grpc/zmq_client.rs | cat -n
echo "== read-only static invariant: count contexts/source of context_length in zmq_client =="
python3 - <<'PY'
from pathlib import Path
p=Path('model_gateway/src/routers/grpc/zmq_client.rs')
text=p.read_text()
for token in ['context_length','max_model_len','EngineCoreReadyResponse','ZmqBackend']:
print(token, text.count(token))
print("contains context_length:", 'context_length' in text)
PYRepository: smg-project/smg
Length of output: 15570
Keep max_context_length from falling back to 0 on SGLang.
get_model_info derives SGLang’s context limit from the handshake ready_response, but SGLang ZMQ workers have no ready response, so consumers receive 0. SGLang’s GetModelInfoResponse normally carries server_args.context_length; use the worker/server config when ready_response.as_ref() is unavailable instead of making 0 look like an unknown/no-limit value.
🤖 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 299 - 307, Update
get_model_info so max_context_length falls back to SGLang’s worker/server
configuration context_length when ready_response.as_ref() is unavailable, rather
than defaulting to 0. Preserve the existing ready_response max_model_len path
for engines that provide a handshake response, and use the established
configuration source already exposed by the engine or worker.
Source: Coding guidelines
| /// Remove an `ipc://` socket file stranded by a crashed predecessor. libzmq | ||
| /// refuses to `bind()` over an existing ipc endpoint file (it fails with | ||
| /// `EADDRINUSE`), so a socket left behind by an earlier process that did not | ||
| /// unbind cleanly would block every future bind at the same path. Only an actual | ||
| /// socket file is removed: a regular file or directory at the path is left in | ||
| /// place and surfaced as an error, so an unexpected collision is never silently | ||
| /// clobbered. A missing path is a no-op. | ||
| /// | ||
| /// This targets the crashed-predecessor case. In the one-worker-per-rank model a | ||
| /// given ipc path has a single binder, so this does not race a live peer. | ||
| async fn remove_stale_ipc_socket(ipc_url: &str) -> WorkerResult<()> { | ||
| let path = ipc_url.strip_prefix("ipc://").unwrap_or(ipc_url); | ||
| let fail = |reason: String| WorkerError::ConnectionFailed { | ||
| url: ipc_url.to_string(), | ||
| reason, | ||
| }; | ||
| let meta = match tokio::fs::symlink_metadata(path).await { | ||
| Ok(meta) => meta, | ||
| Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(()), | ||
| Err(e) => return Err(fail(format!("failed to stat ipc socket {path}: {e}"))), | ||
| }; | ||
| #[cfg(unix)] | ||
| { | ||
| use std::os::unix::fs::FileTypeExt; | ||
| if !meta.file_type().is_socket() { | ||
| return Err(fail(format!( | ||
| "ipc socket path {path} exists but is not a socket; refusing to remove it" | ||
| ))); | ||
| } | ||
| } | ||
| let _ = &meta; | ||
| tokio::fs::remove_file(path) | ||
| .await | ||
| .map_err(|e| fail(format!("failed to remove stale ipc socket {path}: {e}"))) | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/usr/bin/env bash
set -u
fd -a -t f 'Cargo.lock|Cargo.toml' --exec rg -n -C 3 'libzmq|zmq'
rg -n -C 6 'remove_stale_ipc_socket|connect_zmq_backend|ZmqEngineClient::connect' model_gateway/src --glob '*.rs'Repository: smg-project/smg
Length of output: 21159
🏁 Script executed:
#!/usr/bin/env bash
set -u
echo "== dependency declarations =="
rg -n -C 3 'engine-zmq-client|libzmq|zmq' Cargo.lock Cargo.toml model_gateway/Cargo.toml crates/engine_zmq_client/Cargo.toml crates/engine_zmq_client | head -n 220
echo
echo "== engine_zmq_client source outline =="
fd -a -t f . crates/engine_zmq_client/src | sed 's#^\./##' | sort | xargs -r -n 1 ast-grep outline --view compact || true
echo
echo "== connect implementation candidates =="
rg -n -C 12 'pub fn connect|async fn connect|self\.bind|zmq_bind|epilog|ipc|bind\(' crates/engine_zmq_client/src --glob '*.rs'Repository: smg-project/smg
Length of output: 50371
🏁 Script executed:
#!/usr/bin/env bash
set -u
echo "== crates/engine_zmq_client/Cargo.toml =="
cat -n crates/engine_zmq_client/Cargo.toml
echo
echo "== relevant ZmqEngineClient implementation =="
rg -n -C 8 'struct ZmqEngineClient|impl ZmqEngineClient|fn connect|async fn connect|self\.bind|bind_local_sockets|connect_handshake|connect_push_pull' crates/engine_zmq_client/src --glob '*.rs'
echo
echo "== workspace deps around zmq bindings =="
rg -n -C 4 'zmq|libzmq|zmq-sys|sys' Cargo.toml crates --glob '*.toml'Repository: smg-project/smg
Length of output: 27779
🌐 Web query:
zmq-rs zeromq ipc bind EADDRINUSE documentation
💡 Result:
In ZeroMQ, the EADDRINUSE error occurs when attempting to call bind on an IPC (inter-process communication) endpoint that is already being used [1]. When using the IPC transport (e.g., ipc:///tmp/mysocket), ZeroMQ creates a file or a named entity in the operating system's namespace [2][3][4]. If a process attempts to bind to an IPC endpoint that is currently bound or still has a residual file from a previous execution, the bind operation will fail with EADDRINUSE [1][2][4]. Key points regarding EADDRINUSE with IPC: 1. File Persistence: Unlike TCP ports, IPC endpoints often leave files on the filesystem (e.g., in /tmp) [2][3]. If a process crashes or does not clean up, the file remains, potentially causing EADDRINUSE errors upon subsequent restarts. 2. Resolution: Ensure your application cleans up IPC endpoints correctly, or manually remove the stale file if it is known to be orphaned. 3. Behavior: While some transports allow overriding existing bindings (e.g., some ZeroMQ versions or configurations allow a second process to take over), relying on this is generally discouraged, and persistent files are the most common cause of bind failure [2][4]. 4. Rust Specifics: In Rust wrappers like zmq or zmq-rs, the bind method returns a Result type [5][6][7]. If the underlying libzmq bind call fails due to EADDRINUSE, the Rust Result will contain an error reflecting this [5][1]. You should handle this Result and check for the EADDRINUSE error code to implement retry logic or cleanup procedures [5][6]. For more information, refer to the official zmq_ipc(7) and zmq_bind(3) documentation [1][2][3].
Citations:
- 1: http://api.zeromq.org/4-2:zmq-bind
- 2: https://www.mankier.com/7/zmq_ipc
- 3: https://libzmq.readthedocs.io/en/latest/zmq_ipc.html
- 4: https://man.archlinux.org/man/zmq_ipc.7.en.raw
- 5: https://docs.rs/zmq-rs/latest/zmq_rs/struct.Socket.html
- 6: https://docs.rs/zmq/latest/zmq/struct.Socket.html
- 7: https://docs.rs/zmq-rs/latest/zmq_rs/struct.Context.html
🌐 Web query:
zmq-rs zeromq crate version workspace zeromq crates io
💡 Result:
There are two primary Rust crates associated with ZeroMQ, which serve different purposes: 1. zeromq: This is a native, asynchronous Rust implementation of ZeroMQ (often referred to as zmq.rs). It does not require the C libzmq library [1][2]. As of August 4, 2026, the latest version on crates.io is 0.6.0, released on May 4, 2026 [3][4]. 2. zmq: This crate provides high-level Rust bindings to the official C-based libzmq library [5][6]. It is a distinct project from the native zeromq crate [2]. Regarding the concept of a workspace: The native zeromq project (hosted at https://github.com/zeromq/zmq.rs) is developed as a single repository, but it is not commonly referred to as a "workspace" crate in the sense of a multi-crate workspace distribution; rather, it is a single crate project that leverages feature flags (such as tokio-runtime, async-std-runtime, and transport protocols like tcp-transport and ipc-transport) to manage its functionality and dependencies [3][1]. Summary of Key Details: - Native implementation: zeromq (latest 0.6.0) [3]. - C-bindings: zmq (latest versions vary; see crates.io) [5][6]. - Official Repository: https://github.com/zeromq/zmq.rs [3][1].
Citations:
- 1: https://github.com/zeromq/zmq.rs?tab=readme-ov-file
- 2: https://zeromq.org/languages/rust/
- 3: https://crates.io/crates/zeromq
- 4: https://github.com/zeromq/zmq.rs/blob/master/CHANGELOG.md
- 5: https://github.com/erickt/rust-zmq
- 6: https://docs.rs/crate/zmq/latest
🌐 Web query:
crates.io zeromq ipc bind EADDRINUSE implementation
💡 Result:
The ZeroMQ EADDRINUSE error during an ipc:// bind operation generally indicates that the filesystem path or abstract namespace address you are attempting to bind to is already in use by another process or has not been cleaned up properly from a previous session [1][2]. Key technical points regarding IPC binding and EADDRINUSE: 1. Default IPC Behavior: Unlike TCP or inproc, the standard IPC transport in ZeroMQ is designed to allow one process to override another's binding [3][4]. If a second process binds to an endpoint already bound by another, the bind will typically succeed, and the first process will lose its binding [3][4]. 2. Exceptions to the Rule: - Abstract Namespaces: On Linux, if you use the abstract namespace (prefixing your path with @), this override behavior does not apply; if the address is already bound, the bind will fail with an error [3][4]. - OS-Specific Behavior: Implementation details, particularly on non-POSIX systems like Windows, can vary. There are known scenarios where ipc:// binds may fail unexpectedly due to file system permission issues or platform-specific socket limitations [5]. 3. Common Causes: - Stale Files: If a process terminates abruptly, the Unix domain socket file (the IPC path) might remain on the filesystem [2]. Because ZeroMQ attempts to manage these, you may need to manually ensure the file is removed if it persists after process termination. - Resource Leaks: In some cases, child processes may inherit open file descriptors, keeping the socket "in use" even after the parent process has attempted to shut down [2]. - Ephemeral Ports/Conflicts: While less common for IPC than TCP, ensure that your chosen path is not being used or claimed by another application or a previous, "zombie" instance of your own application [2]. Troubleshooting Steps: - Check for Stale Files: Manually check the directory specified in your ipc:///... path and remove any leftover socket files [2]. - Ensure Graceful Shutdown: Ensure that your application closes sockets and destroys the ZeroMQ context (zmq_ctx_destroy()) before exiting to allow the library to clean up its resources properly [2]. - Use Unique Paths: If testing in parallel or in ephemeral environments, use unique pathnames or the wildcard * endpoint [6][4]. If you use the wildcard, use the ZMQ_LAST_ENDPOINT socket option to retrieve the actual generated path so you can manage or unbind it correctly [6][4][7]. For Rust implementations: - If using standard C-bindings (like the zmq crate), EADDRINUSE is handled through the standard error enum, which maps to the underlying libzmq error code [8]. - When using async wrappers (like rzmq or zeromq crates), ensure you are awaiting the bind call properly and checking for returned errors, as the asynchronous nature may delay the exposure of the EADDRINUSE error until the bind is executed on the event loop [9][10][11].
Citations:
- 1: http://api.zeromq.org/4-2:zmq-bind
- 2: https://stackoverflow.com/questions/19159771/recovering-from-zmq-error-zmqerror-address-already-in-use
- 3: https://man.archlinux.org/man/zmq_ipc.7.en.raw
- 4: http://api.zeromq.org/4-2:zmq-ipc
- 5: Problem: ipc connect can fail on Windows, even after bind zeromq/libzmq#4734
- 6: https://libzmq.readthedocs.io/en/latest/zmq_ipc.html
- 7: https://docs.rs/zmq/latest/zmq/struct.Socket.html?search=
- 8: https://github.com/erickt/rust-zmq/blob/master/src/lib.rs
- 9: https://crates.io/crates/zeromq
- 10: https://docs.rs/crate/rzmq/latest
- 11: https://github.com/paddor/omq.rs/blob/main/omq-libzmq/src/socket.rs
🟣 Pre-existing: endpoint ownership is still unclear even if cleanup is needed.
remove_stale_ipc_socket only checks that the path is a Unix socket. That file type does not establish that the socket is stranded by a crashed predecessor. If the endpoint is owned by another live binder, unbinding it during connect_zmq_backend can break that process before this bind has successfully replaced it; use an ownership/liveness check for IPC IPC sockets instead of deleting based on FileTypeExt::is_socket() alone.
🤖 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/worker/worker.rs` around lines 190 - 224, Update
remove_stale_ipc_socket and its connect_zmq_backend call path to verify that the
IPC endpoint is not owned by a live binder before removing it. Replace the
FileTypeExt::is_socket() check alone with an ownership/liveness check that
distinguishes a crashed predecessor from an active endpoint, and preserve errors
for non-socket collisions.
Source: Coding guidelines
| tokio::fs::remove_file(path) | ||
| .await | ||
| .map_err(|e| fail(format!("failed to remove stale ipc socket {path}: {e}"))) | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🟡 Nit: Treat a concurrent unlink as success.
The initial stat treats NotFound as a no-op, but remove_file converts a later NotFound into ConnectionFailed. If another cleanup removes the path between these operations, the endpoint is already clean but the connection attempt fails.
Handle NotFound from remove_file as success.
Proposed fix
- tokio::fs::remove_file(path)
- .await
- .map_err(|e| fail(format!("failed to remove stale ipc socket {path}: {e}")))
+ match tokio::fs::remove_file(path).await {
+ Ok(()) => Ok(()),
+ Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
+ Err(e) => Err(fail(format!(
+ "failed to remove stale ipc socket {path}: {e}"
+ ))),
+ }📝 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.
| tokio::fs::remove_file(path) | |
| .await | |
| .map_err(|e| fail(format!("failed to remove stale ipc socket {path}: {e}"))) | |
| } | |
| match tokio::fs::remove_file(path).await { | |
| Ok(()) => Ok(()), | |
| Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()), | |
| Err(e) => Err(fail(format!( | |
| "failed to remove stale ipc socket {path}: {e}" | |
| ))), | |
| } |
🤖 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/worker/worker.rs` around lines 221 - 224, Update the stale
IPC socket cleanup around remove_file to treat a NotFound error as successful
completion, while preserving fail(...) for all other removal errors.
1 participant
Description
Problem
The direct-ZMQ backend (added in #2000, extended to TokenSpeed in #2036) lets
SMG connect straight to a same-host inference engine over
ipc://, bypassingthe Python gRPC servicer and its serialization + process-hop overhead. Until
now it spoke only the vLLM and TokenSpeed wires — SGLang engines still had to
go through the gRPC servicer even when co-located.
Solution
Teach the direct-ZMQ path to speak SGLang. On this wire SMG plays the SGLang
TokenizerManagerrole itself: it binds the PUSH input and PULL outputipc://sockets, sends already-tokenized generate requests, and receivesbatched token outputs straight off the wire. The SGLang scheduler runs headless
with
--skip-tokenizer-init, so SMG owns tokenization and sampling-paramnormalization end to end.
Changes
protocol/sglang— new msgpack codec for the SGLang tokenizer↔schedulerwire:
request,output,sampling, andtoken_idsencoders. Samplingparams are emitted already-normalized (normalized flag set, greedy resolved,
empty stop lists) because SMG bypasses SGLang's own
normalize().transport— the PUSH/PULL sockets now wait for the engine to attach viaa socket monitor before admitting requests. The PUSH socket returns-to-sender
(rather than buffering) when no peer is connected, so a readiness barrier is
required to avoid dropping the first requests.
connector— output handling generalized to tag-dispatch across enginewires.
routers/grpc/zmq_client.rs— gateway adapter for proto ⇄ SGLangrequest/response translation.
serve.py+_sglang_zmq_launcher.py— launch a headless SGLangscheduler wired to SMG's two sockets, forcing msgpack over the ZMQ wire.
Test Plan
Validated end-to-end against a live SGLang scheduler on a GB300 box:
smg servelaunches the headless scheduler; the router binds theipc://sockets and gates request admission on ZMQ readiness.
DONE ok=True, 95 generated tokens / 95 stream chunks — matching the gRPCpath for the same prompt.
input socket, then flow; no first-request drops.
Unit coverage:
cargo test -p engine-zmq-client(msgpack round-trips for therequest/output/sampling/token-id encoders).
Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspasses