[Doc] Update docs for a better quick start (#814)
* add transfer engine python api * add quick start * add lmcache v1 integration * update trace * fix * modify segment size
This commit is contained in:
Ruoyu Qin
GitHub
parent
cd717c7091
commit
dedfbde5da
BIN
Mooncake-v3.pdf
BIN
Mooncake-v3.pdf
Binary file not shown.
|
|
@ -404,6 +404,7 @@ The general form is `[proto]://[hostname:port]`. For example, the following meta
|
|||
Reclaims all allocated resources and also deletes the global meta data server information.
|
||||
|
||||
## Using Transfer Engine to Your Projects
|
||||
|
||||
### Using C/C++ Interface
|
||||
After compiling Mooncake Store, you can move the compiled static library file `libtransfer_engine.a` and the C header file `transfer_engine_c.h` into your own project. There is no need to reference other files under `src/transfer_engine`.
|
||||
|
||||
|
|
|
|||
|
|
@ -2,6 +2,15 @@
|
|||
|
||||
This document describes how to build Mooncake.
|
||||
|
||||
## PyPI Package
|
||||
Install the Mooncake Transfer Engine package from PyPI, which includes both Mooncake Transfer Engine and Mooncake Store Python bindings:
|
||||
|
||||
```bash
|
||||
pip install mooncake-transfer-engine==0.3.5
|
||||
```
|
||||
|
||||
📦 **Package Details**: [https://pypi.org/project/mooncake-transfer-engine/](https://pypi.org/project/mooncake-transfer-engine/)
|
||||
|
||||
## Automatic
|
||||
|
||||
### Recommended Version
|
||||
|
|
|
|||
|
|
@ -17,7 +17,13 @@ LMCache is an open-source Knowledge Delivery Network (KDN), specifically designe
|
|||
This strategic alliance delivers substantial mutual technological advancements:
|
||||
|
||||
* **LMCache Integrates Mooncake for Enhanced Performance**: LMCache leverages Mooncake’s advanced transfer engine and KVCache storage solutions, significantly enhancing data distribution efficiency and overall system performance.
|
||||
* **Mooncake Integrates LMCache for Enhanced Functionality**: Mooncake now includes LMCache as its KVCache management layer, benefiting from LMCache’s flexible cache control mechanisms and future-oriented functionalities, such as CacheBlend. This integration enables Mooncake to offer more sophisticated features and improved operational flexibility.
|
||||
* **Mooncake Integrates LMCache for Enhanced Functionality**: Mooncake now includes LMCache as its KVCache management layer, benefiting from LMCache's flexible cache control mechanisms and future-oriented functionalities, such as CacheBlend. This integration enables Mooncake to offer more sophisticated features and improved operational flexibility.
|
||||
|
||||
## Getting Started
|
||||
|
||||
For a complete deployment guide with step-by-step instructions, see:
|
||||
|
||||
👉 **[vLLM V1 Disaggregated Serving with Mooncake Store and LMCache](vllm-integration/vllmv1-lmcache-integration.md)**
|
||||
|
||||
## Performance Benchmarking and Results
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,11 @@
|
|||
# vLLM Disaggregated Serving
|
||||
|
||||
::::{toctree}
|
||||
:maxdepth: 1
|
||||
|
||||
vllmv1-lmcache-integration
|
||||
vllm-integration-v0.2
|
||||
vllm-integration-v1
|
||||
::::
|
||||
|
||||
|
||||
|
|
@ -1,13 +1,16 @@
|
|||
# vLLM Disaggregated Serving Demo
|
||||
# vLLM V0 Disaggregated Serving Demo
|
||||
|
||||
## Overview
|
||||
This is the latest version of mooncake-transfer-engine integration doc with the vLLM project based on [PR 10502](https://github.com/vllm-project/vllm/pull/10502) and [PR 10884](https://github.com/vllm-project/vllm/pull/10884) (vllm version: v0.6.4.post1/main) to accelerate KVCache transfer for inter-node disaggregated serving scenario. We have run some experiments to obtain some [preview benchmark results](vllm-benchmark-results-v0.2.md). More benchmark results will be released in due time.
|
||||
This is the latest version of mooncake-transfer-engine integration doc with the vLLM project based on [PR 10502](https://github.com/vllm-project/vllm/pull/10502) and [PR 10884](https://github.com/vllm-project/vllm/pull/10884) (vllm version: v0.6.4.post1/main) to accelerate KVCache transfer for inter-node disaggregated serving scenario. We have run some experiments to obtain some [preview benchmark results](../../../performance/vllm-benchmark-results-v0.2.md). More benchmark results will be released in due time.
|
||||
|
||||
**_Please note that this is still an experimental version and will be modified anytime based on feedback from the vLLM community._**
|
||||
- **Update(Apr 10, 2025)**: We are working on the vLLM v1 integration now. Stay tuned.
|
||||
- **Update(Sep 5, 2025)**: We have released the vLLM v1 integration with Mooncake Store and LMCache. Please refer to [vllmv1-lmcache-integration](vllmv1-lmcache-integration.md) for more details.
|
||||
|
||||
## Installation
|
||||
|
||||
### Prerequisite
|
||||
|
||||
```bash
|
||||
pip3 install mooncake-transfer-engine
|
||||
```
|
||||
|
|
@ -1,4 +1,4 @@
|
|||
# vLLM Disaggregated Serving with MooncakeStore
|
||||
# vLLM V0 Disaggregated Serving with MooncakeStore
|
||||
|
||||
## Overview
|
||||
This is the latest version of the MooncakeStore integration doc with the vLLM project based on [PR 10502](https://github.com/vllm-project/vllm/pull/10502) and [PR 12957](https://github.com/vllm-project/vllm/pull/12957) to support KVCache transfer for intra-node and inter-node disaggregated serving scenario. Benchmark results will be released soon.
|
||||
|
|
@ -13,9 +13,12 @@ Main changes from v0.x to v1:
|
|||
|
||||
**_Please note that this is still an experimental version and will be modified anytime based on feedback from the vLLM community._**
|
||||
- **Update(Apr 10, 2025)**: We are working on the vLLM v1 integration now. Stay tuned.
|
||||
- **Update(Sep 5, 2025)**: We have released the vLLM v1 integration with Mooncake Store and LMCache. Please refer to [vllmv1-lmcache-integration](vllmv1-lmcache-integration.md) for more details.
|
||||
|
||||
## Installation
|
||||
|
||||
### Prerequisite
|
||||
|
||||
```bash
|
||||
pip3 install mooncake-transfer-engine
|
||||
```
|
||||
|
|
@ -0,0 +1,117 @@
|
|||
# vLLM V1 Disaggregated Serving with Mooncake Store and LMCache
|
||||
|
||||
## Overview
|
||||
|
||||
The vLLM v1 version has been released with support for PD disaggregation. The detailed design document can be found [here](https://docs.google.com/document/d/1uPGdbEXksKXeN4Q9nUm9hzotqEjQhYmnpAhidLuAsjk). LMCache immediately implemented the corresponding connector to support storage, transmission, and loading of KVCache, enabling collaborative operation with PD nodes. Mooncake, as LMCache's backend storage engine, has undergone extensive optimizations in usability, performance, and stability. This document explains how to deploy a PD disaggregated serving demo using LMCache + Mooncake.
|
||||
|
||||
## Deployment
|
||||
|
||||
1. First, you need to prepare two GPU-equipped machines, which we will refer to as Machine A and Machine B. Install [vLLM](https://docs.vllm.ai/en/latest/getting_started/quickstart.html), [Mooncake](https://kvcache-ai.github.io/Mooncake/getting_started/build.html) and [LMCache](https://docs.lmcache.ai/getting_started/installation.html) on both Machine A and Machine B. For specific installation instructions, please refer to the official documentation of each repository.
|
||||
|
||||
2. Start the Mooncake Master node on Machine A:
|
||||
```bash
|
||||
mooncake_master -port 50052 -max_threads 64 -metrics_port 9004
|
||||
```
|
||||
|
||||
3. Start the HTTP metadata server (default port 8080) on Machine A for transfer engine metadata:
|
||||
```bash
|
||||
mooncake_http_metadata_server
|
||||
```
|
||||
|
||||
4. Launch the Decoder instance on machine A
|
||||
- Modify the vllm/examples/others/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh file.
|
||||
```diff
|
||||
diff --git a/examples/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh b/examples/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh
|
||||
index 831ef0bb5..a2ff0744c 100644
|
||||
--- a/examples/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh
|
||||
+++ b/examples/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh
|
||||
elif [[ $1 == "decoder" ]]; then
|
||||
# Decoder listens on port 8200
|
||||
- decode_config_file=$SCRIPT_DIR/configs/lmcache-decoder-config.yaml
|
||||
+ decode_config_file=$SCRIPT_DIR/configs/mooncake-decoder-config.yaml
|
||||
|
||||
UCX_TLS=cuda_ipc,cuda_copy,tcp \
|
||||
LMCACHE_CONFIG_FILE=$decode_config_file \
|
||||
LMCACHE_USE_EXPERIMENTAL=True \
|
||||
VLLM_ENABLE_V1_MULTIPROCESSING=1 \
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
CUDA_VISIBLE_DEVICES=1 \
|
||||
```
|
||||
- Add the `mooncake-decoder-config.yaml` file
|
||||
```yaml
|
||||
chunk_size: 256
|
||||
remote_url: "mooncakestore://{IP of Machine A}:50052/"
|
||||
remote_serde: "naive"
|
||||
local_cpu: False
|
||||
max_local_cpu_size: 100
|
||||
|
||||
extra_config:
|
||||
local_hostname: "{IP of Machine A}"
|
||||
metadata_server: "http://{IP of Machine A}:8080/metadata"
|
||||
protocol: "rdma"
|
||||
device_name: "mlx5_0" # Multiple RDMA devices can be specified as comma-separated list
|
||||
master_server_address: "{IP of Machine A}:50052"
|
||||
global_segment_size: 32212254720 # 30GB
|
||||
local_buffer_size: 1073741824 # 1GB
|
||||
transfer_timeout: 1
|
||||
save_chunk_meta: False
|
||||
```
|
||||
|
||||
- Launch the Decoder instance using command
|
||||
```bash
|
||||
bash disagg_vllm_launcher.sh decoder Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4
|
||||
```
|
||||
|
||||
5. Launch the Prefiller instance on machine B
|
||||
- Modify the vllm/examples/others/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh file.
|
||||
```diff
|
||||
diff --git a/examples/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh b/examples/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh
|
||||
index 831ef0bb5..9e5a3f044 100644
|
||||
--- a/examples/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh
|
||||
+++ b/examples/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh
|
||||
@@ -18,12 +18,14 @@ fi
|
||||
|
||||
if [[ $1 == "prefiller" ]]; then
|
||||
# Prefiller listens on port 8100
|
||||
- prefill_config_file=$SCRIPT_DIR/configs/lmcache-prefiller-config.yaml
|
||||
+ prefill_config_file=$SCRIPT_DIR/configs/mooncake-prefiller-config.yaml
|
||||
```
|
||||
|
||||
- Add the `mooncake-prefiller-config.yaml` file
|
||||
```yaml
|
||||
chunk_size: 256
|
||||
remote_url: "mooncakestore://{IP of Machine A}:50052/"
|
||||
remote_serde: "naive"
|
||||
local_cpu: False
|
||||
max_local_cpu_size: 100
|
||||
|
||||
extra_config:
|
||||
local_hostname: "{IP of Machine B}"
|
||||
metadata_server: "http://{IP of Machine A}:8080/metadata"
|
||||
protocol: "rdma"
|
||||
device_name: "mlx5_0" # Multiple RDMA devices can be specified as comma-separated list
|
||||
master_server_address: "{IP of Machine A}:50052"
|
||||
global_segment_size: 32212254720 # 30GB
|
||||
local_buffer_size: 1073741824 # 1GB
|
||||
transfer_timeout: 1
|
||||
save_chunk_meta: False
|
||||
```
|
||||
|
||||
- Launch the Prefiller instance using command
|
||||
```bash
|
||||
bash disagg_vllm_launcher.sh prefiller Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4
|
||||
```
|
||||
|
||||
6. Launch the `disagg_proxy_server` using command
|
||||
|
||||
```bash
|
||||
python3 disagg_proxy_server.py --host localhost --port 9000 --prefiller-host IP_of_Machine_B --prefiller-port 8100 --decoder-host IP_of_Machine_A --decoder-port 8200
|
||||
```
|
||||
|
||||
7. Now we can send the requests to the `disagg_proxy_server` to test PD disaggregated serving.
|
||||
|
||||
## Additional Resources
|
||||
|
||||
* [Mooncake x LMCache: Unite to Pioneer KVCache-Centric LLM Serving System](../../../getting_started/examples/lmcache-integration.md)
|
||||
* [Using Mooncake in LMCache](https://docs.lmcache.ai/kv_cache/storage_backends/mooncake.html)
|
||||
* [Using LMCache in vLLM](https://github.com/vllm-project/vllm/tree/main/examples/others/lmcache)
|
||||
|
|
@ -0,0 +1,219 @@
|
|||
# Quick Start
|
||||
|
||||
This document describes how to quickly start using Mooncake Transfer Engine and Mooncake Store.
|
||||
|
||||
## Installation
|
||||
|
||||
Install the Mooncake Transfer Engine package from PyPI, which includes both Mooncake Transfer Engine and Mooncake Store Python bindings:
|
||||
|
||||
```bash
|
||||
pip install mooncake-transfer-engine==0.3.5
|
||||
```
|
||||
|
||||
📦 **Package Details**: [https://pypi.org/project/mooncake-transfer-engine/](https://pypi.org/project/mooncake-transfer-engine/)
|
||||
|
||||
## Transfer Engine Quick Start
|
||||
|
||||
### Start Transfer Engine Receiver (Server)
|
||||
|
||||
```python
|
||||
import numpy as np
|
||||
import zmq
|
||||
from mooncake.engine import TransferEngine
|
||||
|
||||
def main():
|
||||
# Initialize ZMQ context and socket
|
||||
context = zmq.Context()
|
||||
socket = context.socket(zmq.PUSH)
|
||||
socket.bind("tcp://*:5555") # Bind to port 5555 for buffer info
|
||||
|
||||
HOSTNAME = "localhost" # localhost for simple demo
|
||||
METADATA_SERVER = "P2PHANDSHAKE" # [ETCD_SERVER_URL, P2PHANDSHAKE, ...]
|
||||
PROTOCOL = "tcp" # [rdma, tcp, ...]
|
||||
DEVICE_NAME = "" # auto discovery if empty
|
||||
|
||||
# Initialize server engine
|
||||
server_engine = TransferEngine()
|
||||
server_engine.initialize(
|
||||
HOSTNAME,
|
||||
METADATA_SERVER,
|
||||
PROTOCOL,
|
||||
DEVICE_NAME
|
||||
)
|
||||
session_id = f"{HOSTNAME}:{server_engine.get_rpc_port()}"
|
||||
|
||||
# Allocate memory on server side (1MB buffer)
|
||||
server_buffer = np.zeros(1024 * 1024, dtype=np.uint8)
|
||||
server_ptr = server_buffer.ctypes.data
|
||||
server_len = server_buffer.nbytes
|
||||
|
||||
# Register memory with Mooncake
|
||||
ret_value = server_engine.register_memory(server_ptr, server_len)
|
||||
if ret_value != 0:
|
||||
print("Mooncake memory registration failed.")
|
||||
raise RuntimeError("Mooncake memory registration failed.")
|
||||
|
||||
print(f"Server initialized with session ID: {session_id}")
|
||||
print(f"Server buffer address: {server_ptr}, length: {server_len}")
|
||||
|
||||
# Send buffer info to client
|
||||
buffer_info = {
|
||||
"session_id": session_id,
|
||||
"ptr": server_ptr,
|
||||
"len": server_len
|
||||
}
|
||||
socket.send_json(buffer_info)
|
||||
print("Buffer information sent to client")
|
||||
|
||||
# Keep server running
|
||||
try:
|
||||
while True:
|
||||
input("Press Ctrl+C to exit...")
|
||||
except KeyboardInterrupt:
|
||||
print("\nShutting down server...")
|
||||
finally:
|
||||
# Cleanup
|
||||
ret_value = server_engine.unregister_memory(server_ptr)
|
||||
if ret_value != 0:
|
||||
print("Mooncake memory deregistration failed.")
|
||||
raise RuntimeError("Mooncake memory deregistration failed.")
|
||||
|
||||
socket.close()
|
||||
context.term()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
```
|
||||
|
||||
### Start Transfer Engine Sender (Client)
|
||||
|
||||
```python
|
||||
|
||||
import numpy as np
|
||||
import zmq
|
||||
from mooncake.engine import TransferEngine
|
||||
|
||||
def main():
|
||||
# Initialize ZMQ context and socket
|
||||
context = zmq.Context()
|
||||
socket = context.socket(zmq.PULL)
|
||||
socket.connect(f"tcp://localhost:5555")
|
||||
|
||||
# Wait for buffer info from server
|
||||
print("Waiting for server buffer information...")
|
||||
buffer_info = socket.recv_json()
|
||||
server_session_id = buffer_info["session_id"]
|
||||
server_ptr = buffer_info["ptr"]
|
||||
server_len = buffer_info["len"]
|
||||
print(f"Received server info - Session ID: {server_session_id}")
|
||||
print(f"Server buffer address: {server_ptr}, length: {server_len}")
|
||||
|
||||
# Initialize client engine
|
||||
HOSTNAME = "localhost" # localhost for simple demo
|
||||
METADATA_SERVER = "P2PHANDSHAKE" # [ETCD_SERVER_URL, P2PHANDSHAKE, ...]
|
||||
PROTOCOL = "tcp" # [rdma, tcp, ...]
|
||||
DEVICE_NAME = "" # auto discovery if empty
|
||||
|
||||
client_engine = TransferEngine()
|
||||
client_engine.initialize(
|
||||
HOSTNAME,
|
||||
METADATA_SERVER,
|
||||
PROTOCOL,
|
||||
DEVICE_NAME
|
||||
)
|
||||
session_id = f"{HOSTNAME}:{client_engine.get_rpc_port()}"
|
||||
|
||||
# Allocate and initialize client buffer (1MB)
|
||||
client_buffer = np.ones(1024 * 1024, dtype=np.uint8) # Fill with ones
|
||||
client_ptr = client_buffer.ctypes.data
|
||||
client_len = client_buffer.nbytes
|
||||
|
||||
# Register memory with Mooncake
|
||||
ret_value = client_engine.register_memory(client_ptr, client_len)
|
||||
if ret_value != 0:
|
||||
print("Mooncake memory registration failed.")
|
||||
raise RuntimeError("Mooncake memory registration failed.")
|
||||
|
||||
print(f"Client initialized with session ID: {session_id}")
|
||||
|
||||
# Transfer data from client to server
|
||||
print("Transferring data to server...")
|
||||
for _ in range(10):
|
||||
ret = client_engine.transfer_sync_write(
|
||||
server_session_id,
|
||||
client_ptr,
|
||||
server_ptr,
|
||||
min(client_len, server_len) # Transfer minimum of both lengths
|
||||
)
|
||||
|
||||
if ret >= 0:
|
||||
print("Transfer successful!")
|
||||
else:
|
||||
print("Transfer failed!")
|
||||
|
||||
# Cleanup
|
||||
ret_value = client_engine.unregister_memory(client_ptr)
|
||||
if ret_value != 0:
|
||||
print("Mooncake memory deregistration failed.")
|
||||
raise RuntimeError("Mooncake memory deregistration failed.")
|
||||
|
||||
socket.close()
|
||||
context.term()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
```
|
||||
|
||||
### More Examples and Documentation
|
||||
|
||||
Please refer to the [Transfer Engine Python API](../python-api-reference/transfer-engine.md) and [Transfer Engine](../design/transfer-engine.md) for more examples and documentation.
|
||||
|
||||
## Mooncake Store Quick Start
|
||||
|
||||
### Start Mooncake Store
|
||||
|
||||
Before using the store, start the master and metadata services, which are included in the `mooncake-transfer-engine` package:
|
||||
|
||||
```bash
|
||||
# Start master service (default port 50051)
|
||||
mooncake_master
|
||||
```
|
||||
Start the HTTP metadata server (default port 8080) for transfer engine metadata:
|
||||
```bash
|
||||
# Start HTTP metadata server (default port 8080)
|
||||
mooncake_http_metadata_server
|
||||
```
|
||||
|
||||
### Hello World Example
|
||||
|
||||
```python
|
||||
from mooncake.store import MooncakeDistributedStore
|
||||
|
||||
# 1. Create store instance
|
||||
store = MooncakeDistributedStore()
|
||||
|
||||
# 2. Setup with all required parameters
|
||||
store.setup(
|
||||
"localhost:12345", # Your node's address
|
||||
"http://localhost:8080/metadata", # HTTP metadata server
|
||||
512*1024*1024, # 512MB segment size
|
||||
128*1024*1024, # 128MB local buffer
|
||||
"tcp", # Use TCP (RDMA for high performance)
|
||||
"", # Empty for TCP, specify device for RDMA
|
||||
"localhost:50051" # Master service
|
||||
)
|
||||
|
||||
# 3. Store data
|
||||
store.put("hello_key", b"Hello, Mooncake Store!")
|
||||
|
||||
# 4. Retrieve data
|
||||
data = store.get("hello_key")
|
||||
print(data.decode()) # Output: Hello, Mooncake Store!
|
||||
|
||||
# 5. Clean up
|
||||
store.close()
|
||||
```
|
||||
|
||||
### More Examples and Documentation
|
||||
|
||||
Please refer to the [Mooncake Store Python API](../python-api-reference/mooncake-store.md) and [Mooncake Store](../design/mooncake-store-preview.md) for more examples and documentation.
|
||||
|
|
@ -50,11 +50,11 @@ This repository also hosts its technical report and the open sourced traces.
|
|||
:maxdepth: 2
|
||||
|
||||
getting_started/build
|
||||
getting_started/quick-start
|
||||
getting_started/examples/lmcache-integration
|
||||
getting_started/examples/lmdeploy-integration-v0.9
|
||||
getting_started/examples/sglang-integration-v1
|
||||
getting_started/examples/vllm-integration-v0.2
|
||||
getting_started/examples/vllm-integration-v1
|
||||
getting_started/examples/vllm-integration/index
|
||||
:::
|
||||
|
||||
% Making the most out of Mooncake
|
||||
|
|
@ -72,10 +72,11 @@ performance/allocator_benchmark_result.md
|
|||
% API Documentation
|
||||
|
||||
:::{toctree}
|
||||
:caption: Mooncake Store API
|
||||
:caption: Python API Reference
|
||||
:maxdepth: 1
|
||||
|
||||
mooncake-store-api/python-binding
|
||||
python-api-reference/mooncake-store
|
||||
python-api-reference/transfer-engine
|
||||
:::
|
||||
|
||||
% Explanation of Mooncake internals
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
## PD Disaggregation Performance
|
||||
# PD Disaggregation Performance
|
||||
|
||||
We evaluated the current implementation on two A10 servers. By comparing the performance of a 1P1D configuration with that of two regular (non-disaggregated) instances, we observed that P/D disaggregation achieves approximately 30% lower ITL while maintaining comparable total throughput. This aligns with findings from the Mooncake paper, which highlighted that P/D disaggregation is effective in reducing TBT/ITL under similar throughput conditions—or conversely, in enabling higher throughput under stricter ITL/TBT SLOs.
|
||||
Moreover, we anticipate even greater benefits in larger-scale clusters where both the number of prefill and decode nodes (x and y in xPyD configurations) increase, offering enhanced scheduling flexibility and resource efficiency.
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ The following services need to be started separately:
|
|||
## Quick Start
|
||||
|
||||
### Start Required Services
|
||||
|
||||
Before using the store, start the master and metadata services:
|
||||
|
||||
```bash
|
||||
|
|
@ -0,0 +1,588 @@
|
|||
# Transfer Engine Python API
|
||||
|
||||
## Overview
|
||||
|
||||
The Transfer Engine Python API provides a high-level interface for efficient data transfer between distributed systems using RDMA (Remote Direct Memory Access) and other transport protocols. It enables fast, low-latency data movement between nodes in a cluster.
|
||||
|
||||
For interfaces beyond the Python API (C/C++, Golang, Rust), see [Transfer Engine](../design/transfer-engine.md#using-transfer-engine-to-your-projects).
|
||||
|
||||
## Installation
|
||||
|
||||
Install the Mooncake Transfer Engine package from PyPI, which includes both Mooncake Transfer Engine and Mooncake Store Python bindings:
|
||||
|
||||
```bash
|
||||
pip install mooncake-transfer-engine
|
||||
```
|
||||
|
||||
📦 **Package Details**: [https://pypi.org/project/mooncake-transfer-engine/](https://pypi.org/project/mooncake-transfer-engine/)
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Start Transfer Engine Receiver (Server)
|
||||
|
||||
```python
|
||||
import numpy as np
|
||||
import zmq
|
||||
from mooncake.engine import TransferEngine
|
||||
|
||||
def main():
|
||||
# Initialize ZMQ context and socket
|
||||
context = zmq.Context()
|
||||
socket = context.socket(zmq.PUSH)
|
||||
socket.bind("tcp://*:5555") # Bind to port 5555 for buffer info
|
||||
|
||||
HOSTNAME = "localhost" # localhost for simple demo
|
||||
METADATA_SERVER = "P2PHANDSHAKE" # [ETCD_SERVER_URL, P2PHANDSHAKE, ...]
|
||||
PROTOCOL = "tcp" # [rdma, tcp, ...]
|
||||
DEVICE_NAME = "" # auto discovery if empty
|
||||
|
||||
# Initialize server engine
|
||||
server_engine = TransferEngine()
|
||||
server_engine.initialize(
|
||||
HOSTNAME,
|
||||
METADATA_SERVER,
|
||||
PROTOCOL,
|
||||
DEVICE_NAME
|
||||
)
|
||||
session_id = f"{HOSTNAME}:{server_engine.get_rpc_port()}"
|
||||
|
||||
# Allocate memory on server side (1MB buffer)
|
||||
server_buffer = np.zeros(1024 * 1024, dtype=np.uint8)
|
||||
server_ptr = server_buffer.ctypes.data
|
||||
server_len = server_buffer.nbytes
|
||||
|
||||
# Register memory with Mooncake
|
||||
ret_value = server_engine.register_memory(server_ptr, server_len)
|
||||
if ret_value != 0:
|
||||
print("Mooncake memory registration failed.")
|
||||
raise RuntimeError("Mooncake memory registration failed.")
|
||||
|
||||
print(f"Server initialized with session ID: {session_id}")
|
||||
print(f"Server buffer address: {server_ptr}, length: {server_len}")
|
||||
|
||||
# Send buffer info to client
|
||||
buffer_info = {
|
||||
"session_id": session_id,
|
||||
"ptr": server_ptr,
|
||||
"len": server_len
|
||||
}
|
||||
socket.send_json(buffer_info)
|
||||
print("Buffer information sent to client")
|
||||
|
||||
# Keep server running
|
||||
try:
|
||||
while True:
|
||||
input("Press Ctrl+C to exit...")
|
||||
except KeyboardInterrupt:
|
||||
print("\nShutting down server...")
|
||||
finally:
|
||||
# Cleanup
|
||||
ret_value = server_engine.unregister_memory(server_ptr)
|
||||
if ret_value != 0:
|
||||
print("Mooncake memory deregistration failed.")
|
||||
raise RuntimeError("Mooncake memory deregistration failed.")
|
||||
|
||||
socket.close()
|
||||
context.term()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
```
|
||||
|
||||
### Start Transfer Engine Sender (Client)
|
||||
|
||||
```python
|
||||
|
||||
import numpy as np
|
||||
import zmq
|
||||
from mooncake.engine import TransferEngine
|
||||
|
||||
def main():
|
||||
# Initialize ZMQ context and socket
|
||||
context = zmq.Context()
|
||||
socket = context.socket(zmq.PULL)
|
||||
socket.connect(f"tcp://localhost:5555")
|
||||
|
||||
# Wait for buffer info from server
|
||||
print("Waiting for server buffer information...")
|
||||
buffer_info = socket.recv_json()
|
||||
server_session_id = buffer_info["session_id"]
|
||||
server_ptr = buffer_info["ptr"]
|
||||
server_len = buffer_info["len"]
|
||||
print(f"Received server info - Session ID: {server_session_id}")
|
||||
print(f"Server buffer address: {server_ptr}, length: {server_len}")
|
||||
|
||||
# Initialize client engine
|
||||
HOSTNAME = "localhost" # localhost for simple demo
|
||||
METADATA_SERVER = "P2PHANDSHAKE" # [ETCD_SERVER_URL, P2PHANDSHAKE, ...]
|
||||
PROTOCOL = "tcp" # [rdma, tcp, ...]
|
||||
DEVICE_NAME = "" # auto discovery if empty
|
||||
|
||||
client_engine = TransferEngine()
|
||||
client_engine.initialize(
|
||||
HOSTNAME,
|
||||
METADATA_SERVER,
|
||||
PROTOCOL,
|
||||
DEVICE_NAME
|
||||
)
|
||||
session_id = f"{HOSTNAME}:{client_engine.get_rpc_port()}"
|
||||
|
||||
# Allocate and initialize client buffer (1MB)
|
||||
client_buffer = np.ones(1024 * 1024, dtype=np.uint8) # Fill with ones
|
||||
client_ptr = client_buffer.ctypes.data
|
||||
client_len = client_buffer.nbytes
|
||||
|
||||
# Register memory with Mooncake
|
||||
ret_value = client_engine.register_memory(client_ptr, client_len)
|
||||
if ret_value != 0:
|
||||
print("Mooncake memory registration failed.")
|
||||
raise RuntimeError("Mooncake memory registration failed.")
|
||||
|
||||
print(f"Client initialized with session ID: {session_id}")
|
||||
|
||||
# Transfer data from client to server
|
||||
print("Transferring data to server...")
|
||||
for _ in range(10):
|
||||
ret = client_engine.transfer_sync_write(
|
||||
server_session_id,
|
||||
client_ptr,
|
||||
server_ptr,
|
||||
min(client_len, server_len) # Transfer minimum of both lengths
|
||||
)
|
||||
|
||||
if ret >= 0:
|
||||
print("Transfer successful!")
|
||||
else:
|
||||
print("Transfer failed!")
|
||||
|
||||
# Cleanup
|
||||
ret_value = client_engine.unregister_memory(client_ptr)
|
||||
if ret_value != 0:
|
||||
print("Mooncake memory deregistration failed.")
|
||||
raise RuntimeError("Mooncake memory deregistration failed.")
|
||||
|
||||
socket.close()
|
||||
context.term()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
```
|
||||
|
||||
## API Reference
|
||||
|
||||
### Class: TransferEngine
|
||||
|
||||
The main class that provides all transfer engine functionality.
|
||||
|
||||
### Constructor
|
||||
|
||||
```python
|
||||
TransferEngine()
|
||||
```
|
||||
|
||||
Creates a new TransferEngine instance with default settings.
|
||||
|
||||
### Initialization Methods
|
||||
|
||||
#### initialize()
|
||||
|
||||
```python
|
||||
initialize(local_hostname, metadata_server, protocol, device_name)
|
||||
```
|
||||
|
||||
Initializes the transfer engine with basic configuration.
|
||||
|
||||
**Parameters:**
|
||||
- `local_hostname` (str): The hostname and port of the local server (e.g., "127.0.0.1:12345")
|
||||
- `metadata_server` (str): The metadata server connection string (e.g., "127.0.0.1:2379" or "etcd://127.0.0.1:2379")
|
||||
- `protocol` (str): The transport protocol to use ("rdma", "tcp", etc.)
|
||||
- `device_name` (str): Comma-separated list of device names to filter, or empty string for all devices
|
||||
|
||||
**Returns:**
|
||||
- `int`: 0 on success, negative value on failure
|
||||
|
||||
#### initialize_ext()
|
||||
|
||||
```python
|
||||
initialize_ext(local_hostname, metadata_server, protocol, device_name, metadata_type)
|
||||
```
|
||||
|
||||
Initializes the transfer engine with extended configuration including metadata type specification.
|
||||
|
||||
**Parameters:**
|
||||
- `local_hostname` (str): The hostname and port of the local server
|
||||
- `metadata_server` (str): The metadata server connection string
|
||||
- `protocol` (str): The transport protocol to use
|
||||
- `device_name` (str): Comma-separated list of device names to filter
|
||||
- `metadata_type` (str): The type of metadata server ("etcd", "p2p", etc.)
|
||||
|
||||
**Returns:**
|
||||
- `int`: 0 on success, negative value on failure
|
||||
|
||||
### Network Information
|
||||
|
||||
#### get_rpc_port()
|
||||
|
||||
```python
|
||||
get_rpc_port()
|
||||
```
|
||||
|
||||
Gets the RPC port that the transfer engine is listening on.
|
||||
|
||||
**Returns:**
|
||||
- `int`: The RPC port number
|
||||
|
||||
### Buffer Management
|
||||
|
||||
#### allocate_managed_buffer()
|
||||
|
||||
```python
|
||||
allocate_managed_buffer(length)
|
||||
```
|
||||
|
||||
Allocates a managed buffer of the specified size using a buddy allocation system for efficient memory management.
|
||||
|
||||
**Parameters:**
|
||||
- `length` (int): The size of the buffer to allocate in bytes
|
||||
|
||||
**Returns:**
|
||||
- `int`: The memory address of the allocated buffer as an integer, or 0 on failure
|
||||
|
||||
#### free_managed_buffer()
|
||||
|
||||
```python
|
||||
free_managed_buffer(buffer_addr, length)
|
||||
```
|
||||
|
||||
Frees a previously allocated managed buffer.
|
||||
|
||||
**Parameters:**
|
||||
- `buffer_addr` (int): The memory address of the buffer to free
|
||||
- `length` (int): The size of the buffer in bytes
|
||||
|
||||
**Returns:**
|
||||
- `int`: 0 on success, negative value on failure
|
||||
|
||||
#### get_first_buffer_address()
|
||||
|
||||
```python
|
||||
get_first_buffer_address(segment_name)
|
||||
```
|
||||
|
||||
Gets the address of the first buffer in a specified segment.
|
||||
|
||||
**Parameters:**
|
||||
- `segment_name` (str): The name of the segment
|
||||
|
||||
**Returns:**
|
||||
- `int`: The memory address of the first buffer in the segment
|
||||
|
||||
### Data Transfer Operations
|
||||
|
||||
#### transfer_sync_write()
|
||||
|
||||
```python
|
||||
transfer_sync_write(target_hostname, buffer, peer_buffer_address, length)
|
||||
```
|
||||
|
||||
Performs a synchronous write operation to transfer data from local buffer to remote buffer.
|
||||
|
||||
**Parameters:**
|
||||
- `target_hostname` (str): The hostname of the target server
|
||||
- `buffer` (int): The local buffer address
|
||||
- `peer_buffer_address` (int): The remote buffer address
|
||||
- `length` (int): The number of bytes to transfer
|
||||
|
||||
**Returns:**
|
||||
- `int`: 0 on success, negative value on failure
|
||||
|
||||
#### transfer_sync_read()
|
||||
|
||||
```python
|
||||
transfer_sync_read(target_hostname, buffer, peer_buffer_address, length)
|
||||
```
|
||||
|
||||
Performs a synchronous read operation to transfer data from remote buffer to local buffer.
|
||||
|
||||
**Parameters:**
|
||||
- `target_hostname` (str): The hostname of the target server
|
||||
- `buffer` (int): The local buffer address
|
||||
- `peer_buffer_address` (int): The remote buffer address
|
||||
- `length` (int): The number of bytes to transfer
|
||||
|
||||
**Returns:**
|
||||
- `int`: 0 on success, negative value on failure
|
||||
|
||||
#### transfer_sync()
|
||||
|
||||
```python
|
||||
transfer_sync(target_hostname, buffer, peer_buffer_address, length, opcode)
|
||||
```
|
||||
|
||||
Performs a synchronous transfer operation with specified opcode.
|
||||
|
||||
**Parameters:**
|
||||
- `target_hostname` (str): The hostname of the target server
|
||||
- `buffer` (int): The local buffer address
|
||||
- `peer_buffer_address` (int): The remote buffer address
|
||||
- `length` (int): The number of bytes to transfer
|
||||
- `opcode` (TransferOpcode): The transfer operation type (READ or WRITE)
|
||||
|
||||
**Returns:**
|
||||
- `int`: 0 on success, negative value on failure
|
||||
|
||||
#### transfer_submit_write()
|
||||
|
||||
```python
|
||||
transfer_submit_write(target_hostname, buffer, peer_buffer_address, length)
|
||||
```
|
||||
|
||||
Submits an asynchronous write operation and returns immediately.
|
||||
|
||||
**Parameters:**
|
||||
- `target_hostname` (str): The hostname of the target server
|
||||
- `buffer` (int): The local buffer address
|
||||
- `peer_buffer_address` (int): The remote buffer address
|
||||
- `length` (int): The number of bytes to transfer
|
||||
|
||||
**Returns:**
|
||||
- `int`: Batch ID for tracking the operation, or negative value on failure
|
||||
|
||||
#### transfer_check_status()
|
||||
|
||||
```python
|
||||
transfer_check_status(batch_id)
|
||||
```
|
||||
|
||||
Checks the status of an asynchronous transfer operation.
|
||||
|
||||
**Parameters:**
|
||||
- `batch_id` (int): The batch ID returned from transfer_submit_write()
|
||||
|
||||
**Returns:**
|
||||
- `int`:
|
||||
- 1: Transfer completed successfully
|
||||
- 0: Transfer still in progress
|
||||
- -1: Transfer failed
|
||||
- -2: Transfer timed out
|
||||
|
||||
### Buffer I/O Operations
|
||||
|
||||
#### write_bytes_to_buffer()
|
||||
|
||||
```python
|
||||
write_bytes_to_buffer(dest_address, src_ptr, length)
|
||||
```
|
||||
|
||||
Writes bytes from a Python bytes object to a buffer at the specified address.
|
||||
|
||||
**Parameters:**
|
||||
- `dest_address` (int): The destination buffer address
|
||||
- `src_ptr` (bytes): The source bytes to write
|
||||
- `length` (int): The number of bytes to write
|
||||
|
||||
**Returns:**
|
||||
- `int`: 0 on success, negative value on failure
|
||||
|
||||
#### read_bytes_from_buffer()
|
||||
|
||||
```python
|
||||
read_bytes_from_buffer(source_address, length)
|
||||
```
|
||||
|
||||
Reads bytes from a buffer at the specified address and returns them as a Python bytes object.
|
||||
|
||||
**Parameters:**
|
||||
- `source_address` (int): The source buffer address
|
||||
- `length` (int): The number of bytes to read
|
||||
|
||||
**Returns:**
|
||||
- `bytes`: The bytes read from the buffer
|
||||
|
||||
### Memory Registration (Experimental)
|
||||
|
||||
#### register_memory()
|
||||
|
||||
```python
|
||||
register_memory(buffer_addr, capacity)
|
||||
```
|
||||
|
||||
Registers a memory region for RDMA access (experimental feature).
|
||||
|
||||
**Parameters:**
|
||||
- `buffer_addr` (int): The memory address to register
|
||||
- `capacity` (int): The size of the memory region in bytes
|
||||
|
||||
**Returns:**
|
||||
- `int`: 0 on success, negative value on failure
|
||||
|
||||
#### unregister_memory()
|
||||
|
||||
```python
|
||||
unregister_memory(buffer_addr)
|
||||
```
|
||||
|
||||
Unregisters a previously registered memory region.
|
||||
|
||||
**Parameters:**
|
||||
- `buffer_addr` (int): The memory address to unregister
|
||||
|
||||
**Returns:**
|
||||
- `int`: 0 on success, negative value on failure
|
||||
|
||||
### Enums
|
||||
|
||||
#### TransferOpcode
|
||||
|
||||
```python
|
||||
TransferOpcode.READ # Read operation
|
||||
TransferOpcode.WRITE # Write operation
|
||||
```
|
||||
|
||||
## Environment Variables
|
||||
|
||||
The Transfer Engine respects the following environment variables:
|
||||
|
||||
- `MC_TRANSFER_TIMEOUT`: Sets the transfer timeout in seconds (default: 30)
|
||||
- `MC_METADATA_SERVER`: Default metadata server address
|
||||
- `MC_LEGACY_RPC_PORT_BINDING`: Enables legacy RPC port binding behavior
|
||||
- `MC_TCP_BIND_ADDRESS`: Specifies the TCP bind address
|
||||
- `MC_CUSTOM_TOPO_JSON`: Path to custom topology JSON file
|
||||
- `MC_TE_METRIC`: Enables metrics reporting (set to "1", "true", "yes", or "on")
|
||||
- `MC_TE_METRIC_INTERVAL_SECONDS`: Sets metrics reporting interval in seconds
|
||||
|
||||
## Usage Examples
|
||||
|
||||
### Basic Setup and Data Transfer
|
||||
|
||||
```python
|
||||
from mooncake.engine import TransferEngine
|
||||
import os
|
||||
|
||||
# Create transfer engine instance
|
||||
engine = TransferEngine()
|
||||
|
||||
# Initialize with basic configuration
|
||||
engine.initialize(
|
||||
"127.0.0.1:12345", # local hostname
|
||||
"127.0.0.1:2379", # metadata server
|
||||
"rdma", # transport protocol
|
||||
"" # device name
|
||||
)
|
||||
|
||||
# Allocate and initialize client buffer (1MB)
|
||||
client_buffer = np.ones(1024 * 1024, dtype=np.uint8) # Fill with ones
|
||||
buffer_data = client_buffer.ctypes.data
|
||||
buffer_data_len = client_buffer.nbytes
|
||||
|
||||
# Prepare data
|
||||
data = b"Hello, Transfer Engine!"
|
||||
data_len = len(data)
|
||||
|
||||
engine.register_memory(buffer_data, buffer_data_len)
|
||||
|
||||
# Get Remote Addr from ZMQ or upper-layer inference framework
|
||||
remote_addr = ??
|
||||
|
||||
# Transfer data to remote node
|
||||
ret = engine.transfer_sync_write(
|
||||
"127.0.0.1:12346", # target hostname
|
||||
data, # buffer
|
||||
remote_addr, # peer buffer address
|
||||
data_len # length
|
||||
)
|
||||
|
||||
if ret == 0:
|
||||
print("Data transfer completed successfully")
|
||||
else:
|
||||
print(f"Data transfer failed with code {ret}")
|
||||
|
||||
engine.unregister_memory(data)
|
||||
```
|
||||
|
||||
### Asynchronous Transfer
|
||||
|
||||
```python
|
||||
# Submit asynchronous write
|
||||
batch_id = engine.transfer_submit_write(
|
||||
"127.0.0.1:12346", # target hostname
|
||||
local_addr, # buffer
|
||||
remote_addr, # peer buffer address
|
||||
data_len # length
|
||||
)
|
||||
|
||||
if batch_id < 0:
|
||||
print(f"Failed to submit transfer with code {batch_id}")
|
||||
else:
|
||||
# Poll for completion
|
||||
while True:
|
||||
status = engine.transfer_check_status(batch_id)
|
||||
if status == 1:
|
||||
print("Transfer completed successfully")
|
||||
break
|
||||
elif status == -1:
|
||||
print("Transfer failed")
|
||||
break
|
||||
elif status == -2:
|
||||
print("Transfer timed out")
|
||||
break
|
||||
# Transfer still in progress, continue polling
|
||||
import time
|
||||
time.sleep(0.001) # Small delay to avoid busy waiting
|
||||
```
|
||||
|
||||
### Managed Buffer Allocation
|
||||
|
||||
```python
|
||||
# Allocate managed buffer
|
||||
buffer_size = 1024 * 1024 # 1MB
|
||||
buffer_addr = engine.allocate_managed_buffer(buffer_size)
|
||||
|
||||
if buffer_addr == 0:
|
||||
print("Failed to allocate buffer")
|
||||
else:
|
||||
# Use the buffer
|
||||
test_data = b"Test data for managed buffer"
|
||||
engine.write_bytes_to_buffer(buffer_addr, test_data, len(test_data))
|
||||
|
||||
# Read back
|
||||
read_data = engine.read_bytes_from_buffer(buffer_addr, len(test_data))
|
||||
print(f"Read data: {read_data}")
|
||||
|
||||
# Free the buffer when done
|
||||
engine.free_managed_buffer(buffer_addr, buffer_size)
|
||||
```
|
||||
|
||||
## Error Handling
|
||||
|
||||
All methods return integer status codes:
|
||||
- `0`: Success
|
||||
- Negative values: Error codes indicating various failure conditions
|
||||
|
||||
Common error scenarios:
|
||||
- Network connectivity issues
|
||||
- Invalid buffer addresses
|
||||
- Memory allocation failures
|
||||
- Transfer timeouts
|
||||
- Metadata server connection problems
|
||||
|
||||
## Performance Considerations
|
||||
|
||||
1. **Buffer Reuse**: Reuse allocated buffers when possible to avoid frequent allocation/deallocation overhead
|
||||
2. **Batch Operations**: Use `transfer_submit_write()` and `transfer_check_status()` for better throughput when multiple transfers are needed
|
||||
3. **Memory Alignment**: Ensure buffers are properly aligned for optimal RDMA performance
|
||||
4. **Timeout Configuration**: Adjust `MC_TRANSFER_TIMEOUT` based on your network characteristics and data sizes
|
||||
|
||||
## Thread Safety
|
||||
|
||||
The Transfer Engine Python API is thread-safe for most operations. However, it's recommended to:
|
||||
- Use separate TransferEngine instances for different threads when possible
|
||||
- Avoid concurrent modifications to the same buffer addresses
|
||||
- Use proper synchronization when sharing buffer addresses between threads
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
1. **Initialization Failures**: Check metadata server connectivity and network configuration
|
||||
2. **Transfer Failures**: Verify target hostname is correct and network connectivity is established
|
||||
3. **Memory Issues**: Ensure sufficient system memory and proper buffer alignment
|
||||
4. **Performance Issues**: Check RDMA device configuration and network topology
|
||||
Loading…
Reference in New Issue