diff --git a/docs/design/feature/disaggregated_inference.md b/docs/design/feature/disaggregated_inference.md index 2dae03ff8..57cd39d16 100644 --- a/docs/design/feature/disaggregated_inference.md +++ b/docs/design/feature/disaggregated_inference.md @@ -1,27 +1,35 @@ # Disaggregated Inference for Omni-Modality Models -This guide explains how to configure and use distributed connectors (vllm_omni/distributed/connectors) in vllm-omni for multi-stage pipelines. +This guide explains how to configure and use distributed connectors +(`vllm_omni/distributed/omni_connectors`) in vllm-omni for multi-stage pipelines. -## 1. Overview +Backend-specific setup lives in separate docs: + +- [SharedMemoryConnector](omni_connectors/shared_memory_connector.md) +- [MooncakeConnector](omni_connectors/mooncake_connector.md) +- [YuanrongConnector](omni_connectors/yuanrong_connector.md) + +## Overview Connectors enable data transfer between pipeline stages (e.g., Thinker -> Talker). -Currently supported connectors operate in **D2H2D (Device-to-Host-to-Device)** mode: -1. **SharedMemoryConnector**: Uses system shared memory. -2. **MooncakeConnector**: Uses [Mooncake](https://github.com/kvcache-ai/Mooncake). +Current connectors operate in D2H2D (device to host to device) mode. -* **SharedMemoryConnector (Default)**: Zero-copy (on host), lowest latency. Best for **single-node** deployments. Auto-configured if no connectors are specified. -* **MooncakeConnector**: TCP/RDMA based. Best for **multi-node** distributed deployments. Requires a Mooncake Master service. +## Connector Choices -## 2. API Design +| Use Case | Recommended Connector | Notes | +| :--- | :--- | :--- | +| Single node | SharedMemoryConnector | Auto-configured if no connector is specified. | +| Multi node (Mooncake) | MooncakeConnector | Requires Mooncake Master + metadata server. | +| Multi node (Yuanrong) | YuanrongConnector | Requires Yuanrong Datasystem + etcd. | -The connector system is built around the `OmniConnectorBase` abstraction, which decouples data transport from stage logic. +## Core API -### Core Interface +The connector system is built around `OmniConnectorBase`. ```python class OmniConnectorBase(ABC): @abstractmethod - def put(self, from_stage: str, to_stage: str, request_id: str, data: Any) -> tuple[bool, int, Optional[dict]]: + def put(self, from_stage: str, to_stage: str, put_key: str, data: Any) -> tuple[bool, int, Optional[dict]]: """ Store data. Returns: (success, serialized_size, metadata) @@ -29,145 +37,23 @@ class OmniConnectorBase(ABC): pass @abstractmethod - def get(self, from_stage: str, to_stage: str, request_id: str, metadata: Optional[dict] = None) -> Optional[tuple[Any, int]]: + def get(self, from_stage: str, to_stage: str, get_key: str, metadata: Optional[dict] = None) -> Optional[tuple[Any, int]]: """ Retrieve data. - Args: metadata - Transport-specific handles returned by put() (e.g., SHM name). + Args: metadata - transport-specific handles returned by put() (e.g., SHM name). Returns: (object, serialized_size) """ pass ``` -### Key Concept: Metadata Passing -Unlike a pure key-value store, some connectors (like `SharedMemoryConnector`) generate transient resources (e.g., a shared memory block name) during `put()`. This `metadata` **must be passed via the control plane** (e.g., HTTP headers, queue messages) from the producer stage to the consumer stage so `get()` can locate the data. +### Metadata Passing -## 3. Backends & Use Cases +Some connectors (e.g., SharedMemoryConnector) generate transient resources during `put()`. +This `metadata` must be passed through the control plane so `get()` can locate the data. -### 3.1 SharedMemoryConnector -**Best for:** Single-node, high-performance IPC. +## Configuration Model -* **Mechanism:** - * **Small Payloads (< Threshold)**: Data is serialized and passed directly "inline" within the `metadata` dictionary. This avoids the overhead of creating SHM blocks for tiny messages. - * **Large Payloads (>= Threshold)**: Data is written to a named System Shared Memory block. The block name is returned in `metadata`. -* **Configuration:** - * `shm_threshold_bytes`: Size in bytes to switch from inline to SHM (default: 64KB). - -### 3.2 MooncakeConnector -**Best for:** Multi-node distributed inference. - -* **Mechanism:** Uses Mooncake's distributed KVCache store. - * **Data Plane**: TCP or RDMA for high-bandwidth transfer. - * **Control Plane**: Uses a centralized Mooncake Master and Metadata Server. - * **Keying**: Deterministic keys based on `request_id/from_stage_to_stage`. -* **Requirements**: Requires a running Mooncake Master service. - -## 4. Relationship with vLLM - -vLLM provides specialized distributed mechanisms for specific artifacts: -* **KV Transfer** (`vllm.distributed.kv_transfer`): Optimized for transferring KV caches between prefill and decode instances (using NCCL, Mooncake, etc.). -* **EC Transfer** (`vllm.distributed.ec_transfer`): Optimized for sharing encoder embeddings. -* **Device Communicators** (`vllm.distributed.device_communicators`): Low-level primitives (NCCL, SHM) for Tensor/Pipeline Parallelism. - -`vllm-omni` complements this by introducing a **Generalized Connector Abstraction** (`OmniConnector`) for multimodal pipelines. While vLLM's connectors are artifact-specific, `vllm-omni`: - -1. **Unifies Transport**: Provides a single API (`put`/`get`) to transport *any* stage artifact (Input Embeddings, Hidden States, Audio/Image Tensors, KV Cache, Final Output) between arbitrary pipeline stages (e.g., AudioEncoder -> LLM -> AudioGenerator). -2. **Extends Connectivity**: Enables flexible construction of complex DAGs (Directed Acyclic Graphs) where stages can run in the same process, same node, or across nodes, using the most appropriate backend (SHM, Mooncake, etc.) for each edge. -3. **Wraps & Adapts**: Can internally utilize vLLM's specialized `kv_transfer` for KV paths while using generic transports (SHM/Mooncake) for other data types, presenting a consistent interface to the application layer. - -## 5. Installation (Mooncake) - -If using `MooncakeConnector`, install the library first: - -```bash -# For CUDA-enabled systems (Recommended) -pip install mooncake-transfer-engine - -# For non-CUDA systems -pip install mooncake-transfer-engine-non-cuda -``` - -## 6. Using MooncakeConnector - -### 6.1 Start Mooncake Master - -Start the master service on your primary node: - -```bash -# if you use mooncake SSD storage -mkdir -p ./mc_storage - -mooncake_master \ - --rpc_port=50051 \ - --enable_http_metadata_server=true \ - --http_metadata_server_host=0.0.0.0 \ - --http_metadata_server_port=8080 \ - --metrics_port=9003 \ - --root_fs_dir=./mc_storage/ \ - --cluster_id=mc-local-1 & -``` - -### 6.2 Configuration (YAML) - -Edit your stage config (e.g., `qwen2_5_omni.yaml`). - -**Step 1: Define Connector in Global Runtime** - -```yaml -runtime: - connectors: - connector_of_mooncake: - name: MooncakeConnector - extra: - host: "127.0.0.1" # Local Worker IP - metadata_server: "http://:8080/metadata" - master: ":50051" - segment: 512000000 # 512MB segment - localbuf: 64000000 # 64MB buffer - proto: "tcp" # "tcp" or "rdma" - ``` - -**Mooncake Configuration Parameters:** - -* **host**: The hostname or IP address of the local machine (worker). Mooncake uses this to register itself in the metadata server so other nodes can find it. -* **metadata_server**: The URL of the metadata server. This is used for service discovery and connection establishment (e.g., exchanging QP information for RDMA). -* **master**: The address of the Mooncake Master Server (e.g., `:50051`). This is used for global state management and control plane operations. -* **segment**: The size of the global memory segment in bytes (default: ~512MB). This defines the shared memory region accessible by Mooncake for data transfer. -* **localbuf**: The size of the local buffer in bytes (default: ~64MB). Used for local data buffering during transfer operations. -* **proto**: The transport protocol to use. Options: - * `tcp`: Standard TCP/IP (easier setup, universal compatibility). - * `rdma`: Remote Direct Memory Access (higher performance, requires RDMA-capable hardware). - -For more details, refer to the [Mooncake Repository](https://github.com/kvcache-ai/Mooncake). - - **Step 2: Reference in Stages** - -Explicitly link stages using `input_connectors` and `output_connectors`: - -```yaml -stage_args: - - stage_id: 0 - # ... - output_connectors: - to_stage_1: connector_of_mooncake - - - stage_id: 1 - # ... - input_connectors: - from_stage_0: connector_of_mooncake -``` - -## 7. Using SharedMemoryConnector (Auto-Mode) - -**Best for single-node.** - -The system will automatically create `SharedMemoryConnector`s for any pipeline edge that does not have an explicit connector defined. This is inferred from: -1. `runtime.edges` list in the config. -2. `engine_input_source` dependencies defined in `stage_args`. - -### Threshold Configuration -By default, payloads larger than **64KB** (default threshold) are transferred via shared memory, while smaller ones use the control queue (inline). - -To adjust this threshold (e.g., to 1GB), add the following to your `runtime.connectors`: +Define connectors in runtime: ```yaml runtime: @@ -175,26 +61,46 @@ runtime: connector_of_shared_memory: name: SharedMemoryConnector extra: - shm_threshold_bytes: 1024 # 1KB threshold + shm_threshold_bytes: 65536 ``` -## 8. Summary +Wire stages to connectors: -| Use Case | Recommended Connector | Configuration | -| :--- | :--- | :--- | -| **Single Node** | `SharedMemoryConnector` | **None** (Automatic) or Custom Threshold | -| **Multi Node** | `MooncakeConnector` | Explicit YAML + Mooncake Master | +```yaml +stage_args: + - stage_id: 0 + output_connectors: + to_stage_1: connector_of_shared_memory -## 9. Operational Notes (important) + - stage_id: 1 + input_connectors: + from_stage_0: connector_of_shared_memory +``` -- **Fail-fast config validation**: the loader raises if any expected edge is missing a connector. Define `input_connectors`/`output_connectors` or rely on auto-SHM filling; otherwise startup aborts. -- **Missing payloads halt the stage**: workers expect connector payloads; if metadata or connector config is missing, the stage raises and stops. Verify connector wiring and metadata propagation before production. +If a pipeline edge has no explicit connector, the system auto-creates a +SharedMemoryConnector for that edge. -## 10. Future Roadmap: Device-to-Device (D2D) Transport +## Relationship with vLLM -The current implementations (`SharedMemoryConnector`, `MooncakeConnector`) utilize a **D2H2D (Device-to-Host-to-Device)** data path. Tensors are moved to CPU memory (Host) for transport, which incurs PCIe overhead. +vLLM provides specialized distributed mechanisms for specific artifacts: -As outlined in the design RFC, future versions will introduce **D2D (Device-to-Device)** connectors: +- KV Transfer (`vllm.distributed.kv_transfer`): optimized for KV caches. +- EC Transfer (`vllm.distributed.ec_transfer`): optimized for encoder embeddings. +- Device Communicators (`vllm.distributed.device_communicators`): low-level primitives (NCCL, SHM). -* **Goal**: Direct GPU-to-GPU transfer (via NCCL, UCX, or IPC) to minimize latency for large tensor payloads. -* **Mechanism**: The `OmniConnector` API allows `put()` to initiate a transfer and return a lightweight handle (metadata) via the control plane, while the heavy payload flows directly between devices. +vllm-omni complements this with a generalized connector abstraction: + +1. Unifies transport via a single `put`/`get` API for any stage artifact. +2. Enables DAG-style pipelines across processes or nodes with per-edge transports. +3. Can wrap vLLM-specific transfers for KV paths while keeping a consistent interface. + +## Operational Notes + +- Fail-fast config validation: missing expected edges cause startup failures. +- Missing payloads halt stages: verify connector wiring and metadata propagation. + +## Future Roadmap: D2D Transport + +Current connectors use D2H2D paths. Future versions will introduce direct +device-to-device connectors (NCCL, UCX, IPC) to reduce latency for large +tensor payloads. diff --git a/docs/design/feature/omni_connectors/mooncake_connector.md b/docs/design/feature/omni_connectors/mooncake_connector.md new file mode 100644 index 000000000..fee409ecf --- /dev/null +++ b/docs/design/feature/omni_connectors/mooncake_connector.md @@ -0,0 +1,74 @@ +# MooncakeConnector + +## When to Use + +Best for multi-node distributed inference using Mooncake. + +## Installation + +```bash +# For CUDA-enabled systems (recommended) +pip install mooncake-transfer-engine + +# For non-CUDA systems +pip install mooncake-transfer-engine-non-cuda +``` + +## Start Mooncake Master + +```bash +# If you use Mooncake SSD storage +mkdir -p ./mc_storage + +mooncake_master \ + --rpc_port=50051 \ + --enable_http_metadata_server=true \ + --http_metadata_server_host=0.0.0.0 \ + --http_metadata_server_port=8080 \ + --metrics_port=9003 \ + --root_fs_dir=./mc_storage/ \ + --cluster_id=mc-local-1 & +``` + +## Configuration + +Define the connector in runtime: + +```yaml +runtime: + connectors: + connector_of_mooncake: + name: MooncakeConnector + extra: + host: "127.0.0.1" + metadata_server: "http://:8080/metadata" + master: ":50051" + segment: 512000000 + localbuf: 64000000 + proto: "tcp" +``` + +Wire stages to the connector: + +```yaml +stage_args: + - stage_id: 0 + output_connectors: + to_stage_1: connector_of_mooncake + + - stage_id: 1 + input_connectors: + from_stage_0: connector_of_mooncake +``` + +Parameters: + +- host: local worker IP registered in the metadata server. +- metadata_server: metadata server URL for discovery and setup. +- master: Mooncake Master address. +- segment: global memory segment size in bytes. +- localbuf: local buffer size in bytes. +- proto: transport protocol ("tcp" or "rdma"). + +For more details, refer to the +[Mooncake repository](https://github.com/kvcache-ai/Mooncake). diff --git a/docs/design/feature/omni_connectors/shared_memory_connector.md b/docs/design/feature/omni_connectors/shared_memory_connector.md new file mode 100644 index 000000000..eb91a889a --- /dev/null +++ b/docs/design/feature/omni_connectors/shared_memory_connector.md @@ -0,0 +1,27 @@ +# SharedMemoryConnector + +## When to Use + +Best for single-node deployments where stages run on the same host. It is +auto-configured when no explicit connector is specified for an edge. + +## How It Works + +- Small payloads (< threshold): serialized and passed inline in metadata. +- Large payloads (>= threshold): stored in shared memory; the SHM name is + returned in metadata. + +## Configuration + +```yaml +runtime: + connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 65536 +``` + +## Notes + +- Auto-mode uses SharedMemoryConnector if no connector is declared for an edge. diff --git a/docs/design/feature/omni_connectors/yuanrong_connector.md b/docs/design/feature/omni_connectors/yuanrong_connector.md new file mode 100644 index 000000000..88989d31a --- /dev/null +++ b/docs/design/feature/omni_connectors/yuanrong_connector.md @@ -0,0 +1,101 @@ +# YuanrongConnector + +## When to Use + +Best for multi-node distributed inference using Yuanrong Datasystem. + +## Mechanism + +Uses Yuanrong Datasystem's distributed KV store (`datasystem.kv_client`). + +- Data Plane: TCP or RDMA for high-bandwidth transfer. +- Control Plane: Yuanrong Datasystem workers and etcd. +- Keying: deterministic keys based on `put_key` (often composed as `request_id:fromStage_toStage`). + +## Installation + +```bash +pip install openyuanrong-datasystem +``` + +## Start etcd + +```bash +# Download and install etcd (v3.5.12 or higher) +ETCD_VERSION="v3.5.12" +ETCD_ARCH="linux-arm64" +wget https://github.com/etcd-io/etcd/releases/download/${ETCD_VERSION}/etcd-${ETCD_VERSION}-${ETCD_ARCH}.tar.gz +tar -xvf etcd-${ETCD_VERSION}-${ETCD_ARCH}.tar.gz +cd etcd-${ETCD_VERSION}-${ETCD_ARCH} +sudo cp etcd etcdctl /usr/local/bin/ + +# Start etcd +etcd \ + --name etcd-single \ + --data-dir /tmp/etcd-data \ + --listen-client-urls http://0.0.0.0:2379 \ + --advertise-client-urls http://0.0.0.0:2379 \ + --listen-peer-urls http://0.0.0.0:2380 \ + --initial-advertise-peer-urls http://0.0.0.0:2380 \ + --initial-cluster etcd-single=http://0.0.0.0:2380 & + +# Verify etcd is running +etcdctl --endpoints "127.0.0.1:2379" put key "value" +etcdctl --endpoints "127.0.0.1:2379" get key +``` + +For production environments, refer to the +[official etcd clustering documentation](https://etcd.io/docs/current/op-guide/clustering/). + +## Start Datasystem Worker + +```bash +# Replace ${ETCD_IP} with etcd node IP, ${WORKER_IP} with local node IP +dscli start -w \ + --worker_address "${WORKER_IP}:31501" \ + --etcd_address "${ETCD_IP}:2379" \ + --shared_memory_size_mb 20480 +``` + +To stop the worker: + +```bash +dscli stop --worker_address "${WORKER_IP}:31501" +``` + +## Configuration + +Define the connector in runtime: + +```yaml +runtime: + connectors: + connector_of_yuanrong: + name: YuanrongConnector + extra: + host: "127.0.0.1" + port: 31501 + get_sub_timeout_ms: 1000 +``` + +Wire stages to the connector: + +```yaml +stage_args: + - stage_id: 0 + output_connectors: + to_stage_1: connector_of_yuanrong + + - stage_id: 1 + input_connectors: + from_stage_0: connector_of_yuanrong +``` + +Parameters: + +- host: datasystem worker host. +- port: datasystem worker port. +- get_sub_timeout_ms: get timeout in milliseconds (0 for no timeout). + +For more details, refer to the +[Yuanrong Datasystem repository](https://atomgit.com/openeuler/yuanrong-datasystem). diff --git a/vllm_omni/distributed/__init__.py b/vllm_omni/distributed/__init__.py index dd58b0f22..3a7634219 100644 --- a/vllm_omni/distributed/__init__.py +++ b/vllm_omni/distributed/__init__.py @@ -8,6 +8,7 @@ from .omni_connectors import ( OmniConnectorFactory, OmniTransferConfig, SharedMemoryConnector, + YuanrongConnector, load_omni_transfer_config, ) @@ -20,6 +21,7 @@ __all__ = [ "OmniConnectorFactory", "MooncakeConnector", "SharedMemoryConnector", + "YuanrongConnector", # Utilities "load_omni_transfer_config", ] diff --git a/vllm_omni/distributed/omni_connectors/__init__.py b/vllm_omni/distributed/omni_connectors/__init__.py index b5288798c..fcec32eaa 100644 --- a/vllm_omni/distributed/omni_connectors/__init__.py +++ b/vllm_omni/distributed/omni_connectors/__init__.py @@ -4,6 +4,7 @@ from .connectors.base import OmniConnectorBase from .connectors.mooncake_connector import MooncakeConnector from .connectors.shm_connector import SharedMemoryConnector +from .connectors.yuanrong_connector import YuanrongConnector from .factory import OmniConnectorFactory from .utils.config import ConnectorSpec, OmniTransferConfig from .utils.initialization import ( @@ -26,6 +27,7 @@ __all__ = [ # Specific implementations "MooncakeConnector", "SharedMemoryConnector", + "YuanrongConnector", # Utilities "load_omni_transfer_config", "initialize_connectors_from_config", diff --git a/vllm_omni/distributed/omni_connectors/connectors/mooncake_connector.py b/vllm_omni/distributed/omni_connectors/connectors/mooncake_connector.py index 2e12d3a40..2f969b9bf 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/mooncake_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/mooncake_connector.py @@ -51,7 +51,7 @@ class MooncakeConnector(OmniConnectorBase): def _make_key(self, rid: str, from_stage: str, to_stage: str) -> str: """Generate store key for request between stages.""" - return f"{rid}/{from_stage}_to_{to_stage}" + return f"{rid}/{from_stage}_{to_stage}" def _init_store(self): """Initialize Mooncake store.""" @@ -72,16 +72,14 @@ class MooncakeConnector(OmniConnectorBase): # Use base class serialization methods for consistency - def put( - self, from_stage: str, to_stage: str, request_id: str, data: Any - ) -> tuple[bool, int, dict[str, Any] | None]: + def put(self, from_stage: str, to_stage: str, put_key: str, data: Any) -> tuple[bool, int, dict[str, Any] | None]: if not self.store: logger.error("Store not initialized") return False, 0, None try: serialized_data = self.serialize_obj(data) - key = self._make_key(request_id, from_stage, to_stage) + key = self._make_key(put_key, from_stage, to_stage) self.store.put(key, serialized_data, self.pin) self._metrics["puts"] += 1 @@ -102,7 +100,7 @@ class MooncakeConnector(OmniConnectorBase): return False, 0, None def get( - self, from_stage: str, to_stage: str, request_id: str, metadata: dict[str, Any] | None = None + self, from_stage: str, to_stage: str, get_key: str, metadata: dict[str, Any] | None = None ) -> tuple[Any, int] | None: if not self.store: logger.error("Store not initialized") @@ -110,7 +108,7 @@ class MooncakeConnector(OmniConnectorBase): retries = 20 sleep_s = 0.05 - key = self._make_key(request_id, from_stage, to_stage) + key = self._make_key(get_key, from_stage, to_stage) for attempt in range(retries): try: diff --git a/vllm_omni/distributed/omni_connectors/connectors/yuanrong_connector.py b/vllm_omni/distributed/omni_connectors/connectors/yuanrong_connector.py new file mode 100644 index 000000000..c147d00fd --- /dev/null +++ b/vllm_omni/distributed/omni_connectors/connectors/yuanrong_connector.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + +from ..utils.logging import get_connector_logger +from .base import OmniConnectorBase + +logger = get_connector_logger(__name__) + +try: + from datasystem.kv_client import KVClient, SetParam, WriteMode +except ImportError: + logger.warning("Datasystem not available, YuanrongConnector will not work") + KVClient = None + + +class YuanrongConnector(OmniConnectorBase): + """Datasystem-based distributed connector for OmniConnector.""" + + def __init__(self, config: dict[str, Any]): + if KVClient is None: + raise ImportError("Datasystem not available") + + self.config = config + self.client = None + self.set_param = SetParam() + self.set_param.write_mode = WriteMode.NONE_L2_CACHE_EVICT + self.get_sub_timeout_ms = max(0, int(self.config.get("get_sub_timeout_ms", 1000))) + + self._metrics = { + "puts": 0, + "gets": 0, + "bytes_transferred": 0, + "errors": 0, + "timeouts": 0, + } + + self._init_client() + + def _make_key(self, rid: str, from_stage: str, to_stage: str) -> str: + """Generate key for request between stages.""" + return f"{rid}:{from_stage}_{to_stage}" + + def _init_client(self): + """Initialize Datasystem client.""" + try: + self.host = self.config.get("host", "127.0.0.1") + self.port = int(self.config.get("port", "35001")) + self.client = KVClient(self.host, self.port) + self.client.init() + + logger.info("YuanrongConnector initialized successfully") + except Exception as e: + logger.error("Failed to initialize Datasystem client: %s", e) + raise + + def put(self, from_stage: str, to_stage: str, put_key: str, data: Any) -> tuple[bool, int, dict[str, Any] | None]: + if not self.client: + logger.error("Datasystem client not initialized") + return False, 0, None + + try: + serialized_data = self.serialize_obj(data) + key = self._make_key(put_key, from_stage, to_stage) + self.client.set(key, serialized_data, self.set_param.write_mode) + + self._metrics["puts"] += 1 + self._metrics["bytes_transferred"] += len(serialized_data) + + logger.debug( + "YuanrongConnector: stored %s (%s -> %s) %d bytes", + key, + from_stage, + to_stage, + len(serialized_data), + ) + return True, len(serialized_data), None + + except Exception as exc: + self._metrics["errors"] += 1 + logger.error("YuanrongConnector put failed: %s", exc) + 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 not self.client: + logger.error("Datasystem client not initialized") + return None + + key = self._make_key(get_key, from_stage, to_stage) + try: + raw_list = self.client.get([key], False, self.get_sub_timeout_ms) + raw_data = raw_list[0] if raw_list else None + if raw_data is not None: + data = self.deserialize_obj(raw_data) + self._metrics["gets"] += 1 + payload_size = len(raw_data) + logger.debug( + "YuanrongConnector: retrieved %s (%s -> %s) %d bytes", + key, + from_stage, + to_stage, + payload_size, + ) + return data, payload_size + + except Exception as exc: + self._metrics["timeouts"] += 1 + logger.error("YuanrongConnector get failed: %s", exc) + return None + + def cleanup(self, request_id: str) -> None: + if not self.client: + return + + # Note: Datasystem doesn't have explicit delete, data will be garbage collected + logger.debug("YuanrongConnector: cleanup requested for %s (no-op)", request_id) + + def health(self) -> dict[str, Any]: + if not self.client: + return {"status": "unhealthy", "error": "Datasystem client not initialized"} + + return {"status": "healthy", "host": self.host, "port": self.port, **self._metrics} + + def close(self) -> None: + if not self.client: + return + + self.client = None + logger.info("YuanrongConnector closed") diff --git a/vllm_omni/distributed/omni_connectors/factory.py b/vllm_omni/distributed/omni_connectors/factory.py index ee640d92f..34415f7fd 100644 --- a/vllm_omni/distributed/omni_connectors/factory.py +++ b/vllm_omni/distributed/omni_connectors/factory.py @@ -82,6 +82,19 @@ def _create_shm_connector(config: dict[str, Any]) -> OmniConnectorBase: return SharedMemoryConnector(config) +def _create_yuanrong_connector(config: dict[str, Any]) -> OmniConnectorBase: + try: + from .connectors.yuanrong_connector import YuanrongConnector + except ImportError: + import os + import sys + + sys.path.insert(0, os.path.dirname(os.path.dirname(__file__))) + from omni_connectors.connectors.yuanrong_connector import YuanrongConnector + return YuanrongConnector(config) + + # Register connectors OmniConnectorFactory.register_connector("MooncakeConnector", _create_mooncake_connector) OmniConnectorFactory.register_connector("SharedMemoryConnector", _create_shm_connector) +OmniConnectorFactory.register_connector("YuanrongConnector", _create_yuanrong_connector) diff --git a/vllm_omni/distributed/omni_connectors/utils/config.py b/vllm_omni/distributed/omni_connectors/utils/config.py index 9316c7e11..5556296d9 100644 --- a/vllm_omni/distributed/omni_connectors/utils/config.py +++ b/vllm_omni/distributed/omni_connectors/utils/config.py @@ -13,7 +13,7 @@ logger = get_connector_logger(__name__) class ConnectorSpec: """Specification for a connector instance.""" - name: str # e.g., "MooncakeConnector", "SharedMemoryConnector" + name: str # e.g., "MooncakeConnector", "SharedMemoryConnector", "YuanrongConnector" extra: dict[str, Any] = field(default_factory=dict) # backend-specific config diff --git a/vllm_omni/model_executor/stage_configs/qwen2_5_omni_multiconnector.yaml b/vllm_omni/model_executor/stage_configs/qwen2_5_omni_multiconnector.yaml index f2e4d03dd..aae194d37 100644 --- a/vllm_omni/model_executor/stage_configs/qwen2_5_omni_multiconnector.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen2_5_omni_multiconnector.yaml @@ -117,8 +117,15 @@ runtime: localbuf: 64000000 # 64MB proto: "tcp" + # Yuanrong connector for cross-node/intra-node communication + yuanrong_connector: + name: YuanrongConnector + extra: + host: "127.0.0.1" + port: "35000" + # SharedMemory connector for intra-node communication - #Alternative SHM connector with different threshold + # Alternative SHM connector with different threshold shared_memory_connector: name: SharedMemoryConnector extra: