Ec connector nixl transfer

tests/v1/ec_connector/_fakes.py

@@ -0,0 +1,161 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""In-process fake of the NixlWrapper subset the EC connector uses.
+
+WRITE transfers synchronously memcpy bytes between registered regions
+so end-to-end tests can assert real data arrival without a real NIXL
+runtime.
+"""
+
+from __future__ import annotations
+
+import itertools
+from dataclasses import dataclass
+
+# Global registry so producer fake and consumer fake can "see" each other.
+_AGENTS: dict[str, FakeNixlWrapper] = {}
+
+
+@dataclass
+class _RegEntry:
+    base_addr: int
+    size: int
+
+
+@dataclass
+class _XferDescList:
+    entries: list[tuple[int, int, int]]  # (addr, size, device_id)
+
+
+@dataclass
+class _Dlist:
+    agent_name: str  # owner of the memory
+    descs: _XferDescList
+    handle: int
+
+
+@dataclass
+class _Xfer:
+    op: str
+    local_dlist: _Dlist
+    local_ids: list[int]
+    remote_dlist: _Dlist
+    remote_ids: list[int]
+    notif_msg: bytes
+    state: str = "PROC"
+    error: str | None = None
+
+
+class FakeNixlWrapper:
+    def __init__(self, name: str, config: object | None = None) -> None:
+        self.name = name
+        self._config = config
+        self._registrations: list[_RegEntry] = []
+        self._remote_agents: dict[str, FakeNixlWrapper] = {}
+        self._dlists: dict[int, _Dlist] = {}
+        self._xfers: dict[int, _Xfer] = {}
+        self._dlist_counter = itertools.count(1)
+        self._xfer_counter = itertools.count(1)
+        _AGENTS[name] = self
+
+    # --- agent lifecycle -------------------------------------------------
+
+    def get_agent_metadata(self) -> bytes:
+        # Encode the agent name so add_remote_agent can look it up.
+        return self.name.encode("utf-8")
+
+    def add_remote_agent(self, metadata: bytes) -> str:
+        remote_name = metadata.decode("utf-8")
+        self._remote_agents[remote_name] = _AGENTS[remote_name]
+        return remote_name
+
+    def remove_remote_agent(self, name: str) -> None:
+        self._remote_agents.pop(name, None)
+
+    # --- memory registration --------------------------------------------
+
+    def get_reg_descs(self, data, mem_type):  # noqa: ARG002 — mem_type unused in fake
+        return list(data)  # opaque pass-through
+
+    def register_memory(self, descs, backends=None):  # noqa: ARG002
+        for addr, size, _dev, _lbl in descs:
+            self._registrations.append(_RegEntry(base_addr=addr, size=size))
+
+    def deregister_memory(self, descs) -> None:
+        for addr, size, _dev, _lbl in descs:
+            self._registrations = [
+                r
+                for r in self._registrations
+                if not (r.base_addr == addr and r.size == size)
+            ]
+
+    # --- xfer descriptors + dlists --------------------------------------
+
+    def get_xfer_descs(self, blocks_data, mem_type):  # noqa: ARG002
+        return _XferDescList(entries=list(blocks_data))
+
+    def prep_xfer_dlist(self, agent_name: str, descs) -> int:
+        handle = next(self._dlist_counter)
+        # "NIXL_INIT_AGENT" maps to the local agent in the real NIXL API.
+        owner = self.name if agent_name == "NIXL_INIT_AGENT" else agent_name
+        self._dlists[handle] = _Dlist(agent_name=owner, descs=descs, handle=handle)
+        return handle
+
+    def release_dlist_handle(self, handle: int) -> None:
+        self._dlists.pop(handle, None)
+
+    # --- xfer execution --------------------------------------------------
+
+    def make_prepped_xfer(
+        self,
+        op: str,
+        local_dlist: int,
+        local_ids: list[int],
+        remote_dlist: int,
+        remote_ids: list[int],
+        notif_msg: bytes = b"",
+    ) -> int:
+        xfer_handle = next(self._xfer_counter)
+        self._xfers[xfer_handle] = _Xfer(
+            op=op,
+            local_dlist=self._dlists[local_dlist],
+            local_ids=list(local_ids),
+            remote_dlist=self._dlists[remote_dlist],
+            remote_ids=list(remote_ids),
+            notif_msg=notif_msg,
+        )
+        return xfer_handle
+
+    def transfer(self, xfer_handle: int) -> None:
+        xfer = self._xfers[xfer_handle]
+        try:
+            # WRITE: local (source) → remote (dest).  We do a raw memcpy
+            # via ctypes so tests can observe bytes landing in the dest.
+            import ctypes
+
+            assert len(xfer.local_ids) == len(xfer.remote_ids), "id length mismatch"
+
+            src_entries = xfer.local_dlist.descs.entries
+            dst_entries = xfer.remote_dlist.descs.entries
+
+            for src_i, dst_i in zip(xfer.local_ids, xfer.remote_ids, strict=True):
+                src_addr, src_size, _ = src_entries[src_i]
+                dst_addr, dst_size, _ = dst_entries[dst_i]
+                assert src_size == dst_size, f"size mismatch {src_size} vs {dst_size}"
+                ctypes.memmove(dst_addr, src_addr, src_size)
+
+            xfer.state = "DONE"
+        except Exception as e:  # noqa: BLE001 — mirror real NIXL error surface
+            xfer.state = "ERR"
+            xfer.error = str(e)
+
+    def check_xfer_state(self, xfer_handle: int) -> str:
+        return self._xfers[xfer_handle].state
+
+    def release_xfer_handle(self, xfer_handle: int) -> None:
+        self._xfers.pop(xfer_handle, None)
+
+
+def reset_agents() -> None:
+    """Clear the cross-agent registry between tests."""
+    _AGENTS.clear()

tests/v1/ec_connector/integration/test_ec_zmq_loopback.py

@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""End-to-end test with real ZMQ sockets and FakeNixlWrapper.
+
+Wires one producer ECCPUConnector and one consumer ECCPUConnector
+inside a single process, drives an XferReq through the full pipeline,
+and asserts the consumer's mmap receives the producer's tensor bytes.
+"""
+
+import time
+import uuid
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from tests.v1.ec_connector._fakes import FakeNixlWrapper, reset_agents
+from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorRole
+from vllm.distributed.ec_transfer.ec_connector.cpu_connector import (
+    ECCPUConnector,
+)
+
+
+class _MM:
+    def get_inputs_embeds_size(self):
+        return 16
+
+    model = "m"
+    dtype = torch.bfloat16
+
+
+class _P:
+    world_size = 1
+    rank = 0
+
+
+class _EC:
+    ec_connector = "ECCPUConnector"
+    engine_id = ""
+    ec_role = ""
+    ec_connector_extra_config: dict = {}
+
+    @property
+    def is_ec_producer(self):
+        return self.ec_role in ("ec_producer", "ec_both")
+
+    @property
+    def is_ec_consumer(self):
+        return self.ec_role in ("ec_consumer", "ec_both")
+
+    def get_from_extra_config(self, k, d):
+        return self.ec_connector_extra_config.get(k, d)
+
+
+class _V:
+    def __init__(self, engine_id, role, extra):
+        self.model_config = _MM()
+        self.parallel_config = _P()
+        self.ec_transfer_config = _EC()
+        self.ec_transfer_config.engine_id = engine_id
+        self.ec_transfer_config.ec_role = role
+        self.ec_transfer_config.ec_connector_extra_config = extra or {}
+
+
+def _free_port():
+    import socket
+
+    s = socket.socket()
+    s.bind(("127.0.0.1", 0))
+    p = s.getsockname()[1]
+    s.close()
+    return p
+
+
+class _FakeFeat:
+    def __init__(self, h, L):
+        self.mm_hash = h
+        self.identifier = h
+
+        class _Pos:
+            length: int = 0
+
+        self.mm_position = _Pos()
+        self.mm_position.length = L
+
+
+class _FakeReq:
+    def __init__(self, feats):
+        self.mm_features = feats
+        self.kv_transfer_params = None
+
+
+@pytest.fixture(autouse=True)
+def _reset():
+    reset_agents()
+    yield
+    reset_agents()
+
+
+def _poll_until(pred, timeout=5.0):
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        if pred():
+            return True
+        time.sleep(0.01)
+    return False
+
+
+def test_end_to_end_single_transfer(monkeypatch):
+    port = _free_port()
+    monkeypatch.setattr("vllm.envs.VLLM_EC_SIDE_CHANNEL_HOST", "127.0.0.1")
+    monkeypatch.setattr("vllm.envs.VLLM_EC_SIDE_CHANNEL_PORT", port)
+
+    prod_cfg = _V(f"p-{uuid.uuid4()}", "ec_producer", extra={"num_ec_blocks": 8})
+    cons_cfg = _V(
+        f"c-{uuid.uuid4()}",
+        "ec_consumer",
+        extra={"num_ec_blocks": 8, "default_encoder_node": f"127.0.0.1:{port}"},
+    )
+
+    with patch(
+        "vllm.distributed.ec_transfer.ec_connector.cpu_connector.NixlWrapper",
+        FakeNixlWrapper,
+    ):
+        prod = ECCPUConnector(prod_cfg, ECConnectorRole.SCHEDULER)
+        cons = ECCPUConnector(cons_cfg, ECConnectorRole.SCHEDULER)
+        try:
+            # Producer owns an encoding for mm_hash="h".
+            n = 3
+            bs = prod._block_size_bytes
+            tensor = torch.arange(n * bs, dtype=torch.uint8).contiguous()
+            prod._encodings["h"] = tensor
+
+            # Consumer kicks off the fetch.
+            req = _FakeReq([_FakeFeat("h", L=n)])
+            cons.update_state_after_alloc(req, 0)
+
+            assert _poll_until(lambda: cons.has_cache_item("h") is True), (
+                "consumer never saw transfer complete"
+            )
+
+            class _SO:
+                pass
+
+            meta = cons.build_connector_meta(_SO())
+            assert "h" in meta.mm_hash_to_cpu_blocks
+            block_indices = meta.mm_hash_to_cpu_blocks["h"]
+
+            arrived = cons._region.blocks[block_indices].flatten()
+            assert torch.equal(arrived, tensor.view(n, bs).flatten())
+        finally:
+            cons.shutdown()
+            prod.shutdown()
+
+
+def test_end_to_end_wait_then_drain(monkeypatch):
+    port = _free_port()
+    monkeypatch.setattr("vllm.envs.VLLM_EC_SIDE_CHANNEL_HOST", "127.0.0.1")
+    monkeypatch.setattr("vllm.envs.VLLM_EC_SIDE_CHANNEL_PORT", port)
+
+    prod_cfg = _V(f"p-{uuid.uuid4()}", "ec_producer", extra={"num_ec_blocks": 8})
+    cons_cfg = _V(
+        f"c-{uuid.uuid4()}",
+        "ec_consumer",
+        extra={"num_ec_blocks": 8, "default_encoder_node": f"127.0.0.1:{port}"},
+    )
+
+    with patch(
+        "vllm.distributed.ec_transfer.ec_connector.cpu_connector.NixlWrapper",
+        FakeNixlWrapper,
+    ):
+        prod = ECCPUConnector(prod_cfg, ECConnectorRole.SCHEDULER)
+        cons = ECCPUConnector(cons_cfg, ECConnectorRole.SCHEDULER)
+        try:
+            # Consumer requests BEFORE producer has the data.
+            req = _FakeReq([_FakeFeat("late", L=2)])
+            cons.update_state_after_alloc(req, 0)
+
+            assert _poll_until(lambda: "late" in prod._waiting), (
+                "producer never parked request"
+            )
+
+            # Now populate and drain.
+            n, bs = 2, prod._block_size_bytes
+            tensor = torch.full((n * bs,), 5, dtype=torch.uint8).contiguous()
+            prod._encodings["late"] = tensor
+            prod._drain_waiting("late")
+
+            assert _poll_until(lambda: cons.has_cache_item("late") is True)
+
+            class _SO:
+                pass
+
+            meta = cons.build_connector_meta(_SO())
+            idx = meta.mm_hash_to_cpu_blocks["late"]
+            arrived = cons._region.blocks[idx].flatten()
+            assert torch.all(arrived == 5)
+        finally:
+            cons.shutdown()
+            prod.shutdown()