feat(multimodal): configurable engine-agnostic tensor transport - #1892
Conversation
Promote multimodal tensor transport from TokenSpeed-only, env-only tuning to a first-class, engine-agnostic config surface. - Add `TransportMode` enum (inline|shm|auto) to openai-protocol. - RouterConfig gains `multimodal_tensor_transport` + `multimodal_shm_min_bytes`, exposed via `--multimodal-tensor-transport` / `--multimodal-shm-min-bytes` (CLI, with value_parser), YAML config, and the Python router SDK (RouterArgs + PyRouterConfig). - Per-worker `WorkerSpec` overrides (`multimodal_tensor_transport`, `multimodal_shm_min_bytes`) let co-located and remote workers differ. - Resolution precedence: per-worker override → router config → `SMG_MM_*` env (legacy `SMG_TOKENSPEED_MM_*` kept as deprecated aliases) → default (`inline`, 64 KiB). Router defaults are seeded once at startup (server::startup); the per-request mode drives `shm_enabled` and the resolved threshold rides on `TokenSpeedMultimodalData` so both write paths (encoder input + model-specific tensors) honor it. - Rename the transport resolvers `tokenspeed_* -> mm_*` (in `multimodal/transport.rs`). Behavior change for explicit `shm` users: `shm` now forces SHM whenever SMG can write `/dev/shm` (the operator asserts co-location). Previously `shm` behaved identically to `auto` (both required verified worker sharing), making `shm` redundant. `auto` is unchanged (verifies the worker shares `/dev/shm` via its advertised namespace token). The default is `inline`, so existing deployments are unaffected. Groundwork for engine-neutral SHM (vLLM next). Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
github-actions
Bot
added
documentation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughThis PR adds multimodal transport mode and SHM threshold configuration, threads them through router setup and runtime resolution, and updates multimodal assembly and proto conversion to use the new values. ChangesMultimodal Transport Configuration
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
Suggested labels: Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Clean PR — the config plumbing is thorough (CLI, Python bindings, per-worker overrides, env fallback with legacy aliases) and the resolution precedence is correctly implemented. The OnceLock-based startup seeding in server.rs avoids per-request env reads while keeping the test fallback path sane. One 🟡 nit on silent parse failure in the Python bindings' direct API path.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces configurable multimodal tensor transport modes (inline, shm, auto) and size thresholds (multimodal_shm_min_bytes) across the router, CLI, and per-worker specifications, replacing legacy SMG_TOKENSPEED_MM_* environment variables with SMG_MM_* fallbacks. The review feedback highlights two key issues: first, in disaggregated mode, primary_worker incorrectly returns the prefill worker instead of the encode worker, which bypasses encode-specific overrides; second, invalid transport mode strings passed from Python are silently ignored rather than raising a configuration error.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
model_gateway/src/routers/grpc/multimodal/assemble.rs (1)
521-551: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMissing test coverage for the
shm_min_bytesthreshold boundary.All tests (here and in
serialize.rs/proto_wrapper.rs) pinshm_min_bytesto0, so the newnbytes >= shm_min_bytesgating condition added inserialize_array_as_tokenspeed_tensoris never exercised at/around the threshold. Consider adding a case with a non-zeroshm_min_byteswherenbytesfalls just below and just at the threshold to confirm the SHM/inline split is correct.🤖 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/multimodal/assemble.rs` around lines 521 - 551, Add test coverage for the new shm_min_bytes boundary in the multimodal assembly and serialization tests, since current cases only use zero. Update the tests around pending_tokenspeed_shm_assembly and the serialize_array_as_tokenspeed_tensor path to use a non-zero shm_min_bytes and verify behavior when nbytes is just below and exactly at the threshold. Assert that the SHM vs inline tensor encoding switches correctly at the boundary.
🤖 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/lib.rs`:
- Around line 845-846: The new _Router constructor parameters were inserted in
the middle of the Python argument list, which shifts later positional arguments
and breaks compatibility. Update the constructor definition in the
_Router-related code so multimodal_tensor_transport and multimodal_shm_min_bytes
are appended to the end of the parameter list instead of being placed after
dp_aware, keeping existing positional callers stable.
- Around line 790-794: The multimodal tensor transport parsing in the Python
bindings is silently accepting invalid values by converting parse failures into
None, which causes fallback to env/default behavior. Update the builder path
around multimodal_tensor_transport and config::TransportMode::parse to reject
bad strings explicitly by returning a ConfigError::InvalidValue when parsing
fails, matching the CLI/Rust parsing contract instead of using a silent
fallback.
In `@model_gateway/src/routers/grpc/multimodal/transport.rs`:
- Around line 101-120: The SHM transport resolution currently reads overrides
from the prefill-side worker selection, so pixel tensors can inherit the wrong
mode/threshold when the receiving encode workers differ. Update
resolve_mm_shm_enabled and resolve_mm_shm_min_bytes to use the worker assignment
that actually receives the tensor for pixel values, and fall back conservatively
when the encode workers disagree; keep the existing helper flow around
worker_transport_mode_override, worker_shares_dev_shm, and
worker_shm_min_bytes_override but apply them to the receiving leg rather than
primary_worker().
---
Outside diff comments:
In `@model_gateway/src/routers/grpc/multimodal/assemble.rs`:
- Around line 521-551: Add test coverage for the new shm_min_bytes boundary in
the multimodal assembly and serialization tests, since current cases only use
zero. Update the tests around pending_tokenspeed_shm_assembly and the
serialize_array_as_tokenspeed_tensor path to use a non-zero shm_min_bytes and
verify behavior when nbytes is just below and exactly at the threshold. Assert
that the SHM vs inline tensor encoding switches correctly at the boundary.
🪄 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: ASSERTIVE
Plan: Pro
Run ID: f8ba016b-5f20-4a48-93b3-21e7d30cd36b
📒 Files selected for processing (14)
bindings/python/src/lib.rsbindings/python/src/smg/router_args.pycrates/protocols/src/worker.rsdocs/reference/configuration.mdmodel_gateway/src/config/builder.rsmodel_gateway/src/config/types.rsmodel_gateway/src/main.rsmodel_gateway/src/routers/grpc/epd_encode.rsmodel_gateway/src/routers/grpc/multimodal/assemble.rsmodel_gateway/src/routers/grpc/multimodal/mod.rsmodel_gateway/src/routers/grpc/multimodal/serialize.rsmodel_gateway/src/routers/grpc/multimodal/transport.rsmodel_gateway/src/routers/grpc/proto_wrapper.rsmodel_gateway/src/server.rs
…validation Follow-up on the transport config change addressing bot review plus deprecated-env handling: - transport.rs: warn once when a legacy `SMG_TOKENSPEED_MM_*` env var is read. The aliases stay for backward compat; the warning nudges migration to `SMG_MM_*` so they can be removed in a future major. - transport.rs: `primary_worker` now returns the encode worker in EPD (where the multimodal tensors are actually sent), so per-worker overrides configured on the encode spec are honored instead of the prefill spec's. - bindings/python: reject invalid `multimodal_tensor_transport` strings with a `ConfigError::InvalidValue` instead of silently defaulting, matching the CLI/Rust parsing contract for direct programmatic `RouterArgs` use; and append the new `_Router` constructor params (not mid-list) per the file's documented positional-compatibility convention. - serialize.rs: add a `shm_min_bytes` threshold boundary test (inline below, SHM at/above) so the new gating condition is exercised. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
|
Thanks for the reviews. Addressed in 145613e:
Also added (independent of the review): a one-time deprecation warning when a legacy |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@model_gateway/src/routers/grpc/multimodal/transport.rs`:
- Around line 81-90: The mm_shm_min_bytes_from_env helper is silently ignoring
malformed values by converting parse failures to None, unlike
mm_tensor_transport_mode_from_env which reports invalid input. Update
mm_shm_min_bytes_from_env to detect a parse error from the env value returned by
env_with_deprecated_alias and log an unknown/invalid configuration once (reuse
the same logging pattern used by log_unknown_transport_once or a similar
helper), then fall back to None so the default router config still applies.
🪄 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: ASSERTIVE
Plan: Pro
Run ID: 18207f53-a980-42e7-b015-edc4de64e3ed
📒 Files selected for processing (3)
bindings/python/src/lib.rsmodel_gateway/src/routers/grpc/multimodal/serialize.rsmodel_gateway/src/routers/grpc/multimodal/transport.rs
mm_shm_min_bytes_from_env silently swallowed unparseable values via .ok(), unlike mm_tensor_transport_mode_from_env which warns. A typo like "64k" now logs a one-time warning before falling back to the default, matching the mode reader. Addresses CodeRabbit review. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
|
Addressed in d2ef4bb: |
1 participant
Description
Problem
Multimodal tensor transport (inline vs
/dev/shm) was tunable only viaSMG_TOKENSPEED_MM_*environment variables, TokenSpeed-specific and env-only.There was no CLI/YAML config, no per-worker control (co-located vs remote
workers can't differ), and the knob was framed as TokenSpeed-only even though
the transport is engine-agnostic.
Solution
Promote it to a first-class, engine-agnostic config surface with a clear
resolution precedence:
Router-level defaults are resolved once at startup (
server::startup); theper-request mode drives
shm_enabled, and the resolved size threshold rides onTokenSpeedMultimodalDataso both write paths (encoder input inserialize.rsand model-specific tensors in
proto_wrapper.rs) honor it consistently.Changes
openai-protocol: newTransportModeenum (inline|shm|auto), serde +JsonSchema+FromStr.WorkerSpecgainsmultimodal_tensor_transport+multimodal_shm_min_bytesoverrides.RouterConfig(config/types.rs,builder.rs): the two fields + buildermethods +
Default.main.rs):--multimodal-tensor-transport(value_parserinline|shm|auto) +--multimodal-shm-min-bytes; wired intoto_router_config.transport.rs: resolvers renamedtokenspeed_* → mm_*; precedence logic;per-worker override readers; startup seed
init_mm_transport_defaults.RouterArgsdataclass + argparse +PyRouterConfig(
bindings/python). (Go SDK is a request client — no router-config surface.)Behavior change (explicit
shmonly)shmnow forces SHM whenever SMG can write/dev/shm(operator assertsco-location). Previously
shmwas identical toauto(both required verifiedworker sharing), so
shmwas redundant.autois unchanged. Default isinline, so existing deployments are unaffected.Test Plan
cargo clippy -p smg --all-targets -- -D warnings— clean (lib + tests + main.rs)cargo check -p smg-python— clean (Python binding compiles)cargo +nightly fmt— clean (scoped to the change set)cargo test -p openai-protocol --lib— 77 passcargo test -p smg --lib— 1146 pass (one unrelatedmiddleware::metricsinterner test is a known parallel-test flake; passes in isolation)
cargo test -p smg --bins— new two-path guardmultimodal_transport_flows_into_both_configspassespython3 -m py_compileonrouter_args.py— OKChecklist
cargo +nightly fmtcleancargo clippy -- -D warningscleanto_router_config;to_server_configwraps it) + two-path testvalue_parservalidationDefault+#[serde(default, skip_serializing_if)]Summary by CodeRabbit
inline/shm/auto) andmultimodal_shm_min_bytesacross router config, CLI, Python bindings, and per-worker overrides.SMG_MM_*environment variables (legacySMG_TOKENSPEED_MM_*aliases retained).