anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Core] Cleanup Omni connector runtime and transfer backends (#5096)

Browse Source 
Signed-off-by: spencerr221 <liubingyu62@gmail.com>
This commit is contained in:
Bingyu (Spencer) Liu authored and GitHub committed 2026-07-22 09:44:15 +08:00
1 parent 2cf5f57b3e
commit dcf8b07d39
53 files changed
+349 -448
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/configuration/stage_configs.md
+1 -1
View File
@@ -65,7 +65,7 @@ connectors:
| Connector class | Use case | `extra` keys |
|-----------------|----------|--------------|
| `SharedMemoryConnector` | Same-host KV transfer between stages (default for bundled YAMLs). | `shm_threshold_bytes` (int, default `65536`). |
| `SharedMemoryConnector` | Same-host KV transfer between stages (default for bundled YAMLs). | None. All payloads use shared memory. |
| `MooncakeStoreConnector` | Cross-host KV transfer over TCP. Required for multi-node deployments. | `host`, `metadata_server`, `master`, `segment` (int bytes), `localbuf` (int bytes), `proto` (`"tcp"` / `"rdma"`). |
A stage references a connector by name in its `input_connectors` / `output_connectors`:
docs/design/feature/disaggregated_inference.md
-2
View File
@@ -66,8 +66,6 @@ runtime:
connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
```
Wire stages to connectors:
docs/design/feature/omni_connectors/shared_memory_connector.md
+4 -25
View File
@@ -8,9 +8,7 @@ auto-configured when no explicit connector is specified for an edge.
## How It Works
All payloads are serialized and stored in shared memory (`/dev/shm`); the SHM
segment name is returned in metadata. The configuration exposes a
`shm_threshold_bytes` field for a future inline-vs-SHM split, but the current
implementation always uses shared memory regardless of payload size.
segment name is returned in metadata.
## Configuration
@@ -19,8 +17,6 @@ runtime:
connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
```
## Notes
@@ -166,27 +162,10 @@ This path is mainly for older code paths and is not the preferred mode for the c
### 6. Key Implementation Characteristics
#### 6.1 Threshold Exists, but the Current Code Always Uses SHM
#### 6.1 All Payloads Use Shared Memory
The class keeps a `shm_threshold_bytes` field and still exposes metrics for inline writes. However, the current implementation uses:
```python
if True:
...
```
inside `put()`, which means the current code path always writes to shared memory.
So the design still suggests a future split between:
- small payloads inline
- large payloads in shared memory
but the current behavior is effectively:
- all payloads go through shared memory
This should be documented because it affects real runtime behavior.
`put()` writes every serialized payload to shared memory. The connector has no
inline-payload path or size threshold.
#### 6.2 Cleanup Is Currently Passive
docs/design/feature/ray_based_execution.md
+1 -1
View File
@@ -47,7 +47,7 @@ When running on Ray, the system automatically adapts its communication strategy:
* **Cross-Node**: Recommended to use `MooncakeTransferEngineConnector` (RDMA, fastest) or `MooncakeStoreConnector` (TCP fallback).
* **Same-Node**: Can still use `SharedMemoryConnector` for efficiency, or Ray's native object store (plasma).
* **SHM threshold default differs**: when `worker_backend="ray"`, the SharedMemoryConnector default threshold is set to `sys.maxsize`, which forces payloads to go inline (no SHM). Override `shm_threshold_bytes` in the connector config if you want SHM for Ray runs.
* **SHM payload routing**: `SharedMemoryConnector` always stores payloads in `/dev/shm`.
### 2.4 Internal Helpers
examples/offline_inference/covo_audio/end2end.py
-1
View File
@@ -64,7 +64,6 @@ def main(args):
stage_init_timeout=args.stage_init_timeout,
batch_timeout=args.batch_timeout,
init_timeout=args.init_timeout,
shm_threshold_bytes=args.shm_threshold_bytes,
)
# Stage 0: fused_thinker_talker
examples/offline_inference/mimo_audio/end2end.py
-1
View File
@@ -188,7 +188,6 @@ def main(args):
init_sleep_seconds=args.init_sleep_seconds,
batch_timeout=args.batch_timeout,
init_timeout=args.init_timeout,
shm_threshold_bytes=args.shm_threshold_bytes,
)
thinker_sampling_params = SamplingParams(
examples/offline_inference/step_audio2/end2end.py
-1
View File
@@ -258,7 +258,6 @@ def main(args):
"init_sleep_seconds": args.init_sleep_seconds,
"batch_timeout": args.batch_timeout,
"init_timeout": args.init_timeout,
"shm_threshold_bytes": args.shm_threshold_bytes,
"worker_backend": args.worker_backend,
"ray_address": args.ray_address,
"trust_remote_code": True,
examples/offline_inference/text_to_speech/ming_tts/runner.py
-1
View File
@@ -140,7 +140,6 @@ def build_engine_kwargs(args, stats_log_file: str | None) -> dict:
"stage_init_timeout": args.stage_init_timeout,
"init_timeout": args.init_timeout,
"batch_timeout": args.batch_timeout,
"shm_threshold_bytes": args.shm_threshold_bytes,
"worker_backend": args.worker_backend,
}
if stats_log_file is not None:
tests/config/test_omni_config.py
-2
View File
@@ -302,7 +302,6 @@ def test_from_pipeline_config_maps_orchestrator_cli_overrides():
"omni_dp_size_local": 2,
"omni_lb_policy": "round_robin",
"omni_heartbeat_timeout": 9.5,
"shm_threshold_bytes": 4096,
"batch_timeout": 3,
},
)
@@ -317,7 +316,6 @@ def test_from_pipeline_config_maps_orchestrator_cli_overrides():
assert orchestrator_config.omni_dp_size_local == 2
assert orchestrator_config.omni_lb_policy == "round_robin"
assert orchestrator_config.omni_heartbeat_timeout == 9.5
assert orchestrator_config.shm_threshold_bytes == 4096
assert orchestrator_config.batch_timeout == 3
tests/distributed/omni_connectors/test_adapter_and_flow.py
+2 -3
View File
@@ -142,11 +142,10 @@ def test_recv_no_connector():
def test_shm_connector_flow(mocker: MockerFixture):
"""
Verify the full flow: Send -> Adapter -> Connector -> Adapter -> Recv.
Using real SharedMemoryConnector (inline mode for simplicity).
Using real SharedMemoryConnector and key-addressed SHM metadata.
"""
# 1. Setup Connector
config = {"shm_threshold_bytes": 1024, "inline_small_payloads": True}
connector = SharedMemoryConnector(config)
connector = SharedMemoryConnector({})
connectors_map = {("0", "1"): connector}
# 2. Setup Data
tests/distributed/omni_connectors/test_basic_connectors.py
+47 -10
View File
@@ -60,10 +60,9 @@ def test_ndarray_serialization():
def test_create_shm_connector():
"""Test creating SharedMemoryConnector via Factory."""
spec = ConnectorSpec(name="SharedMemoryConnector", extra={"shm_threshold_bytes": 1024})
spec = ConnectorSpec(name="SharedMemoryConnector")
connector = OmniConnectorFactory.create_connector(spec)
assert isinstance(connector, SharedMemoryConnector)
assert connector.threshold == 1024
def test_create_unknown_connector():
@@ -75,21 +74,20 @@ def test_create_unknown_connector():
@pytest.fixture
def shm_connector():
config = {"shm_threshold_bytes": 100, "inline_small_payloads": True}
return SharedMemoryConnector(config)
connector = SharedMemoryConnector({})
yield connector
connector.close()
def test_put_get_inline(shm_connector):
"""Test inline transfer for small data."""
def test_put_get_small_payload_uses_shm(shm_connector):
"""Small payloads use key-addressed SHM; inline metadata is retired."""
data = {"small": "data"}
success, size, metadata = shm_connector.put("stage_0", "stage_1", "req_1", data)
assert success is True
assert "inline_bytes" in metadata
assert "shm" not in metadata
assert "shm" in metadata
assert "inline_bytes" not in metadata
assert "size" in metadata
assert shm_connector._metrics["inline_writes"] == 1
assert shm_connector._metrics["shm_writes"] == 0
# Retrieve
retrieved_data, ret_size = shm_connector.get("stage_0", "stage_1", "req_1", metadata)
@@ -195,3 +193,42 @@ def test_mooncake_connector_defaults_missing_host_to_detected_ip(monkeypatch: py
assert connector.get_connection_info()["host"] == "10.20.30.40"
finally:
connector.close()
def test_mooncake_get_counts_zmq_timeout_separately(mocker: MockerFixture):
import vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector as mooncake_module
connector = object.__new__(mooncake_module.MooncakeTransferEngineConnector)
connector._closed = False
connector._metrics = {"errors": 0, "timeouts": 0}
connector._resolve_metadata = mocker.Mock(
return_value={"source_host": "127.0.0.1", "source_port": 1234, "data_size": 8}
)
recv_buffer = mocker.Mock()
connector._alloc_recv_buffer = mocker.Mock(return_value=(recv_buffer, 1000))
connector._request_transfer = mocker.Mock(side_effect=mooncake_module.zmq.Again())
assert connector.get("s0", "s1", "req") is None
assert connector._metrics == {"errors": 0, "timeouts": 1}
recv_buffer.release.assert_called_once()
connector._closed = True
def test_mooncake_get_counts_cuda_sync_failure_as_error(mocker: MockerFixture):
import vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector as mooncake_module
connector = object.__new__(mooncake_module.MooncakeTransferEngineConnector)
connector._closed = False
connector._metrics = {"errors": 0, "timeouts": 0}
connector._resolve_metadata = mocker.Mock(
return_value={"source_host": "127.0.0.1", "source_port": 1234, "data_size": 8}
)
recv_buffer = mocker.Mock()
connector._alloc_recv_buffer = mocker.Mock(return_value=(recv_buffer, 1000))
connector._request_transfer = mocker.Mock(return_value=mooncake_module.TRANS_DONE)
connector._sync_receive_buffer_if_needed = mocker.Mock(side_effect=RuntimeError("CUDA sync failed"))
assert connector.get("s0", "s1", "req") is None
assert connector._metrics == {"errors": 1, "timeouts": 0}
recv_buffer.release.assert_called_once()
connector._closed = True
tests/distributed/omni_connectors/test_omni_connector_configs.py
-1
View File
@@ -48,7 +48,6 @@ def test_load_qwen_yaml_configs(yaml_file):
print(f"Testing config load: {yaml_file.name}")
try:
# Attempt to load the config
# default_shm_threshold doesn't matter much for loading correctness, using default
config = load_omni_transfer_config(yaml_file)
assert config is not None, "Config should not be None"
tests/distributed/omni_connectors/test_shm_connector.py
+28 -1
View File
@@ -2,7 +2,10 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for SharedMemoryConnector focusing on TP / CFG / metadata fallback."""
import os
import pytest
import torch
from vllm_omni.distributed.omni_connectors.connectors.shm_connector import (
SharedMemoryConnector,
@@ -13,7 +16,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@pytest.fixture()
def connector():
c = SharedMemoryConnector({"shm_threshold_bytes": 0})
c = SharedMemoryConnector({})
yield c
c.close()
@@ -36,6 +39,30 @@ class TestKeyBasedReadWrite:
assert obj == data
assert rsize == size
assert "test_key_1" not in connector._pending_keys
assert connector._metrics["gets"] == 1
def test_tensor_payload_removes_lock_file(self, connector):
key = "tensor_payload"
payload = torch.ones(2, 2)
ok, _, metadata = connector.put("s0", "s1", key, payload)
assert ok
result = connector.get("s0", "s1", key, metadata=metadata)
assert result is not None
assert torch.equal(result[0], payload)
assert not os.path.exists(f"/dev/shm/shm_{key}_lockfile.lock")
def test_falsey_payload_removes_lock_file(self, connector):
key = "falsey_payload"
ok, _, metadata = connector.put("s0", "s1", key, 0)
assert ok
result = connector.get("s0", "s1", key, metadata=metadata)
assert result is not None
assert result[0] == 0
assert not os.path.exists(f"/dev/shm/shm_{key}_lockfile.lock")
def test_get_nonexistent_key_returns_none(self, connector):
result = connector.get("s0", "s1", "no_such_key_xyz", metadata=None)
tests/e2e/accuracy/test_hunyuan_image3.py
-1
View File
@@ -141,7 +141,6 @@ _DEPLOY_CONFIG = {
"connectors": {
"shared_memory_connector": {
"name": "SharedMemoryConnector",
"extra": {"shm_threshold_bytes": 65536},
},
},
"stages": [
tests/engine/test_async_omni_engine_abort.py
-3
View File
@@ -1,6 +1,5 @@
import asyncio
import os
import sys
from contextlib import ExitStack
import pytest
@@ -67,11 +66,9 @@ async def generate(
@pytest.mark.asyncio
async def test_abort():
with ExitStack() as after:
# Avoid SHM IPC in tests to prevent /dev/shm exhaustion and SIGBUS.
engine = AsyncOmni(
model=model,
stage_configs_path=stage_config,
shm_threshold_bytes=sys.maxsize,
)
after.callback(engine.shutdown)
tests/entrypoints/test_pd_disaggregation.py
+1 -3
View File
@@ -121,9 +121,7 @@ class _FakeStage:
self._in_q = in_q
self._out_q = out_q
def init_stage_worker(
self, model: str, *, is_async=False, shm_threshold_bytes=65536, ctx=None, batch_timeout=10, **kwargs
):
def init_stage_worker(self, model: str, *, is_async=False, ctx=None, batch_timeout=10, **kwargs):
self._proc = _ns(
start=lambda: None,
join=lambda timeout=None: None,
tests/helpers/runtime.py
-2
View File
@@ -2333,7 +2333,6 @@ class OmniRunner:
# production default in AsyncOmniEngine remains 600s; this only
# affects the test runner wrapper.
init_timeout: int = 1800,
shm_threshold_bytes: int = 65536,
log_stats: bool = False,
stage_configs_path: str | None = None,
**kwargs,
@@ -2353,7 +2352,6 @@ class OmniRunner:
stage_init_timeout=stage_init_timeout,
batch_timeout=batch_timeout,
init_timeout=init_timeout,
shm_threshold_bytes=shm_threshold_bytes,
stage_configs_path=stage_configs_path,
**kwargs,
)
tests/worker/test_omni_connector_mixin.py
+2 -27
View File
@@ -110,7 +110,6 @@ class TestMixinAsyncChunkSendRecv(unittest.TestCase):
sender = MixinHost()
sender.init_omni_connectors(
vllm_config=None,
model_config=_make_model_config(stage_id=0, async_chunk=True),
)
sender._omni_connector = connector
@@ -147,7 +146,6 @@ class TestMixinAsyncChunkSendRecv(unittest.TestCase):
sender = MixinHost()
sender.init_omni_connectors(
vllm_config=None,
model_config=_make_model_config(stage_id=0, async_chunk=True),
)
sender._omni_connector = connector
@@ -180,7 +178,6 @@ class TestMixinKVCacheTransfer(unittest.TestCase):
host = MixinHost()
host.init_omni_connectors(
vllm_config=None,
model_config=_make_model_config(),
kv_transfer_manager=mock_kvm,
)
@@ -202,7 +199,6 @@ class TestMixinKVCacheTransfer(unittest.TestCase):
host = MixinHost()
host.init_omni_connectors(
vllm_config=None,
model_config=_make_model_config(),
kv_transfer_manager=mock_kvm,
)
@@ -219,7 +215,6 @@ class TestMixinKVCacheTransfer(unittest.TestCase):
host = MixinHost()
host.init_omni_connectors(
vllm_config=None,
model_config=_make_model_config(),
kv_transfer_manager=mock_kvm,
)
@@ -257,7 +252,6 @@ class TestMixinKVCacheTransfer(unittest.TestCase):
host = MixinHost()
host.init_omni_connectors(
vllm_config=None,
model_config=_make_model_config(),
kv_transfer_manager=mock_kvm,
)
@@ -281,7 +275,6 @@ class TestOmniConnectorOutput(unittest.TestCase):
def test_output_aggregation(self):
host = MixinHost()
host.init_omni_connectors(
vllm_config=None,
model_config=_make_model_config(),
)
@@ -310,7 +303,6 @@ class TestMixinNoConnector(unittest.TestCase):
def test_no_connector(self):
host = MixinHost()
host.init_omni_connectors(
vllm_config=None,
model_config=_make_model_config(),
)
self.assertIsNone(host._omni_connector)
@@ -336,7 +328,6 @@ class TestFinishedLoadReqsDrain(unittest.TestCase):
def test_finished_load_reqs_flow_to_chunk_ready(self):
host = MixinHost()
host.init_omni_connectors(
vllm_config=None,
model_config=_make_model_config(),
)
@@ -389,7 +380,6 @@ class TestFullPayloadSendWithCustomFunc(unittest.TestCase):
host = MixinHost()
host.init_omni_connectors(
vllm_config=None,
model_config=_make_model_config(),
)
host._omni_connector = MockConnector(stage_id=0)
@@ -424,7 +414,6 @@ class TestFullPayloadSendWithCustomFunc(unittest.TestCase):
host = MixinHost()
host.init_omni_connectors(
vllm_config=None,
model_config=_make_model_config(),
)
host._omni_connector = MockConnector(stage_id=0)
@@ -453,7 +442,6 @@ class TestKVSentReqIdsAccumulation(unittest.TestCase):
host = MixinHost()
host.init_omni_connectors(
vllm_config=None,
model_config=_make_model_config(),
kv_transfer_manager=mock_kvm,
)
@@ -488,7 +476,6 @@ class TestChunkStreamCompletedGuard(unittest.TestCase):
def _make_host(self, stage_id: int = 1) -> MixinHost:
host = MixinHost()
host.init_omni_connectors(
vllm_config=None,
model_config=_make_model_config(stage_id=stage_id, async_chunk=True),
)
host._omni_connector = MockConnector(stage_id=stage_id)
@@ -587,7 +574,6 @@ class TestCleanupFinishedRequest(unittest.TestCase):
def _make_host(self, stage_id: int = 1) -> MixinHost:
host = MixinHost()
host.init_omni_connectors(
vllm_config=None,
model_config=_make_model_config(stage_id=stage_id, async_chunk=True),
)
host._omni_connector = MockConnector(stage_id=stage_id)
@@ -771,7 +757,6 @@ class TestSendChunkCachesMapping(unittest.TestCase):
"""send_chunk should cache the internal→external mapping."""
host = MixinHost()
host.init_omni_connectors(
vllm_config=None,
model_config=_make_model_config(stage_id=0, async_chunk=True),
)
host._omni_connector = MockConnector(stage_id=0)
@@ -802,7 +787,6 @@ class TestLocalPayloadCacheLifecycle(unittest.TestCase):
def _make_host(self) -> MixinHost:
host = MixinHost()
host.init_omni_connectors(
vllm_config=None,
model_config=_make_model_config(stage_id=0),
)
host._omni_connector = MockConnector(stage_id=0)
@@ -901,7 +885,6 @@ class TestTPAsyncChunkFanout(unittest.TestCase):
def _make_host(self, rank: int) -> MixinHost:
host = MixinHost()
host.init_omni_connectors(
vllm_config=None,
model_config=_make_model_config(stage_id=2, async_chunk=True, worker_type="gen"),
)
host._omni_connector = MagicMock()
@@ -984,7 +967,6 @@ class TestKVTransferLifecycle(unittest.TestCase):
def _make_host(self) -> MixinHost:
host = MixinHost()
host.init_omni_connectors(
vllm_config=None,
model_config=_make_model_config(stage_id=0),
)
return host
@@ -1039,7 +1021,6 @@ class TestAsyncPayloadLifecycle(unittest.TestCase):
def test_send_side_request_payload_not_cleared_before_payload_is_consumable(self):
host = MixinHost()
host.init_omni_connectors(
vllm_config=None,
model_config=_make_model_config(stage_id=1, async_chunk=True, worker_type="ar"),
)
host._request_ids_mapping["r1"] = "r1"
@@ -1063,7 +1044,6 @@ class TestAsyncPayloadLifecycle(unittest.TestCase):
def test_payload_consumable_ignores_token_horizon_only_updates(self):
host = MixinHost()
host.init_omni_connectors(
vllm_config=None,
model_config=_make_model_config(stage_id=1, async_chunk=True, worker_type="ar"),
)
payload = {
@@ -1084,7 +1064,6 @@ class TestAsyncPayloadLifecycle(unittest.TestCase):
def test_payload_consumable_accepts_decode_embeddings(self):
host = MixinHost()
host.init_omni_connectors(
vllm_config=None,
model_config=_make_model_config(stage_id=1, async_chunk=True, worker_type="ar"),
)
payload = {
@@ -1098,7 +1077,6 @@ class TestAsyncPayloadLifecycle(unittest.TestCase):
def test_ar_metadata_only_followup_chunk_does_not_rewake_request(self):
host = MixinHost()
host.init_omni_connectors(
vllm_config=None,
model_config=_make_model_config(stage_id=1, async_chunk=True, worker_type="ar"),
)
host._omni_connector = MagicMock()
@@ -1138,7 +1116,6 @@ class TestAsyncPayloadLifecycle(unittest.TestCase):
def test_non_ar_recv_does_not_overwrite_unconsumed_staged_chunk(self):
host = MixinHost()
host.init_omni_connectors(
vllm_config=None,
model_config=_make_model_config(stage_id=2, async_chunk=True, worker_type="gen"),
)
host._omni_connector = MagicMock()
@@ -1164,7 +1141,6 @@ class TestAsyncPayloadLifecycle(unittest.TestCase):
def test_non_ar_recv_waits_for_scheduler_handoff_before_fetching_next_chunk(self):
host = MixinHost()
host.init_omni_connectors(
vllm_config=None,
model_config=_make_model_config(stage_id=2, async_chunk=True, worker_type="gen"),
)
host._omni_connector = MagicMock()
@@ -1208,7 +1184,7 @@ class TestAsyncPayloadLifecycle(unittest.TestCase):
class TestRankAwareKVRouting(unittest.TestCase):
def _make_host(self, *, from_tp: int, to_tp: int, local_rank: int) -> MixinHost:
host = MixinHost()
host.init_omni_connectors(vllm_config=None, model_config=_make_model_config(stage_id=1))
host.init_omni_connectors(model_config=_make_model_config(stage_id=1))
host._from_tp = from_tp
host._to_tp = to_tp
host._local_rank = local_rank
@@ -1287,7 +1263,7 @@ class TestConnectorConfigValidation(unittest.TestCase):
model_config.stage_connector_config = {"name": " "}
with self.assertRaisesRegex(RuntimeError, "missing connector name"):
host.init_omni_connectors(vllm_config=None, model_config=model_config)
host.init_omni_connectors(model_config=model_config)
class _FailingConnector:
@@ -1319,7 +1295,6 @@ class TestSendRetry(unittest.TestCase):
def _make_sender(self, connector):
sender = MixinHost()
sender.init_omni_connectors(
vllm_config=None,
model_config=_make_model_config(stage_id=0, async_chunk=True),
)
sender._omni_connector = connector
vllm_omni/config/omni_config.py
-1
View File
@@ -821,7 +821,6 @@ class VllmOmniOrchestratorConfig:
omni_dp_size_local: int = Field(default=1, ge=1)
omni_lb_policy: str = "random"
omni_heartbeat_timeout: float = Field(default=30.0, gt=0.0)
shm_threshold_bytes: int = Field(default=65536, ge=0)
batch_timeout: int = Field(default=10, ge=0)
vllm_omni/deploy/aura_omni.yaml
-1
View File
@@ -12,7 +12,6 @@ connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
codec_streaming: true
connector_get_sleep_s: 0.01
connector_get_max_wait_first_chunk: 3000
vllm_omni/deploy/bagel.yaml
-2
View File
@@ -40,8 +40,6 @@ stages:
connectors:
shared_memory_connector:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
rdma_connector:
name: MooncakeTransferEngineConnector
vllm_omni/deploy/dynin_omni_multiconnector.yaml
-2
View File
@@ -31,8 +31,6 @@ connectors:
# Alternative SHM connector with an explicit threshold.
shared_memory_connector:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536 # 64KB threshold
stages:
- stage_id: 0
vllm_omni/deploy/fish_qwen3_omni.yaml
-1
View File
@@ -9,7 +9,6 @@ connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
codec_streaming: true
connector_get_sleep_s: 0.01
connector_get_max_wait_first_chunk: 3000
vllm_omni/deploy/fish_qwen3_omni_2gpu.yaml
-1
View File
@@ -20,7 +20,6 @@ connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
codec_streaming: true
# Tightened from 10ms to 1ms: chunks become available sooner under
# cross-GPU pipelining, so connector idle waiting dominates if left high.
vllm_omni/deploy/fish_qwen3_omni_high_concurrency_dual_gpu.yaml
-1
View File
@@ -10,7 +10,6 @@ connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
codec_streaming: true
connector_get_sleep_s: 0.001
connector_get_max_wait_first_chunk: 3000
vllm_omni/deploy/fish_qwen3_omni_high_concurrency_single_gpu.yaml
-1
View File
@@ -10,7 +10,6 @@ connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
codec_streaming: true
connector_get_sleep_s: 0.001
connector_get_max_wait_first_chunk: 3000
vllm_omni/deploy/higgs_audio_v2.yaml
-1
View File
@@ -17,7 +17,6 @@ connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
codec_streaming: true
connector_get_sleep_s: 0.01
connector_get_max_wait_first_chunk: 5000
vllm_omni/deploy/higgs_multimodal_qwen3.yaml
-1
View File
@@ -16,7 +16,6 @@ connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
# Stream PCM bytes back to the client as Stage 1 emits each chunk.
codec_streaming: true
# Streaming chunk shape for talker2code2wav_async_chunk:
vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml
-1
View File
@@ -15,7 +15,6 @@ connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
# Stream PCM bytes back to the client as Stage 1 emits each chunk.
codec_streaming: true
# Streaming chunk shape for talker2code2wav_async_chunk:
vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml
-1
View File
@@ -21,7 +21,6 @@ connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
# Stream PCM bytes back to the client as Stage 1 emits each chunk.
codec_streaming: true
# Streaming chunk shape for talker2code2wav_async_chunk:
vllm_omni/deploy/hunyuan_image_3_moe.yaml
-2
View File
@@ -27,8 +27,6 @@ connectors:
memory_pool_device: "npu"
shared_memory_connector:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
stages:
- stage_id: 0
vllm_omni/deploy/indextts2.yaml
-1
View File
@@ -4,7 +4,6 @@ connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
codec_streaming: false
connector_get_sleep_s: 0.01
connector_get_max_wait_first_chunk: 5000
vllm_omni/deploy/indextts2_low_latency.yaml
-1
View File
@@ -17,7 +17,6 @@ connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
codec_streaming: false
connector_get_sleep_s: 0.01
connector_get_max_wait_first_chunk: 5000
vllm_omni/deploy/mimo_audio.yaml
-1
View File
@@ -13,7 +13,6 @@ connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 365536
codec_streaming: true
connector_get_sleep_s: 0.001
connector_get_max_wait_first_chunk: 3000
vllm_omni/deploy/ming_flash_omni_image.yaml
-2
View File
@@ -50,8 +50,6 @@ stages:
connectors:
shared_memory_connector:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536 # 64KB threshold
edges:
- from: 0
vllm_omni/deploy/moss_sound_effect.yaml
-1
View File
@@ -10,7 +10,6 @@ connectors:
shm:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
codec_streaming: true
connector_get_sleep_s: 0.01
connector_get_max_wait_first_chunk: 3000
vllm_omni/deploy/moss_tts.yaml
-1
View File
@@ -8,7 +8,6 @@ connectors:
shm:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
codec_streaming: true
connector_get_sleep_s: 0.01
connector_get_max_wait_first_chunk: 3000
vllm_omni/deploy/moss_tts_80gb.yaml
-1
View File
@@ -11,7 +11,6 @@ connectors:
shm:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
codec_streaming: true
connector_get_sleep_s: 0.01
connector_get_max_wait_first_chunk: 3000
vllm_omni/deploy/moss_tts_local.yaml
-1
View File
@@ -9,7 +9,6 @@ connectors:
shm:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
codec_streaming: true
connector_get_sleep_s: 0.005
connector_get_max_wait_first_chunk: 1000
vllm_omni/deploy/moss_tts_realtime.yaml
-1
View File
@@ -11,7 +11,6 @@ connectors:
shm:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
codec_streaming: true
connector_get_sleep_s: 0.005
connector_get_max_wait_first_chunk: 1000
vllm_omni/deploy/moss_ttsd.yaml
-1
View File
@@ -8,7 +8,6 @@ connectors:
shm:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
codec_streaming: true
connector_get_sleep_s: 0.01
connector_get_max_wait_first_chunk: 3000
vllm_omni/deploy/moss_voice_generator.yaml
-1
View File
@@ -10,7 +10,6 @@ connectors:
shm:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
codec_streaming: true
connector_get_sleep_s: 0.01
connector_get_max_wait_first_chunk: 3000
vllm_omni/deploy/qwen3_tts.yaml
-1
View File
@@ -17,7 +17,6 @@ connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
codec_streaming: true
connector_get_sleep_s: 0.01
connector_get_max_wait_first_chunk: 3000
vllm_omni/deploy/qwen3_tts_high_concurrency.yaml
-1
View File
@@ -15,7 +15,6 @@ connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
codec_streaming: true
connector_get_sleep_s: 0.01
connector_get_max_wait_first_chunk: 3000
vllm_omni/deploy/voxtral_tts.yaml
-1
View File
@@ -10,7 +10,6 @@ connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
codec_streaming: true
connector_get_sleep_s: 0.01
connector_get_max_wait_first_chunk: 3000
vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py
+229 -170
View File
@@ -30,6 +30,9 @@ except ImportError:
# Stale buffer TTL: buffers older than this are automatically reclaimed
# to prevent memory leaks when receiver crashes or gives up.
_BUFFER_TTL_SECONDS = 300 # 5 minutes
_BASE_TRANSFER_TIMEOUT_MS = 30_000
_TRANSFER_TIMEOUT_STEP_BYTES = 100 * 1024 * 1024
_TRANSFER_TIMEOUT_PER_STEP_MS = 5_000
# ZMQ Message constants
TRANS_DONE = b"trans_done"
@@ -119,6 +122,7 @@ class MooncakeTransferEngineConnector(OmniConnectorBase):
self._worker_local = threading.local()
self._last_ttl_check: float = _time_mod.monotonic()
self._sender_endpoints: dict[int, tuple[str, int]] = {}
self._in_flight_buffers: set[str] = set()
self._metrics = {
"puts": 0,
@@ -155,6 +159,7 @@ class MooncakeTransferEngineConnector(OmniConnectorBase):
# --- Memory Pool Configuration ---
self.pool_size = config.get("memory_pool_size", 1024**3) # Default 1GB
self.pool_device = config.get("memory_pool_device", "cpu")
self.buffer_ttl_seconds = float(config.get("buffer_ttl_seconds", _BUFFER_TTL_SECONDS))
# --- Sender Configuration (for receiver to query without metadata) ---
# When receiver doesn't have metadata, it uses these to connect to sender
@@ -175,7 +180,12 @@ class MooncakeTransferEngineConnector(OmniConnectorBase):
self.engine_id = str(uuid.uuid4())
# --- Mooncake Engine Init ---
self._init_engine()
self._init_pool()
self._init_listener(config)
def _init_engine(self) -> None:
"""Initialize the Mooncake engine and discover the local RPC port."""
self.engine = TransferEngine()
# Note: For P2P handshake mode, local_hostname should be just the IP address.
# Mooncake will auto-assign an RPC port, retrievable via get_rpc_port().
@@ -186,7 +196,8 @@ class MooncakeTransferEngineConnector(OmniConnectorBase):
self.rpc_port = self.engine.get_rpc_port()
logger.info(f"MooncakeTransferEngineConnector initialized at {self.host}:{self.rpc_port}")
# --- Pool Allocation & Registration ---
def _init_pool(self) -> None:
"""Allocate and register the memory pool used by TransferEngine."""
logger.info(f"Allocating RDMA Memory Pool: {self.pool_size / 1024**2:.2f} MB on {self.pool_device}")
try:
if self.pool_device == "cpu":
@@ -207,9 +218,8 @@ class MooncakeTransferEngineConnector(OmniConnectorBase):
self.allocator = BufferAllocator(self.pool_size, alignment=4096) # 4KB alignment for safety
# --- State Management & Background Threads ---
# Most fields already initialized at the top of __init__ for teardown
# safety. Only create the real ZMQ context and thread pool here.
def _init_listener(self, config: dict[str, Any]) -> None:
"""Initialize control-plane sockets and start the sender listener."""
self.zmq_ctx = zmq.Context()
self._sender_executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix="mooncake-sender")
@@ -642,44 +652,18 @@ class MooncakeTransferEngineConnector(OmniConnectorBase):
raise RuntimeError("Cannot get data: MooncakeTransferEngineConnector is closed")
get_key = self._make_key(get_key, from_stage, to_stage)
_t0 = _time_mod.perf_counter()
metadata = self._resolve_metadata(get_key, metadata)
if not metadata:
# Path 3: no metadata at all — query default sender
if not self.sender_host or not self.sender_zmq_port or str(self.sender_host).lower() == "auto":
raise RuntimeError(
f"get(metadata=None) requires sender info to be resolved, "
f"but sender_host={self.sender_host!r}, sender_zmq_port={self.sender_zmq_port!r}. "
f"Call update_sender_info(host, port) before using get() without metadata."
)
metadata = self._query_metadata_from_sender(get_key)
if not metadata:
return None
elif "data_size" not in metadata:
# Path 2: partial metadata (host/port only) — query that sender
partial_host = metadata.get("source_host")
partial_port = metadata.get("source_port")
if not partial_host or not partial_port:
logger.warning(
"get(%s): partial metadata missing source_host/source_port, cannot resolve data_size. metadata=%s",
get_key,
metadata,
)
return None
queried = self._query_metadata_at(get_key, str(partial_host), int(partial_port))
if not queried:
return None
metadata = queried
return None
_t1 = _time_mod.perf_counter()
_query_ms = (_t1 - _t0) * 1000
query_ms = (_t1 - _t0) * 1000
src_host = metadata.get("source_host")
src_port = metadata.get("source_port")
data_size = metadata.get("data_size", 0)
is_fast_path = metadata.get("is_fast_path", False)
data_size = int(metadata.get("data_size", 0) or 0)
is_fast_path = bool(metadata.get("is_fast_path", False))
if not src_host or not src_port or str(src_host).lower() == "auto":
logger.error(
f"Invalid metadata for {get_key}: source_host={src_host!r}, "
@@ -691,27 +675,116 @@ class MooncakeTransferEngineConnector(OmniConnectorBase):
logger.warning(f"Skipping get for {get_key}: data_size is 0 (metadata={metadata})")
return None
# 1. Allocate Destination Buffer from Pool
# For non-fast-path (serialized objects), allocate a few extra bytes
# as padding. Mooncake's TCP/memcpy transport may prepend a leading
# protocol byte (0x01) before the payload. The extra padding allows
# the receiver to detect and skip this byte without truncating the
# actual payload. The padding is harmless for fast-path (ManagedBuffer)
# because the caller reads exact-size slices.
_MOONCAKE_TCP_PADDING = 4 if not is_fast_path else 0
alloc_size = data_size + _MOONCAKE_TCP_PADDING
try:
offset = self.allocator.alloc(alloc_size)
recv_buffer = ManagedBuffer(self.allocator, offset, alloc_size, self.pool)
dst_ptr = self.base_ptr + offset
recv_buffer, dst_ptr = self._alloc_recv_buffer(data_size, is_fast_path)
except MemoryError:
logger.error(f"Failed to allocate {alloc_size} bytes in receive pool")
return None
_t2 = _time_mod.perf_counter()
_alloc_ms = (_t2 - _t1) * 1000
alloc_ms = (_t2 - _t1) * 1000
try:
resp = self._request_transfer(
get_key=get_key,
source_host=str(src_host),
source_port=int(src_port),
dst_ptr=dst_ptr,
data_size=data_size,
)
except zmq.Again as e:
self._metrics["timeouts"] += 1
logger.error(f"RDMA Get timed out: {e}", exc_info=True)
recv_buffer.release()
return None
except Exception as e:
self._metrics["errors"] += 1
logger.error(f"RDMA transfer request failed: {e}", exc_info=True)
recv_buffer.release()
return None
# 2. Prepare Handshake
_t3 = _time_mod.perf_counter()
rdma_ms = (_t3 - _t2) * 1000
if resp != TRANS_DONE:
self._metrics["errors"] += 1
logger.error(f"RDMA Get failed: received {resp} instead of TRANS_DONE")
recv_buffer.release()
return None
try:
sync_ms = self._sync_receive_buffer_if_needed()
except Exception as e:
self._metrics["errors"] += 1
logger.error(f"Failed to synchronize Mooncake receive buffer: {e}", exc_info=True)
recv_buffer.release()
return None
try:
return self._decode_result(
get_key=get_key,
recv_buffer=recv_buffer,
data_size=data_size,
is_fast_path=is_fast_path,
query_ms=query_ms,
alloc_ms=alloc_ms,
rdma_ms=rdma_ms,
sync_ms=sync_ms,
start_time=_t0,
)
except Exception as e:
self._metrics["errors"] += 1
logger.error(f"RDMA result decode failed for {get_key}: {e}", exc_info=True)
return None
def _resolve_metadata(self, get_key: str, metadata: dict[str, Any] | None) -> dict[str, Any] | None:
"""Resolve full transfer metadata for a receive request."""
if not metadata:
# Path 3: no metadata at all — query default sender
if not self.sender_host or not self.sender_zmq_port or str(self.sender_host).lower() == "auto":
raise RuntimeError(
f"get(metadata=None) requires sender info to be resolved, "
f"but sender_host={self.sender_host!r}, sender_zmq_port={self.sender_zmq_port!r}. "
f"Call update_sender_info(host, port) before using get() without metadata."
)
return self._query_metadata_from_sender(get_key)
if "data_size" in metadata:
# Path 1: metadata is complete
return metadata
# Path 2: partial metadata (host/port only) — query that sender
partial_host = metadata.get("source_host")
partial_port = metadata.get("source_port")
if not partial_host or not partial_port:
logger.warning(
"get(%s): partial metadata missing source_host/source_port, cannot resolve data_size. metadata=%s",
get_key,
metadata,
)
return None
return self._query_metadata_at(get_key, str(partial_host), int(partial_port))
def _alloc_recv_buffer(self, data_size: int, is_fast_path: bool) -> tuple[ManagedBuffer, int]:
"""Allocate the receive-side buffer and return ``(buffer, dst_ptr)``."""
# For non-fast-path serialized objects, allocate a few extra bytes as
# padding. Mooncake's TCP/memcpy transport may prepend a protocol byte.
padding = 4 if not is_fast_path else 0
alloc_size = data_size + padding
try:
offset = self.allocator.alloc(alloc_size)
except MemoryError:
logger.error(f"Failed to allocate {alloc_size} bytes in receive pool")
raise
recv_buffer = ManagedBuffer(self.allocator, offset, alloc_size, self.pool)
return recv_buffer, self.base_ptr + offset
def _request_transfer(
self,
get_key: str,
source_host: str,
source_port: int,
dst_ptr: int,
data_size: int,
) -> bytes:
"""Ask the sender to write this key into the local receive buffer."""
agent_meta = MooncakeAgentMetadata(
remote_hostname=self.host,
remote_port=self.rpc_port,
@@ -720,109 +793,82 @@ class MooncakeTransferEngineConnector(OmniConnectorBase):
lengths=[data_size],
)
# 3. ZMQ Transaction (uses cached socket to avoid TCP reconnection)
# Timeout scales with data size: base 30s + 5s per 100MB.
# Large RDMA transfers (esp. loopback or cross-NIC) can take tens of
# seconds; if we time out too early the receiver closes and deregisters
# its memory, causing the sender's in-flight write to fail with ret=-1.
_base_timeout_ms = 30000
_size_timeout_ms = max(0, (data_size // (100 * 1024 * 1024))) * 5000
_total_timeout_ms = _base_timeout_ms + _size_timeout_ms
zmq_addr = f"tcp://{src_host}:{src_port}"
req_socket = self._get_req_socket(zmq_addr, timeout_ms=_total_timeout_ms)
size_timeout_ms = (data_size // _TRANSFER_TIMEOUT_STEP_BYTES) * _TRANSFER_TIMEOUT_PER_STEP_MS
total_timeout_ms = _BASE_TRANSFER_TIMEOUT_MS + size_timeout_ms
zmq_addr = f"tcp://{source_host}:{source_port}"
req_socket = self._get_req_socket(zmq_addr, timeout_ms=total_timeout_ms)
try:
req_socket.send(msgspec.msgpack.encode(agent_meta))
resp = req_socket.recv()
_t3 = _time_mod.perf_counter()
_rdma_ms = (_t3 - _t2) * 1000
if resp == TRANS_DONE:
# Success
# Ensure data is visible on GPU
# Note: RDMA write visibility on GPU usually requires some form of fence/sync.
# torch.accelerator.synchronize() is a heavy hammer but safe.
# Ideally Mooncake engine provides a way to poll for completion that guarantees visibility.
# TODO(wzliu): Replace synchronize with cuda event in the future for better performance.
if self.pool.is_cuda:
with torch.cuda.device(self.pool.device):
torch.cuda.current_stream().synchronize()
_t4 = _time_mod.perf_counter()
_sync_ms = (_t4 - _t3) * 1000
if is_fast_path:
# Return ManagedBuffer directly for ALL fast_path data
# (both bytes and tensor payloads). Caller is responsible
# for consuming the data and calling release() afterwards.
# This avoids the expensive to_bytes() copy (~54ms for 115MB).
#
# Usage patterns for callers:
# - Zero-copy read: mv = memoryview(buf.tensor.numpy())
# - Copy to GPU: buf.tensor.to(device); buf.release()
# - Fallback: data = buf.to_bytes(); buf.release()
_t5 = _time_mod.perf_counter()
_total_ms = (_t5 - _t0) * 1000
_mbps = (data_size / 1024 / 1024) / (_total_ms / 1000) if _total_ms > 0 else 0
logger.info(
f"[RDMA GET] {get_key}: query={_query_ms:.1f}ms, alloc={_alloc_ms:.1f}ms, "
f"rdma={_rdma_ms:.1f}ms, sync={_sync_ms:.1f}ms, "
f"total={_total_ms:.1f}ms, {_mbps:.1f} MB/s (fast_path, zero-copy)"
)
self._metrics["gets"] += 1
self._metrics["bytes_transferred"] += data_size
return recv_buffer, data_size
else:
# If it was a serialized object or generic bytes, we assume standard Omni behavior:
# Deserialize and return object. This requires a copy (to_bytes).
# We MUST release the buffer after deserialization.
try:
_t_copy_start = _time_mod.perf_counter()
raw_bytes = recv_buffer.to_bytes()
_t_copy_end = _time_mod.perf_counter()
_copy_ms = (_t_copy_end - _t_copy_start) * 1000
# Try deserializing the first data_size bytes. Mooncake's
# TCP/memcpy transport may prepend a leading protocol byte
# (0x01) before the payload — handle this by retrying from
# offset 1 if the first attempt raises DecodeError.
_t_deser_start = _time_mod.perf_counter()
payload = raw_bytes[:data_size]
try:
val = OmniSerializer.deserialize(payload)
except msgspec.DecodeError:
# Likely a leading Mooncake TCP protocol byte — retry
# from offset 1, using the extra padding bytes if needed.
payload = raw_bytes[1 : data_size + 1]
val = OmniSerializer.deserialize(payload)
_t_deser_end = _time_mod.perf_counter()
_deser_ms = (_t_deser_end - _t_deser_start) * 1000
_total_ms = (_t_deser_end - _t0) * 1000
_mbps = (data_size / 1024 / 1024) / (_total_ms / 1000) if _total_ms > 0 else 0
logger.info(
f"[RDMA GET] {get_key}: query={_query_ms:.1f}ms, alloc={_alloc_ms:.1f}ms, "
f"rdma={_rdma_ms:.1f}ms, sync={_sync_ms:.1f}ms, copy={_copy_ms:.1f}ms, "
f"deser={_deser_ms:.1f}ms, total={_total_ms:.1f}ms, {_mbps:.1f} MB/s"
)
self._metrics["gets"] += 1
self._metrics["bytes_transferred"] += data_size
return val, data_size
finally:
recv_buffer.release()
else:
self._metrics["errors"] += 1
logger.error(f"RDMA Get failed: received {resp} instead of TRANS_DONE")
recv_buffer.release()
return None
except Exception as e:
# Socket may be stuck after timeout; discard it
return req_socket.recv()
except Exception:
# Socket may be stuck after timeout; discard it.
self._invalidate_req_socket(zmq_addr)
self._metrics["timeouts"] += 1
logger.error(f"RDMA Get error: {e}", exc_info=True)
raise
def _sync_receive_buffer_if_needed(self) -> float:
"""Synchronize GPU-visible receive memory until Mooncake exposes a finer fence."""
t0 = _time_mod.perf_counter()
if self.pool.is_cuda:
with torch.cuda.device(self.pool.device):
torch.cuda.current_stream().synchronize()
return (_time_mod.perf_counter() - t0) * 1000
def _decode_result(
self,
get_key: str,
recv_buffer: ManagedBuffer,
data_size: int,
is_fast_path: bool,
query_ms: float,
alloc_ms: float,
rdma_ms: float,
sync_ms: float,
start_time: float,
) -> tuple[Any, int]:
"""Decode a completed transfer result."""
if is_fast_path:
try:
total_ms = (_time_mod.perf_counter() - start_time) * 1000
mbps = (data_size / 1024 / 1024) / (total_ms / 1000) if total_ms > 0 else 0
logger.info(
f"[RDMA GET] {get_key}: query={query_ms:.1f}ms, alloc={alloc_ms:.1f}ms, "
f"rdma={rdma_ms:.1f}ms, sync={sync_ms:.1f}ms, "
f"total={total_ms:.1f}ms, {mbps:.1f} MB/s (fast_path, zero-copy)"
)
self._metrics["gets"] += 1
self._metrics["bytes_transferred"] += data_size
return recv_buffer, data_size
except Exception:
recv_buffer.release()
raise
try:
copy_start = _time_mod.perf_counter()
raw_bytes = recv_buffer.to_bytes()
copy_ms = (_time_mod.perf_counter() - copy_start) * 1000
deser_start = _time_mod.perf_counter()
payload = raw_bytes[:data_size]
try:
value = OmniSerializer.deserialize(payload)
except msgspec.DecodeError:
payload = raw_bytes[1 : data_size + 1]
value = OmniSerializer.deserialize(payload)
deser_ms = (_time_mod.perf_counter() - deser_start) * 1000
total_ms = (_time_mod.perf_counter() - start_time) * 1000
mbps = (data_size / 1024 / 1024) / (total_ms / 1000) if total_ms > 0 else 0
logger.info(
f"[RDMA GET] {get_key}: query={query_ms:.1f}ms, alloc={alloc_ms:.1f}ms, "
f"rdma={rdma_ms:.1f}ms, sync={sync_ms:.1f}ms, copy={copy_ms:.1f}ms, "
f"deser={deser_ms:.1f}ms, total={total_ms:.1f}ms, {mbps:.1f} MB/s"
)
self._metrics["gets"] += 1
self._metrics["bytes_transferred"] += data_size
return value, data_size
finally:
recv_buffer.release()
return None
def cleanup(self, request_id: str, from_stage: str | None = None, to_stage: str | None = None) -> None:
"""Release the producer-side buffer associated with the request.
@@ -847,6 +893,7 @@ class MooncakeTransferEngineConnector(OmniConnectorBase):
with self._local_buffers_lock:
item = self._local_buffers.pop(request_id, None)
if item:
self._in_flight_buffers.discard(request_id)
# item is (src_addrs, lengths, holder, should_release, is_fast_path, created_at)
_, _, holder, should_release, _, _ = item
if should_release and isinstance(holder, ManagedBuffer):
@@ -904,6 +951,7 @@ class MooncakeTransferEngineConnector(OmniConnectorBase):
if should_release and isinstance(holder, ManagedBuffer):
holder.release()
self._local_buffers.clear()
self._in_flight_buffers.clear()
# 5. Close thread-local cached REQ sockets (for the calling thread).
# Sockets cached in *other* threads are not directly accessible here
@@ -945,21 +993,26 @@ class MooncakeTransferEngineConnector(OmniConnectorBase):
# -------------------------------------------------------
def _cleanup_stale_buffers(self) -> None:
"""Reclaim buffers older than ``_BUFFER_TTL_SECONDS``.
"""Reclaim stale buffers not currently used by a transfer.
Prevents permanent memory leaks when a receiver crashes or times out
without ever pulling the data.
TODO(wzliu): In extreme rare case, long transfer time, there might exist
TTL cleanup vs in-flight transfer conflict, which will be handled in the next PR.
"""
if self.buffer_ttl_seconds <= 0:
return
now = _time_mod.monotonic()
with self._local_buffers_lock:
stale_keys = [k for k, v in self._local_buffers.items() if now - v[5] > _BUFFER_TTL_SECONDS]
stale_keys = [
k
for k, v in self._local_buffers.items()
if k not in self._in_flight_buffers and now - v[5] > self.buffer_ttl_seconds
]
for k in stale_keys:
item = self._local_buffers.pop(k)
_, _, holder, should_release, _, _ = item
if should_release and isinstance(holder, ManagedBuffer):
holder.release()
logger.warning(f"TTL expired ({_BUFFER_TTL_SECONDS}s): cleaned up stale buffer for {k}")
logger.warning(f"TTL expired ({self.buffer_ttl_seconds}s): cleaned up stale buffer for {k}")
def _zmq_listener_loop(self):
socket = self.zmq_ctx.socket(zmq.ROUTER)
@@ -1058,30 +1111,36 @@ class MooncakeTransferEngineConnector(OmniConnectorBase):
with self._local_buffers_lock:
item = self._local_buffers.get(meta.request_id)
if item:
self._in_flight_buffers.add(meta.request_id)
if not item:
response_queue.put((identity, TRANS_ERROR))
self._notify_listener(notify_addr)
return
src_addrs, src_lengths, _, _, _, _ = item
remote_session = f"{meta.remote_hostname}:{meta.remote_port}"
try:
src_addrs, src_lengths, _, _, _, _ = item
remote_session = f"{meta.remote_hostname}:{meta.remote_port}"
# RDMA Write
ret = self.engine.batch_transfer_sync_write(remote_session, src_addrs, meta.dst_addrs, src_lengths)
# RDMA Write
ret = self.engine.batch_transfer_sync_write(remote_session, src_addrs, meta.dst_addrs, src_lengths)
if ret == 0:
self.cleanup(meta.request_id)
response_queue.put((identity, TRANS_DONE))
else:
# Keep buffer in _local_buffers so receiver can retry on transient failures.
# Buffer will be cleaned up when: (a) a retry succeeds, (b) close() is called,
# or (c) a future TTL mechanism reclaims stale entries.
logger.warning(
f"RDMA write failed for {meta.request_id} to {remote_session} "
f"(ret={ret}). Buffer retained for retry."
)
response_queue.put((identity, TRANS_ERROR))
if ret == 0:
self.cleanup(meta.request_id)
response_queue.put((identity, TRANS_DONE))
else:
# Keep buffer in _local_buffers so receiver can retry on transient failures.
# Buffer will be cleaned up when: (a) a retry succeeds, (b) close() is called,
# or (c) TTL reclaims it after it is no longer in-flight.
logger.warning(
f"RDMA write failed for {meta.request_id} to {remote_session} "
f"(ret={ret}). Buffer retained for retry."
)
response_queue.put((identity, TRANS_ERROR))
finally:
with self._local_buffers_lock:
self._in_flight_buffers.discard(meta.request_id)
except Exception as e:
logger.error(f"Push failed: {e}")
vllm_omni/distributed/omni_connectors/connectors/shm_connector.py
+28 -54
View File
@@ -27,16 +27,11 @@ class SharedMemoryConnector(OmniConnectorBase):
def __init__(self, config: dict[str, Any]):
self.config = config
self.stage_id = config.get("stage_id", -1)
self.device = config.get("device", "cuda:0")
self.threshold = int(config.get("shm_threshold_bytes", 65536))
self.inline_small_payloads = bool(config.get("inline_small_payloads", False))
self._pending_keys: set[str] = set()
self._metrics = {
"puts": 0,
"gets": 0,
"bytes_transferred": 0,
"shm_writes": 0,
"inline_writes": 0,
}
def put(
@@ -47,33 +42,18 @@ class SharedMemoryConnector(OmniConnectorBase):
data: Any,
) -> tuple[bool, int, dict[str, Any] | None]:
try:
# Always serialize first to check size (and for SHM writing)
# Note: For extremely large objects in "inline" mode (e.g. Ray),
# we might double-serialize if we're not careful, but here we assume
# if it's huge we use SHM, or if Ray, threshold is maxsize.
payload = self.serialize_obj(data)
size = len(payload)
# The legacy async-chunk adapter transfers only the connector key
# across stages; it cannot deliver inline metadata to the receiver.
# Keep key-addressed SHM as the default and only inline payloads
# when the caller explicitly enables the metadata-aware path.
if size >= self.threshold or not self.inline_small_payloads:
lock_file = f"/dev/shm/shm_{put_key}_lockfile.lock"
with open(lock_file, "wb+") as lockf:
fcntl.flock(lockf, fcntl.LOCK_EX)
meta = shm_write_bytes(payload, name=put_key)
fcntl.flock(lockf, fcntl.LOCK_UN)
lock_file = f"/dev/shm/shm_{put_key}_lockfile.lock"
with open(lock_file, "wb+") as lockf:
fcntl.flock(lockf, fcntl.LOCK_EX)
meta = shm_write_bytes(payload, name=put_key)
fcntl.flock(lockf, fcntl.LOCK_UN)
# meta contains {'name': ..., 'size': ...}
metadata = {"shm": meta, "size": size}
self._pending_keys.add(put_key)
self._metrics["shm_writes"] += 1
else:
# Inline small payloads to avoid the mmap + file-lock overhead
# that dominates codec chunk transfers.
metadata = {"inline_bytes": payload, "size": size}
self._metrics["inline_writes"] += 1
# meta contains {'name': ..., 'size': ...}
metadata = {"shm": meta, "size": size}
self._pending_keys.add(put_key)
self._metrics["puts"] += 1
self._metrics["bytes_transferred"] += size
@@ -84,22 +64,26 @@ class SharedMemoryConnector(OmniConnectorBase):
logger.error(f"SharedMemoryConnector put failed for req {put_key}: {e}")
return False, 0, None
def _get_data_with_lock(self, lock_file: str, shm_handle: dict):
obj = None
def _get_data_with_lock(self, lock_file: str, shm_handle: dict[str, Any]) -> tuple[Any, int] | None:
deserialized = False
try:
with open(lock_file, "rb+") as lockf:
fcntl.flock(lockf, fcntl.LOCK_EX)
data_bytes = shm_read_bytes(shm_handle)
fcntl.flock(lockf, fcntl.LOCK_UN)
obj = self.deserialize_obj(data_bytes)
return obj, int(shm_handle.get("size", 0))
result = (obj, int(shm_handle.get("size", 0)))
deserialized = True
return result
except Exception as e:
logger.error(f"SharedMemoryConnector shm get failed for req : {e}")
return None
finally:
# If data has been received, delete lock_file.
if obj and os.path.exists(lock_file):
os.remove(lock_file)
if deserialized:
try:
os.remove(lock_file)
except FileNotFoundError:
pass
def _get_by_key(self, get_key: str) -> tuple[Any, int] | None:
"""Read a SHM segment addressed purely by *get_key*."""
@@ -142,31 +126,21 @@ class SharedMemoryConnector(OmniConnectorBase):
if isinstance(metadata, dict) and get_key in metadata:
metadata = metadata.get(get_key)
if not isinstance(metadata, dict):
return self._get_by_key(get_key)
if "inline_bytes" in metadata:
try:
obj = self.deserialize_obj(metadata["inline_bytes"])
self._pending_keys.discard(get_key)
return obj, int(metadata.get("size", 0))
except Exception as e:
logger.error(f"SharedMemoryConnector inline get failed for req {get_key}: {e}")
return None
if "shm" in metadata:
if isinstance(metadata, dict) and "shm" in metadata:
shm_handle = metadata["shm"]
lock_file = f"/dev/shm/shm_{shm_handle['name']}_lockfile.lock"
result = self._get_data_with_lock(lock_file, shm_handle)
if result is not None:
self._pending_keys.discard(get_key)
return result
else:
# Missing or non-SHM metadata falls back to key-based lookup.
result = self._get_by_key(get_key)
else:
result = self._get_by_key(get_key)
# Metadata is a dict but has no SHM-specific handle (e.g. RDMA-
# style source_host/source_port). Fall back to key-based read.
return self._get_by_key(get_key)
return self._get_by_key(get_key)
if result is not None:
self._metrics["gets"] += 1
return result
def cleanup(self, request_id: str) -> None:
"""Best-effort cleanup of unconsumed SHM segments for *request_id*.
@@ -217,4 +191,4 @@ class SharedMemoryConnector(OmniConnectorBase):
self._pending_keys.clear()
def health(self) -> dict[str, Any]:
return {"status": "healthy", "threshold": self.threshold, **self._metrics}
return {"status": "healthy", **self._metrics}
vllm_omni/distributed/omni_connectors/utils/initialization.py
+4 -18
View File
@@ -4,7 +4,6 @@
"""Utilities for OmniConnector configuration and validation."""
import json
import sys
import warnings
from pathlib import Path
from typing import TYPE_CHECKING, Any
@@ -42,7 +41,6 @@ KV_REPLICA_PORT_STRIDE = 1024
def initialize_connectors_from_config(
config_path: str | Path | None = None,
default_shm_threshold: int = 65536,
purpose: str = "request_forwarding",
caller_stage_id: int | str | None = None,
is_sender: bool | None = None,
@@ -53,7 +51,7 @@ def initialize_connectors_from_config(
Returns:
tuple: (OmniTransferConfig, dict of {(from, to): connector_instance})
"""
transfer_config = load_omni_transfer_config(config_path, default_shm_threshold=default_shm_threshold)
transfer_config = load_omni_transfer_config(config_path)
if not transfer_config:
logger.info("No OmniTransferConfig provided")
@@ -190,7 +188,6 @@ def get_connectors_config_for_stage(transfer_config: OmniTransferConfig | None,
def load_omni_transfer_config(
config_path: str | Path | None = None,
config_dict: dict[str, Any] | None = None,
default_shm_threshold: int = 65536,
) -> OmniTransferConfig | None:
"""Load OmniTransferConfig from file or dict."""
if config_path is None and config_dict is None:
@@ -325,10 +322,7 @@ def load_omni_transfer_config(
expected_edges.add(edge_key)
if edge_key not in connectors:
logger.info(f"Auto-configuring SharedMemoryConnector for edge {edge_key}")
connectors[edge_key] = ConnectorSpec(
name="SharedMemoryConnector",
extra={"shm_threshold_bytes": default_shm_threshold},
)
connectors[edge_key] = ConnectorSpec(name="SharedMemoryConnector")
# Fallback: infer edges from engine_input_source for each stage
for stage_config in stage_args:
@@ -343,9 +337,7 @@ def load_omni_transfer_config(
if edge_key not in connectors:
logger.info(f"Auto-configuring SharedMemoryConnector for edge {edge_key}")
connectors[edge_key] = ConnectorSpec(
name="SharedMemoryConnector", extra={"shm_threshold_bytes": default_shm_threshold}
)
connectors[edge_key] = ConnectorSpec(name="SharedMemoryConnector")
except Exception as e:
logger.warning(f"Failed to auto-configure SHM connectors: {e}")
@@ -369,22 +361,16 @@ def load_omni_transfer_config(
def initialize_orchestrator_connectors(
config_path: str | None, worker_backend: str | None = "multi_process", shm_threshold_bytes: int = 65536
config_path: str | None,
) -> tuple[OmniTransferConfig | None, dict[tuple[str, str], OmniConnectorBase]]:
"""Initialize connectors shared at orchestrator level.
Args:
config_path: The path to the configuration file.
worker_backend: The backend to use for the worker.
Returns:
A tuple containing the OmniTransferConfig and a dictionary of connectors.
"""
if worker_backend == "ray":
default_shm_threshold = sys.maxsize
else:
default_shm_threshold = max(0, shm_threshold_bytes)
transfer_config, connectors = initialize_connectors_from_config(
config_path,
default_shm_threshold=default_shm_threshold,
purpose="request_forwarding",
caller_stage_id="orchestrator",
is_sender=True,
vllm_omni/engine/arg_utils.py
-1
View File
@@ -423,7 +423,6 @@ class OrchestratorArgs:
init_timeout: int = 600
# === Cross-stage Communication ===
shm_threshold_bytes: int = 65536
batch_timeout: int = 10
# === Cluster / Backend ===
vllm_omni/platforms/xpu/stage_configs/voxtral_tts.yaml
-1
View File
@@ -93,7 +93,6 @@ runtime:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
codec_streaming: true
connector_get_sleep_s: 0.01
connector_get_max_wait_first_chunk: 3000
vllm_omni/worker/gpu_ar_model_runner.py
-1
View File
@@ -323,7 +323,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
}
if getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS:
self.init_omni_connectors(
vllm_config=self.vllm_config,
model_config=self.model_config,
kv_transfer_manager=self.kv_transfer_manager,
)
vllm_omni/worker/gpu_generation_model_runner.py
-1
View File
@@ -70,7 +70,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
}
if getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS:
self.init_omni_connectors(
vllm_config=self.vllm_config,
model_config=self.model_config,
)
vllm_omni/worker/omni_connector_model_runner_mixin.py
+2 -85
View File
@@ -17,7 +17,6 @@ import inspect
import os
import threading
from collections import defaultdict, deque
from types import SimpleNamespace
from typing import TYPE_CHECKING, Any
import torch
@@ -36,6 +35,7 @@ from vllm_omni.worker.payload_span import (
_EMBED_SPAN_GROUPS: tuple[tuple[str, str, str], ...] = (("decode", "decode_token_start", "decode_token_end"),)
if TYPE_CHECKING:
from vllm_omni.config.model import OmniModelConfig
from vllm_omni.distributed.omni_connectors.connectors.base import (
OmniConnectorBase,
)
@@ -87,14 +87,12 @@ class OmniConnectorModelRunnerMixin:
def init_omni_connectors(
self,
vllm_config: Any,
model_config: Any,
model_config: OmniModelConfig,
kv_transfer_manager: OmniKVTransferManager | None = None,
) -> None:
"""Initialize connectors and background threads.
Args:
vllm_config: Full vLLM config object.
model_config: Stage-level model config with connector settings.
kv_transfer_manager: Existing KV transfer manager to delegate to.
"""
@@ -1024,31 +1022,6 @@ class OmniConnectorModelRunnerMixin:
self._work_available.set()
return sent_ids
def recv_stage_inputs(self, scheduler_output: Any) -> dict[str, Any] | None:
"""Compatibility wrapper for ``recv_full_payload_inputs``."""
return self.recv_full_payload_inputs(scheduler_output)
def accumulate_batch_output(
self,
req_id: str,
pooler_output: Any,
request: Any,
) -> None:
"""Compatibility wrapper for ``accumulate_full_payload_output``."""
self.accumulate_full_payload_output(req_id, pooler_output, request)
def flush_batch_outputs(self, finished_req_ids: set[str]) -> None:
"""Compatibility wrapper for ``flush_full_payload_outputs``."""
self.flush_full_payload_outputs(finished_req_ids)
def send_stage_outputs(
self,
scheduler_output: Any,
outputs: dict[str, tuple[Any, Any] | Any],
) -> list[str]:
"""Compatibility wrapper for ``send_full_payload_outputs``."""
return self.send_full_payload_outputs(scheduler_output, outputs)
# ------------------------------------------------------------------ #
# Streaming chunk mode (recv_chunk / send_chunk)
# ------------------------------------------------------------------ #
@@ -1505,20 +1478,6 @@ class OmniConnectorModelRunnerMixin:
# Output aggregation
# ------------------------------------------------------------------ #
def _empty_output_with_connector_signals(self) -> Any:
"""Return a minimal ModelRunnerOutput carrying pending connector signals.
Used by early-return paths (e.g. ``num_scheduled_tokens == 0``)
that still need to deliver ``omni_connector_output`` to the
Scheduler so that WAITING_FOR_INPUT / WAITING_FOR_CHUNK
transitions are not lost.
"""
from vllm_omni.outputs import OmniModelRunnerOutput
output = OmniModelRunnerOutput(req_ids=[], req_id_to_index={})
output.omni_connector_output = self.get_omni_connector_output()
return output
def get_omni_connector_output(self) -> OmniConnectorOutput:
"""Collect and reset transfer results for this execute_model cycle.
@@ -2088,48 +2047,6 @@ class OmniConnectorModelRunnerMixin:
# Helpers
# ------------------------------------------------------------------ #
@staticmethod
def _freeze_request_attr(value: Any) -> Any:
if isinstance(value, list):
return list(value)
if isinstance(value, tuple):
return list(value)
if isinstance(value, torch.Tensor):
return value.clone()
raw_list = getattr(value, "_x", None)
if raw_list is not None:
return list(raw_list)
return value
def _snapshot_request_for_send(self, request: Any, external_req_id: str) -> Any:
finished = bool(getattr(request, "is_finished", lambda: False)())
attrs: dict[str, Any] = {}
try:
attrs.update(vars(request))
except TypeError:
pass
for name in (
"request_id",
"req_id",
"external_req_id",
"prompt_token_ids",
"output_token_ids",
"all_token_ids",
"additional_information",
"sampling_params",
"multi_modal_data",
"mm_hashes",
):
if hasattr(request, name):
attrs[name] = self._freeze_request_attr(getattr(request, name))
attrs["external_req_id"] = external_req_id
attrs["_frozen_is_finished"] = finished
snapshot = SimpleNamespace(**attrs)
snapshot.is_finished = lambda: finished
return snapshot
@staticmethod
def _create_connector(model_config: Any) -> OmniConnectorBase | None:
"""Create a connector from model_config, or None if unconfigured."""