diff --git a/docs/design/feature/omni_connectors/mooncake_store_connector.md b/docs/design/feature/omni_connectors/mooncake_store_connector.md index 5812a6abe..2ff116859 100644 --- a/docs/design/feature/omni_connectors/mooncake_store_connector.md +++ b/docs/design/feature/omni_connectors/mooncake_store_connector.md @@ -72,3 +72,241 @@ Parameters: For more details, refer to the [Mooncake repository](https://github.com/kvcache-ai/Mooncake). + +--- + +## Design + +### 1. Overview + +`MooncakeStoreConnector` is the storage-oriented remote connector in `vllm_omni/distributed/omni_connectors`. It is built on top of Mooncake distributed store APIs and provides a uniform `put()` / `get()` abstraction for transferring stage payloads across nodes. + +Compared with `SharedMemoryConnector`, this connector is intended for multi-node deployments. Compared with `MooncakeTransferEngineConnector`, it is simpler and more object-oriented: data is written into a distributed store and fetched back by key, rather than transferred through an explicit remote-memory write protocol. + +Its primary role is to provide a general-purpose remote connector for stage payloads when: + +- the pipeline spans multiple hosts +- a shared-memory connector is not possible +- a simpler remote backend is preferred over the RDMA-oriented transfer engine + +### 2. Relationship with the OmniConnector System + +`MooncakeStoreConnector` implements `OmniConnectorBase`, so it integrates with the same connector lifecycle as the other backends: + +- `OmniConnectorFactory` instantiates it from `ConnectorSpec` +- `load_omni_transfer_config()` maps edge config to the connector spec +- All callers (batch forwarding, chunk transfer, KV transfer, etc.) interact with it through the same `put()` / `get()` contract + +This means the rest of the pipeline does not need store-specific logic. It only relies on the generic connector contract. + +### 3. Design Goals + +The connector is designed around the following goals: + +- **Cross-node object transfer** using a single key-based abstraction +- **Minimal connector-specific control plane**, since the store itself is the shared medium +- **Compatibility with arbitrary Python payloads** through the Omni serializer +- **Operational simplicity** compared with the transfer-engine-based remote connector + +The connector is intentionally not optimized for zero-copy tensor movement or direct remote-memory access. + +### 4. Core Design + +#### 4.1 Object-Oriented Transfer Model + +`MooncakeStoreConnector` treats the transport backend as a distributed object store: + +1. Serialize the Python object into bytes. +2. Generate a unique connector key. +3. Store the bytes in Mooncake. +4. Fetch the bytes from Mooncake on the receiver side. +5. Deserialize the bytes back into the original object. + +This design makes the connector conceptually close to a distributed key-value transport. + +#### 4.2 Key Construction + +The connector uses `OmniConnectorBase._make_key()` to derive the internal store key: + +```text +{request_key}@{from_stage}_{to_stage} +``` + +This adds stage routing information to the logical request key and avoids collisions between different pipeline edges. + +#### 4.3 No Extra Metadata Path + +Unlike `SharedMemoryConnector` and `MooncakeTransferEngineConnector`, this connector returns: + +```python +metadata = None +``` + +from `put()`. + +This is a meaningful design difference: + +- the store itself is the shared rendezvous point +- consumers only need the same key +- no transport handle, address, or side-channel metadata needs to be propagated + +As a result, the control plane is simpler for this connector than for the SHM and RDMA variants. + +### 5. Initialization + +#### 5.1 Required Mooncake Components + +The connector requires the Mooncake Python bindings to expose: + +- `MooncakeDistributedStore` +- `ReplicateConfig` + +If these symbols are unavailable, construction fails immediately with `ImportError`. This keeps startup failures explicit and avoids silent fallback behavior. + +#### 5.2 Store Setup + +During `_init_store()`, the connector: + +1. creates `MooncakeDistributedStore()` +2. calls `store.setup(...)` +3. validates the return code +4. creates a `ReplicateConfig` +5. enables `with_soft_pin = True` + +The setup step is completed during connector construction, not lazily at first use. + +### 6. Put / Get Flow + +#### 6.1 Producer Flow: `put()` + +The producer-side flow is: + +1. Validate that the store has been initialized. +2. Serialize the payload into bytes. +3. Build the internal key with stage routing information. +4. Call `store.put(key, serialized_data, self.pin)`. +5. Update metrics and return success. + +The returned tuple is: + +```python +(True, len(serialized_data), None) +``` + +#### 6.2 Consumer Flow: `get()` + +The consumer-side flow is: + +1. Validate that the store has been initialized. +2. Build the same internal key. +3. Poll `store.get(key)` for up to 20 retries. +4. Sleep 50 ms between retries. +5. If data is found, deserialize it and return `(data, payload_size)`. +6. If all retries are exhausted, record a timeout and return `None`. + +This design gives the connector a bounded waiting model rather than an indefinite blocking get. + +### 7. Integration with Stage Communication + +All callers use the connector through the same `put()` / `get()` contract: + +- the sender calls `put()` to serialize and store the payload +- the receiver calls `get()` to retrieve and deserialize it +- no connector-specific metadata is required, since the store key is the rendezvous point + +Because `put()` returns `metadata=None`, the connector is naturally compatible with callers that do not forward metadata (e.g. polling-based flows). The trade-off is that all payloads incur full serialization and deserialization costs, which makes the connector functional but not the highest-performance option for large payloads such as KV cache blocks. + +### 8. Data Flow in the Pipeline + +The end-to-end transfer model is: + +```mermaid +sequenceDiagram + participant SenderStage + participant MooncakeStoreConnector + participant MooncakeStore + participant ReceiverStage + + SenderStage->>MooncakeStoreConnector: put(from_stage, to_stage, put_key, data) + MooncakeStoreConnector->>MooncakeStoreConnector: serialize object + MooncakeStoreConnector->>MooncakeStore: put(internal_key, bytes) + MooncakeStoreConnector-->>SenderStage: (success, size, None) + + ReceiverStage->>MooncakeStoreConnector: get(from_stage, to_stage, get_key) + MooncakeStoreConnector->>MooncakeStore: get(internal_key) + MooncakeStoreConnector->>MooncakeStoreConnector: deserialize bytes + MooncakeStoreConnector-->>ReceiverStage: (data, size) +``` + +This is a store-mediated remote transfer model rather than a direct peer-to-peer transport model. + +### 9. Strengths and Trade-offs + +#### Strengths + +- Simple conceptual model: store by key, fetch by key. +- No connector-specific metadata handoff is required. +- Works naturally for remote multi-node stage transfer. +- Easy to integrate into the existing connector abstraction. + +#### Trade-offs + +- Full serialization and deserialization are always required. +- Large payloads are more expensive than in direct-memory transports. +- Runtime behavior depends on external Mooncake services being available. +- Cleanup semantics are weaker than request-scoped local buffer management. + +### 10. Important Implementation Characteristics + +#### 10.1 Cleanup Is a No-op + +`cleanup()` only logs a debug message and does not actively delete remote data. The current design assumes Mooncake-side lifecycle management rather than explicit per-request removal. + +This keeps the connector implementation simple but means request-level reclamation is not modeled inside the connector itself. + +#### 10.2 Close Releases the Store Handle + +Unlike `SharedMemoryConnector`, where `close()` is a no-op, `MooncakeStoreConnector.close()` performs a meaningful shutdown: + +1. Calls `self.store.close()` to release the Mooncake store handle. +2. Sets `self.store = None` to prevent further operations. + +This ensures the connector releases its connection to the external Mooncake service on shutdown. Errors during close are logged but do not propagate. + +#### 10.3 Health Output Uses a Shared Schema + +`health()` reports: + +- `host` +- `metadata_server` +- `master` +- metrics + +and also includes placeholder fields such as `protocol`, `pool_device`, and `pool_size` to keep a more uniform shape across connector health outputs. This reflects a system-level consistency choice rather than a store-specific design need. + +#### 10.4 Get Is Retry-Based, Not Event-Driven + +The receiver side polls the store with a short retry loop. This is simple and robust, but it means the connector is latency-sensitive to: + +- store visibility delay +- network jitter +- payload size + +If the payload is not visible within the retry window, the connector reports a timeout. + +### 11. Summary + +`MooncakeStoreConnector` is the remote, store-based connector in the OmniConnector stack. Its design is straightforward: + +- serialize the payload +- store it in Mooncake under a stage-qualified key +- fetch it by the same key on the receiver side +- deserialize it back into the original object + +This connector fills an important gap in the system: + +1. It enables cross-node transfer without requiring shared memory. +2. It keeps the stage communication model uniform. +3. It provides a simpler operational alternative to the transfer-engine-based remote connector. + +Its simplicity is also its main trade-off: it favors a clean object-transfer model over the fast-path and peer-to-peer optimizations implemented by `MooncakeTransferEngineConnector`. diff --git a/docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md b/docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md index 96f93ef9a..798644b96 100644 --- a/docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md +++ b/docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md @@ -205,3 +205,589 @@ for test-specific setup instructions. For more details on the underlying engine, refer to the [Mooncake repository](https://github.com/kvcache-ai/Mooncake). + +--- + +## Design + +### 1. Overview + +`MooncakeTransferEngineConnector` is the high-performance remote connector in `vllm_omni/distributed/omni_connectors`. It is built on top of Mooncake `TransferEngine` and combines: + +- a **direct data plane** for remote memory writes +- a **ZMQ side channel** for metadata lookup, handshake, and completion signaling +- a **managed local memory pool** for both send and receive buffers + +Unlike `MooncakeStoreConnector`, which treats the backend as a distributed store, `MooncakeTransferEngineConnector` is designed as a peer-to-peer transport. Its goal is to move large stage payloads efficiently while still fitting the common `put()` / `get()` API defined by `OmniConnectorBase`. + +It is the most performance-oriented connector in the current OmniConnector family and is intended for large remote payloads such as: + +- KV cache transfer +- stage hidden-state payloads +- streaming chunk payloads +- other binary-heavy inter-stage artifacts + +### 2. Relationship with the OmniConnector System + +`MooncakeTransferEngineConnector` implements the same connector contract as the other backends: + +- `put(from_stage, to_stage, put_key, data)` +- `get(from_stage, to_stage, get_key, metadata=None)` +- `cleanup(request_id, ...)` +- `health()` +- `close()` + +It is integrated into the system through the standard connector plumbing: + +- `OmniConnectorFactory` constructs the connector from `ConnectorSpec` +- `load_omni_transfer_config()` resolves the edge-level connector configuration +- `get_connectors_config_for_stage()` and `resolve_omni_kv_config_for_stage()` inject the connector role +- All callers (batch forwarding, chunk transfer, KV transfer, etc.) interact with it through the same `put()` / `get()` contract + +The key system-level distinction is that this connector is **role-aware**: + +- sender instances expose data and listen for pull requests +- receiver instances allocate buffers and actively pull data from the sender + +### 3. Design Goals + +The connector is designed around four primary goals: + +1. **High-throughput remote transfer** + Avoid store-mediated round trips and write directly into the receiver memory region. + +2. **Fast path for raw payloads** + Support `torch.Tensor`, `bytes`, and `ManagedBuffer` without forcing all traffic through full object serialization. + +3. **Unified connector abstraction** + Preserve the same `put()` / `get()` interface used by the rest of the OmniConnector stack. + +4. **Safe lifecycle management** + Manage allocation, reuse, cleanup, and failure recovery for a registered memory pool. + +### 4. Architecture Overview + +At a high level, the connector is composed of four main subsystems: + +```mermaid +classDiagram + class OmniConnectorBase { + <> + +put(from_stage, to_stage, put_key, data) + +get(from_stage, to_stage, get_key, metadata) + +cleanup(request_id) + +health() + +close() + } + + class MooncakeTransferEngineConnector { + +supports_raw_data: bool + -engine: TransferEngine + -allocator: BufferAllocator + -pool: torch.Tensor + -zmq_ctx: zmq.Context + -_local_buffers: dict + -_sender_executor: ThreadPoolExecutor + -_listener_thread: threading.Thread + +put(...) + +get(...) + +update_sender_info(sender_host, sender_zmq_port) + +get_connection_info() + +cleanup(request_id, from_stage, to_stage) + +close() + } + + class BufferAllocator { + -total_size: int + -alignment: int + -free_blocks: list + +alloc(size) int + +free(offset, size) + } + + class ManagedBuffer { + -allocator: BufferAllocator + -offset: int + -size: int + -pool_tensor: torch.Tensor + +tensor + +as_tensor(dtype, shape) torch.Tensor + +to_bytes() bytes + +release() + } + + class TransferEngine { + +initialize(host, handshake, protocol, device_name) + +register_memory(base_ptr, size) + +batch_transfer_sync_write(remote_session, src_addrs, dst_addrs, lengths) + +unregister_memory(base_ptr) + +get_rpc_port() int + } + + class QueryRequest { + +request_id: str + } + + class QueryResponse { + +request_id: str + +data_size: int + +is_fast_path: bool + } + + class MooncakeAgentMetadata { + +remote_hostname: str + +remote_port: int + +request_id: str + +dst_addrs: list[int] + +lengths: list[int] + } + + OmniConnectorBase <|-- MooncakeTransferEngineConnector + MooncakeTransferEngineConnector *-- BufferAllocator + MooncakeTransferEngineConnector *-- TransferEngine + ManagedBuffer --> BufferAllocator : releases to + ManagedBuffer --> "1" torch.Tensor : views + MooncakeTransferEngineConnector ..> ManagedBuffer : returns / retains + MooncakeTransferEngineConnector ..> QueryRequest : decodes + MooncakeTransferEngineConnector ..> QueryResponse : encodes + MooncakeTransferEngineConnector ..> MooncakeAgentMetadata : exchanges +``` + +#### 4.1 Transfer Engine + +Mooncake `TransferEngine` is responsible for the actual data-plane transfer. It registers local memory and performs synchronous remote writes through: + +```python +batch_transfer_sync_write(...) +``` + +#### 4.2 Managed Memory Pool + +Each connector instance owns a large pre-registered memory pool: + +- CPU pinned memory when `memory_pool_device == "cpu"` +- GPU memory when `memory_pool_device == "cuda"` + +This avoids repeated memory registration and allows each transfer to allocate subranges from one long-lived pool. + +#### 4.3 Buffer Manager + +Two helper classes control local memory ownership: + +- `BufferAllocator` + Manages aligned subrange allocation and free-list merging. + +- `ManagedBuffer` + Represents one live slice of the pool and exposes: + - `.tensor` + - `.as_tensor(dtype, shape)` + - `.to_bytes()` + - `.release()` + +#### 4.4 ZMQ Side Channel + +ZMQ is used for transport coordination, not for the data payload itself. It handles: + +- metadata query from receiver to sender +- pull request submission +- completion or error signaling +- internal notification from worker threads back to the listener thread + +This split makes the control plane lightweight while keeping the bulk payload on the transfer engine data plane. + +### 5. Role Model + +#### 5.1 Sender Role + +A sender connector: + +- accepts `put()` calls +- stores live transfer-ready buffers in `_local_buffers` +- starts a ZMQ listener thread +- responds to metadata queries and pull requests from receivers + +#### 5.2 Receiver Role + +A receiver connector: + +- does not bind the sender-side ZMQ listener +- accepts `get()` calls +- allocates receive buffers from its own pool +- requests metadata or transfer service from the sender + +The role is not inferred dynamically. It is injected by the stage configuration layer: + +- incoming edge for a stage -> `role="receiver"` +- outgoing edge for a stage -> `role="sender"` + +This is important because incorrect role assignment would break initialization semantics. + +#### 5.3 Host Auto-Detection + +The `host` configuration field supports the special value `"auto"`. When set, the connector auto-detects the local IP address that would be used for external communication (via a UDP socket probe to `8.8.8.8`). If that fails, it falls back to hostname resolution, and ultimately to `127.0.0.1`. + +This is useful in environments where the operator does not want to hard-code IP addresses in the connector config. + +#### 5.4 RDMA Device Filtering + +The `device_name` configuration field allows the operator to specify which RDMA NICs to use (comma-separated, e.g. `"mlx5_0,mlx5_1"`). If not set in config, the connector also checks the `RDMA_DEVICE_NAME` environment variable. + +This is important in environments with mixed InfiniBand/RoCE NICs, where not all devices are suitable for the transfer engine. + +### 6. Local Memory Management + +#### 6.1 Memory Pool Registration + +During initialization, the connector: + +1. allocates a large pool tensor +2. records its base pointer +3. registers that memory with Mooncake +4. creates a `BufferAllocator` for subrange management + +This means every later transfer only allocates offsets inside the pre-registered pool rather than registering memory per request. + +#### 6.2 BufferAllocator + +`BufferAllocator` maintains a sorted free list of `(offset, size)` blocks and enforces alignment. Its responsibilities include: + +- aligned allocation +- freeing previously allocated blocks +- adjacent block merging +- double-free detection +- overlap detection to catch corruption + +This is a critical piece of the connector because both sender and receiver depend on long-lived pool reuse. + +#### 6.3 ManagedBuffer + +`ManagedBuffer` is the main fast-path data wrapper. It can: + +- expose the pool slice as a zero-copy 1D `uint8` tensor +- reinterpret that slice as a typed tensor +- copy out the contents as Python `bytes` +- release the slice back to the allocator + +The connector uses `ManagedBuffer` in two different ways: + +- as a send-side holder to keep the pool slice alive +- as a receive-side return type when `is_fast_path=True` + +### 7. Put Flow + +#### 7.1 High-Level Behavior + +`put()` is only valid in sender mode. Its job is to expose a payload for later remote pull by the receiver. + +The high-level flow is: + +1. validate connector state and role +2. convert the input into a pool-backed transferable representation +3. store the transfer metadata in `_local_buffers` +4. return lightweight metadata describing how the receiver can fetch the data + +#### 7.2 Payload Type Handling + +`put()` supports three payload classes: + +**A. `ManagedBuffer`** + +If the buffer belongs to the same pool, the connector can use it directly without copying. This is the most efficient path. + +If the buffer comes from a different pool, the connector falls back to a copy path. + +**B. `torch.Tensor` or `bytes`** + +These payloads are copied into the local pool and marked as fast-path data: + +- no Omni object serialization is required +- receiver can get a `ManagedBuffer` back + +**C. Generic Python object** + +Any other payload is serialized via `OmniSerializer.serialize(...)` and then copied into the pool. + +In this case: + +- `is_fast_path=False` +- the receiver will deserialize back into a Python object + +#### 7.3 Sender Metadata + +The sender returns: + +```python +{ + "source_host": self.host, + "source_port": self.zmq_port, + "data_size": size, + "is_fast_path": is_fast_path, +} +``` + +This metadata is intentionally lightweight. It tells the receiver: + +- where the sender-side control plane lives +- how large the remote transfer will be +- whether the payload should be returned as a `ManagedBuffer` or a deserialized object + +#### 7.4 Sender Buffer Table + +The sender stores each live payload in `_local_buffers` under the stage-qualified key. Each entry contains: + +- source addresses +- lengths +- the holder object +- ownership information (`should_release`) +- `is_fast_path` +- creation time + +This table is the sender-side truth source for both metadata queries and pull requests. + +### 8. Get Flow + +#### 8.1 High-Level Behavior + +`get()` runs on the receiver side and performs four steps: + +1. resolve metadata +2. allocate a destination buffer in the local pool +3. request the sender to write into that destination buffer +4. return either a `ManagedBuffer` or a deserialized object + +#### 8.2 Metadata Resolution Paths + +The `metadata` parameter in `get()` is optional. The connector supports two resolution modes depending on whether the caller supplies it. + +**With metadata** + +When the caller passes metadata, the connector uses it directly. The metadata carries: + +- `source_host` / `source_port` — sender ZMQ endpoint +- `data_size` — payload byte count +- `is_fast_path` — whether the receiver gets a `ManagedBuffer` or a deserialized object + +This mode is suitable when the control plane already forwards the sender's `put()` output to the receiver. + +**Without metadata** + +When `get(metadata=None)` is called, the connector queries the sender over ZMQ to discover the same fields (`data_size`, `is_fast_path`). The caller must first call: + +```python +update_sender_info(sender_host, sender_zmq_port) +``` + +so that the connector knows where to send the query. + +This mode is suitable for polling-based flows (e.g. KV transfer, async chunk transfer) where the receiver does not have metadata from the control plane. + +#### 8.3 Destination Allocation + +Once metadata is resolved, the receiver: + +1. allocates a subrange from its own local pool +2. wraps it in a `ManagedBuffer` +3. builds a `MooncakeAgentMetadata` request containing: + - receiver hostname + - receiver RPC port + - request ID + - destination addresses + - transfer lengths + +This tells the sender exactly where to write the incoming data. + +#### 8.4 Transfer Completion + +The receiver then sends the pull request over ZMQ and waits for: + +- `TRANS_DONE` +- or `TRANS_ERROR` + +If the transfer succeeds: + +- for `is_fast_path=True`, the receiver returns `(ManagedBuffer, size)` +- for `is_fast_path=False`, the receiver copies to bytes, deserializes, releases the buffer, and returns `(object, size)` + +### 9. Sender-Side Listener Design + +#### 9.1 Listener Thread + +In sender mode, the connector starts `_zmq_listener_loop()` after initialization. The listener: + +- binds `tcp://{host}:{zmq_port}` +- receives incoming requests +- uses a poller for socket events and internal notifications +- periodically reclaims stale buffers + +If the bind fails, initialization fails immediately. The code does not silently downgrade the connector role. + +#### 9.2 Worker Thread Pool + +The listener hands work to `_sender_executor` so that the listener thread does not block on transfer work. + +There are two request types: + +- metadata query -> `_handle_query_request(...)` +- transfer request -> `_handle_pull_request(...)` + +#### 9.3 Query Handling + +For metadata queries, the sender looks up the request ID in `_local_buffers` and returns: + +- data size +- fast-path flag + +This supports consumers that only know the sender endpoint but not the original sender metadata. + +#### 9.4 Pull Handling + +For a transfer request, the sender: + +1. locates the source addresses in `_local_buffers` +2. constructs the remote session identifier +3. calls `batch_transfer_sync_write(...)` +4. replies `TRANS_DONE` or `TRANS_ERROR` + +On success, the sender immediately calls `cleanup(meta.request_id)` and frees the producer-side buffer if it owns it. + +This choice is important: it makes the connector effectively a **single-consumer transfer model** for each successful put/get pair. + +### 10. Fast Path Semantics + +This connector explicitly advertises: + +```python +supports_raw_data = True +``` + +That means it can move raw payloads without forcing everything through the Omni object serializer. + +#### Fast Path + +For `torch.Tensor`, `bytes`, or pool-local `ManagedBuffer`: + +- sender returns `is_fast_path=True` +- receiver returns a `ManagedBuffer` +- caller is responsible for calling `release()` + +This avoids an unnecessary copy on the receiver side. + +#### Serialized Object Path + +For arbitrary Python objects: + +- sender serializes the payload +- receiver converts the receive buffer to bytes +- receiver deserializes the object +- receive buffer is released internally + +This preserves a uniform object-oriented API while still allowing optimized raw-data transport when possible. + +### 11. Failure Handling and Cleanup + +#### 11.1 Timeouts and Socket Recovery + +The receiver caches ZMQ REQ sockets per thread, but invalidates them after failures. This avoids reusing sockets that may be stuck in a bad state after timeout or receive errors. + +Timeout is scaled based on payload size: + +- a base timeout +- plus additional time for large payloads + +This is intended to reduce false timeouts for large remote writes. + +#### 11.2 Stale Buffer Reclamation + +The sender periodically reclaims old entries from `_local_buffers` using a TTL policy. This protects the memory pool from permanent leaks if a receiver crashes or never consumes a prepared payload. + +This is a practical recovery mechanism, although the code notes that TTL cleanup can still race with very long-running in-flight transfers. + +#### 11.3 Connector Shutdown + +`close()` is a full resource teardown routine. It: + +- stops the listener thread +- shuts down the worker executor +- releases all pending buffers +- closes cached sockets +- unregisters memory from the engine when supported +- terminates the ZMQ context +- drops the pool reference + +This makes `MooncakeTransferEngineConnector` the most lifecycle-aware connector in the current connector family. + +### 12. Current Implementation Constraints + +The current code documents several important topology constraints. + +#### 12.1 One Sender to One Receiver per Successful Transfer + +After a successful transfer, the sender-side buffer is cleaned up immediately. This means the same prepared payload is not retained for multiple independent receivers. + +#### 12.2 One Receiver to One Active Sender Endpoint + +The receiver only stores one `(sender_host, sender_zmq_port)` pair through `update_sender_info(...)`. So the metadata-query mode is currently single-sender at a time. + +#### 12.3 Explicit Buffer Ownership Matters + +When the connector allocates a pool slice internally, it is responsible for releasing it. When a caller passes an externally owned `ManagedBuffer`, the connector keeps it alive for transfer but does not assume ownership of its eventual release. + +These constraints are consistent with the current implementation and should be treated as design assumptions rather than incidental behavior. + +### 13. Data Flow in the Pipeline + +The end-to-end sender/receiver interaction is: + +```mermaid +sequenceDiagram + participant SenderStage + participant SenderConnector + participant ReceiverConnector + participant ReceiverStage + + SenderStage->>SenderConnector: put(from_stage, to_stage, put_key, data) + SenderConnector->>SenderConnector: place payload in local memory pool + SenderConnector-->>SenderStage: metadata(source_host, source_port, data_size, is_fast_path) + + ReceiverStage->>ReceiverConnector: get(..., metadata) + ReceiverConnector->>ReceiverConnector: allocate destination buffer + ReceiverConnector->>SenderConnector: ZMQ pull request with dst addr + SenderConnector->>ReceiverConnector: TransferEngine remote write + SenderConnector-->>ReceiverConnector: TRANS_DONE + ReceiverConnector-->>ReceiverStage: ManagedBuffer or deserialized object +``` + +For metadata-less polling, the flow simply adds a metadata query step before the pull request. + +### 14. Strengths and Trade-offs + +#### Strengths + +- Best remote-transfer design in the current connector stack for large payloads. +- Supports raw-data fast path. +- Keeps stage communication under the same connector abstraction. +- Includes real lifecycle and memory-pool management. +- Works for both stage payload transfer and KV transfer scenarios. + +#### Trade-offs + +- More complex than the store-based connector. +- Correctness depends on role injection and endpoint coordination. +- Caller must release fast-path receive buffers. +- Current implementation is optimized for single-consumer transfer semantics. + +### 15. Summary + +`MooncakeTransferEngineConnector` is the high-performance peer-to-peer transport in the OmniConnector system. Its design combines: + +- a registered memory pool +- a safe subrange allocator +- a ZMQ control plane +- a Mooncake transfer-engine data plane + +This allows the connector to support both: + +1. a **fast path** for raw tensors and bytes +2. a **generic object path** for arbitrary Python payloads + +Within vLLM-Omni, it is the connector that most directly targets performance-sensitive remote transfer, especially for large payloads and KV cache movement. Its additional complexity is deliberate: it is the connector that turns the generic OmniConnector abstraction into a transport capable of efficient remote memory movement rather than simple object storage. diff --git a/docs/design/feature/omni_connectors/shared_memory_connector.md b/docs/design/feature/omni_connectors/shared_memory_connector.md index eb91a889a..5b35014f2 100644 --- a/docs/design/feature/omni_connectors/shared_memory_connector.md +++ b/docs/design/feature/omni_connectors/shared_memory_connector.md @@ -7,9 +7,10 @@ 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. +All payloads are serialized and stored in shared memory (`/dev/shm`); the SHM +segment name is returned in metadata. The configuration exposes a +`shm_threshold_bytes` field for a future inline-vs-SHM split, but the current +implementation always uses shared memory regardless of payload size. ## Configuration @@ -25,3 +26,234 @@ runtime: ## Notes - Auto-mode uses SharedMemoryConnector if no connector is declared for an edge. + +--- + +## Design + +### 1. Overview + +`SharedMemoryConnector` is the default same-node connector in `vllm_omni/distributed/omni_connectors`. It is designed for stage-to-stage transfer when producer and consumer processes run on the same host and can share `/dev/shm`. + +The connector provides a unified `put()` / `get()` API for arbitrary Python objects while keeping the control plane lightweight: + +- The payload is serialized by the connector. +- The serialized bytes are placed in shared memory. +- The queue/control plane only carries a small metadata handle. + +This makes `SharedMemoryConnector` the simplest connector in the OmniConnector family and the default fallback when an edge does not explicitly configure another backend. + +### 2. Relationship with the OmniConnector System + +`SharedMemoryConnector` implements `OmniConnectorBase`, so it follows the same lifecycle and API contract as the other connectors: + +- `put(from_stage, to_stage, put_key, data)` +- `get(from_stage, to_stage, get_key, metadata=None)` +- `cleanup(request_id)` +- `health()` +- `close()` + +Within the larger system: + +- `load_omni_transfer_config()` automatically fills missing edges with `SharedMemoryConnector`. +- Callers interact with the connector exclusively through the `put()` / `get()` / `cleanup()` contract — the connector does not require caller-specific logic. + +Compared with the remote Mooncake-based connectors, `SharedMemoryConnector` is intentionally minimal and local-only. + +### 3. Design Goals + +The connector is built around the following goals: + +- **Low-friction local transfer** for single-node multi-process pipelines. +- **Unified object semantics** for arbitrary Python payloads. +- **Small control-plane overhead** by passing only metadata through queues. +- **Zero external dependencies** beyond Python shared memory and the existing stage utilities. + +It is not intended to provide cross-node transfer, RDMA, or raw tensor zero-copy semantics across processes. + +### 4. Core Design + +#### 4.1 Serialization Model + +`SharedMemoryConnector` always starts from a Python object and serializes it through the shared Omni serializer: + +```python +payload = self.serialize_obj(data) +``` + +This keeps the connector behavior consistent with the rest of the connector stack: + +- producer code does not need connector-specific serialization logic +- consumer code always receives the original object after deserialization +- the connector can reuse the same serializer used by other backends + +#### 4.2 Shared Memory as the Data Plane + +The actual data plane is a shared-memory segment created by: + +- `shm_write_bytes(...)` +- `shm_read_bytes(...)` + +The connector stores a small metadata object such as: + +```python +{ + "shm": {"name": ..., "size": ...}, + "size": ... +} +``` + +This metadata is passed over the control plane and allows the downstream stage to locate the shared-memory segment. + +#### 4.3 Locking Model + +To avoid races between the producer and consumer, the connector uses a lock file per request: + +```text +/dev/shm/shm_{put_key}_lockfile.lock +``` + +Locking is done with `fcntl.flock`: + +- producer uses `LOCK_EX` +- consumer uses `LOCK_EX` + +Both sides acquire an exclusive lock. This ensures that the shared-memory segment is not read while it is still being written and makes the handoff safer in a multi-process environment. + +### 5. Put / Get Flow + +#### 5.1 Producer Flow: `put()` + +The producer-side flow is: + +1. Serialize the input object to bytes. +2. Compute the payload size. +3. Acquire the per-request lock file. +4. Write the bytes into shared memory. +5. Return lightweight metadata to the caller. + +The returned tuple is: + +```python +(success, serialized_size, metadata) +``` + +where `metadata` contains the shared-memory handle needed by the consumer. + +#### 5.2 Consumer Flow: `get(metadata=...)` + +The primary consumer path is metadata-driven: + +1. Extract the shared-memory handle from `metadata`. +2. Acquire the exclusive lock. +3. Read the raw bytes from shared memory. +4. Deserialize the bytes back into the original Python object. +5. Remove the lock file if it still exists. + +This is the path used by the current stage-to-stage connector flow. + +#### 5.3 Compatibility Flow: `get(metadata=None)` + +The connector also keeps a compatibility path for callers that only know the key: + +1. Attempt to open the shared-memory segment by name via `SharedMemory(name=get_key)`. +2. If the segment exists and has non-zero size, acquire the exclusive lock and read the bytes. +3. Deserialize the bytes and return the object. + +If the segment does not exist or any exception occurs, the call returns `None` immediately. There is no retry loop in this path -- it is a single-attempt open. + +This path is mainly for older code paths and is not the preferred mode for the current connector pipeline. + +### 6. Key Implementation Characteristics + +#### 6.1 Threshold Exists, but the Current Code Always Uses SHM + +The class keeps a `shm_threshold_bytes` field and still exposes metrics for inline writes. However, the current implementation uses: + +```python +if True: + ... +``` + +inside `put()`, which means the current code path always writes to shared memory. + +So the design still suggests a future split between: + +- small payloads inline +- large payloads in shared memory + +but the current behavior is effectively: + +- all payloads go through shared memory + +This should be documented because it affects real runtime behavior. + +#### 6.2 Cleanup Is Currently Passive + +`cleanup()` is currently a no-op. The intended assumption is: + +- the consumer reads the segment +- the underlying shared-memory helpers unlink it + +If the consumer never executes `get()`, the shared-memory segment may remain allocated. This means the connector relies on the normal success path for resource reclamation. + +#### 6.3 Close Is Currently Minimal + +`close()` is also a no-op. There is no connector-owned background thread, socket, or memory pool to tear down, so the lifecycle is simple. The trade-off is that `close()` does not scan or recover leaked shared-memory resources. + +### 7. Data Flow in the Pipeline + +The typical flow with `SharedMemoryConnector` is: + +```mermaid +sequenceDiagram + participant SenderStage + participant SharedMemoryConnector + participant QueueOrControlPlane + participant ReceiverStage + + SenderStage->>SharedMemoryConnector: put(from_stage, to_stage, put_key, data) + SharedMemoryConnector->>SharedMemoryConnector: serialize object + SharedMemoryConnector->>SharedMemoryConnector: write bytes to /dev/shm + SharedMemoryConnector-->>SenderStage: metadata {shm, size} + SenderStage->>QueueOrControlPlane: forward connector metadata + QueueOrControlPlane->>ReceiverStage: task + connector metadata + ReceiverStage->>SharedMemoryConnector: get(from_stage, to_stage, get_key, metadata) + SharedMemoryConnector->>SharedMemoryConnector: read bytes from /dev/shm + SharedMemoryConnector->>SharedMemoryConnector: deserialize object + SharedMemoryConnector-->>ReceiverStage: (data, size) +``` + +This is a classic split-control-plane / data-plane design, but constrained to a single host. + +### 8. Strengths and Trade-offs + +#### Strengths + +- Very simple deployment model. +- No external service dependency. +- Fits naturally into the existing queue-driven orchestration flow. +- Good default for local multi-process pipelines. + +#### Trade-offs + +- Same-node only. +- Full object serialization and deserialization are still required. +- Resource cleanup depends on the normal consumer path. +- Shared memory capacity is limited by host configuration. + +### 9. Summary + +`SharedMemoryConnector` is the baseline local transport for the OmniConnector system. Its design is intentionally straightforward: + +- serialize object +- place bytes in shared memory +- pass metadata through the control plane +- deserialize on the receiving side + +It plays two important roles in vLLM-Omni: + +1. It is the simplest production-ready connector for same-node stage pipelines. +2. It serves as the automatic fallback connector when no explicit edge transport is configured. + +Although the current implementation is deliberately minimal, it provides the foundation for reliable local connector semantics and keeps the stage communication model uniform across the system. diff --git a/docs/design/feature/omni_connectors/yuanrong_connector.md b/docs/design/feature/omni_connectors/yuanrong_connector.md index 88989d31a..12127ba3e 100644 --- a/docs/design/feature/omni_connectors/yuanrong_connector.md +++ b/docs/design/feature/omni_connectors/yuanrong_connector.md @@ -94,8 +94,265 @@ stage_args: Parameters: - host: datasystem worker host. -- port: datasystem worker port. +- port: datasystem worker port (default: `35001` if omitted; the example above uses `31501` to match the worker startup command). - 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). + +--- + +## Design + +### 1. Overview + +`YuanrongConnector` is the Datasystem-based remote connector in `vllm_omni/distributed/omni_connectors`. It uses Yuanrong Datasystem's distributed key-value client as the transport backend and exposes the same `put()` / `get()` interface as the other OmniConnectors. + +Like `MooncakeStoreConnector`, it is a store-oriented remote connector rather than a direct peer-to-peer transport. Its role is to let stage payloads move across nodes through a deterministic key-based storage abstraction while keeping the rest of the pipeline on the common connector API. + +It is intended for deployments that already use Yuanrong Datasystem and want a remote connector that integrates with the existing OmniConnector configuration and orchestration model. + +### 2. Relationship with the OmniConnector System + +`YuanrongConnector` implements `OmniConnectorBase`, so it participates in the same connector lifecycle as the other backends: + +- `OmniConnectorFactory` constructs it from a `ConnectorSpec` +- stage edge configuration is resolved by `load_omni_transfer_config()` +- All callers (batch forwarding, chunk transfer, KV transfer, etc.) interact with it through the same `put()` / `get()` contract + +This means the connector is not exposed directly to stage logic. Stages only interact with the generic connector contract, and the backend choice remains a configuration concern. + +### 3. Design Goals + +The connector is built around the following goals: + +1. **Cross-node payload transfer through Datasystem** + Reuse Yuanrong Datasystem as the remote exchange medium for stage data. + +2. **Uniform object transfer semantics** + Allow arbitrary Python objects to be transmitted through the shared Omni serializer. + +3. **Minimal connector-specific control plane** + Use deterministic keys so that consumers can retrieve data without an extra transport metadata handoff. + +4. **Operational reuse of existing infrastructure** + Fit into environments that already deploy Yuanrong Datasystem workers and etcd. + +The connector is not designed for direct remote-memory writes or tensor-specific fast-path transfer. + +### 4. Core Design + +#### 4.1 Store-Oriented Transfer Model + +`YuanrongConnector` treats the transport backend as a distributed object store: + +1. serialize the Python payload +2. build a deterministic connector key +3. write the serialized bytes into Datasystem +4. read the bytes back on the receiver side +5. deserialize them into the original object + +This is the same broad architectural class as `MooncakeStoreConnector`, but implemented on top of Yuanrong Datasystem APIs instead of Mooncake store APIs. + +#### 4.2 Deterministic Keying + +Unlike the default `_make_key()` in `OmniConnectorBase`, `YuanrongConnector` defines its own key format: + +```text +{request_id}:{from_stage}_{to_stage} +``` + +This has two design implications: + +- the request identifier remains the primary lookup handle +- stage routing information is embedded in the key so that the same logical request ID can safely appear on different edges + +The explicit override also makes the key format easier to align with Datasystem-side debugging and operational inspection. + +#### 4.3 No Extra Metadata Hand-off + +`put()` returns: + +```python +(success, serialized_size, None) +``` + +and does not generate connector-specific metadata. + +This design works because the receiver can reconstruct the exact same key from: + +- `get_key` +- `from_stage` +- `to_stage` + +As a result, the connector does not require a separate side-channel metadata handoff. + +### 5. Initialization + +#### 5.1 Datasystem Client Dependency + +The connector requires the Datasystem Python bindings to expose: + +- `KVClient` +- `SetParam` +- `WriteMode` + +If any of these symbols are unavailable, connector construction fails immediately with `ImportError`. This keeps configuration errors explicit and avoids a partially initialized runtime. + +#### 5.2 Client Setup + +During `_init_client()`, the connector: + +1. reads `host` and `port` +2. creates `KVClient(host, port)` +3. calls `client.init()` + +At construction time it also creates a `SetParam` and fixes: + +```python +self.set_param.write_mode = WriteMode.NONE_L2_CACHE_EVICT +``` + +This means the connector has a stable write policy for all writes and does not currently expose write-mode selection as a higher-level connector option. + +### 6. Put / Get Flow + +#### 6.1 Producer Flow: `put()` + +The producer-side flow is: + +1. verify that the Datasystem client has been initialized +2. serialize the input object with the shared Omni serializer +3. build the Datasystem key using the connector-specific `_make_key()` +4. call `client.set(key, serialized_data, self.set_param.write_mode)` +5. update metrics and return success + +The returned metadata is always `None`, because the Datasystem key itself is the lookup contract between producer and consumer. + +#### 6.2 Consumer Flow: `get()` + +The consumer-side flow is: + +1. verify that the Datasystem client has been initialized +2. rebuild the same key with `from_stage`, `to_stage`, and `get_key` +3. call: + +```python +client.get([key], False, self.get_sub_timeout_ms) +``` + +4. take the first returned element if present +5. deserialize it and return `(data, payload_size)` + +If the returned list is empty or contains no data for the key, `get()` returns `None`. + +### 7. Timeout and Retrieval Semantics + +The connector uses `get_sub_timeout_ms` as its read timeout. Unlike `MooncakeStoreConnector`, which performs an explicit retry loop in Python, `YuanrongConnector` delegates waiting behavior more directly to the Datasystem client call. + +This leads to a slightly different retrieval model: + +- `MooncakeStoreConnector`: retry-oriented polling in connector code +- `YuanrongConnector`: single client call with backend-managed timeout behavior + +From the connector API perspective the result is the same, but operationally the waiting behavior is more dependent on Datasystem client semantics. + +### 8. Integration with Stage Communication + +All callers use the connector through the same `put()` / `get()` contract: + +- the sender calls `put()` to serialize and store the payload +- the receiver calls `get()` to retrieve and deserialize it +- no connector-specific metadata is required, since the Datasystem key is the rendezvous point + +Because `put()` returns `metadata=None`, the connector is naturally compatible with callers that do not forward metadata (e.g. polling-based flows). The trade-off is that all payloads incur full serialization and deserialization costs, and there is no raw tensor fast path, which makes the connector functional but not optimized for the largest payloads. + +### 9. Data Flow in the Pipeline + +The end-to-end transfer model is: + +```mermaid +sequenceDiagram + participant SenderStage + participant YuanrongConnector + participant Datasystem + participant ReceiverStage + + SenderStage->>YuanrongConnector: put(from_stage, to_stage, put_key, data) + YuanrongConnector->>YuanrongConnector: serialize object + YuanrongConnector->>Datasystem: set(key, bytes) + YuanrongConnector-->>SenderStage: (success, size, None) + + ReceiverStage->>YuanrongConnector: get(from_stage, to_stage, get_key) + YuanrongConnector->>Datasystem: get([key], timeout) + YuanrongConnector->>YuanrongConnector: deserialize bytes + YuanrongConnector-->>ReceiverStage: (data, size) +``` + +This is a store-mediated remote connector design with deterministic key lookup and no explicit side-channel metadata exchange. + +### 10. Strengths and Trade-offs + +#### Strengths + +- Reuses existing Yuanrong Datasystem infrastructure. +- Keeps the connector contract simple and uniform. +- Requires no connector-specific metadata handoff. +- Suitable for remote stage transfer in Datasystem-based deployments. + +#### Trade-offs + +- Always pays full serialization and deserialization cost. +- Does not support raw bytes / tensor fast-path semantics. +- Depends on external Datasystem worker availability. +- Retrieval and timeout behavior are partly delegated to the backend client. + +### 11. Important Implementation Characteristics + +#### 11.1 Cleanup Is a No-op + +`cleanup()` only logs a debug message and does not explicitly remove data from Datasystem. The current design assumes backend-side lifecycle or garbage collection rather than request-scoped delete semantics inside the connector. + +This mirrors the same design trade-off seen in other store-based connectors: simplicity over explicit per-request data reclamation. + +#### 11.2 Close Only Releases the Local Client Handle + +`close()` does not perform a remote shutdown. It simply clears the local `client` reference and marks the connector as closed from the process perspective: + +```python +self.client = None +``` + +This is appropriate for a client-based store connector, but it also means that backend resource lifecycle remains outside the connector's control. + +#### 11.3 Error and Timeout Metrics Are Coarse-Grained + +The connector tracks: + +- `puts` +- `gets` +- `bytes_transferred` +- `errors` +- `timeouts` + +In the current code, failed `put()` increments `errors`, while `get()` exceptions increment `timeouts`. This is operationally useful, but it does not distinguish between: + +- backend timeout +- not-found result +- transport failure +- deserialization failure + +So the metrics should be read as high-level health indicators, not detailed root-cause diagnostics. + +### 12. Summary + +`YuanrongConnector` is the Datasystem-backed remote connector in the OmniConnector family. Its design is straightforward: + +- serialize payloads +- store them under a deterministic stage-qualified key +- retrieve them by the same key +- deserialize them on the receiving side + +Within vLLM-Omni, it provides a clean integration point for Yuanrong Datasystem-based deployments while preserving the same connector abstraction used by the rest of the pipeline. + +Its design priorities are simplicity, infrastructure reuse, and API consistency rather than direct-memory transport optimization.