From 6753bbd18c9a1f45cfebcec3eacdd2f1e5ea859b Mon Sep 17 00:00:00 2001 From: Yuan Wu Date: Sun, 13 Sep 2026 16:35:02 +0800 Subject: [PATCH] [Core] Add NIXL omni connector (#6093) Signed-off-by: yuanwu Signed-off-by: Yuan Wu --- docs/.nav.yml | 1 + docs/configuration/environment_variables.md | 7 + .../design/feature/disaggregated_inference.md | 2 + .../feature/omni_connectors/nixl_connector.md | 135 +++ docs/design/index.md | 1 + tests/config/test_environment_variables.py | 2 +- .../sched/test_omni_scheduling_coordinator.py | 14 + .../test_chunk_transfer_adapter.py | 54 +- .../omni_connectors/test_nixl_connector.py | 891 ++++++++++++++ .../test_nixl_connector_native.py | 204 ++++ .../test_orchestrator_kv_sender_info.py | 159 ++- .../test_orchestrator_stage_input_bridge.py | 26 +- tests/worker/test_omni_connector_mixin.py | 192 +++ .../config/environment_variable_inventory.py | 2 + .../core/sched/omni_scheduling_coordinator.py | 2 + vllm_omni/core/sched/output.py | 1 + vllm_omni/distributed/__init__.py | 2 + .../distributed/omni_connectors/__init__.py | 6 + .../connectors/nixl_connector.py | 1044 +++++++++++++++++ .../distributed/omni_connectors/factory.py | 12 + .../omni_connectors/kv_transfer_manager.py | 11 +- .../omni_connector_payload_transport.py | 23 +- .../model_runner/omni_connector_runtime.py | 18 +- .../chunk_transfer_adapter.py | 28 +- .../omni_connectors/utils/config.py | 1 + .../omni_connectors/utils/initialization.py | 152 ++- .../omni_connectors/utils/kv_utils.py | 14 +- vllm_omni/engine/__init__.py | 5 + vllm_omni/engine/orchestrator.py | 22 + vllm_omni/engine/stage_client.py | 2 + vllm_omni/engine/stage_engine_core_client.py | 37 + vllm_omni/engine/stage_engine_core_proc.py | 1 + vllm_omni/engine/stage_init_utils.py | 27 +- 33 files changed, 3032 insertions(+), 66 deletions(-) create mode 100644 docs/design/feature/omni_connectors/nixl_connector.md create mode 100644 tests/distributed/omni_connectors/test_nixl_connector.py create mode 100644 tests/distributed/omni_connectors/test_nixl_connector_native.py create mode 100644 vllm_omni/distributed/omni_connectors/connectors/nixl_connector.py diff --git a/docs/.nav.yml b/docs/.nav.yml index 8dac28d22..beb15bf47 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -136,6 +136,7 @@ nav: - Mooncake Store Connector: design/feature/omni_connectors/mooncake_store_connector.md - Mooncake Transfer Engine Connector: design/feature/omni_connectors/mooncake_transfer_engine_connector.md - Mori Transfer Engine Connector: design/feature/omni_connectors/mori_transfer_engine_connector.md + - NIXL Connector: design/feature/omni_connectors/nixl_connector.md - Shared Memory Connector: design/feature/omni_connectors/shared_memory_connector.md - Yuanrong Store Connector: design/feature/omni_connectors/yuanrong_connector.md - Yuanrong Transfer Engine Connector: design/feature/omni_connectors/yuanrong_transfer_engine_connector.md diff --git a/docs/configuration/environment_variables.md b/docs/configuration/environment_variables.md index 4ddc2939c..3a37b52d2 100644 --- a/docs/configuration/environment_variables.md +++ b/docs/configuration/environment_variables.md @@ -90,6 +90,13 @@ depends on the installed kernels and model path. | `VLLM_VIDEO_ASYNC_CHUNK` | `on` or `off`; default `on` | Streaming video output; read on attribute access | Environment-only setting. Values are trimmed and case-normalized; an invalid value warns once and uses `on`. | Experimental | | `VLLM_VIDEO_AUDIO_DELTA_MODE` | `fast` or `slow`; default `fast` | Streaming video audio deltas; read on attribute access | Environment-only setting. Values are trimmed and case-normalized; an invalid value warns once and uses `fast`. | Experimental | +### NIXL stage transfer + +| Name | Type and default | Applies to and read time | Precedence and invalid values | Lifecycle | +| --- | --- | --- | --- | --- | +| `VLLM_OMNI_NIXL_LEASE_S` | Float seconds; default `3600` | Unclaimed producer payload expiry; connector construction | Nonempty environment value overrides `lease_seconds`; empty falls back to config/default. Invalid floats raise `ValueError`. Claimed READ allocations never expire by time. | Experimental | +| `VLLM_OMNI_NIXL_XFER_TIMEOUT_S` | Float seconds; default `300` | Receiver READ wait; connector construction | Nonempty environment value overrides `transfer_timeout_s`; empty falls back to config/default. Invalid floats raise `ValueError`. Timeout does not cancel DMA or release active allocations. | Experimental | + ### Server storage Storage names use Pydantic's nested-settings delimiter (`__`). They are read diff --git a/docs/design/feature/disaggregated_inference.md b/docs/design/feature/disaggregated_inference.md index ba8c8f9fc..8b94325cb 100644 --- a/docs/design/feature/disaggregated_inference.md +++ b/docs/design/feature/disaggregated_inference.md @@ -9,6 +9,7 @@ Backend-specific setup lives in separate docs: - [MooncakeStoreConnector](omni_connectors/mooncake_store_connector.md) - [MooncakeTransferEngineConnector](omni_connectors/mooncake_transfer_engine_connector.md) - [MoriTransferEngineConnector](omni_connectors/mori_transfer_engine_connector.md) +- [NixlConnector](omni_connectors/nixl_connector.md) - [YuanrongConnector](omni_connectors/yuanrong_connector.md) - [YuanrongTransferEngineConnector](omni_connectors/yuanrong_transfer_engine_connector.md) @@ -25,6 +26,7 @@ Current connectors operate in D2H2D (device to host to device) mode. | Multi node (Mooncake Store) | MooncakeStoreConnector | TCP-based, requires Mooncake Master + metadata server. | | Multi node (Mooncake RDMA) | MooncakeTransferEngineConnector | RDMA/TCP direct transfer with managed memory pool. Fastest. | | Multi node (Mori RDMA) | MoriTransferEngineConnector | RDMA direct transfer via Mori IOEngine. | +| Multi node (NIXL) | NixlConnector | NIXL `READ` over UCX; the only backend covering Intel XPU. | | Multi node (Yuanrong) | YuanrongConnector | Requires Yuanrong Datasystem + etcd. | | Ascend NPU P2P (Yuanrong TE) | YuanrongTransferEngineConnector | Uses Yuanrong TransferEngine directly. Configure NPU device IPv4 and `memory_pool_device: "npu"`. | diff --git a/docs/design/feature/omni_connectors/nixl_connector.md b/docs/design/feature/omni_connectors/nixl_connector.md new file mode 100644 index 000000000..7dfd73946 --- /dev/null +++ b/docs/design/feature/omni_connectors/nixl_connector.md @@ -0,0 +1,135 @@ +# NixlConnector + +## When to Use + +Multi-node or intra-node stage transfer over NIXL, which brokers RDMA/shared-memory +transports through a single API. Useful when the deployment already standardises on +NIXL for KV transfer, or when the target accelerator has no Mooncake/Mori backend -- +NIXL reaches Intel XPU through UCX's Level Zero support. + +## Mechanism + +Uses vLLM's `NixlWrapper` (`vllm.distributed.nixl_utils`) to register the producer's +tensors and let the consumer pull them with a NIXL `READ`. + +- Data Plane: NIXL agent-to-agent `READ`, GPU-to-GPU where the backend allows it. +- Control Plane: a ZMQ ROUTER socket atomically claims the source allocation before + returning descriptors. Forwarded `put()` metadata includes the generation and + endpoint; the consumer claims that generation before submitting any READ. + +Transfer metadata uses schema version 1. Tensor descriptors are grouped by NIXL +memory type (`DRAM`, `VRAM`, and so on), while each descriptor retains its global +tensor index so mixed CPU/accelerator payloads can be reconstructed in order. + +Payloads are not restricted to tensors. A single tensor or a list of tensors is +transferred as-is; a nested structure has its tensor leaves extracted and shipped +alongside a msgpack-encoded skeleton; anything else is msgpack-encoded into one +uint8 tensor. Mapping and sequence structure is preserved. `msgspec.Struct` values, +including `OmniPayloadStruct`, are normalised to mappings so their tensor leaves use +the direct NIXL path rather than CPU object serialisation. + +Each NIXL agent uses strict thread synchronization. A consumer performs one short +metadata query per scheduler poll (10ms by default), which prevents an unavailable +producer from blocking unrelated requests. The regular transfer-completion handshake +keeps the configured transfer timeout. + +## Installation + +CUDA hosts can install the published wheel: + +```bash +pip install nixl +``` + +For Intel XPU, use `docker/Dockerfile.xpu`. Its upstream vLLM XPU base image +provides a matching torch-xpu, UCX with Level Zero support, and NIXL runtime. + +## Configuration + +```yaml +connectors: + nixl_connector: + name: NixlConnector + extra: + host: "auto" + zmq_port: 50061 + backends: ["UCX"] + +stages: + - stage_id: 0 + output_connectors: + to_stage_1: nixl_connector + + - stage_id: 1 + input_connectors: + from_stage_0: nixl_connector +``` + +Parameters: + +- `host`: address the producer binds its handshake socket to (`"auto"` to detect). +- `zmq_port`: handshake port. If a producer omits it, an ephemeral port is opened + for ownership claims and returned in `put()` metadata. The value in the deploy YAML is a + base port. Purpose, replica, rank, and producer-stage offsets are added centrally, + so colocated connectors and tensor-parallel workers receive distinct endpoints + without one config entry per rank. +- `sender_host` / `sender_zmq_port`: consumer-side override naming the producer's + handshake endpoint. Only needed when the consumer cannot learn it from metadata. + In replicated deployments, the producer endpoint is carried with each request and + selected at `get()` time; it is not mutable connector-wide state. + Explicit sender ports are preserved, not recomputed using the receiver replica. + Intermediate NIXL stages retain their incoming receiver configuration and a + separate outgoing bind configuration, so one worker can receive and publish. +- `backends`: NIXL backends to register memory with. Defaults to `["UCX"]`. +- `receive_device`: forces where received tensors land. By default the consumer + keeps the producer's device *type* but uses its own current device of that type. +- `memory_type`: overrides the NIXL memory type, otherwise `DRAM` for CPU tensors + and `VRAM` for accelerator tensors. +- `lease_seconds`: how long a `put()` payload stays registered while waiting to be + read (default 3600). The consumer reports completion, so this only bounds payloads + nobody has claimed. An internal reaper releases only unclaimed registrations. + `VLLM_OMNI_NIXL_LEASE_S` overrides it. +- `transfer_timeout_s`: how long a `get()` waits for its `READ` to complete + (default 300). NIXL 1.3 has no transfer cancellation API, so a timed-out transfer's + buffers and registrations remain owned by the connector until NIXL reports a + terminal state; `close()` waits for that state before releasing them. + `VLLM_OMNI_NIXL_XFER_TIMEOUT_S` overrides it. + +### Source ownership and failure limits + +Claims and expiry use the same lock. Completion identifies both a payload +generation and a unique READ claim; duplicate or stale ACKs cannot release a +different reader or replacement payload. Replacing a claimed key is rejected. +Terminal errors also release claims; timeout, unknown status, and active sibling +transfers retain the complete destination bundle and source claim until terminal. +Zero-element tensors retain their shape/dtype/skeleton positions but never create +DMA descriptors or registrations. + +A lost metadata response, lost completion ACK, or abandoned consumer can retain a +claim indefinitely. NIXL 1.3 cannot prove remote cancellation, so neither TTL nor +`cleanup()` frees those allocations. Producer `close()` rejects new work and +retains its agent, listener and claimed allocations; a later `close()` finishes +teardown after claims drain. Permanently abandoned claims remain until process +exit. This favors memory safety over bounded shutdown/memory usage. Only trusted +peers may access this unauthenticated control plane. Both endpoints must use the +claim-aware protocol; rolling compatibility with older GET_META clients is not +provided. Legacy externally-owned metadata without a generation is accepted only +under the caller's lifetime guarantee, not managed as a leased allocation here. + +## Validation + +The native CUDA smoke test requires NIXL and at least two visible CUDA devices: + +```bash +UCX_RCACHE_MAX_UNRELEASED=1024 pytest -o addopts='' -q -s \ + tests/distributed/omni_connectors/test_nixl_connector_native.py +``` + +It starts independent producer and consumer processes, loads the real UCX backend, +performs the ZMQ `GET_META` handshake and a native NIXL `READ` from GPU 0 to GPU 1, +checks mixed CPU/GPU tensor leaves from an `OmniPayloadStruct`, sends `XFER_DONE`, +and verifies that producer registrations are released. This validates same-host, +cross-GPU operation. Inter-node UCX transport is supported by the design but was not +validated by this test. + +For more details, refer to the [NIXL repository](https://github.com/ai-dynamo/nixl). diff --git a/docs/design/index.md b/docs/design/index.md index 6481aa49b..f70571cb5 100644 --- a/docs/design/index.md +++ b/docs/design/index.md @@ -36,6 +36,7 @@ implementation contract; it is not, by itself, a general support claim. - [Mooncake Store Connector](feature/omni_connectors/mooncake_store_connector.md) - [Mooncake Transfer Engine Connector](feature/omni_connectors/mooncake_transfer_engine_connector.md) - [Mori Transfer Engine Connector](feature/omni_connectors/mori_transfer_engine_connector.md) +- [NIXL Connector](feature/omni_connectors/nixl_connector.md) - [Shared Memory Connector](feature/omni_connectors/shared_memory_connector.md) - [Yuanrong Store Connector](feature/omni_connectors/yuanrong_connector.md) - [Yuanrong Transfer Engine Connector](feature/omni_connectors/yuanrong_transfer_engine_connector.md) diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index 911cfbc80..d958bbf18 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -164,7 +164,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): """Make an inventory expansion an explicit review decision.""" category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values()) assert category_counts == { - EnvironmentVariableCategory.PUBLIC_OMNI: 25, + EnvironmentVariableCategory.PUBLIC_OMNI: 27, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27, EnvironmentVariableCategory.MODEL_SPECIFIC: 63, diff --git a/tests/core/sched/test_omni_scheduling_coordinator.py b/tests/core/sched/test_omni_scheduling_coordinator.py index ae1b3924e..f37aac29f 100644 --- a/tests/core/sched/test_omni_scheduling_coordinator.py +++ b/tests/core/sched/test_omni_scheduling_coordinator.py @@ -58,6 +58,7 @@ def _make_request(req_id: str, status: str = "waiting") -> SimpleNamespace: num_computed_tokens=0, _all_token_ids=[], _output_token_ids=[], + payload_sender_info=None, ) @@ -270,6 +271,19 @@ class TestWaitingForInputTransition(unittest.TestCase): self.assertEqual(len(coord.pending_input_registrations), 1) self.assertEqual(coord.pending_input_registrations[0].request_id, "r1") + def test_pending_input_registration_carries_payload_sender_info(self): + coord = OmniSchedulingCoordinator(stage_id=1) + req = _make_request("r1", status=RequestStatus.WAITING_FOR_INPUT) + req.payload_sender_info = {"host": "10.0.0.1", "zmq_port": 50051} + waiting = MockQueue([req]) + + coord.process_pending_full_payload_inputs(waiting, set()) + + self.assertEqual( + coord.pending_input_registrations[0].payload_sender_info, + {"host": "10.0.0.1", "zmq_port": 50051}, + ) + def test_idle_cycles_retain_received_marker_before_request_appears(self): coord = OmniSchedulingCoordinator(stage_id=1) coord._full_payload_input_received.add("late") diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index a80ded5f0..92c2a7343 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -4,7 +4,7 @@ import threading from collections import deque from types import MethodType, SimpleNamespace -from unittest.mock import Mock, patch +from unittest.mock import Mock, call, patch import pytest import torch @@ -543,7 +543,7 @@ def test_create_connector_config_parsing(monkeypatch, raw_cfg, expected_name, ex assert connector == "ok" assert isinstance(captured["spec"], ConnectorSpec) assert captured["spec"].name == expected_name - assert captured["spec"].extra == expected_extra + assert captured["spec"].extra == {**expected_extra, "stage_id": 0} def test_load_poll(build_adapter): @@ -570,8 +570,51 @@ def test_load_poll(build_adapter): assert request.additional_information == payload assert adapter.get_req_chunk["req-1"] == 1 assert "req-1" not in adapter._finished_load_reqs - assert "req-1" in adapter.upstream_exhausted_requests - assert "req-1" not in adapter._pending_load_reqs + + +def test_load_poll_uses_request_scoped_payload_sender_endpoint(build_adapter): + adapter, connector = build_adapter(stage_id=2, model_mode="ar") + request = _req("req-1", RequestStatus.WAITING, external_req_id="external-1") + request.payload_sender_info = {"host": "10.0.0.1", "zmq_port": 50051} + connector.get.return_value = None + + adapter._poll_single_request(_dequeue_load_entry(adapter, request)) + + connector.get.assert_called_once_with( + "1", + "2", + "req-1_1_0", + {"source_host": "10.0.0.1", "source_port": 50051}, + ) + + +def test_load_poll_keeps_concurrent_payload_sender_endpoints_distinct(build_adapter): + adapter, connector = build_adapter(stage_id=2, model_mode="ar") + first = _req("req-1", RequestStatus.WAITING) + first.payload_sender_info = {"host": "10.0.0.1", "zmq_port": 50051} + second = _req("req-2", RequestStatus.WAITING) + second.payload_sender_info = {"host": "10.0.0.2", "zmq_port": 51051} + connector.get.return_value = None + + second_entry = _dequeue_load_entry(adapter, second) + first_entry = _dequeue_load_entry(adapter, first) + adapter._poll_single_request(second_entry) + adapter._poll_single_request(first_entry) + + assert connector.get.call_args_list == [ + call( + "1", + "2", + "req-2_1_0", + {"source_host": "10.0.0.2", "source_port": 51051}, + ), + call( + "1", + "2", + "req-1_1_0", + {"source_host": "10.0.0.1", "source_port": 50051}, + ), + ] def test_load_async_does_not_requeue_registered_inflight_request(build_adapter): @@ -2770,7 +2813,8 @@ def test_successful_put_keeps_terminal_chunk_sequence(build_adapter, finish_insi scheduler.add_request(final_update) return True, 1, {} - def get_by_key(_from_stage, _to_stage, key): + def get_by_key(_from_stage, _to_stage, key, metadata=None): + assert metadata is None payload = payloads.pop(key, None) return (payload, 1) if payload is not None else None diff --git a/tests/distributed/omni_connectors/test_nixl_connector.py b/tests/distributed/omni_connectors/test_nixl_connector.py new file mode 100644 index 000000000..5b2161eee --- /dev/null +++ b/tests/distributed/omni_connectors/test_nixl_connector.py @@ -0,0 +1,891 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +""" +Unit tests for the NixlConnector metadata handshake. + +NIXL itself is stubbed out, so these tests exercise the ZMQ control plane and +the payload normalisation rather than any actual RDMA transfer. +""" + +import ctypes +import sys +import time +import types + +import msgspec +import pytest +import torch +import zmq + +from vllm_omni.data_entry_keys import HiddenStatesStruct, MetaStruct, OmniPayloadStruct + +pytestmark = [pytest.mark.cpu, pytest.mark.parallel, pytest.mark.core_model] + +PORT = 47431 + + +def _wait_for_metadata(consumer, key, metadata=None, *, timeout=2.0): + """Retry outside the connector's deliberately bounded metadata query.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + resolved = consumer._resolve_metadata(key, metadata) + if resolved is not None: + return resolved + time.sleep(0.01) + pytest.fail(f"Timed out waiting for metadata for {key!r}") + + +@pytest.mark.parametrize("misses", [0, 2]) +def test_wait_for_metadata_retries_until_success(monkeypatch, misses): + expected = {"claim_id": "claim"} + forwarded = {"source_host": "127.0.0.1", "source_port": PORT} + calls = [] + + def resolve(key, metadata): + calls.append((key, metadata)) + return None if len(calls) <= misses else expected + + ticks = iter(range(10)) + monkeypatch.setattr(time, "monotonic", lambda: next(ticks)) + monkeypatch.setattr(time, "sleep", lambda _: None) + consumer = types.SimpleNamespace(_resolve_metadata=resolve) + assert _wait_for_metadata(consumer, "retry", forwarded, timeout=5) is expected + assert calls == [("retry", forwarded)] * (misses + 1) + + +def test_wait_for_metadata_stops_at_deadline(monkeypatch): + calls = [] + ticks = iter([0.0, 0.0, 1.0, 2.0]) + monkeypatch.setattr(time, "monotonic", lambda: next(ticks)) + monkeypatch.setattr(time, "sleep", lambda _: None) + consumer = types.SimpleNamespace(_resolve_metadata=lambda *args: calls.append(args)) + with pytest.raises(pytest.fail.Exception, match="Timed out waiting for metadata for 'missing'"): + _wait_for_metadata(consumer, "missing") + assert calls == [("missing", None)] * 2 + + +@pytest.mark.parametrize("async_chunk", [False, True]) +def test_three_stage_incoming_and_outgoing_endpoints(nixl_connector_cls, async_chunk): + from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec, OmniTransferConfig + from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_connector_spec + from vllm_omni.engine.stage_init_utils import get_stage_connector_spec + + config = OmniTransferConfig( + connectors={ + ("0", "1"): ConnectorSpec(name="NixlConnector", extra={"host": "127.0.0.1", "zmq_port": PORT}), + ("1", "2"): ConnectorSpec(name="NixlConnector", extra={"host": "127.0.0.1", "zmq_port": PORT + 10}), + } + ) + connectors = [] + try: + for stage in range(3): + spec = get_stage_connector_spec(config, stage, async_chunk) + resolved = resolve_connector_spec(ConnectorSpec(**spec), stage_id=stage, role=spec["extra"]["role"]) + connectors.append(nixl_connector_cls(resolved.extra)) + first, middle, last = connectors + assert middle._sender_zmq_port == PORT + assert middle._serving_handshake + assert middle._zmq_port == PORT + 11 + first.put("0", "1", "incoming", torch.ones(1)) + incoming = _wait_for_metadata(middle, "incoming") + assert incoming["sender_zmq_port"] == PORT + assert incoming["generation"] == first._published["incoming"]["generation"] + middle.put("1", "2", "outgoing", torch.ones(1)) + outgoing = _wait_for_metadata(last, "outgoing") + assert outgoing["sender_zmq_port"] == PORT + 11 + assert outgoing["generation"] == middle._published["outgoing"]["generation"] + finally: + for connector in reversed(connectors): + for pending in connector._pending.values(): + pending.claims.clear() # No DMA in this metadata-only probe. + connector.close() + + +@pytest.mark.parametrize("name", ["NixlConnector", "MooncakeTransferEngineConnector"]) +def test_explicit_sender_port_is_not_derived(name): + from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec + from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_connector_spec + + spec = ConnectorSpec(name=name, extra={"sender_zmq_port": 55000, "from_stage": 0}) + resolved = resolve_connector_spec(spec, stage_id=1, role="receiver", replica_id=3, local_rank=2) + assert resolved.extra["sender_zmq_port"] == 55000 + assert spec.extra == {"sender_zmq_port": 55000, "from_stage": 0} + spec.extra["zmq_port"] = "${UNUSED_NIXL_BASE_PORT}" + assert resolve_connector_spec(spec, stage_id=1, role="receiver").extra["sender_zmq_port"] == 55000 + + +@pytest.mark.parametrize("need_recv_cache", [False, True]) +@pytest.mark.parametrize("update_before_init", [False, True]) +def test_manager_receiver_does_not_bind_shared_producer_port(producer, need_recv_cache, update_before_init): + """Payload-only and KV receivers must dial, not bind, the incoming edge.""" + from vllm_omni.distributed.omni_connectors.kv_transfer_manager import ( + OmniKVCacheConfig, + OmniKVTransferManager, + ) + + manager = OmniKVTransferManager( + OmniKVCacheConfig( + connector_config={ + "type": "NixlConnector", + "role": "receiver", + "host": "127.0.0.1", + "zmq_port": PORT, + "backends": ["UCX"], + }, + from_stage="0", + to_stage="1", + stage_id=1, + need_recv_cache=need_recv_cache, + ) + ) + sender_info = {"host": producer.host, "zmq_port": producer._zmq_port} + if update_before_init: + manager.update_sender_info(sender_info) + receiver = manager.connector + try: + assert receiver is not None, "Receiver incorrectly tried to bind the producer's occupied port" + assert receiver._zmq_port is None + assert not receiver._serving_handshake + if not update_before_init: + assert receiver._sender_zmq_port is None + manager.update_sender_info(sender_info) + assert receiver._sender_zmq_port == producer._zmq_port + assert manager.config.connector_config["zmq_port"] == PORT # Do not mutate the shared specification. + success, _, _ = producer.put("0", "1", "manager-receive", torch.ones(1)) + assert success + metadata = _wait_for_metadata(receiver, "manager-receive") + assert metadata["generation"] == producer._published["manager-receive"]["generation"] + receiver._notify_transfer_done("manager-receive", metadata) + assert not producer._pending + finally: + if receiver is not None: + receiver.close() + + +@pytest.mark.parametrize("need_recv_cache", [False, True]) +def test_manager_receiver_preserves_explicit_standalone_sender(producer, need_recv_cache): + from vllm_omni.distributed.omni_connectors.kv_transfer_manager import ( + OmniKVCacheConfig, + OmniKVTransferManager, + ) + + manager = OmniKVTransferManager( + OmniKVCacheConfig( + connector_config={ + "type": "NixlConnector", + "role": "receiver", + "zmq_port": PORT, + "sender_host": producer.host, + "sender_zmq_port": producer._zmq_port, + }, + need_recv_cache=need_recv_cache, + ) + ) + receiver = manager.connector + try: + assert receiver is not None + assert receiver._zmq_port is None and not receiver._serving_handshake + assert (receiver._sender_host, receiver._sender_zmq_port) == (producer.host, producer._zmq_port) + success, _, _ = producer.put("0", "1", "standalone-receive", torch.ones(1)) + assert success + metadata = _wait_for_metadata(receiver, "standalone-receive") + assert metadata["generation"] == producer._published["standalone-receive"]["generation"] + receiver._notify_transfer_done("standalone-receive", metadata) + assert not producer._pending + finally: + if receiver is not None: + receiver.close() + + +def test_middle_stage_advertises_its_outgoing_replica_endpoint(): + from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient + + client = object.__new__(StageEngineCoreClient) + client.stage_id = 1 + client.replica_id = 3 + client.vllm_config = types.SimpleNamespace( + model_config=types.SimpleNamespace( + stage_connector_config={ + "name": "NixlConnector", + "extra": { + "role": "receiver", + "host": "10.0.0.1", + "zmq_port": 47000, + "from_stage": 0, + "outgoing": {"host": "10.0.0.2", "zmq_port": 48000, "from_stage": 1}, + }, + } + ) + ) + assert client._build_payload_sender_info() == {"host": "10.0.0.2", "zmq_port": 51073} + + +@pytest.mark.parametrize("direct", [False, True]) +@pytest.mark.usefixtures("reliable_claim_queries") +def test_claimed_source_survives_expiry_and_cleanup(producer, consumer, direct): + _, _, metadata = producer.put("0", "1", "claimed", torch.ones(1)) + resolved = consumer._resolve_metadata("claimed", metadata if direct else None) + pending = producer._pending["claimed"] + pending.deadline = 0 + producer._cleanup_expired_pending() + assert producer._pending.get("claimed") is pending + producer.cleanup("claimed") + assert producer._pending.get("claimed") is pending + assert producer._agent.registered + consumer._notify_transfer_done("claimed", resolved) + consumer._notify_transfer_done("claimed", resolved) + assert producer._pending == {} + assert producer._agent.registered == [] + + +@pytest.mark.usefixtures("reliable_claim_queries") +def test_claims_are_generation_scoped_and_duplicate_ack_cannot_release_sibling(producer, consumer): + _, _, original = producer.put("0", "1", "owners", torch.ones(1)) + first = consumer._resolve_metadata("owners", original) + second = consumer._resolve_metadata("owners", original) + assert first["claim_id"] != second["claim_id"] + assert not producer.put("0", "1", "owners", torch.zeros(1))[0] + consumer._notify_transfer_done("owners", first) + consumer._notify_transfer_done("owners", first) + assert len(producer._pending["owners"].claims) == 1 + consumer._notify_transfer_done("owners", second) + assert not producer._pending + producer.put("0", "1", "owners", torch.zeros(1)) + consumer._notify_transfer_done("owners", second) + assert "owners" in producer._pending + assert consumer._resolve_metadata("owners", original) is None + + +@pytest.mark.parametrize("get_metadata", [None, {"schema_version": 1, "tensor_specs": [], "descriptor_groups": []}]) +@pytest.mark.usefixtures("reliable_claim_queries") +def test_abandoned_claim_survives_close_and_late_completion(producer, consumer, get_metadata): + _, _, metadata = producer.put("0", "1", "abandoned", torch.ones(1)) + claimed = consumer._resolve_metadata("abandoned", metadata) + producer._pending["abandoned"].deadline = 0 + producer.close() + assert producer._closing and not producer._closed + assert producer._agent.registered + assert producer._listener_thread.is_alive() + with pytest.raises(RuntimeError, match="closed"): + producer.put("0", "1", "new", torch.ones(1)) + metrics = dict(producer._metrics) + with pytest.raises(RuntimeError, match="Cannot get data: NixlConnector is closed"): + producer.get("0", "1", "new", get_metadata) + assert producer._metrics == metrics + assert producer._pending["abandoned"].claims == {claimed["claim_id"]} + assert consumer._resolve_metadata("abandoned", metadata) is None + consumer._notify_transfer_done("abandoned", claimed) + assert not producer._pending + assert not producer._agent.registered + assert producer._listener_thread.is_alive() + producer.close() + assert not producer._agent.registered + assert producer._closed + + +@pytest.fixture +def copying_native_agent(consumer, monkeypatch): + """Strict fake descriptors and actual CPU copies; not native NIXL evidence.""" + calls = [] + + def descriptors(regions, memory_type): + assert regions and all(region[0] > 0 and region[1] > 0 for region in regions) + return regions + + def prepare(operation, local, local_ids, remote, remote_ids): + assert operation == "READ" + assert local_ids == remote_ids == list(range(len(local))) + return list(zip(local, remote, strict=True)) + + def transfer(pairs): + calls.append(pairs) + for destination, source in pairs: + assert destination[1] == source[1] > 0 + ctypes.memmove(destination[0], source[0], source[1]) + + agent = consumer._agent + monkeypatch.setattr(agent, "add_remote_agent", lambda metadata: "producer", raising=False) + monkeypatch.setattr(agent, "get_xfer_descs", descriptors, raising=False) + monkeypatch.setattr(agent, "prep_xfer_dlist", lambda agent, descs: descs, raising=False) + monkeypatch.setattr(agent, "make_prepped_xfer", prepare, raising=False) + monkeypatch.setattr(agent, "transfer", transfer, raising=False) + monkeypatch.setattr(agent, "check_xfer_state", lambda handle: "DONE", raising=False) + for method in ("release_xfer_handle", "release_dlist_handle", "remove_remote_agent"): + monkeypatch.setattr(agent, method, lambda handle: None, raising=False) + return calls + + +@pytest.mark.parametrize("case", ["empty", "all_empty", "mixed", "scalar", "structured_empty", "structured_mixed"]) +def test_zero_byte_leaves_roundtrip_without_native_descriptors(producer, consumer, copying_native_agent, case): + empty = torch.empty((2, 0, 3), dtype=torch.float64) + other = torch.empty((0,), dtype=torch.int64) + scalar = torch.tensor(7) + vector = torch.arange(6, dtype=torch.float32) + payload = { + "empty": empty, + "all_empty": [empty, other], + "mixed": [empty, scalar, other, vector], + "scalar": scalar, + "structured_empty": {"leaves": (empty, [other]), "label": "empty"}, + "structured_mixed": {"leaves": (empty, [scalar, other, vector]), "label": "mixed"}, + }[case] + ok, size, metadata = producer.put("0", "1", "empty-slots", payload) + assert ok + indices = [index for group in metadata["descriptor_groups"] for index in group["tensor_indices"]] + assert sorted(indices) == [i for i, spec in enumerate(metadata["tensor_specs"]) if spec["size"]] + actual, received_size = consumer.get("0", "1", "empty-slots", metadata) + if isinstance(payload, dict): + assert actual["label"] == payload["label"] + torch.testing.assert_close(actual["leaves"], payload["leaves"], rtol=0, atol=0) + else: + torch.testing.assert_close(actual, payload, rtol=0, atol=0) + assert size == received_size + if case in ("empty", "all_empty"): + assert not copying_native_agent + assert not producer._pending and not producer._agent.registered + assert not consumer._agent.registered + + +@pytest.mark.parametrize("direct", [False, True]) +@pytest.mark.parametrize("outcome", ["done", "error", "timeout", "unknown"]) +def test_read_ownership_through_terminal_and_deferred_paths( + producer, consumer, copying_native_agent, monkeypatch, direct, outcome +): + import threading + + _, _, metadata = producer.put("0", "1", "active-read", torch.ones(2)) + entered, proceed = threading.Event(), threading.Event() + state = ["PROC"] + monkeypatch.setattr(consumer._agent, "check_xfer_state", lambda handle: state[0]) + + def wait(handle, key): + entered.set() + assert proceed.wait(5) + if outcome == "done": + state[0] = "DONE" + elif outcome == "error": + state[0] = "ERR" + raise RuntimeError("terminal error") + else: + if outcome == "unknown": + state[0] = "UNKNOWN" + raise TimeoutError("still owned") + + monkeypatch.setattr(consumer, "_wait_for_transfer", wait) + worker = threading.Thread(target=consumer.get, args=("0", "1", "active-read", metadata if direct else None)) + worker.start() + try: + assert entered.wait(5) + pending = producer._pending["active-read"] + pending.deadline = 0 + producer._cleanup_expired_pending() + producer.cleanup("active-read") + assert producer._pending["active-read"] is pending + assert producer._agent.registered + finally: + proceed.set() + worker.join(5) + assert not worker.is_alive() + if outcome in ("timeout", "unknown"): + assert producer._agent.registered + assert consumer._deferred_transfers + consumer._reap_deferred_transfers() + assert producer._agent.registered + state[0] = "DONE" + deadline = time.monotonic() + 5 + while producer._pending and time.monotonic() < deadline: + time.sleep(0.01) + assert not producer._pending + assert not producer._agent.registered + + +class _FakeNixlAgent: + def __init__(self, name, config=None): + self.name = name + self.config = config + self.registered = [] + + def get_reg_descs(self, regions, memory_type): + return ("reg", tuple(regions), memory_type) + + def register_memory(self, descs, backends=None): + self.registered.append(descs) + + def deregister_memory(self, descs): + if descs in self.registered: + self.registered.remove(descs) + + def get_agent_metadata(self): + return f"agent-metadata-{self.name}".encode() + + +@pytest.fixture +def nixl_connector_cls(monkeypatch): + """Import NixlConnector with vLLM's optional NIXL dependency stubbed out.""" + stub = types.ModuleType("vllm.distributed.nixl_utils") + stub.NixlWrapper = _FakeNixlAgent # type: ignore[attr-defined] + config_module = types.ModuleType("test_nixl_config_api") + config_sync_type = types.SimpleNamespace(NIXL_THREAD_SYNC_STRICT="strict") + config_module.nixl_thread_sync_t = config_sync_type # type: ignore[attr-defined] + + def agent_config(**kwargs): + return kwargs + + agent_config.__module__ = config_module.__name__ + config_module.nixl_agent_config = agent_config # type: ignore[attr-defined] + stub.nixl_agent_config = agent_config # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "vllm.distributed.nixl_utils", stub) + monkeypatch.setitem(sys.modules, config_module.__name__, config_module) + sync_type = types.SimpleNamespace(NIXL_THREAD_SYNC_STRICT="wrong-public-enum") + nixl_module = types.ModuleType("nixl") + nixl_module.nixl_thread_sync_t = sync_type # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "nixl", nixl_module) + + from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import NixlConnector + + return NixlConnector + + +@pytest.fixture +def producer(nixl_connector_cls): + connector = nixl_connector_cls({"role": "sender", "host": "127.0.0.1", "zmq_port": PORT}) + yield connector + # These control-plane tests do not submit DMA. Explicitly complete claims + # made by metadata-only probes before tearing down the shared port. + for key, pending in list(connector._pending.items()): + for claim in list(pending.claims): + from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import _XFER_DONE_MSG + + connector._handle_handshake_message( + _XFER_DONE_MSG + + msgspec.msgpack.encode({"key": key, "generation": pending.generation, "claim_id": claim}) + ) + connector.close() + + +@pytest.fixture +def consumer(nixl_connector_cls): + connector = nixl_connector_cls({"role": "receiver", "sender_host": "127.0.0.1", "sender_zmq_port": PORT}) + yield connector + connector.close() + + +@pytest.fixture +def reliable_claim_queries(producer, consumer, monkeypatch): + """Ownership probes need exact claims, not claims retained after lost replies. + + Exercise the real resolver, wire encoding and producer handler synchronously; + discovery tests separately cover bounded queries over real ZMQ sockets. + """ + import uuid + + from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import _GET_META_MSG, _META_NOT_FOUND + + def query(key, host, port, *, generation=None): + assert (host, port) == (producer.host, producer._zmq_port) + request = {"key": key, "generation": generation, "claim_id": uuid.uuid4().hex} + reply = producer._handle_handshake_message(_GET_META_MSG + msgspec.msgpack.encode(request)) + return None if reply == _META_NOT_FOUND else msgspec.msgpack.decode(reply) + + monkeypatch.setattr(consumer, "_query_metadata_at", query) + + +def test_put_publishes_its_handshake_endpoint(producer): + ok, size, metadata = producer.put("0", "1", "req-0", torch.arange(8, dtype=torch.float32)) + + assert ok is True + assert size == 32 + assert metadata["sender_host"] == "127.0.0.1" + assert metadata["sender_zmq_port"] == PORT + + +def test_agent_uses_strict_thread_synchronization(producer): + assert producer._agent.config["sync_mode"] == "strict" + + +def test_handshake_serves_metadata_when_caller_has_none(producer, consumer): + _, _, published = producer.put("0", "1", "req-1", torch.arange(4, dtype=torch.float32)) + + resolved = _wait_for_metadata(consumer, "req-1") + + # msgpack has no tuple type, so region descriptors arrive as lists; get() + # re-tuples them before handing them to NIXL. + assert resolved.pop("claim_id") + assert resolved == msgspec.msgpack.decode(msgspec.msgpack.encode(published)) + assert [tuple(region) for group in resolved["descriptor_groups"] for region in group["regions"]] == [ + region for group in published["descriptor_groups"] for region in group["regions"] + ] + + +def test_legacy_metadata_without_generation_skips_the_handshake(consumer): + """Externally owned metadata without a generation needs no ownership claim.""" + direct = {"schema_version": 1, "kind": "tensors", "tensor_specs": []} + + assert consumer._resolve_metadata("req-2", direct) is direct + + +def test_source_endpoint_metadata_overrides_configured_sender(consumer, monkeypatch): + queried = [] + expected = {"schema_version": 1, "kind": "tensors", "tensor_specs": []} + + def query(get_key, host, port): + queried.append((get_key, host, port)) + return expected + + monkeypatch.setattr(consumer, "_query_metadata_at", query) + + resolved = consumer._resolve_metadata( + "req-tp4-rank3", + {"source_host": "10.0.0.3", "source_port": PORT + 3 * 16}, + ) + + assert resolved is expected + assert queried == [("req-tp4-rank3", "10.0.0.3", PORT + 3 * 16)] + + +def test_unknown_key_is_queried_once(producer, consumer, monkeypatch): + socket = consumer._get_req_socket(f"tcp://127.0.0.1:{PORT}") + send_count = 0 + original_send = socket.send + + def count_send(message): + nonlocal send_count + send_count += 1 + return original_send(message) + + monkeypatch.setattr(socket, "send", count_send) + + assert consumer._resolve_metadata("never-published", None) is None + assert send_count == 1 + assert socket.getsockopt(zmq.RCVTIMEO) == 10 + + +@pytest.mark.usefixtures("reliable_claim_queries") +def test_transfer_done_releases_the_producer_buffer(producer, consumer): + _, _, metadata = producer.put("0", "1", "req-3", torch.arange(4, dtype=torch.float32)) + assert producer._pending and producer._agent.registered + + metadata = consumer._resolve_metadata("req-3", metadata) + consumer._notify_transfer_done("req-3", metadata) + + assert producer._pending == {} + assert producer._published == {} + assert producer._agent.registered == [] + + +def test_direct_metadata_has_ephemeral_ownership_endpoint(nixl_connector_cls): + connector = nixl_connector_cls({}) + try: + assert connector._zmq_ctx is not None + assert connector._listener_thread is not None + _, _, metadata = connector.put("0", "1", "req-4", torch.zeros(2)) + assert metadata["sender_host"] + assert metadata["sender_zmq_port"] > 0 + finally: + connector.close() + + +def test_idle_producer_lease_expires_without_another_put(nixl_connector_cls): + connector = nixl_connector_cls({"role": "sender", "lease_seconds": 0.01}) + try: + connector.put("0", "1", "req-expire", torch.zeros(2)) + + deadline = time.monotonic() + 1.0 + while connector._pending and time.monotonic() < deadline: + time.sleep(0.01) + + assert connector._pending == {} + assert connector._published == {} + assert connector._agent.registered == [] + assert connector._registered_descs == [] + finally: + connector.close() + + +def test_close_stops_lease_reaper(nixl_connector_cls): + connector = nixl_connector_cls({"role": "sender"}) + lease_thread = connector._lease_thread + + connector.close() + + assert lease_thread is not None + assert not lease_thread.is_alive() + assert connector._lease_thread is None + + +def test_deferred_transfer_is_retained_while_active(nixl_connector_cls): + from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import _DeferredTransfer + + connector = nixl_connector_cls({"role": "receiver"}) + released = [] + connector._agent.check_xfer_state = lambda handle: "PROC" + connector._agent.release_xfer_handle = lambda handle: released.append(("handle", handle)) + connector._agent.release_dlist_handle = lambda handle: released.append(("dlist", handle)) + connector._agent.remove_remote_agent = lambda agent: released.append(("agent", agent)) + transfer = _DeferredTransfer( + tensors=[torch.zeros(1)], + registrations=["registration"], + dlists=["local", "remote"], + handles=["transfer"], + remote_agent="producer", + ) + connector._defer_transfer(transfer) + try: + connector._reap_deferred_transfers() + + assert released == [] + assert transfer in connector._deferred_transfers + assert transfer.tensors + finally: + connector._agent.check_xfer_state = lambda handle: "DONE" + connector.close() + + +def test_deferred_transfer_releases_exactly_once_after_done(nixl_connector_cls): + from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import _DeferredTransfer + + connector = nixl_connector_cls({"role": "receiver"}) + released = [] + connector._agent.check_xfer_state = lambda handle: "DONE" + connector._agent.release_xfer_handle = lambda handle: released.append(("handle", handle)) + connector._agent.release_dlist_handle = lambda handle: released.append(("dlist", handle)) + connector._agent.remove_remote_agent = lambda agent: released.append(("agent", agent)) + connector._agent.deregister_memory = lambda descs: released.append(("registration", descs)) + transfer = _DeferredTransfer( + tensors=[torch.zeros(1)], + registrations=["registration"], + dlists=["local", "remote"], + handles=["transfer"], + remote_agent="producer", + ) + connector._defer_transfer(transfer) + try: + connector._reap_deferred_transfers() + connector._reap_deferred_transfers() + + assert released == [ + ("handle", "transfer"), + ("dlist", "local"), + ("dlist", "remote"), + ("agent", "producer"), + ("registration", "registration"), + ] + assert connector._deferred_transfers == [] + assert transfer.tensors == [] + finally: + connector.close() + + +def test_active_sibling_transfer_requires_deferred_ownership(nixl_connector_cls): + connector = nixl_connector_cls({"role": "receiver"}) + connector._agent.check_xfer_state = lambda handle: {"failed": "ERR", "active": "PROC"}[handle] + try: + assert connector._transfers_may_be_active(["failed", "active"]) is True + assert connector._transfers_may_be_active(["failed"]) is False + finally: + connector.close() + + +def test_get_defers_complete_ownership_when_sibling_transfer_is_active(nixl_connector_cls): + connector = nixl_connector_cls({"role": "receiver"}) + connector._agent.add_remote_agent = lambda metadata: "producer" + connector._agent.get_xfer_descs = lambda regions, memory_type: (regions, memory_type) + connector._agent.prep_xfer_dlist = lambda agent, descs: (agent, tuple(descs[0])) + handles = iter(["failed", "active"]) + connector._agent.make_prepped_xfer = lambda *_args: next(handles) + connector._agent.transfer = lambda handle: None + connector._agent.check_xfer_state = lambda handle: {"failed": "ERR", "active": "PROC"}[handle] + released = [] + connector._agent.release_xfer_handle = lambda handle: released.append(("handle", handle)) + connector._agent.release_dlist_handle = lambda handle: released.append(("dlist", handle)) + connector._agent.remove_remote_agent = lambda agent: released.append(("agent", agent)) + connector._agent.deregister_memory = lambda descs: released.append(("registration", descs)) + metadata = { + "schema_version": 1, + "kind": "tensors", + "agent_metadata": b"producer-metadata", + "tensor_specs": [ + {"shape": [1], "dtype": "torch.float32", "device": "cpu", "size": 4}, + {"shape": [1], "dtype": "torch.float32", "device": "cpu", "size": 4}, + ], + "descriptor_groups": [ + {"memory_type": "DRAM", "tensor_indices": [0], "regions": [(1, 4, 0, "")]}, + {"memory_type": "DRAM", "tensor_indices": [1], "regions": [(2, 4, 0, "")]}, + ], + "size": 8, + } + + try: + assert connector.get("0", "1", "req-active-sibling", metadata) is None + assert released == [] + assert len(connector._deferred_transfers) == 1 + deferred = connector._deferred_transfers[0] + assert deferred.handles == ["failed", "active"] + assert len(deferred.registrations) == 2 + assert len(deferred.dlists) == 4 + assert deferred.remote_agent == "producer" + assert len(deferred.tensors) == 2 + finally: + connector._agent.check_xfer_state = lambda handle: "DONE" + connector.close() + + +@pytest.mark.parametrize( + "payload,expected_kind", + [ + (torch.zeros(4), "tensors"), + ([torch.zeros(2), torch.ones(3)], "tensors"), + ({"hidden": torch.zeros(4), "meta": {"token_role_ids": [1, 2]}}, "structured"), + ({"prompt": "a cat", "steps": 8}, "object"), + ], +) +def test_payload_kinds_round_trip_through_the_handshake(producer, consumer, payload, expected_kind): + _, _, published = producer.put("0", "1", "req-kind", payload) + assert published["kind"] == expected_kind + + resolved = _wait_for_metadata(consumer, "req-kind") + + assert resolved["kind"] == expected_kind + assert resolved["tensor_specs"] == published["tensor_specs"] + + +def test_structured_payload_groups_descriptors_by_memory_type(producer, monkeypatch): + monkeypatch.setattr( + producer, + "_resolve_memory_type", + lambda tensor: "DRAM" if tensor.dtype == torch.uint8 else "VRAM", + ) + + _, _, metadata = producer.put( + "0", + "1", + "req-mixed", + {"meta": "value", "hidden": torch.ones(4, dtype=torch.float32)}, + ) + + assert metadata["schema_version"] == 1 + assert metadata["descriptor_groups"] == [ + { + "memory_type": "DRAM", + "tensor_indices": [0], + "regions": metadata["descriptor_groups"][0]["regions"], + }, + { + "memory_type": "VRAM", + "tensor_indices": [1], + "regions": metadata["descriptor_groups"][1]["regions"], + }, + ] + assert len(producer._agent.registered) == 2 + + +def test_omni_payload_struct_uses_structured_tensor_path(producer): + payload = OmniPayloadStruct( + hidden_states=HiddenStatesStruct(output=torch.ones(4, dtype=torch.float32)), + meta=MetaStruct(left_context_size=2), + ) + + _, _, metadata = producer.put("0", "1", "req-struct", payload) + skeleton, tensors = producer._extract_tensor_leaves(payload) + + assert metadata["kind"] == "structured" + assert len(tensors) == 1 + assert torch.equal(tensors[0], payload.hidden_states.output) + assert skeleton["meta"]["left_context_size"] == 2 + + +def test_partial_group_registration_failure_rolls_back(producer, monkeypatch): + monkeypatch.setattr( + producer, + "_resolve_memory_type", + lambda tensor: "DRAM" if tensor.dtype == torch.uint8 else "VRAM", + ) + original_register = producer._agent.register_memory + + def fail_second_group(descs, backends=None): + if descs[2] == "VRAM": + raise RuntimeError("registration failed") + original_register(descs, backends=backends) + + monkeypatch.setattr(producer._agent, "register_memory", fail_second_group) + + success, _, metadata = producer.put( + "0", + "1", + "req-partial", + {"meta": "value", "hidden": torch.ones(4, dtype=torch.float32)}, + ) + + assert success is False + assert metadata is None + assert producer._agent.registered == [] + assert producer._registered_descs == [] + assert "req-partial" not in producer._pending + + +def test_reusing_put_key_releases_previous_registration(producer): + producer.put("0", "1", "req-reuse", torch.zeros(2)) + previous_registration = producer._agent.registered[0] + + producer.put("0", "1", "req-reuse", torch.ones(2)) + + assert previous_registration not in producer._agent.registered + assert len(producer._agent.registered) == 1 + assert len(producer._registered_descs) == 1 + + +def test_expired_snapshot_cannot_claim_replacement(producer): + from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import _PendingPayload + + old = _PendingPayload([torch.zeros(1)], ["old"], 0.0) + replacement = _PendingPayload([torch.ones(1)], ["new"], time.monotonic() + 60) + producer._pending["req-race"] = replacement + + assert producer._take_pending("req-race", expected=old) is None + assert producer._pending["req-race"] is replacement + + +@pytest.mark.parametrize( + "indices", + [ + [0, 0], + [0], + [0, 2], + [-1, 0], + ], +) +def test_descriptor_groups_require_exact_tensor_index_partition(nixl_connector_cls, indices): + metadata = { + "schema_version": 1, + "descriptor_groups": [ + { + "memory_type": "DRAM", + "tensor_indices": indices, + "regions": [(1, 4, 0, "")] * len(indices), + } + ], + } + + with pytest.raises(RuntimeError, match="exact partition"): + nixl_connector_cls._validated_descriptor_groups( + metadata, [{"shape": [1], "dtype": "torch.float32", "size": 4}] * 2 + ) + + +def test_receive_device_ignores_the_producer_index(nixl_connector_cls): + """A producer on cuda:3 must not dictate the consumer's card.""" + connector = nixl_connector_cls({}) + try: + assert connector._resolve_receive_device("cpu") == torch.device("cpu") + assert connector._resolve_receive_device(None) == torch.device("cpu") + finally: + connector.close() + + +def test_receive_device_config_wins(nixl_connector_cls): + connector = nixl_connector_cls({"receive_device": "cpu"}) + try: + assert connector._resolve_receive_device("cuda:3") == torch.device("cpu") + finally: + connector.close() diff --git a/tests/distributed/omni_connectors/test_nixl_connector_native.py b/tests/distributed/omni_connectors/test_nixl_connector_native.py new file mode 100644 index 000000000..e0936a02d --- /dev/null +++ b/tests/distributed/omni_connectors/test_nixl_connector_native.py @@ -0,0 +1,204 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import multiprocessing as mp +import socket +import time +from typing import Any + +import pytest +import torch + +from vllm_omni.data_entry_keys import HiddenStatesStruct, MetaStruct, OmniPayloadStruct +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda, pytest.mark.parallel] + + +def _native_nixl_available() -> bool: + try: + from vllm.distributed.nixl_utils import NixlWrapper + except ImportError: + return False + return NixlWrapper is not None and torch.cuda.is_available() and torch.accelerator.device_count() >= 2 + + +def _free_port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def _producer(port: int, ready: Any, consumed: Any, result: Any, claimed: Any, proceed: Any, wire: Any) -> None: + from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import NixlConnector + + current_omni_platform.set_device(0) + connector = NixlConnector( + { + "role": "sender", + "host": "127.0.0.1", + "zmq_port": port, + "agent_name": "native-smoke-producer", + } + ) + try: + payload = OmniPayloadStruct( + hidden_states=HiddenStatesStruct(output=torch.arange(8, dtype=torch.float32, device="cuda:0")), + meta=MetaStruct(token_role_ids=torch.tensor([1, 2, 3], dtype=torch.int64, device="cpu")), + kv_metadata={"empty": torch.empty((2, 0, 3), device="cuda:0")}, + request_id="native-smoke", + ) + success, size, metadata = connector.put("0", "1", "native-smoke", payload) + result.put(("put", success, size, metadata["kind"] if metadata else None)) + wire.put(metadata) + ready.set() + if not claimed.wait(timeout=30): + raise TimeoutError("consumer did not claim source") + pending = connector._pending["native-smoke"] + pending.deadline = 0 + connector._cleanup_expired_pending() + connector.cleanup("native-smoke") + assert connector._pending["native-smoke"] is pending + assert pending.claims and connector._registered_descs + result.put(("source_retained", len(pending.claims))) + proceed.set() + if not consumed.wait(timeout=30): + raise TimeoutError("consumer did not complete the native NIXL transfer") + deadline = time.monotonic() + 10 + while connector._pending and time.monotonic() < deadline: + time.sleep(0.01) + result.put(("cleanup", len(connector._pending), len(connector._registered_descs))) + except Exception as error: + result.put(("producer_error", repr(error))) + raise + finally: + connector.close() + + +def _consumer( + port: int, ready: Any, consumed: Any, result: Any, claimed: Any, proceed: Any, wire: Any, direct: bool +) -> None: + from vllm_omni.distributed.omni_connectors.kv_transfer_manager import ( + OmniKVCacheConfig, + OmniKVTransferManager, + ) + + if not ready.wait(timeout=30): + raise TimeoutError("producer did not publish native NIXL metadata") + current_omni_platform.set_device(1) + manager = OmniKVTransferManager( + OmniKVCacheConfig( + connector_config={ + "type": "NixlConnector", + "role": "receiver", + "host": "127.0.0.1", + "zmq_port": port, # Shared deployment edge: already bound by producer. + "backends": ["UCX"], + "receive_device": "cuda", + "agent_name": "native-smoke-consumer", + }, + from_stage="0", + to_stage="1", + stage_id=1, + need_recv_cache=direct, # Exercise both payload-only and KV receiver setup. + ) + ) + manager.update_sender_info({"host": "127.0.0.1", "zmq_port": port}) + connector = manager.connector + assert connector is not None, "Receiver must not bind the producer's occupied port" + assert connector._zmq_port is None and not connector._serving_handshake + assert connector._backends == ["UCX"] + try: + original_wait = connector._wait_for_transfer + states = [] + + def wait(handle, key): + states.append(connector._agent.check_xfer_state(handle)) + claimed.set() + if not proceed.wait(timeout=30): + raise TimeoutError("producer did not verify active ownership") + original_wait(handle, key) + + connector._wait_for_transfer = wait + metadata = wire.get(timeout=5) + received = None + deadline = time.monotonic() + 30 + while received is None and time.monotonic() < deadline: + received = connector.get( + "0", + "1", + "native-smoke", + metadata if direct else {"source_host": "127.0.0.1", "source_port": port}, + ) + if received is None: + raise TimeoutError("native NIXL transfer did not complete") + payload, size = received + hidden = payload["hidden_states"]["output"] + token_roles = payload["meta"]["token_role_ids"] + empty = payload["kv_metadata"]["empty"] + assert empty.shape == (2, 0, 3) and str(empty.device) == "cuda:1" + result.put(("native_states", states)) + result.put( + ( + "get", + size, + hidden.cpu().tolist(), + str(hidden.dtype), + list(hidden.shape), + str(hidden.device), + token_roles.cpu().tolist(), + str(token_roles.dtype), + list(token_roles.shape), + str(token_roles.device), + ) + ) + consumed.set() + except Exception as error: + result.put(("consumer_error", repr(error))) + raise + finally: + connector.close() + + +@pytest.mark.skipif(not _native_nixl_available(), reason="requires NIXL and at least two CUDA devices") +@pytest.mark.parametrize("direct", [False, True]) +def test_native_two_process_structured_mixed_device_transfer(direct): + context = mp.get_context("spawn") + ready = context.Event() + consumed = context.Event() + claimed = context.Event() + proceed = context.Event() + wire = context.Queue() + result = context.Queue() + port = _free_port() + producer = context.Process(target=_producer, args=(port, ready, consumed, result, claimed, proceed, wire)) + consumer = context.Process(target=_consumer, args=(port, ready, consumed, result, claimed, proceed, wire, direct)) + + producer.start() + consumer.start() + producer.join(timeout=60) + consumer.join(timeout=60) + if producer.is_alive(): + producer.terminate() + if consumer.is_alive(): + consumer.terminate() + + records = [result.get(timeout=5) for _ in range(5)] + print("Native ownership records:", records) + assert producer.exitcode == 0, records + assert consumer.exitcode == 0, records + assert records[0][0] == "put" + assert records[0][1] is True + assert records[0][2] >= 56 + assert records[0][3] == "structured" + get_record = next(record for record in records if record[0] == "get") + assert get_record[1] == records[0][2] + assert get_record[2:6] == ( + list(range(8)), + "torch.float32", + [8], + "cuda:1", + ) + assert get_record[6:] == ([1, 2, 3], "torch.int64", [3], "cuda:1") + assert next(record for record in records if record[0] == "cleanup") == ("cleanup", 0, 0) + assert next(record for record in records if record[0] == "source_retained") == ("source_retained", 1) diff --git a/tests/engine/test_orchestrator_kv_sender_info.py b/tests/engine/test_orchestrator_kv_sender_info.py index e3f58c129..ca2a26b70 100644 --- a/tests/engine/test_orchestrator_kv_sender_info.py +++ b/tests/engine/test_orchestrator_kv_sender_info.py @@ -1,13 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import asyncio +from collections.abc import Callable from types import SimpleNamespace +from typing import Any import pytest from vllm import SamplingParams +from vllm.v1.engine.core_client import AsyncMPClient, DPLBAsyncMPClient from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker from vllm_omni.engine.messages import OutputMessage from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState -from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient +from vllm_omni.engine.stage_engine_core_client import ( + DPLBStageEngineCoreClient, + StageEngineCoreClient, +) from vllm_omni.engine.stage_pool import StagePool from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -28,7 +37,7 @@ class _DummySenderStage: class _DummyDiffusionStage: stage_type = "diffusion" final_output = True - custom_process_input_func = None + custom_process_input_func: Callable[..., Any] | None = None def __init__(self, engine_input_source=None): self.engine_input_source = engine_input_source or [0] @@ -54,6 +63,97 @@ def _build_sender_pool(stage_id: int, sender_info: dict[str, object]) -> StagePo ) +@pytest.mark.parametrize( + ("client_class", "base_client_class"), + [ + (StageEngineCoreClient, AsyncMPClient), + (DPLBStageEngineCoreClient, DPLBAsyncMPClient), + ], +) +def test_stage_engine_core_client_builds_payload_sender_info_after_base_init( + monkeypatch, client_class, base_client_class +): + vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + omni_kv_config=None, + hf_config=None, + stage_connector_config=None, + ) + ) + metadata = SimpleNamespace( + stage_id=0, + replica_id=0, + stage_type="llm", + model_stage="main", + is_comprehension=False, + requires_multimodal_data=False, + engine_input_source=[], + final_output=False, + final_output_type=None, + default_sampling_params=None, + prompt_transform_func=None, + prompt_expand_func=None, + custom_process_input_func=None, + ) + + def fake_base_init(self, config, *_args, **_kwargs): + self.vllm_config = config + self.resources = SimpleNamespace(engine_dead=False) + + monkeypatch.setattr(base_client_class, "__init__", fake_base_init) + + client = client_class(vllm_config, object, metadata=metadata) + + assert client.vllm_config is vllm_config + assert client.get_payload_sender_info() is None + + +@pytest.mark.parametrize("outgoing", [False, True]) +@pytest.mark.parametrize("base_port", [None, 48000]) +def test_payload_sender_endpoint_matches_resolver_with_unequal_replicas(outgoing, base_port): + from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec + from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_connector_spec + + edge = {"host": "10.0.0.2", "from_stage": 1} + if base_port is not None: + edge["zmq_port"] = base_port + extra = ( + { + "role": "receiver", + "host": "10.0.0.1", + "zmq_port": 47000, + "from_stage": 0, + "outgoing": edge, + } + if outgoing + else {"role": "sender", **edge} + ) + client = object.__new__(StageEngineCoreClient) + client.stage_id = 1 + client.replica_id = 3 + client.vllm_config = SimpleNamespace( + model_config=SimpleNamespace(stage_connector_config={"name": "NixlConnector", "extra": extra}) + ) + producer = resolve_connector_spec( + ConnectorSpec(name="NixlConnector", extra=extra), + stage_id=1, + role=extra["role"], + replica_id=client.replica_id, + ) + consumer = resolve_connector_spec( + ConnectorSpec(name="NixlConnector", extra=edge), stage_id=2, role="receiver", replica_id=7 + ) + + sender_info = client._build_payload_sender_info() + + assert sender_info == { + "host": "10.0.0.2", + "zmq_port": (50051 if base_port is None else base_port) + 3 * 1024 + 1, + } + assert sender_info["zmq_port"] == producer.extra["zmq_port"] + assert sender_info["zmq_port"] != consumer.extra["sender_zmq_port"] + + def test_stage_engine_core_client_builds_kv_sender_info_from_tcp_address(): client = object.__new__(StageEngineCoreClient) client.stage_id = 0 @@ -259,3 +359,58 @@ def test_prewarm_diffusion_attaches_kv_sender_info(): 0: {"host": "10.0.0.3", "zmq_port": 50151}, } assert req_state.stage_submit_ts[1] > 0 + + +def test_prewarm_submits_bound_payload_endpoint_for_concurrent_replicas(): + orchestrator = object.__new__(Orchestrator) + endpoints = {"a": {"host": "10.0.0.2", "zmq_port": 52099}, "b": {"host": "10.0.0.3", "zmq_port": 54147}} + source = SimpleNamespace( + get_bound_client=lambda key: SimpleNamespace(get_payload_sender_info=lambda: endpoints[key]), + get_bound_replica_id=lambda key: {"a": 2, "b": 4}[key], + ) + submitted = {} + + async def submit(key, state, request, **kwargs): + await asyncio.sleep(0) + submitted[key] = request + + target = SimpleNamespace( + stage_type="llm", + stage_client=SimpleNamespace(engine_input_source=[0]), + stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(hf_config=None, max_model_len=64)), + submit_initial=submit, + get_bound_replica_id=lambda key: {"a": 7, "b": 1}[key], + ) + orchestrator.stage_pools = [source, target] + orchestrator._stage_receives_async_chunks = lambda stage: True + orchestrator._record_duplex_stage_submission = lambda *args: None + orchestrator._emit_tx_edge = lambda **kwargs: None + + async def dispatch(callback, **kwargs): + await callback() + return True + + orchestrator._dispatch_or_fail_request = dispatch + + async def run(): + await asyncio.gather( + *[ + orchestrator._prewarm_async_chunk_stages( + key, + SimpleNamespace(prompt_token_ids=[1]), + OrchestratorRequestState( + request_id=key, + prompt={}, + sampling_params_list=[SamplingParams(), SamplingParams()], + final_stage_id=1, + ), + ) + for key in endpoints + ] + ) + + asyncio.run(run()) + from vllm_omni.engine import OmniEngineCoreRequest + + assert all(isinstance(request, OmniEngineCoreRequest) for request in submitted.values()) + assert {key: request.payload_sender_info for key, request in submitted.items()} == endpoints diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py index 68836e2a4..69df71231 100644 --- a/tests/engine/test_orchestrator_stage_input_bridge.py +++ b/tests/engine/test_orchestrator_stage_input_bridge.py @@ -104,9 +104,13 @@ class FakeInputProcessor: class FakePrewarmPool: + """Single-replica pool double that binds requests only on submission.""" + stage_type = "llm" def __init__(self, role: str) -> None: + self.stage_client = SimpleNamespace() + self._bound_request_ids: set[str] = set() self.stage_vllm_config = SimpleNamespace( model_config=SimpleNamespace( max_model_len=64, @@ -115,12 +119,16 @@ class FakePrewarmPool: ) self.submitted: list[Any] = [] - async def submit_initial(self, _request_id, _req_state, request, prompt_text=None): + async def submit_initial(self, request_id, _req_state, request, prompt_text=None): self.submitted.append(request) + self._bound_request_ids.add(request_id) return 0 - def get_bound_replica_id(self, _request_id): - return 0 + def get_bound_replica_id(self, request_id): + return 0 if request_id in self._bound_request_ids else None + + def get_bound_client(self, request_id): + return self.stage_client if self.get_bound_replica_id(request_id) is not None else None def _duplex_stage_port_submission(): @@ -261,11 +269,14 @@ async def test_forward_text_prompt_uses_target_stage_input_processor() -> None: @pytest.mark.asyncio -async def test_async_prewarm_skips_outgoing_only_stage() -> None: +@pytest.mark.parametrize("payload_sender_info", [None, {"host": "10.0.0.2", "zmq_port": 52099}]) +async def test_async_prewarm_skips_outgoing_only_stage(payload_sender_info) -> None: orchestrator = object.__new__(Orchestrator) stage0 = FakePrewarmPool("sender") stage1 = FakePrewarmPool("sender") stage2 = FakePrewarmPool("receiver") + if payload_sender_info is not None: + stage1.stage_client.get_payload_sender_info = MagicMock(return_value=payload_sender_info) orchestrator.stage_pools = [stage0, stage1, stage2] orchestrator._emit_tx_edge = lambda **_kwargs: None orchestrator._record_duplex_stage_submission = MagicMock() @@ -285,7 +296,14 @@ async def test_async_prewarm_skips_outgoing_only_stage() -> None: assert prewarmed is True assert stage1.submitted == [] + assert stage1.get_bound_client("req-prewarm") is None assert len(stage2.submitted) == 1 + assert stage2.get_bound_client("req-prewarm") is stage2.stage_client + assert stage2.submitted[0].payload_sender_info == payload_sender_info + if payload_sender_info is not None: + stage1.stage_client.get_payload_sender_info.assert_called_once_with() + assert stage2.submitted[0].external_req_id == "req-prewarm" + assert stage2.submitted[0].resumable is True assert 1 not in req_state.stage_submit_ts assert 2 in req_state.stage_submit_ts orchestrator._record_duplex_stage_submission.assert_called_once_with( diff --git a/tests/worker/test_omni_connector_mixin.py b/tests/worker/test_omni_connector_mixin.py index 8366e177a..603ca65ff 100644 --- a/tests/worker/test_omni_connector_mixin.py +++ b/tests/worker/test_omni_connector_mixin.py @@ -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.""" diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index fe439457d..fd7617144 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -67,6 +67,8 @@ _PUBLIC_OMNI = ( "VLLM_OMNI_EVENT_DRIVEN_ORCH", "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", "VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", + "VLLM_OMNI_NIXL_LEASE_S", + "VLLM_OMNI_NIXL_XFER_TIMEOUT_S", "VLLM_OMNI_ORCH_MONITOR_PATH", "VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY", "VLLM_OMNI_SERVER_STORAGE__FILE_TTL", diff --git a/vllm_omni/core/sched/omni_scheduling_coordinator.py b/vllm_omni/core/sched/omni_scheduling_coordinator.py index 93f40daa0..3d516f5e3 100644 --- a/vllm_omni/core/sched/omni_scheduling_coordinator.py +++ b/vllm_omni/core/sched/omni_scheduling_coordinator.py @@ -121,6 +121,7 @@ class OmniSchedulingCoordinator: OmniChunkRecvHandle( request_id=request.request_id, external_req_id=getattr(request, "external_req_id", None), + payload_sender_info=getattr(request, "payload_sender_info", None), ) ) elif request.status == RequestStatus.WAITING_FOR_INPUT: @@ -134,6 +135,7 @@ class OmniSchedulingCoordinator: OmniChunkRecvHandle( request_id=request.request_id, external_req_id=getattr(request, "external_req_id", None), + payload_sender_info=getattr(request, "payload_sender_info", None), ) ) if to_remove: diff --git a/vllm_omni/core/sched/output.py b/vllm_omni/core/sched/output.py index b73fbd6aa..deaf44df8 100644 --- a/vllm_omni/core/sched/output.py +++ b/vllm_omni/core/sched/output.py @@ -105,6 +105,7 @@ class OmniChunkRecvHandle: request_id: str external_req_id: str | None = None + payload_sender_info: dict[str, object] | None = None @dataclass diff --git a/vllm_omni/distributed/__init__.py b/vllm_omni/distributed/__init__.py index 3506f02f6..40374cd6d 100644 --- a/vllm_omni/distributed/__init__.py +++ b/vllm_omni/distributed/__init__.py @@ -7,6 +7,7 @@ from .omni_connectors import ( MooncakeStoreConnector, MooncakeTransferEngineConnector, MoriTransferEngineConnector, + NixlConnector, OmniConnectorBase, OmniConnectorFactory, OmniTransferConfig, @@ -27,6 +28,7 @@ __all__ = [ "MooncakeStoreConnector", "MooncakeTransferEngineConnector", "MoriTransferEngineConnector", + "NixlConnector", "SharedMemoryConnector", "YuanrongConnector", "YuanrongTransferEngineConnector", diff --git a/vllm_omni/distributed/omni_connectors/__init__.py b/vllm_omni/distributed/omni_connectors/__init__.py index 0cee596bb..e46142b2a 100644 --- a/vllm_omni/distributed/omni_connectors/__init__.py +++ b/vllm_omni/distributed/omni_connectors/__init__.py @@ -22,6 +22,11 @@ try: from .connectors.mori_transfer_engine_connector import MoriTransferEngineConnector except ImportError: MoriTransferEngineConnector = None # RDMA deps (msgspec/zmq/mori) not installed + +try: + from .connectors.nixl_connector import NixlConnector +except ImportError: + NixlConnector = None # NIXL deps not installed from .factory import OmniConnectorFactory from .utils.config import ConnectorSpec, OmniTransferConfig from .utils.initialization import ( @@ -50,6 +55,7 @@ __all__ = [ "MooncakeStoreConnector", "MooncakeTransferEngineConnector", "MoriTransferEngineConnector", + "NixlConnector", "SharedMemoryConnector", "YuanrongConnector", "YuanrongTransferEngineConnector", diff --git a/vllm_omni/distributed/omni_connectors/connectors/nixl_connector.py b/vllm_omni/distributed/omni_connectors/connectors/nixl_connector.py new file mode 100644 index 000000000..7dda744f6 --- /dev/null +++ b/vllm_omni/distributed/omni_connectors/connectors/nixl_connector.py @@ -0,0 +1,1044 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import math +import os +import socket +import sys +import threading +import time +import uuid +from dataclasses import dataclass, field +from typing import Any + +import msgspec +import torch +import zmq + +from ..utils.logging import get_connector_logger +from .base import OmniConnectorBase + +logger = get_connector_logger(__name__) + +_SCHEMA_VERSION = 1 +_KIND_TENSORS = "tensors" +_KIND_STRUCTURED = "structured" +_KIND_OBJECT = "object" +_INIT_AGENT = "NIXL_INIT_AGENT" +_TENSOR_MARKER = "__nixl_tensor_index__" +_TUPLE_MARKER = "__nixl_tuple__" + +# Handshake control-plane messages. Consumers claim connector-managed payloads +# before READ, even with out-of-band metadata, and acknowledge completion here. +_GET_META_MSG = b"nixl_get_meta" +_XFER_DONE_MSG = b"nixl_xfer_done" +_META_NOT_FOUND = b"nixl_meta_not_found" +_ACK = b"nixl_ack" + +# NIXL has no remote READ cancellation primitive. A closed producer with an +# abandoned claim must keep both its agent and allocations alive until process +# exit; dropping the last Python reference is not a safe cancellation policy. +_RETAINED_PRODUCERS: set[Any] = set() + + +@dataclass +class _DeferredTransfer: + tensors: list[torch.Tensor] + registrations: list[Any] = field(default_factory=list) + dlists: list[Any] = field(default_factory=list) + handles: list[Any] = field(default_factory=list) + remote_agent: Any = None + source_key: str | None = None + source_metadata: dict[str, Any] | None = None + + +@dataclass +class _PendingPayload: + tensors: list[torch.Tensor] + registrations: list[Any] + deadline: float + generation: str = field(default_factory=lambda: uuid.uuid4().hex) + claims: set[str] = field(default_factory=set) + + +class NixlConnector(OmniConnectorBase): + """OmniConnector backed by vLLM's native NIXL wrapper. + + This connector intentionally depends on vLLM's optional NIXL integration + (``vllm.distributed.nixl_utils``). It transfers raw tensor payloads + directly through NIXL READ operations. Non-tensor Python + payloads are serialized with OmniSerializer, packed into a uint8 CPU tensor, + and moved through the same NIXL path. + + A NIXL READ needs the producer's agent metadata and memory descriptors, + which ``put`` returns to its caller. Consumers use the ZMQ control plane to + claim connector-managed payloads before READ and acknowledge completion, + including when metadata is forwarded directly. Producers without a fixed + ``zmq_port`` use an ephemeral listener advertised in that metadata. Callers + passing ``metadata=None`` need a configured or request-specific producer + endpoint to fetch and claim metadata by key. + """ + + supports_raw_data: bool = True + + def __init__(self, config: dict[str, Any]): + self.config = dict(config or {}) + self.stage_id = int(self.config.get("stage_id", 0)) + self._closed = True + self._registered_descs: list[Any] = [] + self._pending: dict[str, _PendingPayload] = {} + self._published: dict[str, dict[str, Any]] = {} + self._state_lock = threading.RLock() + self._remote_agents: list[str] = [] + self._metrics: dict[str, int] = { + "puts": 0, + "gets": 0, + "errors": 0, + "bytes_transferred": 0, + } + + from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config + + if NixlWrapper is None: + raise RuntimeError("NIXL is not available. Install the optional nixl/rixl package to use NixlConnector.") + + backends = self.config.get("backends", ["UCX"]) + self._backends = list(backends) if isinstance(backends, (list, tuple)) else [str(backends)] + if nixl_agent_config is None: + agent_config = None + else: + from nixl import nixl_thread_sync_t + + config_module = sys.modules.get(nixl_agent_config.__module__) + sync_type = getattr(config_module, "nixl_thread_sync_t", nixl_thread_sync_t) + sync_mode = sync_type.NIXL_THREAD_SYNC_STRICT + non_ucx_backends = [backend for backend in self._backends if backend != "UCX"] + if non_ucx_backends: + agent_config = nixl_agent_config( + backends=self._backends, + capture_telemetry=False, + sync_mode=sync_mode, + ) + else: + num_threads = int(self.config.get("num_threads", 4)) + agent_config = nixl_agent_config( + num_threads=num_threads, + capture_telemetry=False, + sync_mode=sync_mode, + ) + + self._agent = NixlWrapper(str(self.config.get("agent_name", uuid.uuid4())), agent_config) + self._receive_device = self._parse_device(self.config.get("receive_device")) + self._default_memory_type = self.config.get("memory_type") + # The lease bounds how long a ``put`` payload stays registered while it + # waits to be read. A consumer stage may legitimately queue for a long + # time (e.g. video diffusion with concurrency above the replica count), + # so the default must be far larger than any plausible queueing delay. + # The environment variables intentionally take precedence over the + # deployment YAML, which ships a fixed 300s lease that is too short for + # long-queueing pipelines and cannot be overridden per deployment. + self._lease_seconds = float( + os.environ.get("VLLM_OMNI_NIXL_LEASE_S") or self.config.get("lease_seconds", 3600.0) + ) + self._transfer_timeout_s = float( + os.environ.get("VLLM_OMNI_NIXL_XFER_TIMEOUT_S") or self.config.get("transfer_timeout_s", 300.0) + ) + self._poll_interval_s = float(self.config.get("poll_interval_s", 0.001)) + self._closed = False + self._init_handshake() + self._lease_wakeup = threading.Event() + self._lease_thread: threading.Thread | None = None + if self._serving_handshake: + self._lease_thread = threading.Thread( + target=self._lease_reaper_loop, + name="nixl-lease-reaper", + daemon=True, + ) + self._lease_thread.start() + self._transfer_wakeup = threading.Event() + self._deferred_transfers: list[_DeferredTransfer] = [] + self._transfer_thread: threading.Thread | None = None + if self._role != "sender": + self._transfer_thread = threading.Thread( + target=self._transfer_reaper_loop, + name="nixl-transfer-reaper", + daemon=True, + ) + self._transfer_thread.start() + + def _init_handshake(self) -> None: + """Set up the ZMQ metadata and payload-ownership control plane. + + Producers listen on a configured or ephemeral port, including for + directly forwarded metadata. Receive-only connectors dial the producer; + intermediate stages can also listen for their outgoing payloads. + """ + role = self.config.get("role") + self._role = str(role).lower() if role else None + # Even out-of-band metadata needs an atomic claim before READ. An + # ephemeral listener supplies that ownership channel when no fixed + # data-plane metadata port was configured. + self._zmq_port = self.config.get("zmq_port", 0 if self._role != "receiver" else None) + self._sender_host = self.config.get("sender_host") + self._sender_zmq_port = self.config.get("sender_zmq_port") + self._handshake_timeout_ms = int(self.config.get("handshake_timeout_ms", 5000)) + self._metadata_query_timeout_ms = int( + self.config.get("metadata_query_timeout_ms", min(self._handshake_timeout_ms, 10)) + ) + self._handshake_max_wait_s = float(self.config.get("handshake_max_wait_s", 60.0)) + self._handshake_retry_s = float(self.config.get("handshake_retry_s", 0.05)) + self._zmq_ctx: zmq.Context | None = None + self._req_local = threading.local() + self._listener_thread: threading.Thread | None = None + self._listener_ready = threading.Event() + self._stop_event = threading.Event() + self._bind_error: BaseException | None = None + self.host: str | None = None + self._serving_handshake = False + + self._handshake_enabled = self._zmq_port is not None or self._sender_host is not None + if not self._handshake_enabled: + return + + self._zmq_ctx = zmq.Context() + # A receiver only dials out, so it never needs a port of its own. + if self._zmq_port is None: + return + + host_value = str(self.config.get("host", "auto")) + self.host = ( + self._get_local_ip() + if host_value.lower() == "auto" or host_value in {"", "*", "0.0.0.0", "::"} + else host_value + ) + self._zmq_port = int(self._zmq_port) + self._listener_thread = threading.Thread( + target=self._handshake_listener_loop, name="nixl-handshake", daemon=True + ) + self._listener_thread.start() + self._listener_ready.wait(timeout=5.0) + if self._bind_error is not None: + raise RuntimeError( + f"NixlConnector failed to bind handshake socket on {self.host}:{self._zmq_port}" + ) from self._bind_error + self._serving_handshake = True + logger.info("NixlConnector handshake listener bound on %s:%s", self.host, self._zmq_port) + + def put( + self, + from_stage: str, + to_stage: str, + put_key: str, + data: Any, + ) -> tuple[bool, int, dict[str, Any] | None]: + if self._closed or getattr(self, "_closing", False): + raise RuntimeError("Cannot put data: NixlConnector is closed") + + try: + self._cleanup_expired_pending() + kind = _KIND_TENSORS + if self._is_tensor_payload(data): + tensors, tensor_specs = self._normalize_tensor_payload(data) + elif self._contains_tensor(data): + skeleton, payload_tensors = self._extract_tensor_leaves(data) + header = self.serialize_obj(skeleton) + header_tensor = torch.frombuffer(header, dtype=torch.uint8).clone().contiguous() + tensors, tensor_specs = self._normalize_tensor_payload([header_tensor, *payload_tensors]) + kind = _KIND_STRUCTURED + else: + payload = self.serialize_obj(data) + tensor = torch.frombuffer(payload, dtype=torch.uint8).clone().contiguous() + tensors, tensor_specs = self._normalize_tensor_payload(tensor) + kind = _KIND_OBJECT + + grouped_tensors: dict[str, list[tuple[int, torch.Tensor]]] = {} + for tensor_index, tensor in enumerate(tensors): + # Preserve spec/skeleton slots, but never register empty DMA regions. + if tensor.numel() == 0: + continue + memory_type = self._resolve_memory_type(tensor) + grouped_tensors.setdefault(memory_type, []).append((tensor_index, tensor)) + + descriptor_groups = [] + registered_descs = [] + try: + for memory_type, indexed_tensors in grouped_tensors.items(): + tensor_indices = [tensor_index for tensor_index, _ in indexed_tensors] + regions = [self._tensor_region(tensor) for _, tensor in indexed_tensors] + reg_descs = self._agent.get_reg_descs(regions, memory_type) + self._agent.register_memory(reg_descs, backends=self._backends) + registered_descs.append(reg_descs) + descriptor_groups.append( + { + "memory_type": memory_type, + "tensor_indices": tensor_indices, + "regions": regions, + } + ) + except Exception: + for reg_descs in registered_descs: + self._safe_call(self._agent.deregister_memory, reg_descs) + raise + with self._state_lock: + self._registered_descs.extend(registered_descs) + + size = sum(spec["size"] for spec in tensor_specs) + metadata = { + "schema_version": _SCHEMA_VERSION, + "kind": kind, + "agent_metadata": self._agent.get_agent_metadata(), + "descriptor_groups": descriptor_groups, + "tensor_specs": tensor_specs, + "size": size, + } + if self._serving_handshake: + metadata["sender_host"] = self.host + metadata["sender_zmq_port"] = self._zmq_port + with self._state_lock: + previous = self._pending.get(put_key) + if previous is not None and previous.claims: + for descs in registered_descs: + self._safe_call(self._agent.deregister_memory, descs) + self._registered_descs.remove(descs) + raise RuntimeError(f"Cannot replace claimed NIXL payload {put_key}") + previous = self._take_pending(put_key) + if previous is not None: + self._release_pending(previous) + pending = _PendingPayload( + tensors=tensors, + registrations=registered_descs, + deadline=time.monotonic() + self._lease_seconds, + ) + metadata["generation"] = pending.generation + self._pending[put_key] = pending + if self._serving_handshake: + self._published[put_key] = metadata + self._lease_wakeup.set() + self._metrics["puts"] += 1 + self._metrics["bytes_transferred"] += size + logger.debug("NixlConnector put %s->%s key=%s size=%d", from_stage, to_stage, put_key, size) + return True, size, metadata + except Exception: + self._metrics["errors"] += 1 + logger.error("NixlConnector put failed for %s", put_key, exc_info=True) + return False, 0, None + + def get( + self, + from_stage: str, + to_stage: str, + get_key: str, + metadata: dict[str, Any] | None = None, + ) -> tuple[Any, int] | None: + if self._closed or getattr(self, "_closing", False): + raise RuntimeError("Cannot get data: NixlConnector is closed") + + remote_agent = None + local_reg_descs_list = [] + dlist_handles = [] + xfer_handles = [] + source_metadata = None + try: + metadata = self._resolve_metadata(get_key, metadata) + if not isinstance(metadata, dict) or metadata.get("schema_version") != _SCHEMA_VERSION: + logger.error("NixlConnector get has invalid metadata for %s", get_key) + return None + + source_metadata = metadata + + tensor_specs = metadata.get("tensor_specs") + if not isinstance(tensor_specs, list): + raise RuntimeError(f"Invalid NIXL metadata for {get_key}: missing tensor_specs") + descriptor_groups = self._validated_descriptor_groups(metadata, tensor_specs) + + local_tensors = [self._allocate_tensor_from_spec(spec, metadata.get("kind")) for spec in tensor_specs] + if descriptor_groups: + remote_agent = self._agent.add_remote_agent(metadata["agent_metadata"]) + self._remote_agents.append(remote_agent) + for descriptor_group in descriptor_groups: + remote_memory_type = descriptor_group["memory_type"] + indexed_regions = list( + zip(descriptor_group["tensor_indices"], descriptor_group["regions"], strict=True) + ) + local_groups: dict[str, list[tuple[int, Any]]] = {} + for tensor_index, remote_region in indexed_regions: + local_memory_type = self._resolve_memory_type(local_tensors[tensor_index]) + local_groups.setdefault(local_memory_type, []).append((tensor_index, remote_region)) + + for local_memory_type, entries in local_groups.items(): + group_tensors = [local_tensors[tensor_index] for tensor_index, _ in entries] + local_regions = [self._tensor_region(tensor) for tensor in group_tensors] + local_reg_descs = self._agent.get_reg_descs(local_regions, local_memory_type) + self._agent.register_memory(local_reg_descs, backends=self._backends) + local_reg_descs_list.append(local_reg_descs) + + remote_regions = [tuple(region)[:3] for _, region in entries] + remote_descs = self._agent.get_xfer_descs(remote_regions, remote_memory_type) + local_descs = self._agent.get_xfer_descs( + [tuple(region)[:3] for region in local_regions], local_memory_type + ) + remote_dlist = self._agent.prep_xfer_dlist(remote_agent, remote_descs) + local_dlist = self._agent.prep_xfer_dlist(_INIT_AGENT, local_descs) + dlist_handles.extend([local_dlist, remote_dlist]) + + desc_ids = list(range(len(entries))) + xfer_handle = self._agent.make_prepped_xfer( + "READ", + local_dlist, + desc_ids, + remote_dlist, + desc_ids, + ) + xfer_handles.append(xfer_handle) + self._agent.transfer(xfer_handle) + + for xfer_handle in xfer_handles: + self._wait_for_transfer(xfer_handle, get_key) + for xfer_handle in xfer_handles: + self._agent.release_xfer_handle(xfer_handle) + xfer_handles.clear() + + size = int(metadata.get("size", sum(spec.get("size", 0) for spec in tensor_specs))) + if metadata.get("kind") == _KIND_OBJECT: + raw = local_tensors[0].detach().cpu().contiguous().view(torch.uint8).numpy().tobytes() + payload = self.deserialize_obj(raw) + elif metadata.get("kind") == _KIND_STRUCTURED: + raw = local_tensors[0].detach().cpu().contiguous().view(torch.uint8).numpy().tobytes() + skeleton = self.deserialize_obj(raw) + payload = self._restore_tensor_leaves(skeleton, local_tensors[1:]) + else: + payload = local_tensors[0] if len(local_tensors) == 1 else local_tensors + self._metrics["gets"] += 1 + self._metrics["bytes_transferred"] += size + logger.debug("NixlConnector get %s->%s key=%s size=%d", from_stage, to_stage, get_key, size) + return payload, size + except Exception: + if xfer_handles and self._transfers_may_be_active(xfer_handles): + self._defer_transfer( + _DeferredTransfer( + tensors=local_tensors, + registrations=local_reg_descs_list, + dlists=dlist_handles, + handles=xfer_handles, + remote_agent=remote_agent, + source_key=get_key, + source_metadata=source_metadata, + ) + ) + local_tensors = [] + local_reg_descs_list = [] + dlist_handles = [] + xfer_handles = [] + remote_agent = None + source_metadata = None + self._metrics["errors"] += 1 + logger.error("NixlConnector get failed for %s", get_key, exc_info=True) + return None + finally: + for xfer_handle in xfer_handles: + self._safe_call(self._agent.release_xfer_handle, xfer_handle) + for dlist_handle in dlist_handles: + self._safe_call(self._agent.release_dlist_handle, dlist_handle) + if remote_agent is not None: + self._safe_call(self._agent.remove_remote_agent, remote_agent) + if remote_agent in self._remote_agents: + self._remote_agents.remove(remote_agent) + for local_reg_descs in local_reg_descs_list: + self._safe_call(self._agent.deregister_memory, local_reg_descs) + if source_metadata is not None: + self._notify_transfer_done(get_key, source_metadata) + + def cleanup(self, request_id: str) -> None: + pending = self._take_pending(request_id) + if pending is None: + return + self._release_pending(pending) + + def _release_pending(self, pending: _PendingPayload) -> None: + for reg_descs in pending.registrations: + self._safe_call(self._agent.deregister_memory, reg_descs) + with self._state_lock: + if reg_descs in self._registered_descs: + self._registered_descs.remove(reg_descs) + + def _take_pending(self, request_id: str, expected: _PendingPayload | None = None) -> _PendingPayload | None: + with self._state_lock: + pending = self._pending.get(request_id) + if pending is None or (expected is not None and pending is not expected): + return None + if pending.claims: + return None + self._published.pop(request_id, None) + return self._pending.pop(request_id) + + def close(self) -> None: + if self._closed: + return + self._closed = True + # Stop accepting puts/gets, but keep serving completion ACKs for outstanding + # remote READs. A later close() call can finish teardown once drained. + for request_id in list(self._pending): + self.cleanup(request_id) + with self._state_lock: + if any(pending.claims for pending in self._pending.values()): + _RETAINED_PRODUCERS.add(self) + self._closed = False + self._closing = True + logger.warning("NIXL close deferred: remote READ claims still own source allocations") + return + _RETAINED_PRODUCERS.discard(self) + self._stop_event.set() + self._lease_wakeup.set() + self._transfer_wakeup.set() + if self._listener_thread is not None: + self._listener_thread.join(timeout=5.0) + self._listener_thread = None + if self._lease_thread is not None: + self._lease_thread.join(timeout=5.0) + self._lease_thread = None + if self._transfer_thread is not None: + self._transfer_thread.join(timeout=5.0) + self._transfer_thread = None + while self._deferred_transfers: + self._reap_deferred_transfers() + if self._deferred_transfers: + time.sleep(max(self._poll_interval_s, 0.01)) + if self._zmq_ctx is not None: + # destroy() rather than term(): REQ sockets live in thread-local + # caches this thread cannot reach, and term() blocks until every + # socket in the context is closed. + self._safe_call(self._zmq_ctx.destroy, 0) + self._zmq_ctx = None + for request_id in list(self._pending): + self.cleanup(request_id) + deferred_agents = { + transfer.remote_agent for transfer in self._deferred_transfers if transfer.remote_agent is not None + } + for agent_name in [agent for agent in self._remote_agents if agent not in deferred_agents]: + self._safe_call(self._agent.remove_remote_agent, agent_name) + self._remote_agents.remove(agent_name) + for reg_descs in list(self._registered_descs): + self._safe_call(self._agent.deregister_memory, reg_descs) + self._registered_descs.clear() + + def health(self) -> dict[str, Any]: + return { + "status": "unhealthy" if self._closed else "healthy", + "pending_requests": len(self._pending), + **self._metrics, + } + + def get_connection_info(self) -> dict[str, Any]: + """Endpoint a consumer needs to reach this producer's handshake socket.""" + return {"host": self.host, "zmq_port": self._zmq_port} + + def update_sender_info(self, sender_host: str, sender_zmq_port: int) -> None: + """Register the producer handshake endpoint on the consumer side. + + Used when the endpoint is only known after both stages have started and + therefore cannot be baked into the connector config. + """ + self._sender_host = sender_host + self._sender_zmq_port = int(sender_zmq_port) + if self._zmq_ctx is None: + self._zmq_ctx = zmq.Context() + self._handshake_enabled = True + + @staticmethod + def _metadata_endpoint(metadata: dict[str, Any] | None) -> tuple[str, int] | None: + if not isinstance(metadata, dict): + return None + host = metadata.get("sender_host") + port = metadata.get("sender_zmq_port") + if not host or not port: + host = metadata.get("source_host") + port = metadata.get("source_port") + if not host or not port: + return None + return str(host), int(port) + + def _resolve_metadata(self, get_key: str, metadata: dict[str, Any] | None) -> dict[str, Any] | None: + """Return complete NIXL transfer metadata for ``get_key``. + + Connector-managed metadata, including directly forwarded ``put`` results, + requires a handshake claim at the supplied or configured producer endpoint. + Only legacy externally owned metadata without a generation bypasses it. + """ + direct = isinstance(metadata, dict) and metadata.get("schema_version") == _SCHEMA_VERSION + if direct and "generation" not in metadata: + # Legacy externally owned metadata has no connector-managed lease. + return metadata + + endpoint = self._metadata_endpoint(metadata) + if endpoint is None and self._sender_host and self._sender_zmq_port: + endpoint = (str(self._sender_host), int(self._sender_zmq_port)) + if endpoint is None: + logger.error( + "NixlConnector get(%s) received no usable metadata and no producer handshake " + "endpoint is configured. Set sender_host/sender_zmq_port on the consumer, " + "pass source_host/source_port for the producer rank, or forward the metadata " + "returned by put().", + get_key, + ) + return None + if self._zmq_ctx is None: + with self._state_lock: + if self._zmq_ctx is None: + self._zmq_ctx = zmq.Context() + + if direct: + return self._query_metadata_at(get_key, *endpoint, generation=metadata["generation"]) + return self._query_metadata_at(get_key, *endpoint) + + def _query_metadata_at( + self, get_key: str, host: str, port: int, *, generation: str | None = None + ) -> dict[str, Any] | None: + """Fetch transfer metadata for ``get_key`` from a producer's ROUTER socket. + + Each call performs one bounded query so a missing key cannot starve + other requests in the shared receive loop. The caller retries later. + """ + zmq_addr = f"tcp://{host}:{port}" + request = _GET_META_MSG + msgspec.msgpack.encode( + {"key": get_key, "generation": generation, "claim_id": uuid.uuid4().hex} + ) + sock = self._get_req_socket(zmq_addr, self._metadata_query_timeout_ms) + try: + sock.send(request) + reply = sock.recv() + except Exception: + self._invalidate_req_socket(zmq_addr) + logger.debug("NixlConnector handshake query to %s failed for %s", zmq_addr, get_key, exc_info=True) + return None + if reply == _META_NOT_FOUND: + return None + return msgspec.msgpack.decode(reply) + + def _notify_transfer_done(self, get_key: str, metadata: dict[str, Any]) -> None: + """Tell the producer its buffer is drained so it can deregister now. + + Without this the producer would hold the registration until the lease + expires, which for long-lived stages means unbounded growth. + """ + endpoint = self._metadata_endpoint(metadata) + if endpoint is None or self._zmq_ctx is None: + return + host, port = endpoint + zmq_addr = f"tcp://{host}:{port}" + sock = self._get_req_socket(zmq_addr) + try: + sock.send( + _XFER_DONE_MSG + + msgspec.msgpack.encode( + {"key": get_key, "generation": metadata.get("generation"), "claim_id": metadata.get("claim_id")} + ) + ) + sock.recv() + except Exception: + self._invalidate_req_socket(zmq_addr) + logger.debug("NixlConnector failed to notify completion for %s", get_key, exc_info=True) + + def _handshake_listener_loop(self) -> None: + router = self._zmq_ctx.socket(zmq.ROUTER) + try: + if self._zmq_port == 0: + self._zmq_port = router.bind_to_random_port(f"tcp://{self.host}") + else: + router.bind(f"tcp://{self.host}:{self._zmq_port}") + except zmq.ZMQError as exc: + logger.error("NixlConnector handshake bind failed on %s:%s: %s", self.host, self._zmq_port, exc) + self._bind_error = exc + self._listener_ready.set() + router.close(linger=0) + return + self._listener_ready.set() + + poller = zmq.Poller() + poller.register(router, zmq.POLLIN) + try: + while not self._stop_event.is_set(): + try: + if not dict(poller.poll(500)): + self._cleanup_expired_pending() + continue + identity, _, payload = router.recv_multipart() + router.send_multipart([identity, b"", self._handle_handshake_message(payload)]) + except zmq.ContextTerminated: + break + except Exception: + logger.debug("NixlConnector handshake listener error", exc_info=True) + finally: + self._safe_call(router.close, 0) + + def _handle_handshake_message(self, payload: bytes) -> bytes: + if payload.startswith(_GET_META_MSG): + request = msgspec.msgpack.decode(payload[len(_GET_META_MSG) :]) + key = request.get("key") + claim = request.get("claim_id") + with self._state_lock: + pending = self._pending.get(key) + metadata = self._published.get(key) + if ( + pending is None + or metadata is None + or not claim + or request.get("generation") not in (None, pending.generation) + or (not pending.claims and time.monotonic() >= pending.deadline) + or self._closed + or getattr(self, "_closing", False) + ): + return _META_NOT_FOUND + # Publish ownership before descriptors can leave this lock. + # Failed/lost replies retain the claim conservatively. + reply = msgspec.msgpack.encode({**metadata, "claim_id": claim}) + pending.claims.add(claim) + return reply + if payload.startswith(_XFER_DONE_MSG): + request = msgspec.msgpack.decode(payload[len(_XFER_DONE_MSG) :]) + key = request.get("key") + with self._state_lock: + pending = self._pending.get(key) + if ( + pending is not None + and request.get("generation") == pending.generation + and request.get("claim_id") in pending.claims + ): + pending.claims.remove(request["claim_id"]) + self.cleanup(key) + return _ACK + logger.warning("NixlConnector handshake received an unknown message") + return _META_NOT_FOUND + + def _get_req_socket(self, zmq_addr: str, timeout_ms: int | None = None) -> zmq.Socket: + """Return a thread-local REQ socket so concurrent calls never interleave.""" + cache: dict[str, zmq.Socket] | None = getattr(self._req_local, "cache", None) + if cache is None: + cache = {} + self._req_local.cache = cache + sock = cache.get(zmq_addr) + if sock is None: + sock = self._zmq_ctx.socket(zmq.REQ) + sock.connect(zmq_addr) + cache[zmq_addr] = sock + timeout_ms = self._handshake_timeout_ms if timeout_ms is None else timeout_ms + sock.setsockopt(zmq.SNDTIMEO, timeout_ms) + sock.setsockopt(zmq.RCVTIMEO, timeout_ms) + return sock + + def _invalidate_req_socket(self, zmq_addr: str) -> None: + cache: dict[str, zmq.Socket] | None = getattr(self._req_local, "cache", None) + if cache is None: + return + sock = cache.pop(zmq_addr, None) + if sock is not None: + self._safe_call(sock.close, 0) + + @staticmethod + def _get_local_ip() -> str: + """Resolve the externally routable local address for the handshake bind.""" + try: + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as probe: + probe.connect(("8.8.8.8", 80)) + return probe.getsockname()[0] + except Exception: + try: + return socket.gethostbyname(socket.gethostname()) + except Exception: + logger.warning("NixlConnector could not resolve a local IP; binding loopback") + return "127.0.0.1" + + def _wait_for_transfer(self, handle: int, request_id: str) -> None: + deadline = time.monotonic() + self._transfer_timeout_s + while True: + state = self._agent.check_xfer_state(handle) + if state == "DONE": + return + if state != "PROC": + raise RuntimeError(f"NIXL transfer for {request_id} failed with state={state}") + if time.monotonic() >= deadline: + raise TimeoutError(f"NIXL transfer for {request_id} timed out") + time.sleep(self._poll_interval_s) + + def _transfers_may_be_active(self, handles: list[Any]) -> bool: + try: + return any(self._agent.check_xfer_state(handle) not in {"DONE", "ERR"} for handle in handles) + except Exception: + return True + + def _cleanup_expired_pending(self) -> None: + now = time.monotonic() + with self._state_lock: + expired = [ + (key, pending, sum(t.numel() * t.element_size() for t in pending.tensors)) + for key, pending in self._pending.items() + if not pending.claims and now >= pending.deadline + ] + for request_id, pending, size in expired: + logger.warning( + "NixlConnector lease expired for request %s after %.0fs; its %d bytes are " + "unclaimed and can safely be reclaimed.", + request_id, + self._lease_seconds, + size, + ) + claimed = self._take_pending(request_id, expected=pending) + if claimed is not None: + self._release_pending(claimed) + + def _lease_reaper_loop(self) -> None: + while not self._stop_event.is_set(): + now = time.monotonic() + with self._state_lock: + next_deadline = min( + (pending.deadline for pending in self._pending.values() if not pending.claims), default=None + ) + if next_deadline is None: + timeout = None + else: + timeout = max(0.0, next_deadline - now) + self._lease_wakeup.wait(timeout=timeout) + self._lease_wakeup.clear() + if not self._stop_event.is_set(): + self._cleanup_expired_pending() + + def _defer_transfer(self, transfer: _DeferredTransfer) -> None: + with self._state_lock: + self._deferred_transfers.append(transfer) + self._transfer_wakeup.set() + + def _transfer_reaper_loop(self) -> None: + while not self._stop_event.is_set(): + self._transfer_wakeup.wait(timeout=max(self._poll_interval_s, 0.01)) + self._transfer_wakeup.clear() + self._reap_deferred_transfers() + + def _reap_deferred_transfers(self) -> None: + with self._state_lock: + transfers = list(self._deferred_transfers) + for transfer in transfers: + try: + states = [self._agent.check_xfer_state(handle) for handle in transfer.handles] + except Exception: + logger.debug("Failed to poll deferred NIXL transfer", exc_info=True) + continue + if any(state not in {"DONE", "ERR"} for state in states): + continue + self._release_deferred_transfer(transfer) + if ( + not transfer.handles + and not transfer.dlists + and not transfer.registrations + and transfer.remote_agent is None + ): + transfer.tensors.clear() + if transfer.source_metadata is not None: + self._notify_transfer_done(transfer.source_key, transfer.source_metadata) + transfer.source_metadata = None + with self._state_lock: + if transfer in self._deferred_transfers: + self._deferred_transfers.remove(transfer) + + def _release_deferred_transfer(self, transfer: _DeferredTransfer) -> None: + transfer.handles = self._release_owned_resources(self._agent.release_xfer_handle, transfer.handles) + transfer.dlists = self._release_owned_resources(self._agent.release_dlist_handle, transfer.dlists) + if transfer.remote_agent is not None: + try: + self._agent.remove_remote_agent(transfer.remote_agent) + except Exception: + logger.debug("Failed to remove deferred NIXL remote agent", exc_info=True) + else: + if transfer.remote_agent in self._remote_agents: + self._remote_agents.remove(transfer.remote_agent) + transfer.remote_agent = None + transfer.registrations = self._release_owned_resources(self._agent.deregister_memory, transfer.registrations) + + @staticmethod + def _release_owned_resources(release: Any, resources: list[Any]) -> list[Any]: + remaining = [] + for resource in resources: + try: + release(resource) + except Exception: + logger.debug("Failed to release deferred NIXL resource", exc_info=True) + remaining.append(resource) + return remaining + + def _resolve_memory_type(self, tensor: torch.Tensor) -> str: + if self._default_memory_type is not None: + return str(self._default_memory_type) + if tensor.device.type == "cpu": + return "DRAM" + return "VRAM" + + @staticmethod + def _validated_descriptor_groups( + metadata: dict[str, Any], tensor_specs: list[dict[str, Any]] + ) -> list[dict[str, Any]]: + # Derive emptiness from shape/dtype rather than the advertised size. + sizes = [] + for spec in tensor_specs: + if not isinstance(spec, dict): + raise RuntimeError("Invalid NIXL metadata: tensor spec must be a mapping") + shape = spec.get("shape") + if not isinstance(shape, list) or any(type(dim) is not int or dim < 0 for dim in shape): + raise RuntimeError("Invalid NIXL metadata: invalid tensor shape") + dtype = getattr(torch, str(spec.get("dtype", "")).removeprefix("torch."), None) + if not isinstance(dtype, torch.dtype): + raise RuntimeError("Invalid NIXL metadata: invalid tensor dtype") + size = math.prod(shape) * dtype.itemsize + if type(spec.get("size")) is not int or spec["size"] != size: + raise RuntimeError("Invalid NIXL metadata: tensor size does not match shape/dtype") + sizes.append(size) + groups = metadata.get("descriptor_groups") + if not isinstance(groups, list): + raise RuntimeError("Invalid NIXL metadata: missing descriptor_groups") + + seen_indices = [] + for group in groups: + if not isinstance(group, dict): + raise RuntimeError("Invalid NIXL metadata: descriptor group must be a mapping") + indices = group.get("tensor_indices") + regions = group.get("regions") + if not isinstance(indices, list) or not isinstance(regions, list) or len(indices) != len(regions): + raise RuntimeError("Invalid NIXL metadata: tensor_indices and regions must have equal lengths") + if not indices: + raise RuntimeError("Invalid NIXL metadata: empty descriptor group") + if not group.get("memory_type"): + raise RuntimeError("Invalid NIXL metadata: descriptor group is missing memory_type") + if any(type(index) is not int or not 0 <= index < len(sizes) for index in indices): + raise RuntimeError("Invalid NIXL metadata: tensor indices must form an exact partition") + for index, region in zip(indices, regions, strict=True): + if ( + not isinstance(region, (list, tuple)) + or len(region) not in (3, 4) + or type(region[0]) is not int + or region[0] <= 0 + or type(region[1]) is not int + or region[1] <= 0 + or region[1] != sizes[index] + or type(region[2]) is not int + or region[2] < 0 + ): + raise RuntimeError("Invalid NIXL metadata: invalid tensor region or byte size") + seen_indices.extend(indices) + + if sorted(seen_indices) != [index for index, size in enumerate(sizes) if size > 0]: + raise RuntimeError("Invalid NIXL metadata: tensor indices must form an exact partition") + return groups + + @staticmethod + def _tensor_region(tensor: torch.Tensor) -> tuple[int, int, int, str]: + device_id = max(tensor.get_device(), 0) if tensor.device.type != "cpu" else 0 + return (tensor.data_ptr(), tensor.numel() * tensor.element_size(), device_id, "") + + @staticmethod + def _is_tensor_payload(payload: Any) -> bool: + return isinstance(payload, torch.Tensor) or ( + isinstance(payload, (list, tuple)) + and bool(payload) + and all(isinstance(item, torch.Tensor) for item in payload) + ) + + @classmethod + def _contains_tensor(cls, payload: Any) -> bool: + if isinstance(payload, torch.Tensor): + return True + if isinstance(payload, msgspec.Struct): + payload = msgspec.structs.asdict(payload) + if isinstance(payload, dict): + return any(cls._contains_tensor(value) for value in payload.values()) + if isinstance(payload, (list, tuple)): + return any(cls._contains_tensor(value) for value in payload) + return False + + @classmethod + def _extract_tensor_leaves(cls, payload: Any) -> tuple[Any, list[torch.Tensor]]: + tensors: list[torch.Tensor] = [] + + def visit(value: Any) -> Any: + if isinstance(value, torch.Tensor): + index = len(tensors) + tensors.append(value) + return {_TENSOR_MARKER: index} + if isinstance(value, msgspec.Struct): + value = msgspec.structs.asdict(value) + if isinstance(value, dict): + return {key: visit(item) for key, item in value.items()} + if isinstance(value, list): + return [visit(item) for item in value] + if isinstance(value, tuple): + return {_TUPLE_MARKER: [visit(item) for item in value]} + return value + + return visit(payload), tensors + + @classmethod + def _restore_tensor_leaves(cls, skeleton: Any, tensors: list[torch.Tensor]) -> Any: + if isinstance(skeleton, dict): + if set(skeleton) == {_TENSOR_MARKER}: + return tensors[int(skeleton[_TENSOR_MARKER])] + if set(skeleton) == {_TUPLE_MARKER}: + return tuple(cls._restore_tensor_leaves(item, tensors) for item in skeleton[_TUPLE_MARKER]) + return {key: cls._restore_tensor_leaves(value, tensors) for key, value in skeleton.items()} + if isinstance(skeleton, list): + return [cls._restore_tensor_leaves(item, tensors) for item in skeleton] + return skeleton + + @staticmethod + def _normalize_tensor_payload(payload: Any) -> tuple[list[torch.Tensor], list[dict[str, Any]]]: + tensors = [payload] if isinstance(payload, torch.Tensor) else list(payload) + normalized: list[torch.Tensor] = [] + specs: list[dict[str, Any]] = [] + for tensor in tensors: + contiguous = tensor.detach().contiguous() + normalized.append(contiguous) + specs.append( + { + "shape": list(contiguous.shape), + "dtype": str(contiguous.dtype), + "device": str(contiguous.device), + "size": contiguous.numel() * contiguous.element_size(), + } + ) + return normalized, specs + + def _allocate_tensor_from_spec(self, spec: dict[str, Any], kind: str | None) -> torch.Tensor: + shape = spec.get("shape") + if not isinstance(shape, list): + raise RuntimeError(f"Invalid NIXL tensor shape: {shape!r}") + dtype_name = str(spec.get("dtype", "")).removeprefix("torch.") + dtype = getattr(torch, dtype_name, None) + if dtype is None: + raise RuntimeError(f"Unsupported NIXL tensor dtype: {spec.get('dtype')!r}") + device = torch.device("cpu") if kind == _KIND_OBJECT else self._resolve_receive_device(spec.get("device")) + return torch.empty(tuple(int(dim) for dim in shape), dtype=dtype, device=device) + + def _resolve_receive_device(self, spec_device: Any) -> torch.device: + if self._receive_device is not None: + return self._receive_device + device = self._parse_device(spec_device) + if device is None or device.type == "cpu": + return torch.device("cpu") + # The producer's device index is meaningless in this process, so keep + # only its type and land the buffer on this stage's own card. + backend = getattr(torch, device.type, None) + index = backend.current_device() if hasattr(backend, "current_device") else 0 + return torch.device(device.type, index) + + @staticmethod + def _parse_device(device_like: Any) -> torch.device | None: + if device_like is None: + return None + try: + return torch.device(device_like) + except Exception as exc: + raise RuntimeError(f"Invalid NIXL receive device: {device_like!r}") from exc + + @staticmethod + def _safe_call(func: Any, *args: Any) -> None: + try: + func(*args) + except Exception: + logger.debug("Ignoring NIXL cleanup failure", exc_info=True) diff --git a/vllm_omni/distributed/omni_connectors/factory.py b/vllm_omni/distributed/omni_connectors/factory.py index c49da1c6c..66dbca079 100644 --- a/vllm_omni/distributed/omni_connectors/factory.py +++ b/vllm_omni/distributed/omni_connectors/factory.py @@ -125,6 +125,17 @@ def _create_mori_transfer_engine_connector(config: dict[str, Any]) -> OmniConnec return MoriTransferEngineConnector(config) +def _create_nixl_connector(config: dict[str, Any]) -> OmniConnectorBase: + try: + from .connectors.nixl_connector import NixlConnector + except ImportError: + import sys + + sys.path.insert(0, os.path.dirname(os.path.dirname(__file__))) + from omni_connectors.connectors.nixl_connector import NixlConnector + return NixlConnector(config) + + # Register connectors OmniConnectorFactory.register_connector("MooncakeStoreConnector", _create_mooncake_store_connector) OmniConnectorFactory.register_connector("MooncakeTransferEngineConnector", _create_mooncake_transfer_engine_connector) @@ -132,5 +143,6 @@ OmniConnectorFactory.register_connector("SharedMemoryConnector", _create_shm_con OmniConnectorFactory.register_connector("YuanrongConnector", _create_yuanrong_connector) OmniConnectorFactory.register_connector("YuanrongTransferEngineConnector", _create_yuanrong_transfer_engine_connector) OmniConnectorFactory.register_connector("MoriTransferEngineConnector", _create_mori_transfer_engine_connector) +OmniConnectorFactory.register_connector("NixlConnector", _create_nixl_connector) # Backward-compatible aliases – will be removed in the future OmniConnectorFactory.register_connector("MooncakeConnector", _create_mooncake_store_connector) diff --git a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py index 6b83349f2..6ab0b36a2 100644 --- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py +++ b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py @@ -564,6 +564,13 @@ class OmniKVTransferManager: # the wrong process. Explicit sender_* YAML values # are still preserved for standalone connector use. + if c_type == "NixlConnector" and c_extra.get("role") == "receiver": + # This manager receives on the shared incoming edge; + # its YAML zmq_port belongs to the producer. NIXL treats + # an explicit zmq_port as a local listener, even for a + # receiver. Keep only the request-scoped sender_* endpoint. + c_extra.pop("zmq_port", None) + logger.info( "Initializing OmniConnector type=%s role=%s", c_type, @@ -922,7 +929,9 @@ class OmniKVTransferManager: The base host/port are also stored so that the receive path can construct per-rank metadata for heterogeneous TP scenarios. """ - if not self.config.need_recv_cache: + # A stage can sit on the receiving end of a payload-only edge, where the + # sender endpoint still has to be applied even though no KV is expected. + if not self.config.need_recv_cache and (self.config.connector_config or {}).get("role") != "receiver": return actual_info = self._resolve_sender_info(sender_info, sender_stage_id=sender_stage_id) diff --git a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py index 25515cfaf..2343315b9 100644 --- a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py +++ b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py @@ -182,14 +182,19 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): from_stage: str, to_stage: str, connector_get_key: str, + metadata: dict[str, Any] | None = None, ) -> Any: """Receive one ordinary non-KV stage payload on the local leader rank only.""" tp_group = self._get_local_tp_group() if tp_group is None or getattr(tp_group, "world_size", 1) <= 1: - return connector.get(from_stage, to_stage, connector_get_key) + if metadata is None: + return connector.get(from_stage, to_stage, connector_get_key) + return connector.get(from_stage, to_stage, connector_get_key, metadata) if not self.is_data_transfer_rank(): return None - return connector.get(from_stage, to_stage, connector_get_key) + if metadata is None: + return connector.get(from_stage, to_stage, connector_get_key) + return connector.get(from_stage, to_stage, connector_get_key, metadata) def _recv_full_payload_result( self, @@ -197,6 +202,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): from_stage: str, to_stage: str, connector_get_key: str, + metadata: dict[str, Any] | None = None, ) -> Any: """Receive one full-payload transfer on the local leader rank only.""" return self._recv_ordinary_stage_result( @@ -204,6 +210,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): from_stage, to_stage, connector_get_key, + metadata, ) def _recv_async_chunk_result( @@ -212,6 +219,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): from_stage: str, to_stage: str, connector_get_key: str, + metadata: dict[str, Any] | None = None, ) -> Any: """Receive one ordinary async chunk on the local leader rank only.""" return self._recv_ordinary_stage_result( @@ -219,6 +227,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): from_stage, to_stage, connector_get_key, + metadata, ) @staticmethod @@ -1027,6 +1036,14 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): chunk_id = self._get_req_chunk[req_id] external_req_id = self._request_ids_mapping.get(req_id, req_id) connector_get_key = f"{external_req_id}_{target_stage_id}_{chunk_id}" + request = self._pending_load_reqs.get(req_id) + sender_info = getattr(request, "payload_sender_info", None) + metadata = None + if isinstance(sender_info, dict): + host = sender_info.get("host") + port = sender_info.get("zmq_port") + if host and port: + metadata = {"source_host": str(host), "source_port": int(port)} if self._async_chunk: result = self._recv_async_chunk_result( @@ -1034,6 +1051,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): str(target_stage_id), str(self._stage_id), connector_get_key, + metadata, ) else: result = self._recv_full_payload_result( @@ -1041,6 +1059,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): str(target_stage_id), str(self._stage_id), connector_get_key, + metadata, ) if result is None: diff --git a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py index 275a2c0af..95a6cac22 100644 --- a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py +++ b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py @@ -19,6 +19,8 @@ from vllm_omni.distributed.omni_connectors.utils.config import ( ConnectorSpec, get_stage_connector_role, ) +from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_connector_spec +from vllm_omni.distributed.omni_connectors.utils.kv_utils import get_local_tp_rank, get_omni_replica_id logger = init_logger("vllm_omni.worker.omni_connector_model_runner_mixin") @@ -186,7 +188,11 @@ class _OmniConnectorRuntimeMixin: getattr(topology, "target_tp_size", None), getattr(topology, "local_rank", None), ) - if all(isinstance(value, int) for value in effective_mapping): + if ( + isinstance(effective_mapping[0], int) + and isinstance(effective_mapping[1], int) + and isinstance(effective_mapping[2], int) + ): rank_cfg = { "from_tp": effective_mapping[0], "to_tp": effective_mapping[1], @@ -521,7 +527,15 @@ class _OmniConnectorRuntimeMixin: elif not isinstance(extra, dict): raise RuntimeError(f"Invalid extra config for connector {name}: expected dict, got {type(extra).__name__}") - spec = ConnectorSpec(name=name, extra=extra) + stage_id = int(getattr(model_config, "stage_id", extra.get("stage_id", 0))) + role = extra.get("role") + spec = resolve_connector_spec( + ConnectorSpec(name=name, extra=extra), + stage_id=stage_id, + role=str(role) if role is not None else None, + local_rank=get_local_tp_rank(), + replica_id=get_omni_replica_id(), + ) try: return OmniConnectorFactory.create_connector(spec) except Exception as exc: diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index 163363931..b801622f9 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -20,6 +20,8 @@ from vllm_omni.data_entry_keys import MetaStruct, OmniPayloadStruct, unflatten_p from ..adapter import construct_next_stage_streaming_input_prompt from ..factory import OmniConnectorFactory from ..utils.config import ConnectorSpec, stage_receives_chunks +from ..utils.initialization import resolve_connector_spec +from ..utils.kv_utils import get_local_tp_rank, get_omni_replica_id from ..utils.logging import get_connector_logger from .base import OmniTransferAdapterBase @@ -47,10 +49,20 @@ class _SenderGeneration: class _LoadEntry: """Identify one receiver registration across queue and I/O boundaries.""" - __slots__ = ("request",) + __slots__ = ("request", "source_metadata") def __init__(self, request: Request) -> None: self.request = request + sender_info = getattr(request, "payload_sender_info", None) + self.source_metadata = None + if isinstance(sender_info, dict): + host = sender_info.get("host") + port = sender_info.get("zmq_port") + if host and port: + self.source_metadata = { + "source_host": str(host), + "source_port": int(port), + } @property def request_id(self) -> str: @@ -280,11 +292,20 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): "extra": getattr(connector_config, "extra", {}), } - connector_specs = ConnectorSpec( + connector_spec = ConnectorSpec( name=connector_config.get("name", "SharedMemoryConnector"), extra=connector_config.get("extra", {}), ) - return OmniConnectorFactory.create_connector(connector_specs) + stage_id = int(getattr(model_config, "stage_id", connector_spec.extra.get("stage_id", 0))) + role = connector_spec.extra.get("role") + resolved_spec = resolve_connector_spec( + connector_spec, + stage_id=stage_id, + role=str(role) if role is not None else None, + local_rank=get_local_tp_rank(), + replica_id=get_omni_replica_id(), + ) + return OmniConnectorFactory.create_connector(resolved_spec) def load_async(self, request: Request): """Register a request for asynchronous chunk retrieval. @@ -484,6 +505,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): str(target_stage_id), str(stage_id), connector_get_key, + entry.source_metadata, ) except Exception as e: logger.error(f"SharedMemoryConnector get failed for req {connector_get_key}: {e}") diff --git a/vllm_omni/distributed/omni_connectors/utils/config.py b/vllm_omni/distributed/omni_connectors/utils/config.py index 1a6d395c2..176a8bf69 100644 --- a/vllm_omni/distributed/omni_connectors/utils/config.py +++ b/vllm_omni/distributed/omni_connectors/utils/config.py @@ -13,6 +13,7 @@ TRANSFER_ENGINE_CONNECTOR_NAMES = frozenset( "MooncakeTransferEngineConnector", "MoriTransferEngineConnector", "YuanrongTransferEngineConnector", + "NixlConnector", } ) diff --git a/vllm_omni/distributed/omni_connectors/utils/initialization.py b/vllm_omni/distributed/omni_connectors/utils/initialization.py index bbbda44e3..d0139b59a 100644 --- a/vllm_omni/distributed/omni_connectors/utils/initialization.py +++ b/vllm_omni/distributed/omni_connectors/utils/initialization.py @@ -4,9 +4,10 @@ """Utilities for OmniConnector configuration and validation.""" import json +import os import warnings from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal from ..factory import OmniConnectorFactory from .config import TRANSFER_ENGINE_CONNECTOR_NAMES, ConnectorSpec, OmniTransferConfig @@ -38,6 +39,91 @@ KV_RANK_PORT_STRIDE = 16 # comfortably sized block per replica for TP-rank and stage offsets. KV_REPLICA_PORT_STRIDE = 1024 +REQUEST_FORWARDING_PORT_OFFSET = 0 +ORCHESTRATOR_PORT_OFFSET = 200 + +ConnectorPurpose = Literal["request_forwarding", "kv_transfer", "orchestrator"] + + +def compute_connector_zmq_port( + base_port: int, + *, + purpose: ConnectorPurpose, + from_stage: int, + local_rank: int = 0, + replica_id: int = 0, +) -> int: + """Return a collision-free connector endpoint for an edge and worker.""" + purpose_offsets = { + "request_forwarding": REQUEST_FORWARDING_PORT_OFFSET, + "kv_transfer": KV_TRANSFER_PORT_OFFSET, + "orchestrator": ORCHESTRATOR_PORT_OFFSET, + } + if purpose not in purpose_offsets: + raise ValueError(f"Unsupported connector purpose: {purpose}") + if not 0 <= from_stage < KV_RANK_PORT_STRIDE: + raise ValueError(f"from_stage must be in [0, {KV_RANK_PORT_STRIDE}), got {from_stage}") + if local_rank < 0: + raise ValueError(f"local_rank must be non-negative, got {local_rank}") + if replica_id < 0: + raise ValueError(f"replica_id must be non-negative, got {replica_id}") + + port = ( + base_port + + purpose_offsets[purpose] + + replica_id * KV_REPLICA_PORT_STRIDE + + local_rank * KV_RANK_PORT_STRIDE + + from_stage + ) + if not 1 <= port <= 65535: + raise ValueError(f"Resolved connector port must be in [1, 65535], got {port}") + return port + + +def resolve_connector_spec( + spec: ConnectorSpec, + *, + stage_id: int, + role: str | None, + purpose: ConnectorPurpose = "request_forwarding", + local_rank: int = 0, + replica_id: int = 0, +) -> ConnectorSpec: + """Copy and resolve a connector spec for one stage worker.""" + extra = dict(spec.extra or {}) + extra["stage_id"] = stage_id + if role is not None: + extra["role"] = role + if spec.name not in TRANSFER_ENGINE_CONNECTOR_NAMES: + return ConnectorSpec(name=spec.name, extra=extra) + + resolved_port = None + if role != "receiver" or "sender_zmq_port" not in extra: + resolved_port = compute_connector_zmq_port( + expand_env_int(extra.get("zmq_port", 50051), "zmq_port"), + purpose=purpose, + from_stage=int(extra.get("from_stage", stage_id)), + local_rank=local_rank, + replica_id=replica_id, + ) + if role == "receiver": + extra.setdefault("sender_host", extra.get("host", "127.0.0.1")) + extra.setdefault("sender_zmq_port", resolved_port) + extra.pop("zmq_port", None) + outgoing = extra.get("outgoing") + if spec.name == "NixlConnector" and isinstance(outgoing, dict): + extra["zmq_port"] = compute_connector_zmq_port( + expand_env_int(outgoing.get("zmq_port", 50051), "zmq_port"), + purpose=purpose, + from_stage=int(outgoing.get("from_stage", stage_id)), + local_rank=local_rank, + replica_id=replica_id, + ) + extra["host"] = outgoing.get("host", "auto") + else: + extra["zmq_port"] = resolved_port + return ConnectorSpec(name=spec.name, extra=extra) + def initialize_connectors_from_config( config_path: str | Path | None = None, @@ -82,52 +168,40 @@ def create_connectors_from_config( Returns: A dictionary of connectors. """ - purpose_port_offsets = { - "request_forwarding": 0, - "kv_transfer": KV_TRANSFER_PORT_OFFSET, - } - port_offset = purpose_port_offsets.get(purpose, 0) - orchestrator_port_offset = 200 - connectors = {} for edge_key, connector_spec in connectors_config.items(): from_stage, to_stage = edge_key try: if connector_spec.name in TRANSFER_ENGINE_CONNECTOR_NAMES: extra = dict(connector_spec.extra) if connector_spec.extra else {} - base_port = expand_env_int(extra.get("zmq_port", 50051), "zmq_port") try: - stage_offset = int(from_stage) - except (TypeError, ValueError): - stage_offset = 0 - - local_rank = 0 if str(caller_stage_id) == "orchestrator" else get_connector_local_rank() - if str(caller_stage_id) == "orchestrator": - rank0_port = base_port + orchestrator_port_offset + stage_offset - else: - rank0_port = base_port + port_offset + stage_offset - adjusted_port = rank0_port + local_rank * KV_RANK_PORT_STRIDE - extra["zmq_port"] = adjusted_port + from_stage_id = int(from_stage) + except (TypeError, ValueError) as exc: + raise ValueError(f"Transfer-engine stage id must be an integer, got {from_stage!r}") from exc + is_orchestrator = str(caller_stage_id) == "orchestrator" + local_rank = 0 if is_orchestrator else get_connector_local_rank() + resolved_purpose: ConnectorPurpose = "orchestrator" if is_orchestrator else purpose # type: ignore[assignment] if is_sender is not None: - extra["role"] = "sender" if is_sender else "receiver" - if not is_sender: - extra.setdefault("sender_host", extra.get("host", "127.0.0.1")) - extra.setdefault("sender_zmq_port", adjusted_port) - elif caller_stage_id is not None: - caller_str = str(caller_stage_id) - if caller_str == from_stage: - extra["role"] = "sender" - elif caller_str == to_stage: - extra["role"] = "receiver" - extra.setdefault("sender_host", extra.get("host", "127.0.0.1")) - extra.setdefault("sender_zmq_port", adjusted_port) - else: - extra["role"] = "sender" + role = "sender" if is_sender else "receiver" + elif caller_stage_id is not None and str(caller_stage_id) == to_stage: + role = "receiver" + elif caller_stage_id is not None and str(caller_stage_id) == from_stage: + role = "sender" else: - extra["role"] = extra.get("role", "auto") - - connector = OmniConnectorFactory.create_connector(ConnectorSpec(name=connector_spec.name, extra=extra)) + role = str(extra.get("role", "auto")) + extra["from_stage"] = from_stage_id + extra["to_stage"] = int(to_stage) + replica_id = max(int(os.environ.get("VLLM_OMNI_REPLICA_ID", "0")), 0) + resolved_spec = resolve_connector_spec( + ConnectorSpec(name=connector_spec.name, extra=extra), + stage_id=int(caller_stage_id) if str(caller_stage_id).isdigit() else from_stage_id, + role=role, + purpose=resolved_purpose, + local_rank=local_rank, + replica_id=replica_id, + ) + connector = OmniConnectorFactory.create_connector(resolved_spec) else: connector = OmniConnectorFactory.create_connector(connector_spec) connectors[edge_key] = connector @@ -174,12 +248,16 @@ def get_connectors_config_for_stage(transfer_config: OmniTransferConfig | None, # Incoming edge → this stage is the receiver extra = dict(spec.extra) if spec.extra else {} extra.setdefault("role", "receiver") + extra["from_stage"] = int(from_stage) + extra["to_stage"] = int(to_stage) stage_connectors_config[f"from_stage_{from_stage}"] = {"spec": {"name": spec.name, "extra": extra}} elif from_stage == target_stage and target_stage == "0": # Outgoing edge for stage 0 — included for async_chunk spec # extraction (omni_stage.py), NOT for connector instantiation. extra = dict(spec.extra) if spec.extra else {} extra.setdefault("role", "sender") + extra["from_stage"] = int(from_stage) + extra["to_stage"] = int(to_stage) stage_connectors_config[f"to_stage_{to_stage}"] = {"spec": {"name": spec.name, "extra": extra}} return stage_connectors_config diff --git a/vllm_omni/distributed/omni_connectors/utils/kv_utils.py b/vllm_omni/distributed/omni_connectors/utils/kv_utils.py index b7f51460c..4be32ba7e 100644 --- a/vllm_omni/distributed/omni_connectors/utils/kv_utils.py +++ b/vllm_omni/distributed/omni_connectors/utils/kv_utils.py @@ -16,7 +16,7 @@ from vllm.distributed.parallel_state import ( ) from vllm.logger import init_logger -from .initialization import KV_RANK_PORT_STRIDE, KV_REPLICA_PORT_STRIDE, KV_TRANSFER_PORT_OFFSET +from .initialization import compute_connector_zmq_port logger = init_logger(__name__) @@ -123,12 +123,12 @@ def kv_zmq_port( """ replica = get_omni_replica_id() if replica_id is None else max(int(replica_id), 0) - return ( - base_port - + KV_TRANSFER_PORT_OFFSET - + replica * KV_REPLICA_PORT_STRIDE - + local_rank * KV_RANK_PORT_STRIDE - + from_stage + return compute_connector_zmq_port( + base_port, + purpose="kv_transfer", + from_stage=from_stage, + local_rank=local_rank, + replica_id=replica, ) diff --git a/vllm_omni/engine/__init__.py b/vllm_omni/engine/__init__.py index e72ce9fed..a96c8027f 100644 --- a/vllm_omni/engine/__init__.py +++ b/vllm_omni/engine/__init__.py @@ -79,6 +79,7 @@ class OmniEngineCoreRequest(EngineCoreRequest): # GPUModelRunner.model_intermediate_buffer instead of using the deprecated # additional_information request transport. model_intermediate_buffer: dict[str, Any] | None = None + payload_sender_info: dict[str, Any] | None = None @classmethod def from_request( @@ -88,6 +89,7 @@ class OmniEngineCoreRequest(EngineCoreRequest): prompt_embeds: torch.Tensor | None = None, additional_information: AdditionalInformationPayload | None = None, model_intermediate_buffer: dict[str, Any] | None = None, + payload_sender_info: dict[str, Any] | None = None, ) -> "OmniEngineCoreRequest": """Clone an EngineCoreRequest into an OmniEngineCoreRequest with optional payload overrides.""" @@ -97,6 +99,8 @@ class OmniEngineCoreRequest(EngineCoreRequest): additional_information = getattr(request, "additional_information", None) if model_intermediate_buffer is None: model_intermediate_buffer = getattr(request, "model_intermediate_buffer", None) + if payload_sender_info is None: + payload_sender_info = getattr(request, "payload_sender_info", None) return cls( request_id=request.request_id, @@ -121,6 +125,7 @@ class OmniEngineCoreRequest(EngineCoreRequest): abort_immediately=request.abort_immediately, additional_information=additional_information, model_intermediate_buffer=model_intermediate_buffer, + payload_sender_info=payload_sender_info, ) diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 299e1dd07..498a9dc89 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -2148,6 +2148,7 @@ class Orchestrator: *, mm_features: list | None = None, resumable: bool = False, + payload_sender_info: dict[str, Any] | None = None, ) -> Any: next_pool = self.stage_pools[next_stage_id] if self._next_stage_input_is_tokens(next_input): @@ -2160,6 +2161,7 @@ class Orchestrator: resumable=resumable, ) request.external_req_id = request.request_id + request.payload_sender_info = payload_sender_info return request processor = self._get_stage_input_processor(next_stage_id) @@ -2181,6 +2183,7 @@ class Orchestrator: ) request = self._upgrade_processed_stage_request(request, next_input) request.external_req_id = req_id + request.payload_sender_info = payload_sender_info return request @staticmethod @@ -2829,6 +2832,7 @@ class Orchestrator: params=params, mm_features=mm_features, resumable=next_stage_resumable, + payload_sender_info=self._build_payload_sender_info(src_stage_id, request_id=req_id), ) if already_submitted: @@ -2967,6 +2971,7 @@ class Orchestrator: model_config=next_pool.stage_vllm_config.model_config, resumable=downstream_resumable, ) + request.payload_sender_info = self._build_payload_sender_info(next_stage_id - 1, request_id=request_id) request.external_req_id = request.request_id submitted = await self._dispatch_or_fail_request( lambda: next_pool.submit_initial( @@ -3043,6 +3048,23 @@ class Orchestrator: return sender_infos or None + def _build_payload_sender_info( + self, + sender_stage_id: int, + *, + request_id: str | None = None, + ) -> dict[str, Any] | None: + if sender_stage_id < 0 or sender_stage_id >= len(self.stage_pools): + return None + sender_pool = self.stage_pools[sender_stage_id] + sender_stage = sender_pool.get_bound_client(request_id) if request_id is not None else None + if sender_stage is None: + sender_stage = sender_pool.stage_client + get_sender_info = getattr(sender_stage, "get_payload_sender_info", None) + if not callable(get_sender_info): + return None + return get_sender_info() + # ---- Shutdown / lifecycle ---- def _shutdown_stages(self) -> None: diff --git a/vllm_omni/engine/stage_client.py b/vllm_omni/engine/stage_client.py index 2a3c26177..7ee949146 100644 --- a/vllm_omni/engine/stage_client.py +++ b/vllm_omni/engine/stage_client.py @@ -94,6 +94,8 @@ class StagePoolLLMClient(StagePoolClient, Protocol): kv_transfer_port_offset: int = ..., ) -> dict[str, Any] | None: ... + def get_payload_sender_info(self) -> dict[str, Any] | None: ... + class StagePoolDiffusionClient(StagePoolClient, Protocol): """Pool-facing API for diffusion stages.""" diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index 87baaed5f..5fdd74912 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Stage Engine Core Client for vLLM-Omni multi-stage runtime. @@ -210,6 +213,7 @@ class StageEngineCoreClientBase(StageClientBase): ) raise + self._payload_sender_info = self._build_payload_sender_info() self._initialize_kv_sender_endpoint() logger.info( @@ -291,6 +295,39 @@ class StageEngineCoreClientBase(StageClientBase): return None return connector_config + def _build_payload_sender_info(self) -> dict[str, Any] | None: + model_config = getattr(self.vllm_config, "model_config", None) + connector_config = getattr(model_config, "stage_connector_config", None) + if not isinstance(connector_config, dict): + return None + extra = connector_config.get("extra") + if not isinstance(extra, dict): + return None + if isinstance(extra.get("outgoing"), dict): + extra = extra["outgoing"] + elif extra.get("role") != "sender": + return None + base_port = extra.get("zmq_port", 50051) + if base_port is None: + return None + sender_host = self._resolve_sender_host_from_config(extra) + if sender_host is None: + return None + from vllm_omni.distributed.omni_connectors.utils.initialization import compute_connector_zmq_port + + return { + "host": sender_host, + "zmq_port": compute_connector_zmq_port( + int(os.path.expandvars(str(base_port))), + purpose="request_forwarding", + from_stage=int(extra.get("from_stage", self.stage_id)), + replica_id=self.replica_id, + ), + } + + def get_payload_sender_info(self) -> dict[str, Any] | None: + return dict(self._payload_sender_info) if self._payload_sender_info is not None else None + def _resolve_sender_host_from_config(self, connector_config: dict[str, Any]) -> str | None: host = connector_config.get("sender_host") or connector_config.get("host") if host in {None, "", "auto", "*", "0.0.0.0", "::"}: diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index e483e1bd4..a52f58925 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -88,6 +88,7 @@ class StageEngineCoreProc(EngineCoreProc): scheduler_request, current_wave = super().preprocess_add_request(request) scheduler_request.additional_information = request.additional_information scheduler_request.external_req_id = getattr(request, "external_req_id", request.request_id) + scheduler_request.payload_sender_info = request.payload_sender_info return scheduler_request, current_wave @staticmethod diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index de97f7aa3..9668dfc99 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -50,6 +50,9 @@ from vllm_omni.config.omni_config import ( ) from vllm_omni.config.stage_config import StageType from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.distributed.omni_connectors.utils.config import ( + TRANSFER_ENGINE_CONNECTOR_NAMES, +) from vllm_omni.engine.arg_utils import OmniEngineArgs from vllm_omni.entrypoints.stage_utils import _to_dict, set_stage_devices from vllm_omni.entrypoints.utils import filter_dataclass_kwargs @@ -1882,17 +1885,37 @@ def get_stage_connector_spec( stage_connectors_cfg = get_stage_connector_config(omni_transfer_config, stage_id) for cfg in stage_connectors_cfg.values(): - return dict(cfg.get("spec", {})) + connector_spec = dict(cfg.get("spec", {})) + if connector_spec.get("name") == "NixlConnector": + # A middle worker uses one connector for both directions. Preserve + # the incoming edge and carry the outgoing bind config separately. + for (source, target), outgoing in getattr(omni_transfer_config, "connectors", {}).items(): + if source == str(stage_id): + if outgoing.name != "NixlConnector": + raise ValueError("A NIXL middle stage requires NIXL on its outgoing edge") + extra = dict(connector_spec.get("extra", {})) + extra["outgoing"] = {**(outgoing.extra or {}), "from_stage": int(source), "to_stage": int(target)} + connector_spec["extra"] = extra + break + if connector_spec.get("name") not in TRANSFER_ENGINE_CONNECTOR_NAMES: + extra = dict(connector_spec.get("extra", {})) + extra.pop("from_stage", None) + extra.pop("to_stage", None) + connector_spec["extra"] = extra + return connector_spec # A producer does not consume connector data itself. Keep its connector # for both async-chunk and terminal full-payload sends, but mark it # sender-only so the scheduler does not park orchestrator-provided inputs # waiting for an upstream payload. target_stage = str(stage_id) - for (from_stage, _to_stage), spec in getattr(omni_transfer_config, "connectors", {}).items(): + for (from_stage, to_stage), spec in getattr(omni_transfer_config, "connectors", {}).items(): if from_stage == target_stage: extra = dict(spec.extra or {}) extra.setdefault("role", "sender") + if spec.name in TRANSFER_ENGINE_CONNECTOR_NAMES: + extra["from_stage"] = int(from_stage) + extra["to_stage"] = int(to_stage) return {"name": spec.name, "extra": extra} return {}