[Core] Cleanup Omni connector runtime and transfer backends

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`:

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:

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
 

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
 

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

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(

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,

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:

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
 
 

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

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 unregister_memory(self, base_ptr):
         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

@@ -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

@@ -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 @@
 
 @pytest.fixture()
 def connector():
-    c = SharedMemoryConnector({"shm_threshold_bytes": 0})
+    c = SharedMemoryConnector({})
     yield c
     c.close()
 
@@ -36,6 +39,30 @@ def test_put_then_get_by_key(self, connector):
         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

@@ -141,7 +141,6 @@ def _default_ar_dit_devices() -> tuple[str, str]:
     "connectors": {
         "shared_memory_connector": {
             "name": "SharedMemoryConnector",
-            "extra": {"shm_threshold_bytes": 65536},
         },
     },
     "stages": [

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)
 

tests/entrypoints/test_pd_disaggregation.py

@@ -121,9 +121,7 @@ def attach_queues(self, in_q, out_q):
         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

@@ -2333,7 +2333,6 @@ def __init__(
         # 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 @@ def __init__(
             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,
         )