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

[Core] Add NIXL omni connector (#6093)

Browse Source 
Signed-off-by: yuanwu <yuan.wu@intel.com>
Signed-off-by: Yuan Wu <yuan.wu@intel.com>
This commit is contained in:
Yuan Wu authored and GitHub committed 2026-09-13 16:35:02 +08:00
1 parent 01a2f93256
commit 6753bbd18c
33 files changed
+3032 -66
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/worker/test_omni_connector_mixin.py
+192
View File
@@ -10,6 +10,7 @@ from __future__ import annotations
import time
import unittest
from copy import deepcopy
from types import SimpleNamespace
from typing import Any
from unittest.mock import MagicMock, patch
@@ -20,6 +21,10 @@ import torch
from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
OmniKVTransferManager,
)
from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
from vllm_omni.distributed.omni_connectors.utils.initialization import (
resolve_connector_spec,
)
from vllm_omni.outputs import OmniConnectorOutput
from vllm_omni.worker.omni_connector_model_runner_mixin import (
OmniConnectorModelRunnerMixin,
@@ -564,6 +569,55 @@ class TestChunkStreamCompletedGuard(unittest.TestCase):
host.shutdown_omni_connectors()
def test_poll_uses_request_scoped_payload_sender_endpoint(self):
host = self._make_host(stage_id=1)
host._omni_connector.get = MagicMock(return_value=None)
req = _make_request("req-1", "ext-req-1")
req.payload_sender_info = {"host": "10.0.0.1", "zmq_port": 50051}
host.register_chunk_recv(req)
host._poll_single_request("req-1")
host._omni_connector.get.assert_called_once_with(
"0",
"1",
"ext-req-1_0_0",
{"source_host": "10.0.0.1", "source_port": 50051},
)
host.shutdown_omni_connectors()
def test_concurrent_requests_keep_distinct_payload_sender_endpoints(self):
host = self._make_host(stage_id=1)
host._omni_connector.get = MagicMock(return_value=None)
first = _make_request("req-1", "ext-req-1")
first.payload_sender_info = {"host": "10.0.0.1", "zmq_port": 50051}
second = _make_request("req-2", "ext-req-2")
second.payload_sender_info = {"host": "10.0.0.2", "zmq_port": 51051}
host.register_chunk_recv(first)
host.register_chunk_recv(second)
host._poll_single_request("req-2")
host._poll_single_request("req-1")
self.assertEqual(
host._omni_connector.get.call_args_list,
[
unittest.mock.call(
"0",
"1",
"ext-req-2_0_0",
{"source_host": "10.0.0.2", "source_port": 51051},
),
unittest.mock.call(
"0",
"1",
"ext-req-1_0_0",
{"source_host": "10.0.0.1", "source_port": 50051},
),
],
)
host.shutdown_omni_connectors()
def test_finish_sentinel_populates_completed_set(self):
"""Receiving is_finished=True adds to _chunk_stream_completed."""
host = self._make_host(stage_id=1)
@@ -1381,6 +1435,144 @@ class TestConnectorConfigValidation(unittest.TestCase):
host.init_omni_connectors(model_config=model_config)
class TestRankAwareHandshakePort(unittest.TestCase):
"""A sender must bind the port the orchestrator advertises for its rank."""
@staticmethod
def _resolve(name, extra, local_rank=0):
spec = resolve_connector_spec(
ConnectorSpec(name=name, extra=extra),
stage_id=int(extra.get("stage_id", 0)),
role=extra.get("role"),
local_rank=local_rank,
)
return spec.extra
def test_sender_port_matches_the_orchestrator_formula(self):
for local_rank in (0, 1, 3):
with self.subTest(local_rank=local_rank):
extra = {"zmq_port": 50071, "stage_id": 2, "role": "sender"}
resolved = self._resolve("NixlConnector", extra, local_rank=local_rank)
self.assertEqual(resolved["zmq_port"], 50071 + 2 + local_rank * 16)
self.assertEqual(extra["zmq_port"], 50071, "the source config must not be mutated")
def test_receiver_gets_sender_endpoint_without_a_bind_port(self):
receiver = self._resolve(
"NixlConnector",
{"zmq_port": 50071, "stage_id": 3, "from_stage": 2, "role": "receiver"},
local_rank=1,
)
self.assertNotIn("zmq_port", receiver)
self.assertEqual(receiver["sender_zmq_port"], 50071 + 2 + 16)
def test_handshake_less_connector_keeps_its_config(self):
shm = self._resolve("SharedMemoryConnector", {"zmq_port": 50071, "role": "sender"}, local_rank=1)
self.assertEqual(shm["zmq_port"], 50071)
def test_unresolvable_port_is_rejected(self):
extra = {"zmq_port": "${MISSING_PORT_VAR}", "stage_id": 0, "role": "sender"}
with self.assertRaises(ValueError):
self._resolve("NixlConnector", extra)
@pytest.mark.parametrize(
("extra", "expected"),
[
(
{"role": "sender", "from_stage": 2, "zmq_port": 50071},
{"zmq_port": 52153},
),
(
{"role": "receiver", "from_stage": 1, "zmq_port": 50071, "host": "producer"},
{"sender_zmq_port": 52152, "sender_host": "producer"},
),
(
{
"role": "receiver",
"from_stage": 1,
"zmq_port": 50071,
"sender_zmq_port": 49000,
"sender_host": "bound-producer",
"outgoing": {"from_stage": 2, "zmq_port": 51071, "host": "local-worker"},
},
{"sender_zmq_port": 49000, "sender_host": "bound-producer", "zmq_port": 53153, "host": "local-worker"},
),
],
)
def test_split_runtime_factory_resolves_nixl_endpoints(extra, expected):
"""The runner must pass the resolved spec, not the raw config, to the factory."""
from vllm_omni.distributed.omni_connectors.model_runner import omni_connector_runtime as runtime
original = deepcopy(extra)
config = SimpleNamespace(stage_id=2, stage_connector_config={"name": "NixlConnector", "extra": extra})
with (
patch.object(runtime, "get_local_tp_rank", return_value=2, create=True),
patch.object(runtime, "get_omni_replica_id", return_value=2, create=True),
patch.object(runtime.OmniConnectorFactory, "create_connector") as factory,
):
assert MixinHost._create_connector(config) is factory.return_value
spec = factory.call_args.args[0]
assert spec.name == "NixlConnector"
assert spec.extra["stage_id"] == 2
for key, value in expected.items():
assert spec.extra[key] == value
if extra["role"] == "receiver" and "outgoing" not in extra:
assert "zmq_port" not in spec.extra
assert extra == original
@pytest.mark.parametrize("async_chunk", [False, True])
def test_split_payload_endpoint_retry_cleanup_and_reuse(async_chunk):
host = MixinHost()
host.init_omni_connectors(_make_model_config(async_chunk=async_chunk))
host._stage_id = 1
host._omni_connector = MagicMock()
host._omni_connector.get.return_value = None
first = _make_request("req-1", "ext-req-1")
first.payload_sender_info = {"host": "producer-a", "zmq_port": "50051"}
second = _make_request("req-2", "ext-req-2")
second.payload_sender_info = {"host": "producer-b", "zmq_port": 51051}
try:
host.register_chunk_recv(first)
host.register_chunk_recv(second)
for req_id in ("req-1", "req-2", "req-1"):
assert not host._poll_single_request(req_id)
assert host._omni_connector.get.call_args_list == [
unittest.mock.call("0", "1", "ext-req-1_0_0", {"source_host": "producer-a", "source_port": 50051}),
unittest.mock.call("0", "1", "ext-req-2_0_0", {"source_host": "producer-b", "source_port": 51051}),
unittest.mock.call("0", "1", "ext-req-1_0_0", {"source_host": "producer-a", "source_port": 50051}),
]
host.cleanup_finished_request("req-1")
assert "req-1" not in host._pending_load_reqs
assert host._pending_load_reqs["req-2"] is second
host.register_chunk_recv(_make_request("req-1", "reused"))
assert not host._poll_single_request("req-1")
host._omni_connector.get.assert_called_with("0", "1", "reused_0_0")
finally:
host.shutdown_omni_connectors()
@pytest.mark.parametrize("method", ["_recv_full_payload_result", "_recv_async_chunk_result"])
@pytest.mark.parametrize("rank", [0, 1])
def test_split_payload_sender_metadata_respects_tp_leader(method, rank):
host = MixinHost()
connector = MagicMock()
metadata = {"source_host": "producer", "source_port": 50051}
with (
patch.object(host, "_get_local_tp_group", return_value=_FakeTPGroup(world_size=2, rank_in_group=rank)),
patch.object(host, "is_data_transfer_rank", return_value=rank == 0),
):
result = getattr(host, method)(connector, "0", "1", "req_0_0", metadata)
if rank == 0:
connector.get.assert_called_once_with("0", "1", "req_0_0", metadata)
assert result is connector.get.return_value
else:
connector.get.assert_not_called()
assert result is None
class _FailingConnector:
"""Connector whose put() fails a configurable number of times."""