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

Signed-off-by: spencerr221 <liubingyu62@gmail.com>

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

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