Python API 与工具概览

Python API 与工具概览

上一篇 | 目录 | 下一篇

概述#

Mooncake 通过 mooncake-transfer-engine Python 包（即 mooncake-wheel）为上层推理框架提供 Python API。该包封装了三大核心 C++ 模块的 pybind11 绑定，并附带若干 CLI 工具和辅助服务。本文将深入解析 Python 包的结构、各层绑定接口、关键 Python 类以及运维工具链。

1. Python 包结构#

mooncake-wheel/mooncake/ 目录的完整文件布局如下：

mooncake-wheel/mooncake/
├── __init__.py                        # 包入口 (向后兼容导入)
├── engine.so                          # [编译产物] TransferEngine pybind11 模块
├── store.so                           # [编译产物] MooncakeDistributedStore pybind11 模块
├── ep_<torch_version>.so              # [编译产物] Expert Parallelism pybind11 模块
├── ep.py                              # EP 模块版本适配加载器
├── mooncake_config.py                 # MooncakeConfig 配置解析
├── mooncake_connector_v1.py           # vLLM v1 KVConnector 实现
├── mooncake_ep_buffer.py              # EP Buffer 管理
├── mooncake_store_service.py          # Store REST API 服务
├── http_metadata_server.py            # HTTP 元数据服务器
├── vllm_v1_proxy_server.py            # vLLM v1 Proxy Server
├── transfer_engine_topology_dump.py   # 拓扑信息转储工具
├── cli.py                             # mooncake_master CLI 入口
├── cli_client.py                      # mooncake_client CLI 入口
├── cli_bench.py                       # transfer_engine_bench CLI 入口
└── README.md                          # 使用说明

1.1 Python 绑定架构#

2. Pybind11 绑定架构#

Mooncake 的三大核心 C++ 模块通过 pybind11 暴露给 Python，源码位于 mooncake-integration/ 目录下。

2.1 TransferEngine 绑定 (engine.so)#

源文件: mooncake-integration/transfer_engine/transfer_engine_py.cpp / transfer_engine_py.h

TransferEngine 是最底层的数据传输引擎，提供 RDMA/TCP 内存注册与远程读写能力。

核心类 - TransferEngine#

from mooncake.engine import TransferEngine

engine = TransferEngine()
engine.initialize(local_hostname, metadata_server, protocol, device_name)

主要方法#

辅助类型#

• TransferOpcode - 传输操作码枚举：Read, Write
• TransferNotify - 传输通知结构：name, msg

2.2 Store 绑定 (store.so)#

源文件: mooncake-integration/store/store_py.cpp

MooncakeDistributedStore 是分布式对象存储的 Python 接口，支持原始字节和 PyTorch Tensor 的存取。

核心类 - MooncakeDistributedStore#

from mooncake.store import MooncakeDistributedStore

store = MooncakeDistributedStore()
store.setup(
    local_hostname="192.168.0.1",
    metadata_server="192.168.0.1:2379",
    global_segment_size=3355443200,  # 3.125 GiB
    local_buffer_size=1073741824,    # 1 GiB
    protocol="rdma",
    rdma_devices="erdma_0",
    master_server_addr="192.168.0.1:50051"
)

主要方法#

基础 KV 操作

Tensor 操作

Tensor Parallelism 支持

零拷贝直接访问

辅助类型#

• ReplicateConfig - 副本配置：replica_num, with_soft_pin, preferred_segments, prefer_alloc_in_same_node
• ReplicaStatus - 副本状态枚举：UNDEFINED, INITIALIZED, PROCESSING, COMPLETE, REMOVED, FAILED
• BufferHandle - 缓冲区句柄：ptr(), size(), 支持 Python Buffer Protocol
• MooncakeHostMemAllocator - 主机内存分配器：alloc(size), free(ptr)

2.3 Expert Parallelism 绑定 (ep.so)#

源文件: mooncake-integration/ep/ep_py.cpp

EP 模块提供 MoE 模型的 Expert Parallelism 通信后端，作为 torch.distributed 的自定义 Backend 注册。

版本适配机制#

ep.py 根据当前 PyTorch 版本动态加载对应的编译产物：

# ep.py 核心逻辑
import importlib
torch_version = re.match(r"\d+(?:\.\d+)*", torch.__version__).group()
version_suffix = "_" + torch_version.replace(".", "_")
backend_module = importlib.import_module("mooncake.ep" + version_suffix)

例如 PyTorch 2.5.0 会加载 mooncake.ep_2_5_0 模块。

核心类与函数#

Buffer 类方法#

3. 关键 Python 类#

3.1 MooncakeConfig#

源文件: mooncake-wheel/mooncake/mooncake_config.py

配置类，支持从 JSON 文件和环境变量加载：

@dataclass
class MooncakeConfig:
    local_hostname: str            # 本机主机名/IP
    metadata_server: str           # 元数据服务地址
    global_segment_size: int       # 全局段大小 (默认 3.125 GiB)
    local_buffer_size: int         # 本地缓冲区大小 (默认 1 GiB)
    protocol: str                  # 传输协议 (rdma/tcp)
    device_name: Optional[str]     # RDMA 设备名
    master_server_address: str     # Master Server 地址

加载方式：

相关环境变量：

3.2 MooncakeStoreService#

源文件: mooncake-wheel/mooncake/mooncake_store_service.py

提供 REST API 封装的 Store 服务，基于 aiohttp：

service = MooncakeStoreService(config_path="mooncake.json")
await service.start_store_service()
await service.start_http_service(port=8080)

REST API 端点：

3.3 KVBootstrapServer#

源文件: mooncake-wheel/mooncake/http_metadata_server.py

轻量级 HTTP 元数据服务器，可替代 etcd 用于存储 TransferEngine 的连接元数据：

server = KVBootstrapServer(port=8080, host="0.0.0.0")
server.run()  # 后台线程启动

支持的操作：GET /metadata?key=xxx, PUT /metadata?key=xxx, DELETE /metadata?key=xxx

4. CLI 工具#

Mooncake 提供三个 CLI 工具，通过 setup.py 的 console_scripts 入口点注册。

4.1 mooncake_master#

入口: mooncake-wheel/mooncake/cli.py

Master 守护进程，负责 Store 服务的集群管理与协调。

mooncake_master --port 50051

内部实际调用打包在 wheel 中的预编译二进制文件 mooncake_master。

4.2 mooncake_client#

入口: mooncake-wheel/mooncake/cli_client.py

Store 客户端工具，用于与 Store 服务交互。

mooncake_client [options]

4.3 transfer_engine_bench (mooncake_bench)#

入口: mooncake-wheel/mooncake/cli_bench.py

TransferEngine 性能基准测试工具。

transfer_engine_bench [options]

4.4 工具链概览#

5. HTTP 元数据服务器#

5.1 概述#

http_metadata_server.py 实现了一个轻量级的 HTTP 元数据存储服务，用于替代 etcd/redis 作为 TransferEngine 的元数据后端。它特别适合快速部署和测试场景，无需安装和维护外部依赖。

5.2 架构设计#

class KVBootstrapServer:
    """
    基于 aiohttp 的 HTTP 元数据服务器
    - 线程安全的内存 KV 存储
    - 支持 GET/PUT/DELETE 操作
    - 后台线程运行，不阻塞主程序
    """

5.3 服务状态枚举#

class KVPoll(Enum):
    Failed = 0           # 服务失败
    Bootstrapping = 1    # 启动中
    WaitingForInput = 2  # 等待输入
    Transferring = 3     # 传输中
    Success = 4          # 运行正常

5.4 使用方式#

独立启动：

python -m mooncake.http_metadata_server --port 8080 --host 0.0.0.0

编程使用：

from mooncake.http_metadata_server import KVBootstrapServer

server = KVBootstrapServer(port=8080)
server.run()

# 检查服务状态
status = server.poll()  # KVPoll.Success

# 优雅关闭
server.close()

5.5 元数据后端选择指南#

参考资源#

• 包源码: mooncake-wheel/mooncake/
• TransferEngine 绑定: mooncake-integration/transfer_engine/transfer_engine_py.cpp
• Store 绑定: mooncake-integration/store/store_py.cpp
• EP 绑定: mooncake-integration/ep/ep_py.cpp
• 配置说明: mooncake-wheel/mooncake/mooncake_config.py

上一篇 | 目录 | 下一篇