From 21a18a082475bf055e5006900ae19ea7b77bb672 Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Mon, 13 Jul 2026 09:57:40 -0700 Subject: [PATCH 1/7] feat(multimodal): wire vLLM RDMA multimodal tensor pull 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> --- crates/grpc_client/proto/vllm_engine.proto | 5 ++ grpc_servicer/smg_grpc_servicer/mm_rdma.py | 5 ++ .../smg_grpc_servicer/vllm/servicer.py | 29 ++++-- grpc_servicer/tests/test_mm_rdma.py | 88 +++++++++++++++++++ grpc_servicer/tests/test_vllm_mm_rdma.py | 86 ++++++++++++++++++ .../src/routers/grpc/multimodal/assemble.rs | 8 +- .../src/routers/grpc/multimodal/transport.rs | 65 ++++++++++++++ .../src/routers/grpc/proto_wrapper.rs | 17 ++++ 8 files changed, 295 insertions(+), 8 deletions(-) create mode 100644 grpc_servicer/tests/test_mm_rdma.py create mode 100644 grpc_servicer/tests/test_vllm_mm_rdma.py diff --git a/crates/grpc_client/proto/vllm_engine.proto b/crates/grpc_client/proto/vllm_engine.proto index fd637f635..c7497bf40 100644 --- a/crates/grpc_client/proto/vllm_engine.proto +++ b/crates/grpc_client/proto/vllm_engine.proto @@ -360,6 +360,11 @@ message GetServerInfoResponse { // the router can verify a shared /dev/shm before using the SHM tensor // transport under `auto`. Empty when it can't be determined. string shm_namespace_id = 10; + + // Whether this worker can pull RDMA (NIXL) multimodal tensor payloads. The + // router only emits a `remote` payload when this is true, so an older worker + // that cannot consume it stays on the inline/SHM path. + bool supports_rdma_pull = 11; } // ===================== diff --git a/grpc_servicer/smg_grpc_servicer/mm_rdma.py b/grpc_servicer/smg_grpc_servicer/mm_rdma.py index 44b7b17c4..481dab93f 100644 --- a/grpc_servicer/smg_grpc_servicer/mm_rdma.py +++ b/grpc_servicer/smg_grpc_servicer/mm_rdma.py @@ -169,6 +169,11 @@ def __init__( self._nixl_agent = None self._landing_free = None + @property + def ready(self) -> bool: + """True once the NIXL agent + landing pool initialized (RDMA pulls possible).""" + return self._nixl_agent is not None and self._landing_free is not None + def _ensure_remote_ready(self, ip: str, port: int, remote, room: int) -> None: """One-time metadata handshake per gateway listener.""" key = (ip, port) diff --git a/grpc_servicer/smg_grpc_servicer/vllm/servicer.py b/grpc_servicer/smg_grpc_servicer/vllm/servicer.py index 1b1449518..d64521f95 100755 --- a/grpc_servicer/smg_grpc_servicer/vllm/servicer.py +++ b/grpc_servicer/smg_grpc_servicer/vllm/servicer.py @@ -9,6 +9,8 @@ import hashlib import itertools import json +import os +import socket import time from collections.abc import AsyncGenerator, AsyncIterator from datetime import datetime, timezone @@ -38,6 +40,7 @@ from vllm.sampling_params import RequestOutputKind, StructuredOutputsParams from smg_grpc_servicer import mm_shm +from smg_grpc_servicer.mm_rdma import RdmaPixelPuller from smg_grpc_servicer.tokenizer_bundle import CHUNK_SIZE, build_tokenizer_zip from smg_grpc_servicer.vllm.kv_events import ( endpoint_for_rank, @@ -69,13 +72,14 @@ def _filtered_sampling_defaults(params: dict | None) -> dict: # Proto dtype string → torch dtype _PROTO_DTYPE_MAP: dict[str, torch.dtype] = { "float32": torch.float32, + "bfloat16": torch.bfloat16, "int64": torch.int64, "uint32": torch.uint32, } -def _tensor_from_proto(td: vllm_engine_pb2.TensorData) -> torch.Tensor: - """Deserialize a TensorData proto message into a torch.Tensor.""" +def _inline_tensor_from_proto(td: vllm_engine_pb2.TensorData) -> torch.Tensor: + """Deserialize an inline/SHM TensorData proto message into a torch.Tensor.""" torch_dtype = _PROTO_DTYPE_MAP.get(td.dtype) if torch_dtype is None: raise ValueError(f"Unsupported proto tensor dtype: {td.dtype!r}") @@ -151,8 +155,22 @@ def __init__(self, async_llm: EngineClient, start_time: float): # Resolve KV-event publishing config from the engine. Non-None only when # vLLM was started with --kv-events-config enabling the ZMQ publisher. self._kv_events_config = resolve_kv_events_config(async_llm) + # No-op unless SMG_MM_PIXEL_RDMA is set; a unique agent name avoids NIXL + # metadata collisions since vLLM has no bootstrap host/port. + self._rdma_pixel_puller = RdmaPixelPuller( + agent_name=f"smg-vllm-{socket.gethostname()}-{os.getpid()}", + log_prefix="vLLM RDMA", + ) logger.info("VllmEngineServicer initialized") + def _tensor_from_proto(self, td: vllm_engine_pb2.TensorData) -> torch.Tensor: + """Deserialize a TensorData proto; RDMA `remote` payloads are pulled via NIXL.""" + if td.WhichOneof("payload") == "remote": + return self._rdma_pixel_puller.feature_from_remote( + td, explicit_room=None, cast_to=self.engine.model_config.dtype + ) + return _inline_tensor_from_proto(td) + async def Generate( self, request: vllm_engine_pb2.GenerateRequest, @@ -490,6 +508,7 @@ async def GetServerInfo( 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, ) async def GetLoads( @@ -628,12 +647,12 @@ def mm_key(key: str) -> str: return "pixel_values_videos" return key - # Deserialize all tensors from proto + # 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), } for key, td in mm_proto.model_specific_tensors.items(): - hf_dict[mm_key(key)] = _tensor_from_proto(td) + hf_dict[mm_key(key)] = self._tensor_from_proto(td) # Cast floating-point tensors to model dtype (e.g. bfloat16). # This mirrors _postprocess_output in multimodal/processing/context.py diff --git a/grpc_servicer/tests/test_mm_rdma.py b/grpc_servicer/tests/test_mm_rdma.py new file mode 100644 index 000000000..2be9f19d8 --- /dev/null +++ b/grpc_servicer/tests/test_mm_rdma.py @@ -0,0 +1,88 @@ +"""Engine-free tests for the shared RDMA pixel puller wire format. + +Covers the `SMGRDMA1` descriptor parse (shared by the TokenSpeed and vLLM +pullers) and the disabled-by-default readiness signal. No torch / nixl / vLLM +required — the module's heavy deps are imported lazily inside the pull path. + +Run with: pytest grpc_servicer/tests/test_mm_rdma.py -v +""" + +import importlib.util +import sys +import types +from pathlib import Path + +import pytest + +_MODULE_PATH = Path(__file__).parents[1] / "smg_grpc_servicer" / "mm_rdma.py" +_spec = importlib.util.spec_from_file_location("smg_mm_rdma_under_test", _MODULE_PATH) +mm_rdma = importlib.util.module_from_spec(_spec) +# Register before exec so dataclass annotation resolution can find the module. +sys.modules[_spec.name] = mm_rdma +_spec.loader.exec_module(mm_rdma) + + +def _descriptor(addr: int, gen: int, room: int, port: int, ip: str) -> bytes: + return ( + b"SMGRDMA1" + + addr.to_bytes(8, "little") + + gen.to_bytes(8, "little") + + room.to_bytes(8, "little", signed=True) + + port.to_bytes(2, "little") + + ip.encode() + ) + + +def _td(descriptor: bytes): + """Minimal stand-in for a proto TensorData carrying a remote descriptor.""" + return types.SimpleNamespace(remote=types.SimpleNamespace(descriptor=descriptor)) + + +class TestParseDescriptor: + def test_round_trips_all_fields(self): + desc = _descriptor(0xDEADBEEF, 7, 12345, 18515, "172.16.1.80") + parsed = mm_rdma._parse_descriptor(_td(desc), explicit_room=None) + assert parsed.remote_addr == 0xDEADBEEF + assert parsed.expected_gen == 7 + assert parsed.room == 12345 + assert parsed.port == 18515 + assert parsed.ip == "172.16.1.80" + + def test_matching_explicit_room_accepted(self): + desc = _descriptor(0x1000, 1, 99, 18515, "127.0.0.1") + parsed = mm_rdma._parse_descriptor(_td(desc), explicit_room=99) + assert parsed.room == 99 + + def test_room_mismatch_rejected(self): + desc = _descriptor(0x1000, 1, 99, 18515, "127.0.0.1") + with pytest.raises(ValueError, match="room mismatch"): + mm_rdma._parse_descriptor(_td(desc), explicit_room=100) + + def test_too_short_rejected(self): + with pytest.raises(ValueError, match="too short"): + mm_rdma._parse_descriptor(_td(b"SMGRDMA1\x00\x00"), explicit_room=None) + + def test_legacy_descriptor_needs_explicit_room(self): + # A pre-SMGRDMA1 descriptor (no magic) carries no room inline. + legacy = ( + (0x2000).to_bytes(8, "little") + + (5).to_bytes(8, "little") + + (18515).to_bytes(2, "little") + + b"127.0.0.1" + ) + parsed = mm_rdma._parse_descriptor(_td(legacy), explicit_room=42) + assert parsed.room == 42 + assert parsed.remote_addr == 0x2000 + assert parsed.expected_gen == 5 + with pytest.raises(ValueError, match="legacy remote descriptor lacks room"): + mm_rdma._parse_descriptor(_td(legacy), explicit_room=None) + + +class TestReadiness: + def test_disabled_when_env_unset(self, monkeypatch): + monkeypatch.delenv("SMG_MM_PIXEL_RDMA", raising=False) + puller = mm_rdma.RdmaPixelPuller(agent_name="test-agent", log_prefix="test") + assert puller.ready is False + + def test_default_gateway_agent_name(self): + assert mm_rdma.DEFAULT_GATEWAY_AGENT_NAME == "smg-gateway-encode" diff --git a/grpc_servicer/tests/test_vllm_mm_rdma.py b/grpc_servicer/tests/test_vllm_mm_rdma.py new file mode 100644 index 000000000..a754b5bf8 --- /dev/null +++ b/grpc_servicer/tests/test_vllm_mm_rdma.py @@ -0,0 +1,86 @@ +"""Integration tests for the vLLM servicer's RDMA multimodal wiring. + +Exercises the servicer's tensor-deserialization seam: `remote` (RDMA) payloads +route to the NIXL puller with the model dtype, inline payloads deserialize +directly, and bfloat16 is a supported wire dtype. Skipped where vLLM isn't +installed (the servicer module imports vLLM at import time). + +Run with: pytest grpc_servicer/tests/test_vllm_mm_rdma.py -v +""" + +import struct +import types + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("vllm") +pytest.importorskip("smg_grpc_proto") + +from smg_grpc_proto import vllm_engine_pb2 # noqa: E402 +from smg_grpc_proto.generated import common_pb2 # noqa: E402 +from smg_grpc_servicer.vllm import servicer as vllm_servicer # noqa: E402 + + +class _FakePuller: + def __init__(self, ready: bool = True): + self.ready = ready + self.calls = [] + + def feature_from_remote(self, td, *, explicit_room, cast_to): + self.calls.append((td, explicit_room, cast_to)) + return torch.zeros(1, dtype=cast_to) + + +def _servicer(puller, dtype): + """Servicer instance bypassing __init__ (no real engine needed).""" + s = vllm_servicer.VllmEngineServicer.__new__(vllm_servicer.VllmEngineServicer) + s._rdma_pixel_puller = puller + s.engine = types.SimpleNamespace(model_config=types.SimpleNamespace(dtype=dtype)) + return s + + +def test_bfloat16_is_a_supported_wire_dtype(): + assert vllm_servicer._PROTO_DTYPE_MAP["bfloat16"] is torch.bfloat16 + + +def test_remote_payload_routes_to_puller_with_model_dtype(): + puller = _FakePuller() + s = _servicer(puller, torch.bfloat16) + td = vllm_engine_pb2.TensorData( + shape=[1], + dtype="float32", + remote=common_pb2.RemoteTensorHandle(transport="nixl", descriptor=b"x", nbytes=4), + ) + + out = s._tensor_from_proto(td) + + assert len(puller.calls) == 1 + called_td, explicit_room, cast_to = puller.calls[0] + assert called_td is td + assert explicit_room is None # vLLM has no EPD bootstrap_room + assert cast_to is torch.bfloat16 + assert out.dtype is torch.bfloat16 + + +def test_inline_bfloat16_payload_deserializes(): + s = _servicer(_FakePuller(), torch.bfloat16) + # Two bf16 1.0 values (0x3F80 little-endian), carried inline. + payload = (0x3F80).to_bytes(2, "little") * 2 + td = vllm_engine_pb2.TensorData(shape=[2], dtype="bfloat16", inline=payload) + + out = s._tensor_from_proto(td) + + assert out.dtype is torch.bfloat16 + assert list(out.shape) == [2] + assert out.tolist() == [1.0, 1.0] + + +def test_inline_payload_does_not_touch_puller(): + puller = _FakePuller() + s = _servicer(puller, torch.float32) + td = vllm_engine_pb2.TensorData(shape=[1], dtype="float32", inline=struct.pack(") -> usi worker_shm_min_bytes_override(workers).unwrap_or_else(|| mm_transport_defaults().shm_min_bytes) } +/// Resolve whether vLLM `pixel_values` may use the RDMA lane for this request: the +/// resolved transport mode is `rdma`, the gateway exporter is up, and the worker +/// advertises it can pull. The capability gate keeps SMG from emitting a `remote` +/// payload to a worker that would reject it. +pub(super) fn resolve_mm_rdma_enabled(workers: Option<&WorkerSelection>) -> bool { + let mode = + worker_transport_mode_override(workers).unwrap_or_else(|| mm_transport_defaults().mode); + mode == TransportMode::Rdma + && mm_rdma_exporter().is_some() + && worker_supports_rdma_pull(workers) +} + +/// Whether the request's worker advertises RDMA-pull support via the +/// `supports_rdma_pull` label (lifted from vLLM `GetServerInfo`). Missing or +/// non-`"true"` reads as no. +fn worker_supports_rdma_pull(workers: Option<&WorkerSelection>) -> bool { + primary_worker(workers).is_some_and(|worker| { + worker + .metadata() + .spec + .labels + .get("supports_rdma_pull") + .map(String::as_str) + == Some("true") + }) +} + // ===================== RDMA pixel lane ===================== // // The gateway owns all RDMA *policy*: it decides whether the lane is on (a @@ -604,6 +631,44 @@ mod tests { assert_eq!(clamp_pool_slots(64, 32 * 1024 * 1024), 64); } + fn single_worker_with_labels(pairs: &[(&str, &str)]) -> WorkerSelection { + use crate::worker::BasicWorkerBuilder; + let labels = pairs + .iter() + .map(|(k, v)| ((*k).to_string(), (*v).to_string())) + .collect(); + let worker = BasicWorkerBuilder::new("http://test:8000") + .labels(labels) + .build(); + WorkerSelection::Single { + worker: Arc::new(worker), + } + } + + #[test] + fn worker_supports_rdma_pull_reads_label() { + assert!(worker_supports_rdma_pull(Some(&single_worker_with_labels( + &[("supports_rdma_pull", "true",)] + )))); + assert!(!worker_supports_rdma_pull(Some( + &single_worker_with_labels(&[("supports_rdma_pull", "false")]) + ))); + // Missing label (older worker) reads as no. + assert!(!worker_supports_rdma_pull(Some( + &single_worker_with_labels(&[]) + ))); + assert!(!worker_supports_rdma_pull(None)); + } + + #[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))); + } + #[test] #[cfg(target_os = "linux")] fn local_shm_namespace_id_resolves_on_linux() { diff --git a/model_gateway/src/routers/grpc/proto_wrapper.rs b/model_gateway/src/routers/grpc/proto_wrapper.rs index d8b4deb09..6a0b81203 100644 --- a/model_gateway/src/routers/grpc/proto_wrapper.rs +++ b/model_gateway/src/routers/grpc/proto_wrapper.rs @@ -2331,4 +2331,21 @@ mod tests { let image = vllm_mm_data(common::Modality::Image).into_proto(); assert_eq!(image.modality, common::Modality::Image as i32); } + + #[test] + fn vllm_rdma_enabled_without_exporter_stays_inline() { + // With rdma_enabled=true but no gateway exporter (never built in the stub + // test build), pixel_values must fall back to inline rather than emit an + // unfulfillable `remote` payload. + let mut data = vllm_mm_data(common::Modality::Image); + data.rdma_enabled = true; + let proto = data.into_proto(); + assert!( + matches!( + proto.pixel_values.unwrap().payload, + Some(vllm::tensor_data::Payload::Inline(_)) + ), + "pixel_values must stay inline when the RDMA exporter is absent" + ); + } } From fd8a446556c7c1cb232913686078fb22973bea33 Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Mon, 13 Jul 2026 13:55:31 -0700 Subject: [PATCH 2/7] fix(multimodal): harden vLLM RDMA transport per review - 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> --- crates/grpc_client/python/pyproject.toml | 2 +- grpc_servicer/pyproject.toml | 4 +++- grpc_servicer/smg_grpc_servicer/mm_rdma.py | 10 +++++++++- .../smg_grpc_servicer/vllm/servicer.py | 11 ++++++++-- grpc_servicer/tests/test_mm_rdma.py | 20 ++++++++++++++++++- grpc_servicer/tests/test_vllm_mm_rdma.py | 3 ++- .../src/routers/grpc/multimodal/transport.rs | 4 +++- 7 files changed, 46 insertions(+), 8 deletions(-) diff --git a/crates/grpc_client/python/pyproject.toml b/crates/grpc_client/python/pyproject.toml index 3892134c9..037762b91 100644 --- a/crates/grpc_client/python/pyproject.toml +++ b/crates/grpc_client/python/pyproject.toml @@ -8,7 +8,7 @@ build-backend = "setuptools.build_meta" [project] name = "smg-grpc-proto" -version = "0.4.14" +version = "0.4.15" description = "SMG gRPC proto definitions for vLLM, TRT-LLM, MLX, TokenSpeed, and SGLang" requires-python = ">=3.10" dependencies = [ diff --git a/grpc_servicer/pyproject.toml b/grpc_servicer/pyproject.toml index 023e75a6b..4247c2077 100644 --- a/grpc_servicer/pyproject.toml +++ b/grpc_servicer/pyproject.toml @@ -8,7 +8,9 @@ version = "0.6.0" description = "SMG gRPC servicer implementations for LLM inference engines (vLLM, MLX, TokenSpeed, SGLang)" requires-python = ">=3.10" dependencies = [ - "smg-grpc-proto>=0.4.13", + # >=0.4.15 ships GetServerInfoResponse.supports_rdma_pull (field 11); the vLLM + # servicer sets it, so older stubs would raise on the unknown keyword. + "smg-grpc-proto>=0.4.15", "grpcio>=1.81.1", "grpcio-reflection>=1.81.1", "grpcio-health-checking>=1.81.1", diff --git a/grpc_servicer/smg_grpc_servicer/mm_rdma.py b/grpc_servicer/smg_grpc_servicer/mm_rdma.py index 481dab93f..bbe8abb23 100644 --- a/grpc_servicer/smg_grpc_servicer/mm_rdma.py +++ b/grpc_servicer/smg_grpc_servicer/mm_rdma.py @@ -20,6 +20,14 @@ _LOCAL_IP_CACHE: dict[str, bool] = {} +def _rdma_enabled_from_env() -> bool: + """Whether the RDMA lane is on: the first-class `SMG_MM_TENSOR_TRANSPORT=rdma` + (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" + + @dataclass(frozen=True) class _RemotePixelDescriptor: remote_addr: int @@ -121,7 +129,7 @@ def __init__( self._landing_slot_bytes = 0 self._landing_free = None - if os.environ.get("SMG_MM_PIXEL_RDMA") not in ("1", "true"): + if not _rdma_enabled_from_env(): return try: diff --git a/grpc_servicer/smg_grpc_servicer/vllm/servicer.py b/grpc_servicer/smg_grpc_servicer/vllm/servicer.py index d64521f95..0cb40adc1 100755 --- a/grpc_servicer/smg_grpc_servicer/vllm/servicer.py +++ b/grpc_servicer/smg_grpc_servicer/vllm/servicer.py @@ -72,6 +72,7 @@ def _filtered_sampling_defaults(params: dict | None) -> dict: # Proto dtype string → torch dtype _PROTO_DTYPE_MAP: dict[str, torch.dtype] = { "float32": torch.float32, + "float16": torch.float16, "bfloat16": torch.bfloat16, "int64": torch.int64, "uint32": torch.uint32, @@ -212,8 +213,14 @@ async def Generate( if has_preprocessed_mm and input_type == "tokenized": # Preprocessed multimodal from Rust router. # Token IDs already have expanded placeholders; tensors are - # ready for the model. Bypass the renderer entirely. - prompt = self._build_preprocessed_mm_inputs(request.tokenized, request.mm_inputs) + # ready for the model. Bypass the renderer entirely. Offloaded to + # a thread: the RDMA (remote) tensor pull does blocking NIXL waits + # that must not stall the asyncio event loop. + prompt = await asyncio.to_thread( + self._build_preprocessed_mm_inputs, + request.tokenized, + request.mm_inputs, + ) prompt["arrival_time"] = arrival_time elif input_type == "tokenized": prompt: TokensPrompt = {"prompt_token_ids": list(request.tokenized.input_ids)} diff --git a/grpc_servicer/tests/test_mm_rdma.py b/grpc_servicer/tests/test_mm_rdma.py index 2be9f19d8..8f3964cb1 100644 --- a/grpc_servicer/tests/test_mm_rdma.py +++ b/grpc_servicer/tests/test_mm_rdma.py @@ -78,11 +78,29 @@ def test_legacy_descriptor_needs_explicit_room(self): mm_rdma._parse_descriptor(_td(legacy), explicit_room=None) -class TestReadiness: +class TestEnvGate: def test_disabled_when_env_unset(self, monkeypatch): monkeypatch.delenv("SMG_MM_PIXEL_RDMA", raising=False) + monkeypatch.delenv("SMG_MM_TENSOR_TRANSPORT", raising=False) + assert mm_rdma._rdma_enabled_from_env() is False puller = mm_rdma.RdmaPixelPuller(agent_name="test-agent", log_prefix="test") assert puller.ready is False + def test_enabled_by_legacy_flag(self, monkeypatch): + monkeypatch.setenv("SMG_MM_PIXEL_RDMA", "1") + monkeypatch.delenv("SMG_MM_TENSOR_TRANSPORT", raising=False) + assert mm_rdma._rdma_enabled_from_env() is True + + def test_enabled_by_transport_mode(self, monkeypatch): + # The first-class transport switch (matching the gateway) enables the lane. + monkeypatch.delenv("SMG_MM_PIXEL_RDMA", raising=False) + monkeypatch.setenv("SMG_MM_TENSOR_TRANSPORT", "rdma") + assert mm_rdma._rdma_enabled_from_env() is True + + def test_other_transport_mode_stays_disabled(self, monkeypatch): + monkeypatch.delenv("SMG_MM_PIXEL_RDMA", raising=False) + monkeypatch.setenv("SMG_MM_TENSOR_TRANSPORT", "shm") + assert mm_rdma._rdma_enabled_from_env() is False + def test_default_gateway_agent_name(self): assert mm_rdma.DEFAULT_GATEWAY_AGENT_NAME == "smg-gateway-encode" diff --git a/grpc_servicer/tests/test_vllm_mm_rdma.py b/grpc_servicer/tests/test_vllm_mm_rdma.py index a754b5bf8..9418eb49d 100644 --- a/grpc_servicer/tests/test_vllm_mm_rdma.py +++ b/grpc_servicer/tests/test_vllm_mm_rdma.py @@ -40,8 +40,9 @@ def _servicer(puller, dtype): return s -def test_bfloat16_is_a_supported_wire_dtype(): +def test_half_precision_wire_dtypes_supported(): assert vllm_servicer._PROTO_DTYPE_MAP["bfloat16"] is torch.bfloat16 + assert vllm_servicer._PROTO_DTYPE_MAP["float16"] is torch.float16 def test_remote_payload_routes_to_puller_with_model_dtype(): diff --git a/model_gateway/src/routers/grpc/multimodal/transport.rs b/model_gateway/src/routers/grpc/multimodal/transport.rs index 4d38b41e6..12fd4a2be 100644 --- a/model_gateway/src/routers/grpc/multimodal/transport.rs +++ b/model_gateway/src/routers/grpc/multimodal/transport.rs @@ -157,7 +157,9 @@ pub(super) fn resolve_mm_shm_min_bytes(workers: Option<&WorkerSelection>) -> usi /// Resolve whether vLLM `pixel_values` may use the RDMA lane for this request: the /// resolved transport mode is `rdma`, the gateway exporter is up, and the worker /// advertises it can pull. The capability gate keeps SMG from emitting a `remote` -/// payload to a worker that would reject it. +/// payload to a worker that would reject it. The exporter is a process-wide +/// resource (NIXL agent + arena), so a per-worker mode override can gate RDMA off +/// but cannot turn it on without the router-level transport mode / env. pub(super) fn resolve_mm_rdma_enabled(workers: Option<&WorkerSelection>) -> bool { let mode = worker_transport_mode_override(workers).unwrap_or_else(|| mm_transport_defaults().mode); From 70317178db97c855181ec640ee24b94da4273f2b Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Mon, 13 Jul 2026 15:31:05 -0700 Subject: [PATCH 3/7] fix(multimodal): gate vLLM RDMA emit on exporter presence, not request mode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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> --- .../src/routers/grpc/multimodal/transport.rs | 23 +++++++++---------- 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/model_gateway/src/routers/grpc/multimodal/transport.rs b/model_gateway/src/routers/grpc/multimodal/transport.rs index 12fd4a2be..80dfe1aba 100644 --- a/model_gateway/src/routers/grpc/multimodal/transport.rs +++ b/model_gateway/src/routers/grpc/multimodal/transport.rs @@ -155,17 +155,17 @@ pub(super) fn resolve_mm_shm_min_bytes(workers: Option<&WorkerSelection>) -> usi } /// Resolve whether vLLM `pixel_values` may use the RDMA lane for this request: the -/// resolved transport mode is `rdma`, the gateway exporter is up, and the worker -/// advertises it can pull. The capability gate keeps SMG from emitting a `remote` -/// payload to a worker that would reject it. The exporter is a process-wide -/// resource (NIXL agent + arena), so a per-worker mode override can gate RDMA off -/// but cannot turn it on without the router-level transport mode / env. +/// gateway exporter is up and the worker advertises it can pull. The capability +/// gate keeps SMG from emitting a `remote` payload to a worker that would reject +/// it, and is the per-worker opt-out (a worker that shouldn't RDMA reports +/// `supports_rdma_pull=false`). +/// +/// The exporter is a process-wide resource built once the RDMA lane is enabled +/// (router transport mode `rdma`, `SMG_MM_TENSOR_TRANSPORT=rdma`, or the legacy +/// `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 { - let mode = - worker_transport_mode_override(workers).unwrap_or_else(|| mm_transport_defaults().mode); - mode == TransportMode::Rdma - && mm_rdma_exporter().is_some() - && worker_supports_rdma_pull(workers) + mm_rdma_exporter().is_some() && worker_supports_rdma_pull(workers) } /// Whether the request's worker advertises RDMA-pull support via the @@ -666,8 +666,7 @@ mod tests { 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")]); + let workers = single_worker_with_labels(&[("supports_rdma_pull", "true")]); assert!(!resolve_mm_rdma_enabled(Some(&workers))); } From 7374bfc425501f2b3c2033a184fbcd4b513cdc13 Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Mon, 13 Jul 2026 16:34:52 -0700 Subject: [PATCH 4/7] test(multimodal): add vLLM RDMA multimodal e2e test 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> --- .../chat_completions/test_multimodal_rdma.py | 117 ++++++++++++++++++ 1 file changed, 117 insertions(+) create mode 100644 e2e_test/chat_completions/test_multimodal_rdma.py diff --git a/e2e_test/chat_completions/test_multimodal_rdma.py b/e2e_test/chat_completions/test_multimodal_rdma.py new file mode 100644 index 000000000..d90d86895 --- /dev/null +++ b/e2e_test/chat_completions/test_multimodal_rdma.py @@ -0,0 +1,117 @@ +"""vLLM RDMA multimodal transport E2E test. + +Exercises the full RDMA (NIXL) pixel lane end to end: the gateway stages +``pixel_values`` into a pre-registered arena and the vLLM worker pulls them with a +one-sided READ instead of receiving them inline. Asserts the model still +understands the image AND — critically — that the pixels actually travelled over +the remote path, so a silent fall-back to inline cannot pass this test. + +Opt-in via ``SMG_E2E_MM_RDMA=1`` because it needs more than a stock CI box: + - the gateway binary built with ``--features mm-rdma`` and ``libnixl`` on + ``LD_LIBRARY_PATH``, and + - a NIXL/UCX-capable host (loopback RoCE is fine). +Without those the RDMA lane safely degrades to inline, which would (correctly) +fail the remote-path assertion, so the class skips unless explicitly enabled. + +Usage: + SMG_E2E_MM_RDMA=1 pytest e2e_test/chat_completions/test_multimodal_rdma.py -v +""" + +from __future__ import annotations + +import base64 +import logging +import os +from pathlib import Path + +import httpx +import pytest + +logger = logging.getLogger(__name__) + +FIXTURES_DIR = Path(__file__).parent.parent / "fixtures" / "images" +DOG_IMAGE_PATH = FIXTURES_DIR / "dog.jpg" # Black labrador puppy + +pytestmark = pytest.mark.skipif( + os.environ.get("SMG_E2E_MM_RDMA", "").lower() not in ("1", "true", "yes"), + reason="RDMA e2e needs a mm-rdma gateway build + NIXL; set SMG_E2E_MM_RDMA=1 to run", +) + + +def _image_to_base64_url(path: Path) -> str: + data = base64.b64encode(path.read_bytes()).decode("utf-8") + return f"data:image/jpeg;base64,{data}" + + +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 + + +@pytest.fixture(scope="class", autouse=True) +def _rdma_env(): + """Turn the RDMA lane on for both the gateway and the vLLM worker, which + inherit this process's environment when launched locally. Restored on teardown. + + - ``SMG_MM_TENSOR_TRANSPORT=rdma``: the first-class transport switch (also set + as a gateway CLI flag) that enables the worker-side puller. + - ``SMG_MM_PIXEL_RDMA=1``: the legacy puller switch, belt-and-suspenders. + - ``SMG_RDMA_LISTEN_IP=127.0.0.1``: the gateway's NIXL listener (loopback). + """ + with pytest.MonkeyPatch.context() as mp: + mp.setenv("SMG_MM_TENSOR_TRANSPORT", "rdma") + mp.setenv("SMG_MM_PIXEL_RDMA", "1") + mp.setenv("SMG_RDMA_LISTEN_IP", "127.0.0.1") + yield + + +@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"]) +@pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True) +class TestMultimodalRdmaQwen3VL: + """vLLM multimodal over the RDMA pixel lane via gRPC.""" + + def test_single_image_uses_rdma(self, model, setup_backend): + _, _, client, gateway = setup_backend + + response = client.chat.completions.create( + model=model, + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What animal is in this image?"}, + { + "type": "image_url", + "image_url": {"url": _image_to_base64_url(DOG_IMAGE_PATH)}, + }, + ], + } + ], + temperature=0, + max_tokens=100, + ) + + text = response.choices[0].message.content + assert any(k in text.lower() for k in ["dog", "puppy", "labrador"]), ( + f"Expected dog-related content over the RDMA transport, got: {text}" + ) + logger.info("RDMA multimodal response: %s", text) + + # Critical: the request must have actually used the remote (RDMA) path. + # A silent inline fall-back would still answer correctly, so correctness + # alone does not prove the lane worked — the transport metric does. + remote_bytes = _remote_pixel_bytes(gateway.metrics_url) + assert remote_bytes > 0, ( + 'expected smg_mm_tensor_bytes_total{path="remote"} > 0; the pixels ' + "silently fell back to inline (RDMA lane not exercised)" + ) + logger.info("RDMA remote pixel bytes: %s", remote_bytes) From 928e68638b3961f78ed7e69d4555661a937377e8 Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Tue, 14 Jul 2026 05:50:35 -0700 Subject: [PATCH 5/7] test(multimodal): actually exercise RDMA in the vLLM e2e test 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> --- .../chat_completions/test_multimodal_rdma.py | 22 +++++-------------- 1 file changed, 6 insertions(+), 16 deletions(-) diff --git a/e2e_test/chat_completions/test_multimodal_rdma.py b/e2e_test/chat_completions/test_multimodal_rdma.py index d90d86895..72783b5fe 100644 --- a/e2e_test/chat_completions/test_multimodal_rdma.py +++ b/e2e_test/chat_completions/test_multimodal_rdma.py @@ -2,26 +2,21 @@ Exercises the full RDMA (NIXL) pixel lane end to end: the gateway stages ``pixel_values`` into a pre-registered arena and the vLLM worker pulls them with a -one-sided READ instead of receiving them inline. Asserts the model still -understands the image AND — critically — that the pixels actually travelled over -the remote path, so a silent fall-back to inline cannot pass this test. +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. -Opt-in via ``SMG_E2E_MM_RDMA=1`` because it needs more than a stock CI box: - - the gateway binary built with ``--features mm-rdma`` and ``libnixl`` on - ``LD_LIBRARY_PATH``, and - - a NIXL/UCX-capable host (loopback RoCE is fine). -Without those the RDMA lane safely degrades to inline, which would (correctly) -fail the remote-path assertion, so the class skips unless explicitly enabled. +Requires a gateway built with ``--features mm-rdma`` and a NIXL/UCX-capable host. Usage: - SMG_E2E_MM_RDMA=1 pytest e2e_test/chat_completions/test_multimodal_rdma.py -v + pytest e2e_test/chat_completions/test_multimodal_rdma.py -v """ from __future__ import annotations import base64 import logging -import os from pathlib import Path import httpx @@ -32,11 +27,6 @@ FIXTURES_DIR = Path(__file__).parent.parent / "fixtures" / "images" DOG_IMAGE_PATH = FIXTURES_DIR / "dog.jpg" # Black labrador puppy -pytestmark = pytest.mark.skipif( - os.environ.get("SMG_E2E_MM_RDMA", "").lower() not in ("1", "true", "yes"), - reason="RDMA e2e needs a mm-rdma gateway build + NIXL; set SMG_E2E_MM_RDMA=1 to run", -) - def _image_to_base64_url(path: Path) -> str: data = base64.b64encode(path.read_bytes()).decode("utf-8") From 294608dd3f9c621b63dd60633211662b22f2d4df Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Tue, 14 Jul 2026 18:49:55 -0700 Subject: [PATCH 6/7] fix(multimodal): address RDMA e2e review comments - 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 --- .../chat_completions/test_multimodal_rdma.py | 5 ++-- e2e_test/infra/worker_pool.py | 26 ++++++++++++++----- .../smg_grpc_servicer/vllm/servicer.py | 5 ++-- 3 files changed, 26 insertions(+), 10 deletions(-) diff --git a/e2e_test/chat_completions/test_multimodal_rdma.py b/e2e_test/chat_completions/test_multimodal_rdma.py index 72783b5fe..41332113f 100644 --- a/e2e_test/chat_completions/test_multimodal_rdma.py +++ b/e2e_test/chat_completions/test_multimodal_rdma.py @@ -35,9 +35,10 @@ def _image_to_base64_url(path: Path) -> str: 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 + resp = httpx.get(f"{metrics_url}/metrics", timeout=10) + resp.raise_for_status() total = 0.0 - for line in text.splitlines(): + for line in resp.text.splitlines(): if line.startswith("smg_mm_tensor_bytes_total") and 'path="remote"' in line: total += float(line.rsplit(" ", 1)[1]) return total diff --git a/e2e_test/infra/worker_pool.py b/e2e_test/infra/worker_pool.py index 75f58c4a3..439ece763 100644 --- a/e2e_test/infra/worker_pool.py +++ b/e2e_test/infra/worker_pool.py @@ -25,6 +25,7 @@ import atexit import logging +import os import threading from .constants import DEFAULT_STARTUP_TIMEOUT, ConnectionMode, WorkerType @@ -33,12 +34,24 @@ logger = logging.getLogger(__name__) -# Key is (engine, model_id, mode, worker_type, count, gpus, extra_engine_args). -# ``count`` is part of the key because a class asking for count=2 after a -# count=1 class on the same backend would otherwise reuse a 1-worker entry and -# run with the wrong topology; ``gpus``/``extra_engine_args`` likewise change -# the launched topology (e.g. --data-parallel-size). -_PoolKey = tuple[str, str, ConnectionMode, WorkerType, int, int | None, tuple[str, ...] | None] +# Key is (engine, model_id, mode, worker_type, count, gpus, extra_engine_args, +# rdma_env). ``count``/``gpus``/``extra_engine_args`` change the launched +# topology, so a class asking for a different one must not reuse a cached entry. +# ``rdma_env`` is included because workers inherit the launcher's environment and +# the RDMA lane is env-gated, so an RDMA-enabled worker must not be reused for an +# inline test (or vice versa). +_PoolKey = tuple[ + str, str, ConnectionMode, WorkerType, int, int | None, tuple[str, ...] | None, tuple[str, ...] +] + + +def _rdma_env_signature() -> tuple[str, ...]: + """RDMA-lane env inherited by workers at launch; part of the pool key.""" + return ( + os.environ.get("SMG_MM_TENSOR_TRANSPORT", ""), + os.environ.get("SMG_MM_PIXEL_RDMA", ""), + os.environ.get("SMG_RDMA_LISTEN_IP", ""), + ) class WorkerPool: @@ -123,6 +136,7 @@ def acquire( count, gpus, tuple(extra_engine_args) if extra_engine_args else None, + _rdma_env_signature(), ) if self._key == key and all(w.is_alive() for w in self._workers): diff --git a/grpc_servicer/smg_grpc_servicer/vllm/servicer.py b/grpc_servicer/smg_grpc_servicer/vllm/servicer.py index 0cb40adc1..f52319ab3 100755 --- a/grpc_servicer/smg_grpc_servicer/vllm/servicer.py +++ b/grpc_servicer/smg_grpc_servicer/vllm/servicer.py @@ -156,8 +156,9 @@ def __init__(self, async_llm: EngineClient, start_time: float): # Resolve KV-event publishing config from the engine. Non-None only when # vLLM was started with --kv-events-config enabling the ZMQ publisher. self._kv_events_config = resolve_kv_events_config(async_llm) - # No-op unless SMG_MM_PIXEL_RDMA is set; a unique agent name avoids NIXL - # metadata collisions since vLLM has no bootstrap host/port. + # No-op unless the RDMA lane is on (SMG_MM_TENSOR_TRANSPORT=rdma or legacy + # SMG_MM_PIXEL_RDMA); a unique agent name avoids NIXL metadata collisions + # since vLLM has no bootstrap host/port. self._rdma_pixel_puller = RdmaPixelPuller( agent_name=f"smg-vllm-{socket.gethostname()}-{os.getpid()}", log_prefix="vLLM RDMA", From e0e4cf4eb8d33fcf2feb046d59f27a36ca1fb078 Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Wed, 15 Jul 2026 10:40:45 -0700 Subject: [PATCH 7/7] ci(mm-rdma): dedicated native wheel for the RDMA e2e 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 --- .github/workflows/e2e-gpu-job.yml | 7 +++- .github/workflows/pr-test-rust.yml | 63 +++++++++++++++++++++++++++++- bindings/python/Cargo.toml | 1 + scripts/ci_build_wheel_mm_rdma.sh | 48 +++++++++++++++++++++++ 4 files changed, 117 insertions(+), 2 deletions(-) create mode 100755 scripts/ci_build_wheel_mm_rdma.sh diff --git a/.github/workflows/e2e-gpu-job.yml b/.github/workflows/e2e-gpu-job.yml index b5abed138..917b61424 100644 --- a/.github/workflows/e2e-gpu-job.yml +++ b/.github/workflows/e2e-gpu-job.yml @@ -54,6 +54,11 @@ on: type: string default: "nixl" description: "KV transfer backend for vLLM PD workers: nixl or mooncake" + wheel_artifact: + required: false + type: string + default: "smg-wheel" + description: "Name of the wheel artifact to install (e.g. smg-wheel-mm-rdma)" jobs: run: @@ -92,7 +97,7 @@ jobs: - name: Download wheel artifact uses: actions/download-artifact@v8 with: - name: smg-wheel + name: ${{ inputs.wheel_artifact }} path: wheel/ - name: Download WASM test fixtures diff --git a/.github/workflows/pr-test-rust.yml b/.github/workflows/pr-test-rust.yml index 4e5d11dee..026fa9991 100644 --- a/.github/workflows/pr-test-rust.yml +++ b/.github/workflows/pr-test-rust.yml @@ -212,6 +212,37 @@ jobs: python3 -c "from smg.smg_rs import Router; print('Rust extension: OK')" python3 -m smg.launch_router --help > /dev/null && echo "Entry point: OK" + build-wheel-mm-rdma: + # Dedicated mm-rdma wheel for the RDMA e2e, built natively on the bare Ubuntu + # runner (GNU libstdc++). The shared wheel is cross-compiled with + # `maturin --zig`, which ships LLVM libc++ and can't link nixl-sys's + # GNU-libstdc++ C++ stub; building it separately here keeps that off the + # shared wheel so it can't break the other e2e lanes. Runs on the same runner + # image as the GPU e2e lanes, so the native wheel loads there unchanged. + needs: detect-changes + if: >- + always() + && !cancelled() + && (github.event_name != 'pull_request' + || (needs.detect-changes.result == 'success' + && (needs.detect-changes.outputs.common == 'true' + || needs.detect-changes.outputs.chat-completions == 'true'))) + runs-on: k8s-runner-cpu + permissions: + contents: read + steps: + - uses: actions/checkout@v7 + - name: Setup Rust + uses: ./.github/actions/setup-rust + - name: Build mm-rdma wheel + run: bash scripts/ci_build_wheel_mm_rdma.sh + - name: Upload mm-rdma wheel artifact + uses: actions/upload-artifact@v7 + with: + name: smg-wheel-mm-rdma + path: bindings/python/dist/*.whl + retention-days: 1 + python-unit-tests: needs: build-wheel runs-on: k8s-runner-cpu @@ -508,6 +539,9 @@ jobs: - engine: vllm timeout: 24 test_timeout: 18 + # The RDMA multimodal test needs the dedicated mm-rdma wheel + a NIXL + # runner; it runs in e2e-1gpu-mm-rdma, not this shared-wheel lane. + test_filter: "--ignore=e2e_test/chat_completions/test_multimodal_rdma.py" - engine: trtllm timeout: 32 test_timeout: 18 @@ -529,6 +563,31 @@ jobs: timeout: ${{ matrix.timeout }} test_timeout: ${{ matrix.test_timeout }} test_dirs: ${{ matrix.test_dirs || 'e2e_test/chat_completions' }} + test_filter: ${{ matrix.test_filter || '' }} + secrets: inherit + + e2e-1gpu-mm-rdma: + # Isolated RDMA multimodal e2e: installs the dedicated mm-rdma wheel and runs + # only test_multimodal_rdma.py against a vLLM worker on a NIXL-capable runner. + needs: [build-wheel, build-wheel-mm-rdma, detect-changes] + if: >- + always() + && !cancelled() + && needs.build-wheel.result == 'success' + && needs.build-wheel-mm-rdma.result == 'success' + && (github.event_name != 'pull_request' + || (needs.detect-changes.result == 'success' + && (needs.detect-changes.outputs.common == 'true' + || needs.detect-changes.outputs.chat-completions == 'true'))) + uses: ./.github/workflows/e2e-gpu-job.yml + with: + engine: vllm + gpu_tier: "1" + runner: 1-gpu-h100 + timeout: 24 + test_timeout: 18 + test_dirs: e2e_test/chat_completions/test_multimodal_rdma.py + wheel_artifact: smg-wheel-mm-rdma secrets: inherit e2e-1gpu-completions: @@ -1037,7 +1096,7 @@ jobs: path: benchmark_go_bindings/ finish: - needs: [pre-commit, python-lint, grpc-proto-build-check, build-wheel, python-unit-tests, unit-tests, benchmarks, e2e-1gpu-chat, e2e-1gpu-completions, e2e-1gpu-embeddings, e2e-1gpu-gateway, e2e-1gpu-responses, e2e-2gpu-pd, e2e-4gpu-chat, e2e-4gpu-gateway, e2e-4gpu-epd, e2e-vendor, go-unit-tests, go-bindings-e2e] + needs: [pre-commit, python-lint, grpc-proto-build-check, build-wheel, build-wheel-mm-rdma, python-unit-tests, unit-tests, benchmarks, e2e-1gpu-chat, e2e-1gpu-mm-rdma, e2e-1gpu-completions, e2e-1gpu-embeddings, e2e-1gpu-gateway, e2e-1gpu-responses, e2e-2gpu-pd, e2e-4gpu-chat, e2e-4gpu-gateway, e2e-4gpu-epd, e2e-vendor, go-unit-tests, go-bindings-e2e] if: always() runs-on: k8s-runner-cpu permissions: {} @@ -1048,10 +1107,12 @@ jobs: "${{ needs.python-lint.result }}" == "failure" || \ "${{ needs.grpc-proto-build-check.result }}" == "failure" || \ "${{ needs.build-wheel.result }}" == "failure" || \ + "${{ needs.build-wheel-mm-rdma.result }}" == "failure" || \ "${{ needs.python-unit-tests.result }}" == "failure" || \ "${{ needs.unit-tests.result }}" == "failure" || \ "${{ needs.benchmarks.result }}" == "failure" || \ "${{ needs.e2e-1gpu-chat.result }}" == "failure" || \ + "${{ needs.e2e-1gpu-mm-rdma.result }}" == "failure" || \ "${{ needs.e2e-1gpu-completions.result }}" == "failure" || \ "${{ needs.e2e-1gpu-embeddings.result }}" == "failure" || \ "${{ needs.e2e-1gpu-gateway.result }}" == "failure" || \ diff --git a/bindings/python/Cargo.toml b/bindings/python/Cargo.toml index 18fde4a7b..dd43f1b58 100644 --- a/bindings/python/Cargo.toml +++ b/bindings/python/Cargo.toml @@ -36,6 +36,7 @@ workspace = true default = ["pyo3/extension-module"] opencv-video = ["smg/opencv-video"] vendored-openssl = ["smg/vendored-openssl"] +mm-rdma = ["smg/mm-rdma"] [profile.ci] inherits = "release" diff --git a/scripts/ci_build_wheel_mm_rdma.sh b/scripts/ci_build_wheel_mm_rdma.sh new file mode 100755 index 000000000..b646024f4 --- /dev/null +++ b/scripts/ci_build_wheel_mm_rdma.sh @@ -0,0 +1,48 @@ +#!/bin/bash +# Build the mm-rdma smg wheel for the RDMA e2e, natively on the bare Ubuntu +# runner (GNU libstdc++). +# +# Why a separate build (not the shared ci_build_wheel.sh): the shared wheel is +# cross-compiled with `maturin --zig`, which ships LLVM libc++. nixl-sys (pulled +# by the mm-rdma feature) compiles a C++ stub that references GNU libstdc++ +# (std::cerr / _ZSt4cerr), so it must be linked against GNU libstdc++ or the +# wheel fails to import ("undefined symbol: _ZSt4cerr"). Building natively with +# the runner's GNU toolchain links the C++ stub cleanly, and the wheel runs on +# the GPU e2e runner unchanged (same runner image, same libstdc++.so.6). +# +# Assumes ./.github/actions/setup-rust already ran (Rust + build-essential/g++ + +# protoc). This wheel is CI-only (installed by the RDMA e2e job), so a plain +# linux tag is fine. +set -euxo pipefail + +# libclang for nixl-sys bindgen; setup-rust installs build-essential (g++ -> +# libstdc++) but not clang. +export DEBIAN_FRONTEND=noninteractive +if command -v sudo >/dev/null 2>&1; then + sudo apt-get update + sudo apt-get install -y --no-install-recommends libclang-dev +else + apt-get update + apt-get install -y --no-install-recommends libclang-dev +fi + +# setup-rust adds cargo to GITHUB_PATH; source for good measure (and local runs). +if [ -f "$HOME/.cargo/env" ]; then + source "$HOME/.cargo/env" +fi +export RUSTC_WRAPPER="${RUSTC_WRAPPER:-sccache}" + +python3 -m pip install --upgrade pip maturin + +echo "Building mm-rdma wheel (native, GNU libstdc++)..." +cd bindings/python +# abi3 wheel (pyo3 abi3-py38) -> one build works for all 3.8+. +# --manylinux off: CI-only wheel; skip auditwheel's manylinux policy and +# dynamically link libstdc++.so.6 (present on the GPU e2e runner). +maturin build \ + --profile ci \ + --features vendored-openssl,mm-rdma \ + --manylinux off \ + --out dist +echo "mm-rdma wheel: OK" +ls -lh dist/