diff --git a/docs/design/feature/disaggregated_inference.md b/docs/design/feature/disaggregated_inference.md index 83c35ac01..755d6aeea 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) - [YuanrongConnector](omni_connectors/yuanrong_connector.md) +- [YuanrongTransferEngineConnector](omni_connectors/yuanrong_transfer_engine_connector.md) ## Overview @@ -23,6 +24,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 (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"`. | ## Core API diff --git a/docs/design/feature/omni_connectors/yuanrong_transfer_engine_connector.md b/docs/design/feature/omni_connectors/yuanrong_transfer_engine_connector.md new file mode 100644 index 000000000..21d5f1841 --- /dev/null +++ b/docs/design/feature/omni_connectors/yuanrong_transfer_engine_connector.md @@ -0,0 +1,211 @@ +# YuanrongTransferEngineConnector + +## When to Use + +Use `YuanrongTransferEngineConnector` for high-performance peer-to-peer KV cache +transfer on Ascend NPU. It is different from `YuanrongConnector`: + +- `YuanrongConnector` uses Yuanrong Datasystem as a distributed KV store and + requires Datasystem workers plus etcd. +- `YuanrongTransferEngineConnector` uses Yuanrong TransferEngine directly. It + registers local memory in each stage worker and lets the receiver pull data + from the sender with TransferEngine. + +The implementation currently lives under +`vllm_omni/platforms/npu/omni_connectors/` because the supported backend is +Ascend NPU only. The connector is still registered with the generic +OmniConnector factory under the name `YuanrongTransferEngineConnector`. + +For Ascend P2P transfer, the connector should use NPU memory as its transfer +pool. The Mooncake-style CPU pool is not supported for this connector. + +## Prerequisites + +Install Yuanrong Datasystem/TransferEngine Python bindings in the runtime +environment. The packaged install is: + +```bash +pip install openyuanrong-datasystem +``` + +To build from source, follow the upstream project instructions: + +- Source repository: [openeuler/yuanrong-datasystem](https://atomgit.com/openeuler/yuanrong-datasystem) +- Installation guide: [Yuanrong Datasystem Linux installation](https://pages.openeuler.openatom.cn/openyuanrong-datasystem/docs/zh-cn/latest/installation/installation_linux.html) + +The Ascend driver, CANN runtime, HCCL/HCCP, and NPU device network must also be +configured on every machine participating in the transfer. + +## Check NPU Device IPv4 + +Yuanrong TransferEngine with `protocol: "ascend"` uses the NPU device network, +not only the host management IP. Each participating NPU must have an IPv4 HCCN +device IP. + +Check the device IP, link, and health state: + +```bash +for i in {0..7}; do + echo "===== NPU $i =====" + hccn_tool -i "$i" -ip -g + hccn_tool -i "$i" -link -g + hccn_tool -i "$i" -net_health -g +done +``` + +Expected output for each NPU should include an IPv4 address and healthy link: + +```text +ipaddr:10.86.1.80 +netmask:255.255.255.0 +link status: UP +net health status: Success +``` + +If the output says `Get ipconf failed, because no ip was preset there!`, the NPU +device IP is not configured and Ascend P2P will fail. If the address is IPv6, +configure IPv4 for this connector path or validate that the underlying Yuanrong +TransferEngine build supports the IPv6 HCCN path before using it. + +## Check Device-to-Device Connectivity + +After collecting the target-stage NPU IPv4 addresses, ping them from each +source-stage NPU: + +```bash +# Example: AR uses NPU 0..3 and DiT uses NPU 4..7. +# Replace the destination IPs with the values reported by hccn_tool -i 4..7 -ip -g. +for src in 0 1 2 3; do + for dst_ip in 10.86.1.78 10.86.1.79 10.86.1.80 10.86.1.81; do + echo "===== NPU $src -> $dst_ip =====" + hccn_tool -i "$src" -ping -g address "$dst_ip" || true + done +done +``` + +The expected result is `0.00% packet loss` for every required source/target pair. +This verifies basic HCCN IPv4 connectivity. TransferEngine can still fail later +if RDMA/QP/HCCL setup is broken, but failed ping means the device network must be +fixed first. + +## YAML Configuration + +For Ascend P2P, configure the connector with `protocol: "ascend"` and an NPU +memory pool: + +```yaml +runtime: + enabled: true + defaults: + window_size: -1 + max_inflight: 1 + connectors: + yuanrong_te_connector: + name: YuanrongTransferEngineConnector + extra: + host: "auto" + zmq_port: 50051 + rpc_port: "auto" + protocol: "ascend" + device_name: "auto" + memory_pool_size: 1073741824 + memory_pool_device: "npu" + edges: + - from: 0 + to: 1 + window_size: -1 +``` + +Important fields: + +| Parameter | Recommended Value | Notes | +|---|---|---| +| `host` | `"auto"` | Advertises a routable local host IP for ZMQ metadata exchange. | +| `zmq_port` | `50051` | Base ZMQ port. The runtime applies stage/rank offsets. | +| `rpc_port` | `"auto"` | Lets TransferEngine choose an available RPC port. | +| `protocol` | `"ascend"` | Uses the Ascend TransferEngine backend. | +| `device_name` | `"auto"` | Each TP worker resolves to its local logical NPU, for example `npu:0`, `npu:1`, etc. | +| `memory_pool_device` | `"npu"` | Required for the Ascend fast path. This registers NPU memory with TransferEngine. | +| `memory_pool_size` | `1073741824` | 1 GiB per worker is a practical starting point for KV-cache transfer validation. Increase if the transfer pool is exhausted. | + +`memory_pool_device: "cpu"` is rejected for +`YuanrongTransferEngineConnector` with `protocol: "ascend"`. CPU pool registration +can lead to Ascend P2P failures such as HCCL RA init errors, QP timeouts, or +FFTS/SDMA runtime errors. `memory_pool_device: "cpu"` is common for +`MooncakeTransferEngineConnector`, but Yuanrong TE on Ascend only supports +`"npu"`. + +## 4-Card AR-to-DiT TP Example + +For AR 4-way TP sending KV cache to DiT 4-way TP, set the stage TP topology in +both stage `omni_kv_config` blocks: + +```yaml +stage_args: + - stage_id: 0 + runtime: + devices: "0,1,2,3" + engine_args: + tensor_parallel_size: 4 + omni_kv_config: + need_send_cache: true + rank_mapping: + from_tp: 4 + to_tp: 4 + output_connectors: + to_stage_1: yuanrong_te_connector + + - stage_id: 1 + runtime: + devices: "4,5,6,7" + engine_args: + parallel_config: + tensor_parallel_size: 4 + omni_kv_config: + need_recv_cache: true + rank_mapping: + from_tp: 4 + to_tp: 4 + input_connectors: + from_stage_0: yuanrong_te_connector +``` + +In this topology each DiT rank receives the corresponding AR rank shard: + +```text +DiT rank 0 <- AR rank 0 +DiT rank 1 <- AR rank 1 +DiT rank 2 <- AR rank 2 +DiT rank 3 <- AR rank 3 +``` + +## Validate a Successful Run + +Search the inference log for Yuanrong TE receive-side messages: + +```bash +grep -E "\[YR GET\]|Successfully received KV cache|transfer_engine get failed" inference.log +``` + +A successful Ascend fast-path transfer contains lines like: + +```text +[YR GET] ... fast_path, zero-copy +Successfully received KV cache ... across 1 key(s) +Applied CFG KV caches +[KV Reuse] ... +``` + +Sender-side `KV transfer OK` only means that the sender registered the payload +and metadata. The actual P2P data transfer happens on the receiver side during +`[YR GET]`. + +## Troubleshooting + +| Symptom | Likely Cause | Action | +|---|---|---| +| `IPv4 device IP not found` | NPU HCCN IPv4 is missing. | Run `hccn_tool -i -ip -g` and configure IPv4 for every participating NPU. | +| `P2PCommInitRootInfo failed`, `ra init failed`, `RA_QP_STATUS_TIMEOUT` | Ascend P2P/RDMA/QP setup failed. | First verify device IPv4, link health, and device-to-device ping. Then test a minimal Yuanrong TE pair outside vLLM. | +| `fftsplus sdma error`, `context is abort` | Runtime SDMA/FFTS failure, often after an invalid or failed P2P transfer. | Use `memory_pool_device: "npu"` for Ascend TE, reduce `memory_pool_size` if OOM occurs, and isolate the failing NPU pair with a Yuanrong TE demo. | +| Sender logs `KV transfer OK` but receiver never logs `[YR GET]` | Sender `put()` succeeded, receiver P2P read failed or timed out. | Check receiver logs and search for `transfer_engine get failed`. | +| Pool allocation fails or OOM | Transfer pool is too small or too large for available NPU memory. | Start with 1 GiB (`1073741824`), reduce to 512 MiB for validation, or increase for larger KV payloads. | diff --git a/tests/distributed/omni_connectors/test_mooncake_transfer_engine_buffer.py b/tests/distributed/omni_connectors/test_mooncake_transfer_engine_buffer.py index 55351efcb..3e9b97d4e 100755 --- a/tests/distributed/omni_connectors/test_mooncake_transfer_engine_buffer.py +++ b/tests/distributed/omni_connectors/test_mooncake_transfer_engine_buffer.py @@ -15,8 +15,9 @@ from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_c BufferAllocator, ManagedBuffer, ) +from vllm_omni.distributed.omni_connectors.kv_transfer_manager import KVCacheTransferData -# All tests in this file are pure-CPU unit tests for the memory allocator. +# Most tests in this file are pure-CPU unit tests; CUDA smoke tests are skip-guarded. pytestmark = [pytest.mark.cpu, pytest.mark.parallel, pytest.mark.core_model] @@ -225,3 +226,70 @@ class TestManagedBuffer: # Space should be reusable new_offset = self.allocator.alloc(128) assert new_offset == offset + + def test_to_bytes_cpu_pool(self): + """Smoke-test Mooncake's ManagedBuffer byte extraction on CPU pools.""" + expected = bytes(range(64)) + offset = self.allocator.alloc(len(expected)) + buf = ManagedBuffer(self.allocator, offset, len(expected), self.pool) + self.pool[offset : offset + len(expected)] = torch.tensor(list(expected), dtype=torch.uint8) + + try: + assert buf.to_bytes() == expected + finally: + buf.release() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for GPU pool smoke test") + def test_to_bytes_cuda_pool(self): + """Smoke-test Mooncake's existing CUDA pool D2H byte path.""" + expected = bytes(range(64)) + allocator = BufferAllocator(total_size=4096, alignment=64) + pool = torch.zeros(4096, dtype=torch.uint8, device="cuda:0") + offset = allocator.alloc(len(expected)) + buf = ManagedBuffer(allocator, offset, len(expected), pool) + pool[offset : offset + len(expected)].copy_(torch.tensor(list(expected), dtype=torch.uint8, device="cuda:0")) + + try: + assert buf.to_bytes() == expected + finally: + buf.release() + + +def _make_kv_payload(device: str | torch.device = "cpu") -> tuple[KVCacheTransferData, torch.Tensor, torch.Tensor]: + key_tensor = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4) + value_tensor = key_tensor + 100 + payload = KVCacheTransferData( + request_id="moon-smoke", + layer_blocks={"key_cache": [key_tensor], "value_cache": [value_tensor]}, + block_ids=[1], + metadata={"seq_len": 3}, + ) + return payload, key_tensor, value_tensor + + +class TestMooncakePackedPayloadSmoke: + """Smoke tests for Mooncake raw payload helpers without real Mooncake/RDMA.""" + + def test_load_header_from_cpu_tensor(self): + payload, _, _ = _make_kv_payload() + raw = payload.to_bytes() + packed = torch.tensor(list(raw), dtype=torch.uint8) + header, data_start = KVCacheTransferData._load_header_from_tensor(packed) + + assert header["rid"] == "moon-smoke" + assert data_start == 4 + int.from_bytes(raw[:4], "big") + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for packed CUDA smoke test") + def test_from_bytes_device_cuda_round_trip(self): + payload, key_tensor, value_tensor = _make_kv_payload(device="cuda:0") + packed = payload.to_gpu_tensor() + + header, data_start = KVCacheTransferData._load_header_from_tensor(packed) + decoded = KVCacheTransferData.from_bytes_device(packed) + + assert header["rid"] == "moon-smoke" + assert data_start == 4 + int.from_bytes(payload.to_bytes()[:4], "big") + assert decoded["layer_blocks"]["key_cache"][0].is_cuda + assert decoded["layer_blocks"]["value_cache"][0].is_cuda + assert torch.equal(decoded["layer_blocks"]["key_cache"][0].cpu(), key_tensor.cpu()) + assert torch.equal(decoded["layer_blocks"]["value_cache"][0].cpu(), value_tensor.cpu()) diff --git a/vllm_omni/deploy/hunyuan_image3.yaml b/vllm_omni/deploy/hunyuan_image3.yaml index 20d72304b..a6e2add17 100644 --- a/vllm_omni/deploy/hunyuan_image3.yaml +++ b/vllm_omni/deploy/hunyuan_image3.yaml @@ -15,6 +15,16 @@ connectors: device_name: "" memory_pool_size: 4294967296 memory_pool_device: "cpu" + yuanrong_te_connector: + name: YuanrongTransferEngineConnector + extra: + host: "auto" + zmq_port: 50051 + rpc_port: "auto" + protocol: "ascend" + device_name: "auto" + memory_pool_size: 1073741824 + memory_pool_device: "npu" shared_memory_connector: name: SharedMemoryConnector extra: diff --git a/vllm_omni/distributed/__init__.py b/vllm_omni/distributed/__init__.py index f43a833f8..7af13a322 100644 --- a/vllm_omni/distributed/__init__.py +++ b/vllm_omni/distributed/__init__.py @@ -11,6 +11,7 @@ from .omni_connectors import ( OmniTransferConfig, SharedMemoryConnector, YuanrongConnector, + YuanrongTransferEngineConnector, load_omni_transfer_config, ) @@ -26,6 +27,7 @@ __all__ = [ "MooncakeTransferEngineConnector", "SharedMemoryConnector", "YuanrongConnector", + "YuanrongTransferEngineConnector", # Utilities "load_omni_transfer_config", ] diff --git a/vllm_omni/distributed/omni_connectors/__init__.py b/vllm_omni/distributed/omni_connectors/__init__.py index 383267928..98e3a5c36 100644 --- a/vllm_omni/distributed/omni_connectors/__init__.py +++ b/vllm_omni/distributed/omni_connectors/__init__.py @@ -6,6 +6,13 @@ from .connectors.mooncake_store_connector import MooncakeStoreConnector from .connectors.shm_connector import SharedMemoryConnector from .connectors.yuanrong_connector import YuanrongConnector +try: + from vllm_omni.platforms.npu.omni_connectors.yuanrong_transfer_engine_connector import ( + YuanrongTransferEngineConnector, + ) +except ImportError: + YuanrongTransferEngineConnector = None + try: from .connectors.mooncake_transfer_engine_connector import MooncakeTransferEngineConnector except ImportError: @@ -39,6 +46,7 @@ __all__ = [ "MooncakeTransferEngineConnector", "SharedMemoryConnector", "YuanrongConnector", + "YuanrongTransferEngineConnector", # Utilities "load_omni_transfer_config", "initialize_connectors_from_config", diff --git a/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py b/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py index 4b21ed0ec..1c058c7f0 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py @@ -16,6 +16,7 @@ import torch import zmq from ..utils.logging import get_connector_logger +from ..utils.memory_pool import BufferAllocator, ManagedBuffer from ..utils.serialization import OmniSerializer from .base import OmniConnectorBase @@ -66,160 +67,6 @@ class MooncakeAgentMetadata: lengths: list[int] -class BufferAllocator: - """ - Manages the allocation of memory segments within the registered pool. - Thread-safe implementation using a simple free list. - """ - - def __init__(self, total_size: int, alignment: int = 4096): - self.total_size = total_size - self.alignment = alignment - self.lock = threading.Lock() - # Free list: [(start, size), ...] sorted by start - self.free_blocks = [(0, total_size)] - - def alloc(self, size: int) -> int: - """ - Allocates a block of 'size' bytes. - Returns the starting offset. - """ - # Align size upwards - aligned_size = (size + self.alignment - 1) // self.alignment * self.alignment - - with self.lock: - for i, (start, block_size) in enumerate(self.free_blocks): - if block_size >= aligned_size: - # Found a block - new_start = start + aligned_size - new_size = block_size - aligned_size - - if new_size > 0: - self.free_blocks[i] = (new_start, new_size) - else: - self.free_blocks.pop(i) - return start - - raise MemoryError(f"Out of memory in buffer pool. Requested {size} bytes (aligned {aligned_size}).") - - def free(self, offset: int, size: int): - """ - Frees a previously allocated block. - """ - aligned_size = (size + self.alignment - 1) // self.alignment * self.alignment - - with self.lock: - # Check for double-free and corruption - for start, length in self.free_blocks: - # Case 1: Exact match = double free, safe to ignore - if offset == start and aligned_size == length: - logger.warning(f"Double free detected at offset {offset}, size {aligned_size}. Ignoring.") - return - # Case 2: Block is fully contained within an existing free block = also double free - # This happens when the block was freed and then merged with adjacent blocks - if offset >= start and offset + aligned_size <= start + length: - logger.warning( - f"Double free detected: block {offset}-{offset + aligned_size} " - f"is already within free block {start}-{start + length}. Ignoring." - ) - return - # Case 3: Partial overlap (but not fully contained) = memory corruption - if not (offset + aligned_size <= start or start + length <= offset): - raise RuntimeError( - f"Memory corruption detected: freeing {offset}-{offset + aligned_size} " - f"partially overlaps with free block {start}-{start + length}" - ) - - self.free_blocks.append((offset, aligned_size)) - self.free_blocks.sort() # Sort by offset - - # Merge adjacent blocks - i = 0 - while i < len(self.free_blocks) - 1: - curr_start, curr_size = self.free_blocks[i] - next_start, next_size = self.free_blocks[i + 1] - - if curr_start + curr_size == next_start: - self.free_blocks[i] = (curr_start, curr_size + next_size) - self.free_blocks.pop(i + 1) - else: - i += 1 - - -class ManagedBuffer: - """ - A temporary view into the global memory pool. - Must be kept alive while the data view is being used. - """ - - def __init__(self, allocator: BufferAllocator, offset: int, size: int, pool_tensor: torch.Tensor): - self.allocator = allocator - self.offset = offset - self.size = size - self.pool_tensor = pool_tensor - self._released = False - - def release(self): - """Explicitly release the buffer back to the pool.""" - if not self._released: - self.allocator.free(self.offset, self.size) - self._released = True - - def __del__(self): - self.release() - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - self.release() - - @property - def tensor(self) -> torch.Tensor: - """ - Returns a 1D uint8 zero-copy view of the buffer. - """ - return self.pool_tensor[self.offset : self.offset + self.size] - - def as_tensor(self, dtype: torch.dtype, shape: tuple) -> torch.Tensor: - """ - Returns a typed, shaped zero-copy view. - Validates size, shape, and alignment. - """ - itemsize = torch.tensor([], dtype=dtype).element_size() - - # Calculate expected size - expected_bytes = itemsize - for dim in shape: - if dim < 0: - raise ValueError("Dynamic dimension (-1) is not supported in as_tensor") - expected_bytes *= dim - - if expected_bytes != self.size: - raise ValueError( - f"Shape {shape} with dtype {dtype} requires {expected_bytes} bytes, but buffer size is {self.size}" - ) - - # Check alignment (offset must be divisible by itemsize) - if self.offset % itemsize != 0: - raise RuntimeError(f"Buffer offset {self.offset} is not aligned for dtype {dtype} (itemsize {itemsize})") - - raw_view = self.tensor - # view() requires contiguous memory, slice of contiguous tensor is contiguous - typed_view = raw_view.view(dtype) - return typed_view.reshape(shape) - - def to_bytes(self) -> bytes: - """ - Returns a copy of the data as python bytes. - Performs D2H copy if pool is on GPU. - """ - t = self.tensor - if t.is_cuda: - t = t.cpu() - return t.numpy().tobytes() - - class MooncakeTransferEngineConnector(OmniConnectorBase): """ OmniConnector implementation using Mooncake Transfer Engine with a managed memory pool. diff --git a/vllm_omni/distributed/omni_connectors/connectors/yuanrong_transfer_engine_connector.py b/vllm_omni/distributed/omni_connectors/connectors/yuanrong_transfer_engine_connector.py new file mode 100644 index 000000000..3254743c5 --- /dev/null +++ b/vllm_omni/distributed/omni_connectors/connectors/yuanrong_transfer_engine_connector.py @@ -0,0 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Compatibility import for the NPU-only Yuanrong TransferEngine connector.""" + +from vllm_omni.platforms.npu.omni_connectors.yuanrong_transfer_engine_connector import ( + YuanrongTransferEngineConnector, +) + +__all__ = ["YuanrongTransferEngineConnector"] diff --git a/vllm_omni/distributed/omni_connectors/factory.py b/vllm_omni/distributed/omni_connectors/factory.py index a680da953..7ee2dcde3 100644 --- a/vllm_omni/distributed/omni_connectors/factory.py +++ b/vllm_omni/distributed/omni_connectors/factory.py @@ -91,6 +91,18 @@ def _create_yuanrong_connector(config: dict[str, Any]) -> OmniConnectorBase: return YuanrongConnector(config) +def _create_yuanrong_transfer_engine_connector(config: dict[str, Any]) -> OmniConnectorBase: + try: + from vllm_omni.platforms.npu.omni_connectors import YuanrongTransferEngineConnector + except ImportError as exc: + raise ImportError( + "YuanrongTransferEngineConnector is only available in the NPU platform " + "environment. Install the Ascend/Yuanrong runtime dependencies before " + "using this connector." + ) from exc + return YuanrongTransferEngineConnector(config) + + def _create_mooncake_transfer_engine_connector(config: dict[str, Any]) -> OmniConnectorBase: try: from .connectors.mooncake_transfer_engine_connector import MooncakeTransferEngineConnector @@ -107,5 +119,6 @@ OmniConnectorFactory.register_connector("MooncakeStoreConnector", _create_moonca OmniConnectorFactory.register_connector("MooncakeTransferEngineConnector", _create_mooncake_transfer_engine_connector) OmniConnectorFactory.register_connector("SharedMemoryConnector", _create_shm_connector) OmniConnectorFactory.register_connector("YuanrongConnector", _create_yuanrong_connector) +OmniConnectorFactory.register_connector("YuanrongTransferEngineConnector", _create_yuanrong_transfer_engine_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 241a6070d..372c2cb3d 100644 --- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py +++ b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py @@ -13,7 +13,8 @@ import torch from vllm.logger import init_logger from .factory import OmniConnectorFactory -from .utils.config import ConnectorSpec +from .utils.config import TRANSFER_ENGINE_CONNECTOR_NAMES, ConnectorSpec +from .utils.env import expand_env_int from .utils.initialization import KV_RANK_PORT_STRIDE from .utils.kv_utils import ( KVTPTopology, @@ -106,7 +107,7 @@ class KVCacheTransferData: t = tensor.detach().contiguous() if cpu: t = t.cpu() - elif device is None and t.is_cuda: + elif device is None and getattr(t.device, "type", "cpu") != "cpu": device = t.device nbytes = t.numel() * t.element_size() tensors_desc.append( @@ -143,10 +144,10 @@ class KVCacheTransferData: return b"".join([self._build_header_bytes(tensors_desc)] + chunks) def to_gpu_tensor(self) -> torch.Tensor: - """Convert to a packed GPU tensor for raw-data connectors.""" + """Convert to a packed device tensor for raw-data connectors.""" tensors_desc, chunks, data_offset, device = self._build_tensors_desc(cpu=False) if device is None: - raise RuntimeError("No CUDA tensors found, use to_bytes() instead") + raise RuntimeError("No device tensors found, use to_bytes() instead") header_prefix = self._build_header_bytes(tensors_desc) output = torch.empty(len(header_prefix) + data_offset, dtype=torch.uint8, device=device) output[: len(header_prefix)].copy_(torch.frombuffer(bytearray(header_prefix), dtype=torch.uint8)) @@ -169,19 +170,19 @@ class KVCacheTransferData: return json.loads(bytes(raw_mv[4 : 4 + header_len])), raw_mv[4 + header_len :] @staticmethod - def _load_header_from_tensor(gpu_tensor: torch.Tensor) -> tuple[dict[str, Any], int]: - if gpu_tensor.dtype != torch.uint8 or gpu_tensor.dim() != 1: - raise ValueError("Packed GPU KV payload must be a 1-D uint8 tensor") + def _load_header_from_tensor(tensor: torch.Tensor) -> tuple[dict[str, Any], int]: + if tensor.dtype != torch.uint8 or tensor.dim() != 1: + raise ValueError("Packed device KV payload must be a 1-D uint8 tensor") - total_bytes = int(gpu_tensor.numel()) + total_bytes = int(tensor.numel()) if total_bytes < 4: raise ValueError("Corrupted KV payload: missing 4-byte header length") - header_len = struct.unpack(">I", gpu_tensor[:4].cpu().numpy().tobytes())[0] + header_len = struct.unpack(">I", tensor[:4].cpu().numpy().tobytes())[0] if header_len > total_bytes - 4: raise ValueError(f"Corrupted KV payload: header_len={header_len} exceeds buffer size={total_bytes}") - header_bytes = gpu_tensor[4 : 4 + header_len].cpu().numpy().tobytes() + header_bytes = tensor[4 : 4 + header_len].cpu().numpy().tobytes() return json.loads(header_bytes), 4 + header_len @staticmethod @@ -260,17 +261,22 @@ class KVCacheTransferData: return KVCacheTransferData._populate_caches(header, _get) @staticmethod - def from_bytes_gpu(gpu_tensor: torch.Tensor) -> dict[str, Any]: - """Reconstruct KV cache data from a packed GPU tensor.""" - header, data_start = KVCacheTransferData._load_header_from_tensor(gpu_tensor) - data_len = int(gpu_tensor.numel()) - data_start + def from_bytes_device(tensor: torch.Tensor) -> dict[str, Any]: + """Reconstruct KV cache data from a packed device tensor.""" + header, data_start = KVCacheTransferData._load_header_from_tensor(tensor) + data_len = int(tensor.numel()) - data_start def _get(info: dict) -> torch.Tensor: offset, nbytes = KVCacheTransferData._validate_tensor_span(info["n"], info, data_len) - return gpu_tensor[data_start + offset : data_start + offset + nbytes].clone() + return tensor[data_start + offset : data_start + offset + nbytes].clone() return KVCacheTransferData._populate_caches(header, _get) + @staticmethod + def from_bytes_gpu(tensor: torch.Tensor) -> dict[str, Any]: + """Compatibility alias for callers using the old GPU-specific name.""" + return KVCacheTransferData.from_bytes_device(tensor) + class OmniKVTransferManager: """Unified management for OmniConnector and KV cache transfer. @@ -405,8 +411,8 @@ class OmniKVTransferManager: if cfg and (c_type := cfg.get("type")): try: c_extra = {k: v for k, v in cfg.items() if k != "type"} - if c_type == "MooncakeTransferEngineConnector": - base_port = c_extra.get("zmq_port", 50051) + if c_type in TRANSFER_ENGINE_CONNECTOR_NAMES: + base_port = expand_env_int(c_extra.get("zmq_port", 50051), "zmq_port") c_extra["from_stage"] = ( str(self.config.from_stage) if self.config.from_stage is not None else "0" ) @@ -429,8 +435,13 @@ class OmniKVTransferManager: c_extra["zmq_port"] = zmq_port elif self.config.need_recv_cache: c_extra["role"] = "receiver" - c_extra.setdefault("sender_host", c_extra.get("host", "127.0.0.1")) - c_extra.setdefault("sender_zmq_port", zmq_port) + # Receiver-side sender endpoints are request-scoped. + # They are attached by the orchestrator as + # kv_sender_info and applied in update_sender_info(). + # Do not derive sender_host from this receiver's + # local YAML host; on multi-node runs that points at + # the wrong process. Explicit sender_* YAML values + # are still preserved for standalone connector use. logger.info( "Initializing OmniConnector type=%s role=%s", @@ -1058,8 +1069,8 @@ class OmniKVTransferManager: managed_buffer = raw_data try: buf_tensor = raw_data.tensor - if buf_tensor.is_cuda: - data = KVCacheTransferData.from_bytes_gpu(buf_tensor) + if getattr(buf_tensor.device, "type", "cpu") != "cpu": + data = KVCacheTransferData.from_bytes_device(buf_tensor) raw_data.release() managed_buffer = None else: @@ -1075,7 +1086,10 @@ class OmniKVTransferManager: elif isinstance(raw_data, (bytes, bytearray)): data = KVCacheTransferData.from_bytes(raw_data) elif isinstance(raw_data, torch.Tensor) and raw_data.dtype == torch.uint8 and raw_data.dim() == 1: - data = KVCacheTransferData.from_bytes(raw_data.cpu().numpy().tobytes()) + if getattr(raw_data.device, "type", "cpu") != "cpu": + data = KVCacheTransferData.from_bytes_device(raw_data) + else: + data = KVCacheTransferData.from_bytes(raw_data.numpy().tobytes()) else: data = raw_data diff --git a/vllm_omni/distributed/omni_connectors/utils/config.py b/vllm_omni/distributed/omni_connectors/utils/config.py index f18e52c4d..abe6e6a20 100644 --- a/vllm_omni/distributed/omni_connectors/utils/config.py +++ b/vllm_omni/distributed/omni_connectors/utils/config.py @@ -8,6 +8,13 @@ from .logging import get_connector_logger logger = get_connector_logger(__name__) +TRANSFER_ENGINE_CONNECTOR_NAMES = frozenset( + { + "MooncakeTransferEngineConnector", + "YuanrongTransferEngineConnector", + } +) + @dataclass class ConnectorSpec: diff --git a/vllm_omni/distributed/omni_connectors/utils/env.py b/vllm_omni/distributed/omni_connectors/utils/env.py new file mode 100644 index 000000000..ee8b324ea --- /dev/null +++ b/vllm_omni/distributed/omni_connectors/utils/env.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import os +import re +from typing import Any + +AUTO_HOST_VALUES = {"", "auto", "*", "0.0.0.0", "::"} +AUTO_DEVICE_VALUES = {"", "auto"} + +_UNEXPANDED_ENV_VAR_PATTERN = re.compile(r"\$\{([^}]+)\}|\$([A-Za-z_][A-Za-z0-9_]*)|%([^%]+)%") + + +class EnvVarExpansionError(ValueError): + pass + + +def _unexpanded_env_vars(value: str) -> list[str]: + names: list[str] = [] + for match in _UNEXPANDED_ENV_VAR_PATTERN.finditer(value): + name = next(group for group in match.groups() if group) + if name not in names: + names.append(name) + return names + + +def expand_env_value(value: Any, *, strict: bool = False, field_name: str = "value") -> Any: + if not isinstance(value, str): + return value + + expanded = os.path.expandvars(value) + if strict: + unresolved = _unexpanded_env_vars(expanded) + if unresolved: + names = ", ".join(unresolved) + raise EnvVarExpansionError( + f"{field_name} references unset environment variable(s): {names}. Original value: {value!r}" + ) + return expanded + + +def expand_env_int(value: Any, field_name: str) -> int: + expanded = expand_env_value(value, strict=True, field_name=field_name) + try: + return int(expanded) + except (TypeError, ValueError) as exc: + raise ValueError(f"{field_name} must be an integer after environment expansion, got {expanded!r}") from exc diff --git a/vllm_omni/distributed/omni_connectors/utils/initialization.py b/vllm_omni/distributed/omni_connectors/utils/initialization.py index a8cafba61..f7a464809 100644 --- a/vllm_omni/distributed/omni_connectors/utils/initialization.py +++ b/vllm_omni/distributed/omni_connectors/utils/initialization.py @@ -10,7 +10,9 @@ from pathlib import Path from typing import TYPE_CHECKING, Any from ..factory import OmniConnectorFactory -from .config import ConnectorSpec, OmniTransferConfig +from .config import TRANSFER_ENGINE_CONNECTOR_NAMES, ConnectorSpec, OmniTransferConfig +from .env import expand_env_int +from .local_rank import get_connector_local_rank from .logging import get_connector_logger if TYPE_CHECKING: @@ -93,18 +95,20 @@ def create_connectors_from_config( for edge_key, connector_spec in connectors_config.items(): from_stage, to_stage = edge_key try: - if connector_spec.name == "MooncakeTransferEngineConnector": + if connector_spec.name in TRANSFER_ENGINE_CONNECTOR_NAMES: extra = dict(connector_spec.extra) if connector_spec.extra else {} - base_port = extra.get("zmq_port", 50051) + 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": - adjusted_port = base_port + orchestrator_port_offset + stage_offset + rank0_port = base_port + orchestrator_port_offset + stage_offset else: - adjusted_port = base_port + port_offset + stage_offset + rank0_port = base_port + port_offset + stage_offset + adjusted_port = rank0_port + local_rank * KV_RANK_PORT_STRIDE extra["zmq_port"] = adjusted_port if is_sender is not None: diff --git a/vllm_omni/distributed/omni_connectors/utils/local_rank.py b/vllm_omni/distributed/omni_connectors/utils/local_rank.py new file mode 100644 index 000000000..4c2b49e8f --- /dev/null +++ b/vllm_omni/distributed/omni_connectors/utils/local_rank.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import os + +_LOCAL_RANK_ENV_NAMES = ( + "LOCAL_RANK", + "VLLM_LOCAL_RANK", + "OMPI_COMM_WORLD_LOCAL_RANK", + "MV2_COMM_WORLD_LOCAL_RANK", +) +_GLOBAL_RANK_ENV_NAMES = ("RANK", "RANK_ID") + + +def _visible_device_count() -> int | None: + visible = os.environ.get("ASCEND_RT_VISIBLE_DEVICES") or os.environ.get("CUDA_VISIBLE_DEVICES") + if not visible: + return None + devices = [item.strip() for item in visible.split(",") if item.strip()] + return len(devices) or None + + +def _get_tensor_model_parallel_rank() -> int: + from vllm.distributed.parallel_state import get_tensor_model_parallel_rank + + return int(get_tensor_model_parallel_rank()) + + +def get_connector_local_rank() -> int: + try: + return _get_tensor_model_parallel_rank() + except Exception: + pass + + for env_name in _LOCAL_RANK_ENV_NAMES: + try: + return int(os.environ[env_name]) + except (KeyError, TypeError, ValueError): + pass + + visible_count = _visible_device_count() + for env_name in _GLOBAL_RANK_ENV_NAMES: + try: + rank = int(os.environ[env_name]) + except (KeyError, TypeError, ValueError): + continue + return rank % visible_count if visible_count else rank + + return 0 diff --git a/vllm_omni/distributed/omni_connectors/utils/memory_pool.py b/vllm_omni/distributed/omni_connectors/utils/memory_pool.py new file mode 100644 index 000000000..2e2e7ba00 --- /dev/null +++ b/vllm_omni/distributed/omni_connectors/utils/memory_pool.py @@ -0,0 +1,153 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import threading + +import torch + +from .logging import get_connector_logger + +logger = get_connector_logger(__name__) + + +class BufferAllocator: + """ + Manages the allocation of memory segments within a registered pool. + Thread-safe implementation using a simple free list. + """ + + def __init__(self, total_size: int, alignment: int = 4096): + self.total_size = total_size + self.alignment = alignment + self.lock = threading.Lock() + self.free_blocks = [(0, total_size)] + + def alloc(self, size: int) -> int: + """ + Allocates a block of *size* bytes and returns the starting offset. + """ + aligned_size = (size + self.alignment - 1) // self.alignment * self.alignment + + with self.lock: + for i, (start, block_size) in enumerate(self.free_blocks): + if block_size >= aligned_size: + new_start = start + aligned_size + new_size = block_size - aligned_size + + if new_size > 0: + self.free_blocks[i] = (new_start, new_size) + else: + self.free_blocks.pop(i) + return start + + raise MemoryError(f"Out of memory in buffer pool. Requested {size} bytes (aligned {aligned_size}).") + + def free(self, offset: int, size: int) -> None: + """ + Frees a previously allocated block. + """ + aligned_size = (size + self.alignment - 1) // self.alignment * self.alignment + + with self.lock: + for start, length in self.free_blocks: + if offset == start and aligned_size == length: + logger.warning("Double free detected at offset %s, size %s. Ignoring.", offset, aligned_size) + return + if offset >= start and offset + aligned_size <= start + length: + logger.warning( + "Double free detected: block %s-%s is already within free block %s-%s. Ignoring.", + offset, + offset + aligned_size, + start, + start + length, + ) + return + if not (offset + aligned_size <= start or start + length <= offset): + raise RuntimeError( + f"Memory corruption detected: freeing {offset}-{offset + aligned_size} " + f"partially overlaps with free block {start}-{start + length}" + ) + + self.free_blocks.append((offset, aligned_size)) + self.free_blocks.sort() + + i = 0 + while i < len(self.free_blocks) - 1: + curr_start, curr_size = self.free_blocks[i] + next_start, next_size = self.free_blocks[i + 1] + + if curr_start + curr_size == next_start: + self.free_blocks[i] = (curr_start, curr_size + next_size) + self.free_blocks.pop(i + 1) + else: + i += 1 + + +class ManagedBuffer: + """ + A temporary view into a global memory pool. + Must be kept alive while the data view is being used. + """ + + def __init__(self, allocator: BufferAllocator, offset: int, size: int, pool_tensor: torch.Tensor): + self.allocator = allocator + self.offset = offset + self.size = size + self.pool_tensor = pool_tensor + self._released = False + + def release(self) -> None: + """Explicitly release the buffer back to the pool.""" + if not self._released: + self.allocator.free(self.offset, self.size) + self._released = True + + def __del__(self): + self.release() + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.release() + + @property + def tensor(self) -> torch.Tensor: + """ + Returns a 1D uint8 zero-copy view of the buffer. + """ + return self.pool_tensor[self.offset : self.offset + self.size] + + def as_tensor(self, dtype: torch.dtype, shape: tuple) -> torch.Tensor: + """ + Returns a typed, shaped zero-copy view. + Validates size, shape, and alignment. + """ + itemsize = torch.tensor([], dtype=dtype).element_size() + + expected_bytes = itemsize + for dim in shape: + if dim < 0: + raise ValueError("Dynamic dimension (-1) is not supported in as_tensor") + expected_bytes *= dim + + if expected_bytes != self.size: + raise ValueError( + f"Shape {shape} with dtype {dtype} requires {expected_bytes} bytes, but buffer size is {self.size}" + ) + + if self.offset % itemsize != 0: + raise RuntimeError(f"Buffer offset {self.offset} is not aligned for dtype {dtype} (itemsize {itemsize})") + + typed_view = self.tensor.view(dtype) + return typed_view.reshape(shape) + + def to_bytes(self) -> bytes: + """ + Returns a copy of the data as Python bytes. + Performs D2H copy if the pool is on a device. + """ + t = self.tensor + if getattr(t.device, "type", "cpu") != "cpu": + t = t.cpu() + return t.numpy().tobytes() diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index e9e2944a7..84bf72863 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -8,6 +8,7 @@ from __future__ import annotations import inspect import multiprocessing.connection +import os import socket import threading import weakref @@ -20,6 +21,9 @@ from vllm.v1.engine import EngineCoreRequest from vllm.v1.engine.core_client import AsyncMPClient, DPLBAsyncMPClient from vllm.v1.engine.exceptions import EngineDeadError +from vllm_omni.distributed.omni_connectors.utils.config import ( + TRANSFER_ENGINE_CONNECTOR_NAMES, +) from vllm_omni.distributed.omni_connectors.utils.initialization import ( KV_TRANSFER_PORT_OFFSET, ) @@ -331,11 +335,13 @@ class StageEngineCoreClientBase(StageClientBase): if sender_host is not None: self._kv_sender_host = sender_host + connector_type = connector_config.get("type") sender_port = connector_config.get("sender_zmq_port") - if sender_port is None: + if connector_type in TRANSFER_ENGINE_CONNECTOR_NAMES or sender_port is None: base_port = connector_config.get("zmq_port") if base_port is None: return + base_port = os.path.expandvars(str(base_port)) omni_kv_config = getattr(self, "_omni_kv_config", None) from_stage = self.stage_id diff --git a/vllm_omni/platforms/npu/omni_connectors/__init__.py b/vllm_omni/platforms/npu/omni_connectors/__init__.py new file mode 100644 index 000000000..e411df850 --- /dev/null +++ b/vllm_omni/platforms/npu/omni_connectors/__init__.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +try: + from vllm_omni.platforms.npu.omni_connectors.yuanrong_transfer_engine_connector import ( + YuanrongTransferEngineConnector, + ) +except ImportError: + YuanrongTransferEngineConnector = None + +__all__ = ["YuanrongTransferEngineConnector"] diff --git a/vllm_omni/platforms/npu/omni_connectors/yuanrong_transfer_engine_connector.py b/vllm_omni/platforms/npu/omni_connectors/yuanrong_transfer_engine_connector.py new file mode 100644 index 000000000..649d4d45c --- /dev/null +++ b/vllm_omni/platforms/npu/omni_connectors/yuanrong_transfer_engine_connector.py @@ -0,0 +1,862 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import queue +import socket +import threading +import time as _time_mod +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from typing import Any + +import msgspec +import torch +import zmq + +from vllm_omni.distributed.omni_connectors.connectors.base import OmniConnectorBase +from vllm_omni.distributed.omni_connectors.utils.env import ( + AUTO_DEVICE_VALUES, + AUTO_HOST_VALUES, + expand_env_int, + expand_env_value, +) +from vllm_omni.distributed.omni_connectors.utils.local_rank import get_connector_local_rank +from vllm_omni.distributed.omni_connectors.utils.logging import get_connector_logger +from vllm_omni.distributed.omni_connectors.utils.memory_pool import BufferAllocator, ManagedBuffer +from vllm_omni.distributed.omni_connectors.utils.serialization import OmniSerializer + +logger = get_connector_logger(__name__) + +try: + from yr.datasystem import TransferEngine +except ImportError: + TransferEngine = None + +_BUFFER_TTL_SECONDS = 300 + +QUERY_INFO = b"query_info" +CLEANUP_KEY = b"cleanup_key" +ACK_OK = b"ack_ok" +ACK_ERROR = b"ack_error" +INFO_NOT_FOUND = b"info_not_found" + + +@dataclass +class QueryRequest: + """Request to query metadata for a specific key.""" + + request_id: str + + +@dataclass +class CleanupRequest: + """Response containing metadata for a request.""" + + request_id: str + + +@dataclass +class QueryResponse: + request_id: str + rpc_port: int + data_size: int + is_fast_path: bool + src_addrs: list[int] + lengths: list[int] + + +def _resolve_device_name(configured_device: Any, protocol: str) -> str: + device_name = str(expand_env_value(configured_device or "auto")).strip() + if device_name.lower() not in AUTO_DEVICE_VALUES: + return device_name + if protocol.lower() in {"ascend", "hccl", "npu"}: + return f"npu:{get_connector_local_rank()}" + return "" + + +def _resolve_pool_device(configured_device: Any) -> str: + pool_device = str(expand_env_value(configured_device or "npu")).strip().lower() + if pool_device in AUTO_DEVICE_VALUES: + return "npu" + if not (pool_device == "npu" or pool_device.startswith("npu:")): + raise ValueError( + f"YuanrongTransferEngineConnector requires an NPU memory pool. Got memory_pool_device={pool_device!r}." + ) + return pool_device + + +class YuanrongTransferEngineConnector(OmniConnectorBase): + """Pull-based connector backed by Yuanrong transfer_engine.""" + + supports_raw_data: bool = True + + def __init__(self, config: dict[str, Any]): + if TransferEngine is None: + raise ImportError("Yuanrong is not available.") + + self._closed = False + self._bind_error: Exception | None = None + self._stop_event = threading.Event() + self._listener_thread: threading.Thread | None = None + self._listener_ready = threading.Event() + self._control_executor: ThreadPoolExecutor | None = None + self._local_buffers: dict[str, Any] = {} + self._local_buffers_lock = threading.Lock() + self._req_local = threading.local() + self._worker_local = threading.local() + self._last_ttl_check: float = _time_mod.monotonic() + + self._metrics = { + "puts": 0, + "gets": 0, + "bytes_transferred": 0, + "errors": 0, + "timeouts": 0, + } + + self.config = {key: expand_env_value(value) for key, value in config.items()} + config = self.config + + raw_host_config = config.get("host") + host_config = "auto" if raw_host_config is None else str(raw_host_config).strip() + self.host = self._get_local_ip() if host_config.lower() in AUTO_HOST_VALUES else host_config + self.zmq_port = self._resolve_port(config.get("zmq_port"), self.host, "zmq_port") + self.rpc_port = self._resolve_port(config.get("rpc_port"), self.host, "rpc_port") + self.protocol = str(config.get("protocol", "ascend")) + self.device_name = _resolve_device_name(config.get("device_name", "auto"), self.protocol) + self.pool_size = int(config.get("memory_pool_size", 1024**3)) + self.pool_device = _resolve_pool_device(config.get("memory_pool_device", "npu")) + self.sender_host = config.get("sender_host") + sender_zmq_port = config.get("sender_zmq_port") + self.sender_zmq_port = self._resolve_optional_port(sender_zmq_port, "sender_zmq_port") + + role = str(config.get("role", "sender")).lower() + if role not in {"sender", "receiver"}: + raise ValueError( + f"Invalid role={role!r} for YuanrongTransferEngineConnector. Expected 'sender' or 'receiver'." + ) + self.can_put = role == "sender" + + self.engine = TransferEngine() + local_endpoint = f"{self.host}:{self.rpc_port}" + self._ensure_result_ok( + self.engine.initialize(local_endpoint, self.protocol, self.device_name), + f"initialize transfer engine on {local_endpoint}", + ) + resolved_rpc_port = int(self.engine.get_rpc_port()) + if resolved_rpc_port > 0: + self.rpc_port = resolved_rpc_port + + logger.info( + "YuanrongTransferEngineConnector initialized at %s:%s (protocol=%s, device=%s)", + self.host, + self.rpc_port, + self.protocol, + self.device_name, + ) + + try: + self.pool = torch.empty(self.pool_size, dtype=torch.uint8, device=self.pool_device) + self.base_ptr = int(self.pool.data_ptr()) + self._ensure_result_ok( + self.engine.register_memory(self.base_ptr, self.pool_size), + f"register memory pool ({self.pool_size} bytes)", + ) + except Exception as exc: + logger.error("Failed to allocate/register transfer_engine pool: %s", exc) + raise + + self.allocator = BufferAllocator(self.pool_size, alignment=4096) + self.zmq_ctx = zmq.Context() + self._control_executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix="yuanrong-transfer-control") + + if self.can_put: + self._listener_thread = threading.Thread(target=self._zmq_listener_loop, daemon=True) + self._listener_thread.start() + self._listener_ready.wait(timeout=1.0) + if self._bind_error is not None: + raise RuntimeError( + f"YuanrongTransferEngineConnector failed to bind ZMQ on " + f"{self.host}:{self.zmq_port}: {self._bind_error}" + ) from self._bind_error + else: + logger.info( + "YuanrongTransferEngineConnector started in receiver mode (sender_host=%r, sender_zmq_port=%r)", + self.sender_host, + self.sender_zmq_port, + ) + + @staticmethod + def _ensure_result_ok(result: Any, action: str) -> None: + if result.is_error(): + raise RuntimeError(f"{action} failed: {result.to_string()}") + + @staticmethod + def _resolve_optional_port(configured_port: Any, field_name: str) -> int | None: + configured_port = expand_env_value(configured_port, strict=True, field_name=field_name) + if configured_port is None or str(configured_port).strip().lower() in {"", "auto"}: + return None + return expand_env_int(configured_port, field_name) + + @classmethod + def _resolve_port(cls, configured_port: Any, host: str, field_name: str) -> int: + resolved_port = cls._resolve_optional_port(configured_port, field_name) + if resolved_port is not None: + return resolved_port + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind((host, 0)) + return int(sock.getsockname()[1]) + + @staticmethod + def _get_local_ip() -> str: + try: + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock: + sock.connect(("8.8.8.8", 80)) + return sock.getsockname()[0] + except Exception: + try: + return socket.gethostbyname(socket.gethostname()) + except Exception: + return "127.0.0.1" + + @staticmethod + def _sync_tensor_device(tensor: torch.Tensor | None) -> None: + if tensor is None: + return + device = getattr(tensor, "device", None) + if device is None or getattr(device, "type", "cpu") == "cpu": + return + + backend = getattr(torch, device.type, None) + if backend is None: + return + + current_stream = getattr(backend, "current_stream", None) + if callable(current_stream): + try: + current_stream(device).synchronize() + return + except TypeError: + try: + current_stream().synchronize() + return + except Exception: + pass + except Exception: + pass + + synchronize = getattr(backend, "synchronize", None) + if callable(synchronize): + try: + synchronize(device) + except TypeError: + synchronize() + + def get_connection_info(self) -> dict[str, Any]: + return { + "host": self.host, + "zmq_port": self.zmq_port, + "rpc_port": self.rpc_port, + "can_put": self.can_put, + } + + def update_sender_info( + self, + sender_host: str, + sender_zmq_port: int, + ) -> None: + self.sender_host = sender_host + self.sender_zmq_port = int(sender_zmq_port) + + def _get_req_socket(self, zmq_addr: str, timeout_ms: int) -> zmq.Socket: + cache: dict[str, zmq.Socket] = getattr(self._req_local, "cache", None) # type: ignore[assignment] + 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 + + 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] = getattr(self._req_local, "cache", None) # type: ignore[assignment] + if cache is None: + return + sock = cache.pop(zmq_addr, None) + if sock is not None: + try: + sock.close(linger=0) + except Exception: + pass + + def put(self, from_stage: str, to_stage: str, put_key: str, data: Any) -> tuple[bool, int, dict[str, Any] | None]: + if self._closed: + raise RuntimeError("Cannot put data: YuanrongTransferEngineConnector is closed") + + if not self.can_put: + self._metrics["errors"] += 1 + logger.warning("Rejecting put for %s: connector is in receiver mode", put_key) + return False, 0, None + + put_key = self._make_key(put_key, from_stage, to_stage) + + try: + src_addr = 0 + size = 0 + holder = None + is_fast_path = True + should_release_holder = False + + if isinstance(data, bytes) and len(data) == 0: + self._metrics["errors"] += 1 + logger.warning("Rejecting put for %s: empty bytes payload", put_key) + return False, 0, None + if isinstance(data, torch.Tensor) and data.nbytes == 0: + self._metrics["errors"] += 1 + logger.warning("Rejecting put for %s: zero-size tensor", put_key) + return False, 0, None + + if not isinstance(data, (ManagedBuffer, torch.Tensor, bytes)): + data = OmniSerializer.serialize(data) + is_fast_path = False + + if isinstance(data, ManagedBuffer): + if data.pool_tensor.data_ptr() != self.pool.data_ptr(): + logger.warning("ManagedBuffer from different pool detected. Falling back to copy path.") + data = data.tensor.contiguous() + else: + src_addr = self.base_ptr + data.offset + size = data.size + holder = data + should_release_holder = False + + if isinstance(data, (torch.Tensor, bytes)): + if isinstance(data, torch.Tensor): + size = data.nbytes + tensor_data = data + else: + size = len(data) + tensor_data = torch.frombuffer(data, dtype=torch.uint8) + + try: + offset = self.allocator.alloc(size) + holder = ManagedBuffer(self.allocator, offset, size, self.pool) + should_release_holder = True + except MemoryError: + self._metrics["errors"] += 1 + logger.error("Pool exhausted, cannot put data size %s", size) + return False, 0, None + + try: + dst_tensor = holder.tensor + if isinstance(data, torch.Tensor): + if not data.is_contiguous(): + data = data.contiguous() + src_view = data.view(torch.uint8).flatten() + dst_tensor.copy_(src_view, non_blocking=src_view.device != dst_tensor.device) + self._sync_tensor_device(src_view) + self._sync_tensor_device(dst_tensor) + else: + dst_tensor.copy_(tensor_data) + self._sync_tensor_device(dst_tensor) + except Exception as exc: + holder.release() + self._metrics["errors"] += 1 + logger.error("Failed to copy data to pool: %s", exc) + return False, 0, None + + src_addr = self.base_ptr + offset + + if size <= 0: + if should_release_holder and isinstance(holder, ManagedBuffer): + holder.release() + self._metrics["errors"] += 1 + logger.warning("Rejecting put for %s: final payload size is %s", put_key, size) + return False, 0, None + + with self._local_buffers_lock: + old_item = self._local_buffers.pop(put_key, None) + if old_item: + _, _, old_holder, old_should_release, _, _ = old_item + if old_should_release and isinstance(old_holder, ManagedBuffer): + old_holder.release() + self._local_buffers[put_key] = ( + [src_addr], + [size], + holder, + should_release_holder, + is_fast_path, + _time_mod.monotonic(), + ) + + metadata = { + "source_host": self.host, + "source_port": self.zmq_port, + "source_rpc_port": self.rpc_port, + "source_addrs": [src_addr], + "lengths": [size], + "data_size": size, + "is_fast_path": is_fast_path, + } + + self._metrics["puts"] += 1 + self._metrics["bytes_transferred"] += size + return True, size, metadata + except Exception as exc: + self._metrics["errors"] += 1 + logger.error("transfer_engine put failed for %s: %s", put_key, exc, exc_info=True) + return False, 0, None + + def _query_metadata_at(self, get_key: str, host: str, port: int) -> dict[str, Any] | None: + zmq_addr = f"tcp://{host}:{port}" + req_socket = self._get_req_socket(zmq_addr, timeout_ms=5000) + + try: + query = QueryRequest(request_id=get_key) + req_socket.send(QUERY_INFO + msgspec.msgpack.encode(query)) + resp = req_socket.recv() + if resp == INFO_NOT_FOUND: + return None + + query_resp = msgspec.msgpack.decode(resp, type=QueryResponse) + return { + "source_host": host, + "source_port": port, + "source_rpc_port": query_resp.rpc_port, + "source_addrs": query_resp.src_addrs, + "lengths": query_resp.lengths, + "data_size": query_resp.data_size, + "is_fast_path": query_resp.is_fast_path, + } + except Exception as exc: + self._invalidate_req_socket(zmq_addr) + logger.debug("Failed to query metadata at %s for %s: %s", zmq_addr, get_key, exc) + return None + + @classmethod + def _buffer_to_bytes(cls, buffer: ManagedBuffer) -> bytes: + tensor = buffer.tensor + cls._sync_tensor_device(tensor) + if getattr(tensor.device, "type", "cpu") != "cpu": + tensor = tensor.cpu() + return tensor.numpy().tobytes() + + def _request_sender_cleanup(self, request_id: str, source_host: str | None, source_port: int | None) -> bool: + if not source_host or not source_port: + return False + + zmq_addr = f"tcp://{source_host}:{source_port}" + req_socket = self._get_req_socket(zmq_addr, timeout_ms=3000) + + try: + payload = CleanupRequest(request_id=request_id) + req_socket.send(CLEANUP_KEY + msgspec.msgpack.encode(payload)) + return req_socket.recv() == ACK_OK + except Exception as exc: + self._invalidate_req_socket(zmq_addr) + logger.warning("Failed to cleanup sender buffer for %s: %s", request_id, exc) + return False + + 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: + raise RuntimeError("Cannot get data: YuanrongTransferEngineConnector is closed") + + get_key = self._make_key(get_key, from_stage, to_stage) + t0 = _time_mod.perf_counter() + + if not metadata: + if not self.sender_host or self.sender_zmq_port is None or str(self.sender_host).lower() == "auto": + self._metrics["errors"] += 1 + raise RuntimeError( + "get(metadata=None) requires sender_host and sender_zmq_port to be resolved. " + "Call update_sender_info(host, port) before polling." + ) + metadata = self._query_metadata_at(get_key, str(self.sender_host), int(self.sender_zmq_port)) + if not metadata: + return None + elif not (metadata.get("source_host") and metadata.get("source_rpc_port") and metadata.get("source_addrs")): + partial_host = metadata.get("source_host") + partial_port = metadata.get("source_port") + if not partial_host or not partial_port: + self._metrics["errors"] += 1 + logger.warning( + "get(%s): partial metadata missing source_host/source_port. metadata=%s", + get_key, + metadata, + ) + return None + queried = self._query_metadata_at(get_key, str(partial_host), int(partial_port)) + if not queried: + return None + metadata = queried + + t1 = _time_mod.perf_counter() + query_ms = (t1 - t0) * 1000 + + source_host = metadata.get("source_host") + source_port = metadata.get("source_port") + source_rpc_port = metadata.get("source_rpc_port") + source_addrs = list(metadata.get("source_addrs") or []) + lengths = list(metadata.get("lengths") or []) + data_size = int(metadata.get("data_size", 0)) + is_fast_path = bool(metadata.get("is_fast_path", False)) + + if not source_host or not source_rpc_port or str(source_host).lower() == "auto": + self._metrics["errors"] += 1 + logger.error( + "Invalid metadata for %s: source_host=%r, source_rpc_port=%r", + get_key, + source_host, + source_rpc_port, + ) + return None + if not source_addrs: + self._metrics["errors"] += 1 + logger.error("Invalid metadata for %s: missing source_addrs", get_key) + return None + if not lengths: + lengths = [data_size] + if len(source_addrs) != len(lengths): + self._metrics["errors"] += 1 + logger.error( + "Invalid metadata for %s: source_addrs/lengths mismatch (%s vs %s)", + get_key, + len(source_addrs), + len(lengths), + ) + return None + + total_size = sum(int(length) for length in lengths) + if data_size <= 0: + data_size = total_size + if data_size <= 0: + self._metrics["errors"] += 1 + logger.warning("Skipping get for %s: data_size is 0", get_key) + return None + + try: + offset = self.allocator.alloc(data_size) + recv_buffer = ManagedBuffer(self.allocator, offset, data_size, self.pool) + base_dst_ptr = self.base_ptr + offset + except MemoryError: + self._metrics["errors"] += 1 + logger.error("Failed to allocate %s bytes in receive pool", data_size) + return None + + dst_addrs: list[int] = [] + cursor = 0 + for length in lengths: + dst_addrs.append(base_dst_ptr + cursor) + cursor += int(length) + if cursor != data_size: + recv_buffer.release() + self._metrics["errors"] += 1 + logger.error("Length sum mismatch for %s: lengths_sum=%s, data_size=%s", get_key, cursor, data_size) + return None + + t2 = _time_mod.perf_counter() + alloc_ms = (t2 - t1) * 1000 + target_hostname = f"{source_host}:{int(source_rpc_port)}" + + try: + result = self.engine.batch_transfer_sync_read(target_hostname, dst_addrs, source_addrs, lengths) + self._ensure_result_ok(result, f"read from {target_hostname}") + self._sync_tensor_device(recv_buffer.tensor) + + t3 = _time_mod.perf_counter() + read_ms = (t3 - t2) * 1000 + + cleanup_ok = self._request_sender_cleanup(get_key, source_host, int(source_port) if source_port else None) + if not cleanup_ok: + logger.warning("Sender-side cleanup not confirmed for %s", get_key) + + if is_fast_path: + total_ms = (_time_mod.perf_counter() - t0) * 1000 + mbps = (data_size / 1024 / 1024) / (total_ms / 1000) if total_ms > 0 else 0 + logger.info( + "[YR GET] %s: query=%.1fms, alloc=%.1fms, read=%.1fms, total=%.1fms, %.1f MB/s " + "(fast_path, zero-copy)", + get_key, + query_ms, + alloc_ms, + read_ms, + total_ms, + mbps, + ) + self._metrics["gets"] += 1 + self._metrics["bytes_transferred"] += data_size + return recv_buffer, data_size + + try: + copy_start = _time_mod.perf_counter() + raw_bytes = self._buffer_to_bytes(recv_buffer) + copy_end = _time_mod.perf_counter() + value = OmniSerializer.deserialize(raw_bytes) + end = _time_mod.perf_counter() + total_ms = (end - t0) * 1000 + mbps = (data_size / 1024 / 1024) / (total_ms / 1000) if total_ms > 0 else 0 + logger.info( + "[YR GET] %s: query=%.1fms, alloc=%.1fms, read=%.1fms, copy=%.1fms, total=%.1fms, %.1f MB/s", + get_key, + query_ms, + alloc_ms, + read_ms, + (copy_end - copy_start) * 1000, + total_ms, + mbps, + ) + self._metrics["gets"] += 1 + self._metrics["bytes_transferred"] += data_size + return value, data_size + finally: + recv_buffer.release() + except Exception as exc: + self._metrics["errors"] += 1 + self._metrics["timeouts"] += 1 + logger.error("transfer_engine get failed for %s: %s", get_key, exc, exc_info=True) + recv_buffer.release() + return None + + def cleanup(self, request_id: str, from_stage: str | None = None, to_stage: str | None = None) -> None: + if (from_stage is None) != (to_stage is None): + raise ValueError( + "cleanup() requires both from_stage and to_stage, or neither. " + f"Got from_stage={from_stage!r}, to_stage={to_stage!r}" + ) + if from_stage is not None and to_stage is not None: + request_id = self._make_key(request_id, from_stage, to_stage) + + with self._local_buffers_lock: + item = self._local_buffers.pop(request_id, None) + if not item: + return + _, _, holder, should_release, _, _ = item + if should_release and isinstance(holder, ManagedBuffer): + holder.release() + + def health(self) -> dict[str, Any]: + if self._closed: + return {"status": "unhealthy", "error": "Connector is closed"} + + return { + "status": "healthy", + "host": self.host, + "zmq_port": self.zmq_port, + "rpc_port": self.rpc_port, + "protocol": self.protocol, + "pool_device": self.pool_device, + "pool_size": self.pool_size, + **self._metrics, + } + + def close(self) -> None: + if getattr(self, "_closed", True): + return + self._closed = True + + self._stop_event.set() + + if self._listener_thread is not None and self._listener_thread.is_alive(): + self._listener_thread.join(timeout=2.0) + + if self._control_executor is not None: + self._control_executor.shutdown(wait=True, cancel_futures=False) + + with self._local_buffers_lock: + for _, item in list(self._local_buffers.items()): + _, _, holder, should_release, _, _ = item + if should_release and isinstance(holder, ManagedBuffer): + holder.release() + self._local_buffers.clear() + + cache: dict[str, zmq.Socket] = getattr(self._req_local, "cache", None) # type: ignore[assignment] + if cache: + for sock in cache.values(): + try: + sock.close(linger=0) + except Exception: + pass + cache.clear() + + try: + if hasattr(self, "engine") and hasattr(self.engine, "unregister_memory"): + result = self.engine.unregister_memory(self.base_ptr) + if result.is_error(): + logger.warning("Failed to unregister memory: %s", result.to_string()) + except Exception as exc: + logger.warning("Failed to unregister memory: %s", exc) + + try: + if hasattr(self, "engine") and hasattr(self.engine, "finalize"): + result = self.engine.finalize() + if result.is_error(): + logger.warning("Failed to finalize transfer engine: %s", result.to_string()) + except Exception as exc: + logger.warning("Failed to finalize transfer engine: %s", exc) + + try: + if hasattr(self, "zmq_ctx"): + self.zmq_ctx.term() + except Exception as exc: + logger.warning("Failed to terminate ZMQ context: %s", exc) + + self.pool = None + + def _cleanup_stale_buffers(self) -> None: + now = _time_mod.monotonic() + with self._local_buffers_lock: + stale_keys = [k for k, v in self._local_buffers.items() if now - v[5] > _BUFFER_TTL_SECONDS] + for key in stale_keys: + item = self._local_buffers.pop(key) + _, _, holder, should_release, _, _ = item + if should_release and isinstance(holder, ManagedBuffer): + holder.release() + logger.warning("TTL expired (%ss): cleaned up stale buffer for %s", _BUFFER_TTL_SECONDS, key) + + def _zmq_listener_loop(self) -> None: + socket_obj = self.zmq_ctx.socket(zmq.ROUTER) + try: + socket_obj.bind(f"tcp://{self.host}:{self.zmq_port}") + except zmq.ZMQError as exc: + logger.error("ZMQ bind failed on %s:%s: %s", self.host, self.zmq_port, exc) + self.can_put = False + self._bind_error = exc + self._listener_ready.set() + return + + self._listener_ready.set() + + notify_addr = f"inproc://notify-{id(self)}" + notify_recv = self.zmq_ctx.socket(zmq.PULL) + notify_recv.bind(notify_addr) + + poller = zmq.Poller() + poller.register(socket_obj, zmq.POLLIN) + poller.register(notify_recv, zmq.POLLIN) + response_queue: queue.Queue = queue.Queue() + + try: + while not self._stop_event.is_set(): + try: + events = dict(poller.poll(1000)) + + if notify_recv in events: + while True: + try: + notify_recv.recv(zmq.NOBLOCK) + except zmq.Again: + break + + while True: + try: + identity, response = response_queue.get_nowait() + socket_obj.send_multipart([identity, b"", response]) + except queue.Empty: + break + + now = _time_mod.monotonic() + if now - self._last_ttl_check >= 10.0: + self._last_ttl_check = now + self._cleanup_stale_buffers() + + if socket_obj in events: + frames = socket_obj.recv_multipart() + if len(frames) >= 2: + self._control_executor.submit( + self._handle_control_request, + response_queue, + notify_addr, + frames[0], + frames[-1], + ) + except zmq.ContextTerminated: + break + except Exception: + logger.debug("Listener loop error", exc_info=True) + finally: + try: + notify_recv.close(linger=0) + socket_obj.close(linger=0) + except Exception: + pass + + def _handle_control_request(self, response_queue: queue.Queue, notify_addr: str, identity, payload: bytes) -> None: + try: + if payload.startswith(QUERY_INFO): + self._handle_query_request(response_queue, notify_addr, identity, payload[len(QUERY_INFO) :]) + return + if payload.startswith(CLEANUP_KEY): + self._handle_cleanup_request(response_queue, notify_addr, identity, payload[len(CLEANUP_KEY) :]) + return + + response_queue.put((identity, ACK_ERROR)) + except Exception as exc: + logger.error("Control request failed: %s", exc) + response_queue.put((identity, ACK_ERROR)) + + self._notify_listener(notify_addr) + + def _handle_query_request(self, response_queue: queue.Queue, notify_addr: str, identity, payload: bytes) -> None: + try: + query = msgspec.msgpack.decode(payload, type=QueryRequest) + with self._local_buffers_lock: + item = self._local_buffers.get(query.request_id) + + if not item: + response_queue.put((identity, INFO_NOT_FOUND)) + else: + src_addrs, src_lengths, _, _, is_fast_path, _ = item + resp = QueryResponse( + request_id=query.request_id, + rpc_port=self.rpc_port, + data_size=sum(int(length) for length in src_lengths), + is_fast_path=is_fast_path, + src_addrs=src_addrs, + lengths=src_lengths, + ) + response_queue.put((identity, msgspec.msgpack.encode(resp))) + except Exception as exc: + logger.error("Query request failed: %s", exc) + response_queue.put((identity, INFO_NOT_FOUND)) + + self._notify_listener(notify_addr) + + def _handle_cleanup_request(self, response_queue: queue.Queue, notify_addr: str, identity, payload: bytes) -> None: + try: + request = msgspec.msgpack.decode(payload, type=CleanupRequest) + self.cleanup(request.request_id) + response_queue.put((identity, ACK_OK)) + except Exception as exc: + logger.error("Cleanup request failed: %s", exc) + response_queue.put((identity, ACK_ERROR)) + + self._notify_listener(notify_addr) + + def _notify_listener(self, notify_addr: str) -> None: + try: + local = self._worker_local + sock = getattr(local, "notify_socket", None) + cached_addr = getattr(local, "notify_addr", None) + if sock is None or cached_addr != notify_addr: + if sock is not None: + sock.close(linger=0) + sock = self.zmq_ctx.socket(zmq.PUSH) + sock.connect(notify_addr) + local.notify_socket = sock + local.notify_addr = notify_addr + sock.send(b"", zmq.NOBLOCK) + except Exception: + local.notify_socket = None + local.notify_addr = None