From dcf8b07d39d553e34338c997e5497a0a503cfdc4 Mon Sep 17 00:00:00 2001 From: "Bingyu (Spencer) Liu" Date: Wed, 22 Jul 2026 09:44:15 +0800 Subject: [PATCH] [Core] Cleanup Omni connector runtime and transfer backends (#5096) Signed-off-by: spencerr221 --- docs/configuration/stage_configs.md | 2 +- .../design/feature/disaggregated_inference.md | 2 - .../shared_memory_connector.md | 29 +- docs/design/feature/ray_based_execution.md | 2 +- .../offline_inference/covo_audio/end2end.py | 1 - .../offline_inference/mimo_audio/end2end.py | 1 - .../offline_inference/step_audio2/end2end.py | 1 - .../text_to_speech/ming_tts/runner.py | 1 - tests/config/test_omni_config.py | 2 - .../omni_connectors/test_adapter_and_flow.py | 5 +- .../omni_connectors/test_basic_connectors.py | 57 ++- .../test_omni_connector_configs.py | 1 - .../omni_connectors/test_shm_connector.py | 29 +- tests/e2e/accuracy/test_hunyuan_image3.py | 1 - tests/engine/test_async_omni_engine_abort.py | 3 - tests/entrypoints/test_pd_disaggregation.py | 4 +- tests/helpers/runtime.py | 2 - tests/worker/test_omni_connector_mixin.py | 29 +- vllm_omni/config/omni_config.py | 1 - vllm_omni/deploy/aura_omni.yaml | 1 - vllm_omni/deploy/bagel.yaml | 2 - .../deploy/dynin_omni_multiconnector.yaml | 2 - vllm_omni/deploy/fish_qwen3_omni.yaml | 1 - vllm_omni/deploy/fish_qwen3_omni_2gpu.yaml | 1 - ..._qwen3_omni_high_concurrency_dual_gpu.yaml | 1 - ...wen3_omni_high_concurrency_single_gpu.yaml | 1 - vllm_omni/deploy/higgs_audio_v2.yaml | 1 - vllm_omni/deploy/higgs_multimodal_qwen3.yaml | 1 - ...iggs_multimodal_qwen3_high_throughput.yaml | 1 - .../higgs_multimodal_qwen3_low_latency.yaml | 1 - vllm_omni/deploy/hunyuan_image_3_moe.yaml | 2 - vllm_omni/deploy/indextts2.yaml | 1 - vllm_omni/deploy/indextts2_low_latency.yaml | 1 - vllm_omni/deploy/mimo_audio.yaml | 1 - vllm_omni/deploy/ming_flash_omni_image.yaml | 2 - vllm_omni/deploy/moss_sound_effect.yaml | 1 - vllm_omni/deploy/moss_tts.yaml | 1 - vllm_omni/deploy/moss_tts_80gb.yaml | 1 - vllm_omni/deploy/moss_tts_local.yaml | 1 - vllm_omni/deploy/moss_tts_realtime.yaml | 1 - vllm_omni/deploy/moss_ttsd.yaml | 1 - vllm_omni/deploy/moss_voice_generator.yaml | 1 - vllm_omni/deploy/qwen3_tts.yaml | 1 - .../deploy/qwen3_tts_high_concurrency.yaml | 1 - vllm_omni/deploy/voxtral_tts.yaml | 1 - .../mooncake_transfer_engine_connector.py | 399 ++++++++++-------- .../connectors/shm_connector.py | 82 ++-- .../omni_connectors/utils/initialization.py | 22 +- vllm_omni/engine/arg_utils.py | 1 - .../xpu/stage_configs/voxtral_tts.yaml | 1 - vllm_omni/worker/gpu_ar_model_runner.py | 1 - .../worker/gpu_generation_model_runner.py | 1 - .../omni_connector_model_runner_mixin.py | 87 +--- 53 files changed, 349 insertions(+), 448 deletions(-) diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index 1aa4ca1b3..e49487b72 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -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`: diff --git a/docs/design/feature/disaggregated_inference.md b/docs/design/feature/disaggregated_inference.md index 15d9e09df..c31086c0d 100644 --- a/docs/design/feature/disaggregated_inference.md +++ b/docs/design/feature/disaggregated_inference.md @@ -66,8 +66,6 @@ runtime: connectors: connector_of_shared_memory: name: SharedMemoryConnector - extra: - shm_threshold_bytes: 65536 ``` Wire stages to connectors: diff --git a/docs/design/feature/omni_connectors/shared_memory_connector.md b/docs/design/feature/omni_connectors/shared_memory_connector.md index 5b35014f2..62d77cd6d 100644 --- a/docs/design/feature/omni_connectors/shared_memory_connector.md +++ b/docs/design/feature/omni_connectors/shared_memory_connector.md @@ -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 diff --git a/docs/design/feature/ray_based_execution.md b/docs/design/feature/ray_based_execution.md index ae10d661f..4f4caa642 100644 --- a/docs/design/feature/ray_based_execution.md +++ b/docs/design/feature/ray_based_execution.md @@ -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 diff --git a/examples/offline_inference/covo_audio/end2end.py b/examples/offline_inference/covo_audio/end2end.py index ae7e5c6bb..ce728373d 100644 --- a/examples/offline_inference/covo_audio/end2end.py +++ b/examples/offline_inference/covo_audio/end2end.py @@ -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 diff --git a/examples/offline_inference/mimo_audio/end2end.py b/examples/offline_inference/mimo_audio/end2end.py index 139351aab..4ba8f483a 100644 --- a/examples/offline_inference/mimo_audio/end2end.py +++ b/examples/offline_inference/mimo_audio/end2end.py @@ -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( diff --git a/examples/offline_inference/step_audio2/end2end.py b/examples/offline_inference/step_audio2/end2end.py index 8cc7e2a26..e0d1836a8 100644 --- a/examples/offline_inference/step_audio2/end2end.py +++ b/examples/offline_inference/step_audio2/end2end.py @@ -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, diff --git a/examples/offline_inference/text_to_speech/ming_tts/runner.py b/examples/offline_inference/text_to_speech/ming_tts/runner.py index 35852ce04..6f9d75740 100644 --- a/examples/offline_inference/text_to_speech/ming_tts/runner.py +++ b/examples/offline_inference/text_to_speech/ming_tts/runner.py @@ -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: diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 23e53ea3d..1da71b21d 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -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 diff --git a/tests/distributed/omni_connectors/test_adapter_and_flow.py b/tests/distributed/omni_connectors/test_adapter_and_flow.py index 548d16eb6..92dc1d7c0 100644 --- a/tests/distributed/omni_connectors/test_adapter_and_flow.py +++ b/tests/distributed/omni_connectors/test_adapter_and_flow.py @@ -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 diff --git a/tests/distributed/omni_connectors/test_basic_connectors.py b/tests/distributed/omni_connectors/test_basic_connectors.py index 1518e98b1..787d5594b 100644 --- a/tests/distributed/omni_connectors/test_basic_connectors.py +++ b/tests/distributed/omni_connectors/test_basic_connectors.py @@ -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 diff --git a/tests/distributed/omni_connectors/test_omni_connector_configs.py b/tests/distributed/omni_connectors/test_omni_connector_configs.py index 3dc80b57b..82f13c5cd 100644 --- a/tests/distributed/omni_connectors/test_omni_connector_configs.py +++ b/tests/distributed/omni_connectors/test_omni_connector_configs.py @@ -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" diff --git a/tests/distributed/omni_connectors/test_shm_connector.py b/tests/distributed/omni_connectors/test_shm_connector.py index 572376aef..bba6002a0 100644 --- a/tests/distributed/omni_connectors/test_shm_connector.py +++ b/tests/distributed/omni_connectors/test_shm_connector.py @@ -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) diff --git a/tests/e2e/accuracy/test_hunyuan_image3.py b/tests/e2e/accuracy/test_hunyuan_image3.py index 2214b734d..e2d84b13d 100644 --- a/tests/e2e/accuracy/test_hunyuan_image3.py +++ b/tests/e2e/accuracy/test_hunyuan_image3.py @@ -141,7 +141,6 @@ _DEPLOY_CONFIG = { "connectors": { "shared_memory_connector": { "name": "SharedMemoryConnector", - "extra": {"shm_threshold_bytes": 65536}, }, }, "stages": [ diff --git a/tests/engine/test_async_omni_engine_abort.py b/tests/engine/test_async_omni_engine_abort.py index 794475388..dfe00732f 100644 --- a/tests/engine/test_async_omni_engine_abort.py +++ b/tests/engine/test_async_omni_engine_abort.py @@ -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) diff --git a/tests/entrypoints/test_pd_disaggregation.py b/tests/entrypoints/test_pd_disaggregation.py index 5ffabfbf2..57e484af1 100644 --- a/tests/entrypoints/test_pd_disaggregation.py +++ b/tests/entrypoints/test_pd_disaggregation.py @@ -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, diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index 263f34da2..7c19c643d 100644 --- a/tests/helpers/runtime.py +++ b/tests/helpers/runtime.py @@ -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, ) diff --git a/tests/worker/test_omni_connector_mixin.py b/tests/worker/test_omni_connector_mixin.py index b8eb899fa..87cf4ba73 100644 --- a/tests/worker/test_omni_connector_mixin.py +++ b/tests/worker/test_omni_connector_mixin.py @@ -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 diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 4c2154401..be89e1cae 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -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) diff --git a/vllm_omni/deploy/aura_omni.yaml b/vllm_omni/deploy/aura_omni.yaml index 486e44b57..cf6826fdd 100644 --- a/vllm_omni/deploy/aura_omni.yaml +++ b/vllm_omni/deploy/aura_omni.yaml @@ -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 diff --git a/vllm_omni/deploy/bagel.yaml b/vllm_omni/deploy/bagel.yaml index 8de6f9305..644e2c3c0 100644 --- a/vllm_omni/deploy/bagel.yaml +++ b/vllm_omni/deploy/bagel.yaml @@ -40,8 +40,6 @@ stages: connectors: shared_memory_connector: name: SharedMemoryConnector - extra: - shm_threshold_bytes: 65536 rdma_connector: name: MooncakeTransferEngineConnector diff --git a/vllm_omni/deploy/dynin_omni_multiconnector.yaml b/vllm_omni/deploy/dynin_omni_multiconnector.yaml index c863ec3d5..5abb4c5a3 100644 --- a/vllm_omni/deploy/dynin_omni_multiconnector.yaml +++ b/vllm_omni/deploy/dynin_omni_multiconnector.yaml @@ -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 diff --git a/vllm_omni/deploy/fish_qwen3_omni.yaml b/vllm_omni/deploy/fish_qwen3_omni.yaml index 539b37775..f10ed1656 100644 --- a/vllm_omni/deploy/fish_qwen3_omni.yaml +++ b/vllm_omni/deploy/fish_qwen3_omni.yaml @@ -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 diff --git a/vllm_omni/deploy/fish_qwen3_omni_2gpu.yaml b/vllm_omni/deploy/fish_qwen3_omni_2gpu.yaml index 9fec34d69..360b72974 100644 --- a/vllm_omni/deploy/fish_qwen3_omni_2gpu.yaml +++ b/vllm_omni/deploy/fish_qwen3_omni_2gpu.yaml @@ -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. diff --git a/vllm_omni/deploy/fish_qwen3_omni_high_concurrency_dual_gpu.yaml b/vllm_omni/deploy/fish_qwen3_omni_high_concurrency_dual_gpu.yaml index a65609c6c..57ea6e757 100644 --- a/vllm_omni/deploy/fish_qwen3_omni_high_concurrency_dual_gpu.yaml +++ b/vllm_omni/deploy/fish_qwen3_omni_high_concurrency_dual_gpu.yaml @@ -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 diff --git a/vllm_omni/deploy/fish_qwen3_omni_high_concurrency_single_gpu.yaml b/vllm_omni/deploy/fish_qwen3_omni_high_concurrency_single_gpu.yaml index d771e4163..0c9771b18 100644 --- a/vllm_omni/deploy/fish_qwen3_omni_high_concurrency_single_gpu.yaml +++ b/vllm_omni/deploy/fish_qwen3_omni_high_concurrency_single_gpu.yaml @@ -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 diff --git a/vllm_omni/deploy/higgs_audio_v2.yaml b/vllm_omni/deploy/higgs_audio_v2.yaml index 07e25d5c1..9c181400b 100644 --- a/vllm_omni/deploy/higgs_audio_v2.yaml +++ b/vllm_omni/deploy/higgs_audio_v2.yaml @@ -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 diff --git a/vllm_omni/deploy/higgs_multimodal_qwen3.yaml b/vllm_omni/deploy/higgs_multimodal_qwen3.yaml index 5fc44a295..da299415b 100644 --- a/vllm_omni/deploy/higgs_multimodal_qwen3.yaml +++ b/vllm_omni/deploy/higgs_multimodal_qwen3.yaml @@ -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: diff --git a/vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml b/vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml index 86427be13..8a751483c 100644 --- a/vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml +++ b/vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml @@ -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: diff --git a/vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml b/vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml index 8210db4a8..0d6ae8218 100644 --- a/vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml +++ b/vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml @@ -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: diff --git a/vllm_omni/deploy/hunyuan_image_3_moe.yaml b/vllm_omni/deploy/hunyuan_image_3_moe.yaml index b661674c7..4ae4520a7 100644 --- a/vllm_omni/deploy/hunyuan_image_3_moe.yaml +++ b/vllm_omni/deploy/hunyuan_image_3_moe.yaml @@ -27,8 +27,6 @@ connectors: memory_pool_device: "npu" shared_memory_connector: name: SharedMemoryConnector - extra: - shm_threshold_bytes: 65536 stages: - stage_id: 0 diff --git a/vllm_omni/deploy/indextts2.yaml b/vllm_omni/deploy/indextts2.yaml index 75eb60188..c5fefcf30 100644 --- a/vllm_omni/deploy/indextts2.yaml +++ b/vllm_omni/deploy/indextts2.yaml @@ -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 diff --git a/vllm_omni/deploy/indextts2_low_latency.yaml b/vllm_omni/deploy/indextts2_low_latency.yaml index 621116ab2..678144a2b 100644 --- a/vllm_omni/deploy/indextts2_low_latency.yaml +++ b/vllm_omni/deploy/indextts2_low_latency.yaml @@ -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 diff --git a/vllm_omni/deploy/mimo_audio.yaml b/vllm_omni/deploy/mimo_audio.yaml index a50632964..daeedec99 100644 --- a/vllm_omni/deploy/mimo_audio.yaml +++ b/vllm_omni/deploy/mimo_audio.yaml @@ -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 diff --git a/vllm_omni/deploy/ming_flash_omni_image.yaml b/vllm_omni/deploy/ming_flash_omni_image.yaml index 9f7c6a34d..a02bd07e5 100644 --- a/vllm_omni/deploy/ming_flash_omni_image.yaml +++ b/vllm_omni/deploy/ming_flash_omni_image.yaml @@ -50,8 +50,6 @@ stages: connectors: shared_memory_connector: name: SharedMemoryConnector - extra: - shm_threshold_bytes: 65536 # 64KB threshold edges: - from: 0 diff --git a/vllm_omni/deploy/moss_sound_effect.yaml b/vllm_omni/deploy/moss_sound_effect.yaml index 0fe616290..1dcf0d42f 100644 --- a/vllm_omni/deploy/moss_sound_effect.yaml +++ b/vllm_omni/deploy/moss_sound_effect.yaml @@ -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 diff --git a/vllm_omni/deploy/moss_tts.yaml b/vllm_omni/deploy/moss_tts.yaml index b1549d0d0..92bd2f3ba 100644 --- a/vllm_omni/deploy/moss_tts.yaml +++ b/vllm_omni/deploy/moss_tts.yaml @@ -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 diff --git a/vllm_omni/deploy/moss_tts_80gb.yaml b/vllm_omni/deploy/moss_tts_80gb.yaml index e61d4cf5d..57dcbc1df 100644 --- a/vllm_omni/deploy/moss_tts_80gb.yaml +++ b/vllm_omni/deploy/moss_tts_80gb.yaml @@ -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 diff --git a/vllm_omni/deploy/moss_tts_local.yaml b/vllm_omni/deploy/moss_tts_local.yaml index 1455ccc09..ab66df300 100644 --- a/vllm_omni/deploy/moss_tts_local.yaml +++ b/vllm_omni/deploy/moss_tts_local.yaml @@ -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 diff --git a/vllm_omni/deploy/moss_tts_realtime.yaml b/vllm_omni/deploy/moss_tts_realtime.yaml index b7909661a..ca63161e2 100644 --- a/vllm_omni/deploy/moss_tts_realtime.yaml +++ b/vllm_omni/deploy/moss_tts_realtime.yaml @@ -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 diff --git a/vllm_omni/deploy/moss_ttsd.yaml b/vllm_omni/deploy/moss_ttsd.yaml index f0b5f6eae..3ffba8678 100644 --- a/vllm_omni/deploy/moss_ttsd.yaml +++ b/vllm_omni/deploy/moss_ttsd.yaml @@ -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 diff --git a/vllm_omni/deploy/moss_voice_generator.yaml b/vllm_omni/deploy/moss_voice_generator.yaml index 010f04205..fa3cd91af 100644 --- a/vllm_omni/deploy/moss_voice_generator.yaml +++ b/vllm_omni/deploy/moss_voice_generator.yaml @@ -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 diff --git a/vllm_omni/deploy/qwen3_tts.yaml b/vllm_omni/deploy/qwen3_tts.yaml index f27da1b5b..a12c920ca 100644 --- a/vllm_omni/deploy/qwen3_tts.yaml +++ b/vllm_omni/deploy/qwen3_tts.yaml @@ -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 diff --git a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml index 939396840..847383955 100644 --- a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml +++ b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml @@ -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 diff --git a/vllm_omni/deploy/voxtral_tts.yaml b/vllm_omni/deploy/voxtral_tts.yaml index 929daddb1..d024c3900 100644 --- a/vllm_omni/deploy/voxtral_tts.yaml +++ b/vllm_omni/deploy/voxtral_tts.yaml @@ -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 diff --git a/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py b/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py index 34d19062a..2d0c1f11a 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py @@ -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}") diff --git a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py index 5de0d54db..3a51795c4 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py @@ -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} diff --git a/vllm_omni/distributed/omni_connectors/utils/initialization.py b/vllm_omni/distributed/omni_connectors/utils/initialization.py index f7a464809..bbbda44e3 100644 --- a/vllm_omni/distributed/omni_connectors/utils/initialization.py +++ b/vllm_omni/distributed/omni_connectors/utils/initialization.py @@ -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, diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index bdf43711c..19db55c63 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -423,7 +423,6 @@ class OrchestratorArgs: init_timeout: int = 600 # === Cross-stage Communication === - shm_threshold_bytes: int = 65536 batch_timeout: int = 10 # === Cluster / Backend === diff --git a/vllm_omni/platforms/xpu/stage_configs/voxtral_tts.yaml b/vllm_omni/platforms/xpu/stage_configs/voxtral_tts.yaml index 0820ab632..cf7937e52 100644 --- a/vllm_omni/platforms/xpu/stage_configs/voxtral_tts.yaml +++ b/vllm_omni/platforms/xpu/stage_configs/voxtral_tts.yaml @@ -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 diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 6cd2b3566..411796d37 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -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, ) diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index ec41900b1..bcb647605 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -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, ) diff --git a/vllm_omni/worker/omni_connector_model_runner_mixin.py b/vllm_omni/worker/omni_connector_model_runner_mixin.py index b173f81ac..f9f13cbe7 100644 --- a/vllm_omni/worker/omni_connector_model_runner_mixin.py +++ b/vllm_omni/worker/omni_connector_model_runner_mixin.py @@ -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."""