feat(multimodal): wire vLLM RDMA multimodal tensor pull - #1916
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds RDMA capability advertisement and worker-based gateway gating, then enables the vLLM servicer to pull remote multimodal tensors while retaining inline/SHM deserialization. It also adds unit, integration, and opt-in end-to-end coverage. ChangesRDMA multimodal transport
Estimated code review effort: 3 (Moderate) | ~30 minutes Sequence Diagram(s)sequenceDiagram
participant Router
participant Worker
participant VllmEngineServicer
participant RdmaPixelPuller
participant GatewayMetrics
Router->>Worker: read supports_rdma_pull capability
Router->>VllmEngineServicer: send multimodal TensorData
VllmEngineServicer->>RdmaPixelPuller: pull remote tensor
RdmaPixelPuller-->>VllmEngineServicer: return tensor
VllmEngineServicer-->>Router: complete image request
Router->>GatewayMetrics: record remote pixel bytes
Possibly related PRs
Suggested labels: Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
github-actions
Bot
added
grpc
![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 support for RDMA (NIXL) multimodal tensor payloads in the vLLM engine servicer and updates the model gateway to resolve and utilize the RDMA transport lane when supported by the worker. The feedback highlights two critical improvements: offloading synchronous, blocking RDMA network operations to a separate thread using asyncio.to_thread to prevent freezing the asyncio event loop, and adding a mapping for float16 in _PROTO_DTYPE_MAP to avoid runtime deserialization errors.
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.
| # Deserialize all tensors from proto (pixel_values may arrive over RDMA). | ||
| hf_dict: dict[str, torch.Tensor] = { | ||
| mm_key("pixel_values"): _tensor_from_proto(mm_proto.pixel_values), | ||
| mm_key("pixel_values"): self._tensor_from_proto(mm_proto.pixel_values), |
There was a problem hiding this comment.
The _tensor_from_proto method performs synchronous, blocking RDMA network operations via RdmaPixelPuller.feature_from_remote (which includes a busy-wait loop and synchronous socket/NIXL operations). Since _build_preprocessed_mm_inputs is called synchronously on the main thread of the asyncio event loop in Generate (line 216), this will completely block the event loop during RDMA pulls, freezing all concurrent requests and health checks.
To prevent blocking the event loop, consider offloading the call to _build_preprocessed_mm_inputs to a separate thread using asyncio.to_thread in Generate:
prompt = await asyncio.to_thread(
self._build_preprocessed_mm_inputs,
request.tokenized,
request.mm_inputs,
)References
- Avoid using
asyncio.to_threadfor synchronous I/O or CPU-intensive tasks if the operation is not on the hot path and does not block concurrent requests.
| # Proto dtype string → torch dtype | ||
| _PROTO_DTYPE_MAP: dict[str, torch.dtype] = { | ||
| "float32": torch.float32, | ||
| "bfloat16": torch.bfloat16, |
There was a problem hiding this comment.
The _PROTO_DTYPE_MAP is missing a mapping for float16 (half precision). Since the gateway can resolve and export multimodal tensors in float16 (e.g., via worker labels or environment configuration), the servicer should support deserializing it to prevent runtime ValueError exceptions.
| "bfloat16": torch.bfloat16, | |
| "float16": torch.float16, | |
| "bfloat16": torch.bfloat16, |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b3ccc79ccc
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
|
Hi @slin1237, this PR has merge conflicts that must be resolved before it can be merged. Please rebase your branch: git fetch origin main
git rebase origin/main
# resolve any conflicts, then:
git push --force-with-lease |
mergify
Bot
added
the
needs-rebase
b3ccc79 to
e0eb21f
Compare
mergify
Bot
removed
the
needs-rebase
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e0eb21f3ec
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| kv_engine_id=kv_engine_id, | ||
| data_parallel_size=parallel.data_parallel_size, | ||
| shm_namespace_id=mm_shm.shm_namespace_id(), | ||
| supports_rdma_pull=self._rdma_pixel_puller.ready, |
There was a problem hiding this comment.
This now passes supports_rdma_pull into the generated GetServerInfoResponse, but grpc_servicer/pyproject.toml still allows smg-grpc-proto>=0.4.13 and the proto package version was not bumped even though field 11 was added. In installs that satisfy the servicer with the previously published proto stubs, this keyword is unknown and GetServerInfo raises at runtime, so vLLM worker discovery/health fails before registration. Please publish/bump the proto package and tighten the servicer dependency, or set the field only when the installed stubs expose it.
Useful? React with 👍 / 👎.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@grpc_servicer/smg_grpc_servicer/vllm/servicer.py`:
- Around line 166-172: Move the blocking remote branch of _tensor_from_proto out
of the asyncio event loop by making the preprocessed multimodal path
asynchronous and awaiting an asyncio.to_thread or executor call to
RdmaPixelPuller.feature_from_remote. Update _build_preprocessed_mm_inputs and
its callers, including Generate, to propagate the await while keeping inline
tensor deserialization synchronous and unchanged.
In `@model_gateway/src/routers/grpc/multimodal/transport.rs`:
- Around line 663-670: Update resolve_mm_rdma_enabled_requires_exporter to
configure the worker’s multimodal_tensor_transport as Some(TransportMode::Rdma)
rather than relying on the "_mode" label. Preserve the supports_rdma_pull label
and assertion so the test reaches and validates the exporter gate.
🪄 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: 3e2551a6-884a-40b3-807c-eec1c989fea8
📒 Files selected for processing (8)
crates/grpc_client/proto/vllm_engine.protogrpc_servicer/smg_grpc_servicer/mm_rdma.pygrpc_servicer/smg_grpc_servicer/vllm/servicer.pygrpc_servicer/tests/test_mm_rdma.pygrpc_servicer/tests/test_vllm_mm_rdma.pymodel_gateway/src/routers/grpc/multimodal/assemble.rsmodel_gateway/src/routers/grpc/multimodal/transport.rsmodel_gateway/src/routers/grpc/proto_wrapper.rs
| #[test] | ||
| fn resolve_mm_rdma_enabled_requires_exporter() { | ||
| // Even a capable worker cannot RDMA without the gateway exporter, which is | ||
| // never built in the default (stub) test build -> the gate stays off. | ||
| let workers = | ||
| single_worker_with_labels(&[("supports_rdma_pull", "true"), ("_mode", "unused")]); | ||
| assert!(!resolve_mm_rdma_enabled(Some(&workers))); | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
rg -n 'fn mm_transport_defaults' -A 15 model_gateway/src/routers/grpc/multimodal
rg -n 'multimodal_tensor_transport' model_gateway/src/routers/grpc/multimodal model_gateway/src/core -A3 -B3
rg -n 'fn labels' model_gateway/src/core/worker.rs -A 10Repository: lightseekorg/smg
Length of output: 2411
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '50,90p' model_gateway/src/routers/grpc/multimodal/transport.rs
sed -n '350,390p' model_gateway/src/routers/grpc/multimodal/transport.rs
rg -n 'single_worker_with_labels|worker_transport_mode_override|multimodal_tensor_transport|supports_rdma_pull' model_gateway/src -A 6 -B 6Repository: lightseekorg/smg
Length of output: 21616
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n 'enum TransportMode|impl Default for TransportMode|derive\(.*Default.*\)' model_gateway/src -A 20 -B 5
rg -n 'struct BasicWorkerBuilder|impl BasicWorkerBuilder|multimodal_tensor_transport' model_gateway/src -A 8 -B 8Repository: lightseekorg/smg
Length of output: 50372
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '1,140p' model_gateway/src/routers/grpc/multimodal/transport.rs
rg -n 'struct BasicWorkerBuilder|impl BasicWorkerBuilder|multimodal_tensor_transport|labels' model_gateway/src/worker.rs model_gateway/src -A 6 -B 6Repository: lightseekorg/smg
Length of output: 50373
🏁 Script executed:
#!/bin/bash
set -euo pipefail
fd -a 'worker.rs' model_gateway/src
fd -a 'mod.rs' model_gateway/src/worker
rg -n 'struct BasicWorkerBuilder|impl BasicWorkerBuilder|fn labels\(|multimodal_tensor_transport|_mode' model_gateway/src/worker.rs model_gateway/src/worker -A 12 -B 12Repository: lightseekorg/smg
Length of output: 50373
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n 'multimodal_tensor_transport|_mode' model_gateway/src/worker/builder.rs model_gateway/src/worker -A 6 -B 6
sed -n '1,220p' model_gateway/src/worker/builder.rsRepository: lightseekorg/smg
Length of output: 50373
Make this test force Rdma mode
_mode has no effect here; BasicWorkerBuilder only stores labels in spec.labels, while resolve_mm_rdma_enabled() reads spec.multimodal_tensor_transport. With the default inline mode, this assertion can pass without ever reaching the exporter gate. Set the worker spec to Some(TransportMode::Rdma) instead.
🤖 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/transport.rs` around lines 663 -
670, Update resolve_mm_rdma_enabled_requires_exporter to configure the worker’s
multimodal_tensor_transport as Some(TransportMode::Rdma) rather than relying on
the "_mode" label. Preserve the supports_rdma_pull label and assertion so the
test reaches and validates the exporter gate.
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
grpc_servicer/smg_grpc_servicer/vllm/servicer.py (1)
159-164: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winDocument both RDMA enablement switches.
The adjacent comment says initialization is a no-op unless
SMG_MM_PIXEL_RDMAis set, but_rdma_enabled_from_env()also enables it forSMG_MM_TENSOR_TRANSPORT=rdma. Update the comment so operators do not miss the first-class configuration.🤖 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 `@grpc_servicer/smg_grpc_servicer/vllm/servicer.py` around lines 159 - 164, The comment above _rdma_pixel_puller must document both RDMA enablement switches: SMG_MM_PIXEL_RDMA and SMG_MM_TENSOR_TRANSPORT=rdma. Update only the comment so it accurately states initialization is a no-op unless either configuration enables RDMA.
🤖 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.
Outside diff comments:
In `@grpc_servicer/smg_grpc_servicer/vllm/servicer.py`:
- Around line 159-164: The comment above _rdma_pixel_puller must document both
RDMA enablement switches: SMG_MM_PIXEL_RDMA and SMG_MM_TENSOR_TRANSPORT=rdma.
Update only the comment so it accurately states initialization is a no-op unless
either configuration enables RDMA.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: b80a97de-6796-47ec-9ffb-906441c4f8c1
📒 Files selected for processing (7)
crates/grpc_client/python/pyproject.tomlgrpc_servicer/pyproject.tomlgrpc_servicer/smg_grpc_servicer/mm_rdma.pygrpc_servicer/smg_grpc_servicer/vllm/servicer.pygrpc_servicer/tests/test_mm_rdma.pygrpc_servicer/tests/test_vllm_mm_rdma.pymodel_gateway/src/routers/grpc/multimodal/transport.rs
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: fca7eb88f2
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b268a95662
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| /// `SMG_MM_PIXEL_RDMA`). Gating on its presence — rather than the per-request | ||
| /// transport mode — matches TokenSpeed and honors every way the lane is enabled. | ||
| pub(super) fn resolve_mm_rdma_enabled(workers: Option<&WorkerSelection>) -> bool { | ||
| mm_rdma_exporter().is_some() && worker_supports_rdma_pull(workers) |
There was a problem hiding this comment.
In deployments where the router default is rdma but a specific worker is configured with WorkerSpec.multimodal_tensor_transport = inline or shm, this gate still enables vLLM RDMA whenever the exporter exists and the worker advertises supports_rdma_pull. The config/docs state that per-worker multimodal_tensor_transport overrides the router-level transport, and vllm_tensor_payload prefers RDMA before SHM, so the worker override is ignored and the router emits remote payloads even for workers explicitly forced to another transport.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
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/transport.rs (1)
665-670: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winMake the RDMA-gating test independent of global exporter state.
mm_rdma_exporter()is a process-wideOnceLockbacked by env-derived config, so this test only covers the default stub path. Split out a helper that takesexporter_availableand assert bothfalseandtruecases.🤖 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/transport.rs` around lines 665 - 670, Update the RDMA gating test around resolve_mm_rdma_enabled to avoid reading process-wide mm_rdma_exporter state. Extract or use a helper that accepts exporter_available explicitly, then assert RDMA remains disabled when it is false and enabled for a capable worker when it is true.
🤖 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.
Outside diff comments:
In `@model_gateway/src/routers/grpc/multimodal/transport.rs`:
- Around line 665-670: Update the RDMA gating test around
resolve_mm_rdma_enabled to avoid reading process-wide mm_rdma_exporter state.
Extract or use a helper that accepts exporter_available explicitly, then assert
RDMA remains disabled when it is false and enabled for a capable worker when it
is true.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 88b26d8e-49cd-47bc-a96d-0f3731188d96
📒 Files selected for processing (1)
model_gateway/src/routers/grpc/multimodal/transport.rs
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3e0ee0b7a5
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| mp.setenv("SMG_MM_TENSOR_TRANSPORT", "rdma") | ||
| mp.setenv("SMG_MM_PIXEL_RDMA", "1") | ||
| mp.setenv("SMG_RDMA_LISTEN_IP", "127.0.0.1") |
There was a problem hiding this comment.
When this opt-in test runs after another Qwen3-VL gRPC class in the same pytest session, setup_backend can reuse the session-cached vLLM worker because the worker-pool key does not include environment variables. In that scenario these RDMA env vars are applied only to the current process/gateway launch, not to the already-running worker, so the worker keeps advertising supports_rdma_pull=false and the remote-byte assertion fails even though the gateway code is working. Evict/bypass the pool for this test or make the RDMA env part of the worker cache key.
Useful? React with 👍 / 👎.
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 `@e2e_test/chat_completions/test_multimodal_rdma.py`:
- Around line 46-53: Harden _remote_pixel_bytes by calling raise_for_status() on
the /metrics response before parsing, then replace manual line splitting with
prometheus_client.parser.text_string_to_metric_families and sum samples from
smg_mm_tensor_bytes_total whose path label is "remote". Preserve the float total
result while allowing HTTP and malformed-metrics errors to surface directly.
🪄 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: 663045aa-e2e6-44f6-8b6a-168b440d1e3a
📒 Files selected for processing (1)
e2e_test/chat_completions/test_multimodal_rdma.py
| def _remote_pixel_bytes(metrics_url: str) -> float: | ||
| """Sum ``smg_mm_tensor_bytes_total`` samples on the RDMA (remote) path.""" | ||
| text = httpx.get(f"{metrics_url}/metrics", timeout=10).text | ||
| total = 0.0 | ||
| for line in text.splitlines(): | ||
| if line.startswith("smg_mm_tensor_bytes_total") and 'path="remote"' in line: | ||
| total += float(line.rsplit(" ", 1)[1]) | ||
| return total |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win
Harden _remote_pixel_bytes parsing.
Manual line-splitting assumes each matching sample line has no trailing timestamp field, and the response is never checked for HTTP errors — an unreachable /metrics endpoint would silently produce total == 0.0, making the eventual assertion failure ("fell back to inline") misleading about the actual root cause.
♻️ Proposed fix using prometheus_client's parser
def _remote_pixel_bytes(metrics_url: str) -> float:
"""Sum ``smg_mm_tensor_bytes_total`` samples on the RDMA (remote) path."""
- text = httpx.get(f"{metrics_url}/metrics", timeout=10).text
- total = 0.0
- for line in text.splitlines():
- if line.startswith("smg_mm_tensor_bytes_total") and 'path="remote"' in line:
- total += float(line.rsplit(" ", 1)[1])
- return total
+ from prometheus_client.parser import text_string_to_metric_families
+
+ response = httpx.get(f"{metrics_url}/metrics", timeout=10)
+ response.raise_for_status()
+ total = 0.0
+ for family in text_string_to_metric_families(response.text):
+ if family.name != "smg_mm_tensor_bytes_total":
+ continue
+ for sample in family.samples:
+ if sample.labels.get("path") == "remote":
+ total += sample.value
+ return totalConfirmed via web search that prometheus_client.parser.text_string_to_metric_families is the documented API for parsing Prometheus text-format responses, so this is a drop-in, well-supported replacement for the manual parsing.
📝 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.
| def _remote_pixel_bytes(metrics_url: str) -> float: | |
| """Sum ``smg_mm_tensor_bytes_total`` samples on the RDMA (remote) path.""" | |
| text = httpx.get(f"{metrics_url}/metrics", timeout=10).text | |
| total = 0.0 | |
| for line in text.splitlines(): | |
| if line.startswith("smg_mm_tensor_bytes_total") and 'path="remote"' in line: | |
| total += float(line.rsplit(" ", 1)[1]) | |
| return total | |
| def _remote_pixel_bytes(metrics_url: str) -> float: | |
| """Sum ``smg_mm_tensor_bytes_total`` samples on the RDMA (remote) path.""" | |
| from prometheus_client.parser import text_string_to_metric_families | |
| response = httpx.get(f"{metrics_url}/metrics", timeout=10) | |
| response.raise_for_status() | |
| total = 0.0 | |
| for family in text_string_to_metric_families(response.text): | |
| if family.name != "smg_mm_tensor_bytes_total": | |
| continue | |
| for sample in family.samples: | |
| if sample.labels.get("path") == "remote": | |
| total += sample.value | |
| return total |
🧰 Tools
🪛 ast-grep (0.44.1)
[warning] 47-47: Request-controlled URL passed to httpx; validate against an allowlist to prevent SSRF.
Context: httpx.get(f"{metrics_url}/metrics", timeout=10)
Note: [CWE-918] Server-Side Request Forgery (SSRF).
(avoid-ssrf)
🤖 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 `@e2e_test/chat_completions/test_multimodal_rdma.py` around lines 46 - 53,
Harden _remote_pixel_bytes by calling raise_for_status() on the /metrics
response before parsing, then replace manual line splitting with
prometheus_client.parser.text_string_to_metric_families and sum samples from
smg_mm_tensor_bytes_total whose path label is "remote". Preserve the float total
result while allowing HTTP and malformed-metrics errors to surface directly.
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 `@e2e_test/chat_completions/test_multimodal_rdma.py`:
- Around line 5-10: Restore explicit capability gating for the RDMA-only
multimodal test, using the existing SMG_E2E_MM_RDMA environment-based skip or an
equivalent check before test execution. Keep the test enabled only when the
required RDMA gateway feature and NIXL/UCX support are explicitly available.
🪄 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: 371f6cfe-2fec-4cf7-860b-5fab3edfa1f5
📒 Files selected for processing (1)
e2e_test/chat_completions/test_multimodal_rdma.py
| one-sided READ instead of receiving them inline. Verifies both that the model | ||
| still understands the image over RDMA AND that the pixels actually travelled the | ||
| remote path (``smg_mm_tensor_bytes_total{path="remote"}`` grew), so a silent | ||
| fall-back to inline fails the test rather than passing on correctness alone. | ||
|
|
||
| Requires a gateway built with ``--features mm-rdma`` and a NIXL/UCX-capable host. |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Restore explicit opt-in gating for this RDMA-only test.
Removing the SMG_E2E_MM_RDMA skip causes pytest to run the test on hosts without the required gateway feature or NIXL/UCX support, turning an optional validation into an unconditional failure. Restore the previous environment-based skip or an equivalent capability gate.
🤖 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 `@e2e_test/chat_completions/test_multimodal_rdma.py` around lines 5 - 10,
Restore explicit capability gating for the RDMA-only multimodal test, using the
existing SMG_E2E_MM_RDMA environment-based skip or an equivalent check before
test execution. Keep the test enabled only when the required RDMA gateway
feature and NIXL/UCX support are explicitly available.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2d298d10f5
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| /// `SMG_MM_PIXEL_RDMA`). Gating on its presence — rather than the per-request | ||
| /// transport mode — matches TokenSpeed and honors every way the lane is enabled. | ||
| pub(super) fn resolve_mm_rdma_enabled(workers: Option<&WorkerSelection>) -> bool { | ||
| mm_rdma_exporter().is_some() && worker_supports_rdma_pull(workers) |
There was a problem hiding this comment.
In deployments that still have SMG_MM_PIXEL_RDMA=1 in the gateway environment, this gate enables vLLM RDMA solely because the exporter exists and the worker advertises pull support, even if the router was explicitly started with --multimodal-tensor-transport inline (or config mode inline). init_mm_transport_defaults gives router config precedence over env, so bypassing the resolved mode here defeats the explicit opt-out and can force unexpected remote payloads; check the resolved transport mode before falling back to the legacy env switch.
Useful? React with 👍 / 👎.
2d298d1 to
da6b6db
Compare
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: da6b6db4b4
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| self._rdma_pixel_puller = RdmaPixelPuller( | ||
| agent_name=f"smg-vllm-{socket.gethostname()}-{os.getpid()}", | ||
| log_prefix="vLLM RDMA", | ||
| ) |
There was a problem hiding this comment.
When a vLLM worker process survives a gateway restart, this long-lived puller keeps RdmaPixelPuller._rdma_md_ready entries keyed only by (ip, port). Because the gateway RDMA listener defaults to the same port and uses a fixed agent name, the next gateway instance at the same address is treated as already handshaken, so no fresh fetch_remote_metadata/check_remote_metadata runs before the READ; that can leave the worker using stale NIXL metadata/rkeys and fail or hang until timeout. Include a gateway instance token in the descriptor/cache key or refresh metadata when the exporter restarts.
Useful? React with 👍 / 👎.
da6b6db to
d248695
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9364ee5d02
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| /// `SMG_MM_PIXEL_RDMA`). Gating on its presence — rather than the per-request | ||
| /// transport mode — matches TokenSpeed and honors every way the lane is enabled. | ||
| pub(super) fn resolve_mm_rdma_enabled(workers: Option<&WorkerSelection>) -> bool { | ||
| mm_rdma_exporter().is_some() && worker_supports_rdma_pull(workers) |
There was a problem hiding this comment.
When RDMA is configured but the selected vLLM worker lacks supports_rdma_pull=true (older worker, failed NIXL init, or explicit non-RDMA worker), this evaluates mm_rdma_exporter() before checking the worker capability. In a real mm-rdma build that can allocate/register the default 2 GiB gateway arena and start NIXL even though the second predicate immediately disables RDMA for the request; check worker_supports_rdma_pull(workers) first so unsupported workers stay on inline/SHM without consuming RDMA resources.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e36f5084ef
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| set -euo pipefail | ||
|
|
||
| # libclang for nixl-sys bindgen; the container already ships gcc/g++ (libstdc++). | ||
| dnf install -y clang-devel >/dev/null |
There was a problem hiding this comment.
In the new manylinux RDMA wheel job, this script only adds clang-devel before running maturin, but the workspace build still compiles smg-grpc-client, whose crates/grpc_client/build.rs invokes tonic_prost_build::compile_protos. The regular Rust setup installs a protobuf compiler for that step, and the manylinux image documentation only guarantees its listed wheel-building tools, so this dedicated job can fail before uploading smg-wheel-mm-rdma; add a suitable protoc/protobuf-compiler install or otherwise put a recent protoc on PATH here.
Useful? React with 👍 / 👎.
github-actions
Bot
removed
python-bindings
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c7452f913d
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| @pytest.mark.engine("vllm") | ||
| @pytest.mark.gpu(1) | ||
| @pytest.mark.e2e | ||
| @pytest.mark.model("Qwen/Qwen3-VL-8B-Instruct") | ||
| @pytest.mark.gateway(extra_args=["--multimodal-tensor-transport", "rdma"]) |
There was a problem hiding this comment.
This new vLLM/gpu1/e2e test lives under e2e_test/chat_completions, so it is collected by the normal e2e-1gpu-chat vLLM job (pr-test-rust.yml runs that directory via e2e-gpu-job.yml). That job downloads the standard smg-wheel, and scripts/ci_build_wheel.sh builds it with only vendored-openssl, not mm-rdma; in that default stub build RDMA exports fall back inline, so this test's remote_bytes > 0 assertion fails on regular chat CI even though the app still answers correctly. Please either skip/filter this test unless the RDMA wheel is installed or move it to a dedicated mm-rdma CI lane.
Useful? React with 👍 / 👎.
| (matching the gateway) or the legacy `SMG_MM_PIXEL_RDMA` flag.""" | ||
| if os.environ.get("SMG_MM_PIXEL_RDMA") in ("1", "true"): | ||
| return True | ||
| return os.environ.get("SMG_MM_TENSOR_TRANSPORT", "").strip().lower() == "rdma" |
There was a problem hiding this comment.
When a deployment enables RDMA with only the supported deprecated alias SMG_TOKENSPEED_MM_TENSOR_TRANSPORT=rdma, the gateway still treats the lane as enabled because transport.rs reads that alias, but this worker-side predicate returns false unless SMG_MM_TENSOR_TRANSPORT or SMG_MM_PIXEL_RDMA is set. The vLLM worker then advertises supports_rdma_pull=false, so the gateway never emits remote payloads for a configuration it otherwise accepts; include the same alias here to keep gateway and worker policy aligned.
Useful? React with 👍 / 👎.
github-actions
Bot
added
python-bindings
|
Hi @slin1237, this PR has merge conflicts that must be resolved before it can be merged. Please rebase your branch: git fetch origin main
git rebase origin/main
# resolve any conflicts, then:
git push --force-with-lease |
mergify
Bot
added
the
needs-rebase
Complete the engine-neutral RDMA transport: vLLM workers can now pull `pixel_values` over NIXL instead of receiving them inline, gated on a capability handshake so the gateway never emits an unfulfillable payload. Gateway (Rust): - `resolve_mm_rdma_enabled` enables the vLLM RDMA lane only when the resolved transport mode is `rdma`, the exporter is up, and the worker advertises `supports_rdma_pull`. `assemble_vllm` uses it (replacing the hard-coded off). - `worker_supports_rdma_pull` reads the capability label auto-lifted from vLLM `GetServerInfo`. Proto: - Add `bool supports_rdma_pull` to vLLM `GetServerInfoResponse`. Worker (Python): - The vLLM servicer builds an `RdmaPixelPuller` (no-op unless SMG_MM_PIXEL_RDMA is set) with a host+pid agent name; `_tensor_from_proto` pulls `remote` payloads via NIXL cast to the model dtype; `GetServerInfo` reports `supports_rdma_pull` only when the puller initialized. Add `bfloat16` to the proto dtype map. Tests: - Rust: capability-gate + label reads; vLLM emit degrades to inline without an exporter. - Python: engine-free descriptor parse (round-trip, room mismatch, legacy) and readiness; vLLM servicer remote-routing + dtype map (skipped without vLLM). Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
- Offload the preprocessed-multimodal build to a thread in Generate: the RDMA (remote) tensor pull does blocking NIXL waits that would otherwise stall the asyncio event loop and delay concurrent RPCs. (Gemini, CodeRabbit) - Add float16 to the vLLM proto dtype map so half-precision tensors deserialize instead of raising. (Gemini) - Enable the vLLM puller from the first-class SMG_MM_TENSOR_TRANSPORT=rdma (matching the gateway), not only the legacy SMG_MM_PIXEL_RDMA; otherwise a worker started with the documented transport flag advertises supports_rdma_pull=false and never receives remote payloads. (Codex) - Bump smg-grpc-proto to 0.4.15 (adds GetServerInfoResponse.supports_rdma_pull) and require it in the servicer so an older stub can't raise on the unknown keyword at GetServerInfo. (Codex) - Document that the RDMA exporter is process-wide, so a per-worker mode override can gate RDMA off but cannot turn it on. (Codex, CodeRabbit) Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
…t mode resolve_mm_rdma_enabled required the per-request transport mode to be `rdma`, but the exporter is built whenever the lane is enabled — including via the legacy `SMG_MM_PIXEL_RDMA` flag, which leaves the mode at the default `inline`. That combination built the exporter and let the worker advertise supports_rdma_pull=true, yet the gateway never emitted `remote` payloads. Gate purely on the exporter being present plus the worker's capability label, matching TokenSpeed (which emits whenever the exporter is up) and honoring every way the lane is enabled. The capability label is the per-worker opt-out. Drop the bogus `_mode` label from the test (mode is no longer read; the exporter gate is now exercised directly). (Codex, CodeRabbit) Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
Add an end-to-end test that drives an image request through a gateway + vLLM
worker configured for the RDMA (NIXL) pixel lane and asserts both that the model
understands the image and that the pixels actually travelled over the remote
path (scraping smg_mm_tensor_bytes_total{path="remote"}), so a silent inline
fall-back cannot pass.
Modeled on the existing TestMultimodalQwen3VL: engine=vllm, gpu=1, Qwen3-VL over
gRPC, transport set via the --multimodal-tensor-transport gateway flag plus the
worker/gateway RDMA env. Opt-in via SMG_E2E_MM_RDMA=1 since it needs a gateway
built with --features mm-rdma and a NIXL-capable host; it safely skips otherwise.
Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
Drop the invented gating from the RDMA e2e test and make it verify the lane for
real. It no longer skips behind a fabricated SMG_E2E_MM_RDMA flag or a
/workers supports_rdma_pull probe; instead it runs with the real SMG transport
config (--multimodal-tensor-transport rdma + the SMG_MM_TENSOR_TRANSPORT /
SMG_MM_PIXEL_RDMA / SMG_RDMA_LISTEN_IP env) and asserts both that the model
understands the image AND that pixels actually travelled the remote path
(smg_mm_tensor_bytes_total{path="remote"} > 0), so a silent inline fall-back
fails rather than passing on correctness alone.
Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
- worker pool: include the RDMA-lane env in the cache key so the RDMA e2e spawns its own worker instead of reusing an inline-only cached one (the test would otherwise fail even when the gateway is correct). - e2e metrics: raise_for_status() in _remote_pixel_bytes so an unreachable /metrics surfaces a real error instead of a misleading inline-fallback. - vllm servicer: correct the RDMA-puller comment to mention SMG_MM_TENSOR_TRANSPORT. Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
7a72763 to
4ef0b5e
Compare
mergify
Bot
removed
the
needs-rebase
The RDMA multimodal e2e needs a gateway built with --features mm-rdma, but the shared wheel is cross-compiled with maturin --zig (LLVM libc++), which can't link nixl-sys's GNU-libstdc++ C++ stub. Build a dedicated mm-rdma wheel natively on the bare Ubuntu CPU runner (GNU libstdc++, same runner image as the GPU e2e lanes), upload it as smg-wheel-mm-rdma, and run test_multimodal_rdma.py in a dedicated e2e-1gpu-mm-rdma job via e2e-gpu-job's new wheel_artifact input. The shared wheel and other e2e lanes are untouched; the RDMA test is excluded from the shared vLLM lane. Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
4ef0b5e to
e0e4cf4
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e0e4cf4eb8
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| && (github.event_name != 'pull_request' | ||
| || (needs.detect-changes.result == 'success' | ||
| && (needs.detect-changes.outputs.common == 'true' | ||
| || needs.detect-changes.outputs.chat-completions == 'true'))) |
There was a problem hiding this comment.
For PRs that touch only the RDMA implementation (for example crates/mm_rdma/**) or the new scripts/ci_build_wheel_mm_rdma.sh, the detect-changes filters in this workflow set neither common nor chat-completions because those paths are not listed there. This new job and the e2e-1gpu-mm-rdma job therefore skip exactly the mm-rdma wheel/e2e validation those changes need, so a broken RDMA crate or build script can merge without exercising the dedicated lane; include these paths in the gate or key it off rust-ci as well.
Useful? React with 👍 / 👎.
|
This pull request has been automatically marked as stale because it has not had any activity within 14 days. It will be automatically closed if no further activity occurs within 16 days. Leave a comment if you feel this pull request should remain open. Thank you! |
1 participant
Description
Problem
#1908 extracted the RDMA mechanics into
smg-mm-rdmaand #1915 made RDMA a first-class transport in the shared payload path — but vLLM emit was deliberately held off (rdma_enabled=false): vLLM workers had no puller and no way to advertise they could consume aremotepayload, so emitting one would hard-fail.Solution
Close the loop for vLLM: give the vLLM servicer a NIXL puller, advertise the capability via
GetServerInfo, and gate the gateway's emit on that capability. TokenSpeed/EPD are unchanged (they already consumeremote).Changes
Proto
bool supports_rdma_pull = 11to vLLMGetServerInfoResponse.Gateway (Rust)
resolve_mm_rdma_enabled(workers)enables the vLLM RDMA lane only when the resolved transport mode isrdma, the gateway exporter is up, and the worker advertisessupports_rdma_pull.assemble_vllmnow uses it (was hard-codedfalse).worker_supports_rdma_pullreads thesupports_rdma_pulllabel, whichclient.rs::flat_labelsauto-lifts fromGetServerInfo— zero extra discovery wiring.Worker (Python)
RdmaPixelPuller(no-op unlessSMG_MM_PIXEL_RDMAis set) with a uniquesmg-vllm-{host}-{pid}agent name (vLLM has no EPD bootstrap host/port)._tensor_from_protoroutesremotepayloads through the puller cast to the model dtype;GetServerInforeportssupports_rdma_pullonly when the puller actually initialized.bfloat16added to the proto dtype map.mm_shm.py'sremoteraise stays as a defensive backstop.Test Plan
The
SMGRDMA1wire format is unchanged; TokenSpeed/EPD paths are untouched.Added tests
transport.rs):worker_supports_rdma_pulllabel reads (true / false / missing / no worker);resolve_mm_rdma_enabledstays off without an exporter.proto_wrapper.rs): vLLMinto_protowithrdma_enabled=truedegrades to inline when no exporter is present (never emits an unfulfillableremote).test_mm_rdma.py, engine-free):SMGRDMA1descriptor parse — round-trip, room mismatch, too-short, legacy fallback — plus disabled-by-default readiness. Runs everywhere.test_vllm_mm_rdma.py): vLLM servicer remote-routing (calls the puller withexplicit_room=None+ model dtype), inline bf16 deserialize, and the dtype map. Skipped where vLLM isn't installed.Verification run
cargo clippy -p smgdefault and--features mm-rdma -- -D warningsclean;cargo testtransport + proto_wrapper pass;cargo +nightly fmt --all --checkclean.ruff check/ruff format --checkclean;pytest grpc_servicer/tests/test_mm_rdma.py test_vllm_mm_rdma.py→ 7 passed, 1 skipped (no local vLLM);codespellclean.Live NIXL/vLLM loopback validation on the GB300 box (
SMG_MM_TENSOR_TRANSPORT=rdma,SMG_MM_PIXEL_RDMA=1, a real Qwen3-VL worker) is the last step and will be run before merge once the stack lands.Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit
supports_rdma_pull) and enabled RDMA-based multimodalpixel_valuestransfers when RDMA is configured and the selected worker supports it.bfloat16), and a new end-to-end RDMA multimodal test using metrics.