diff --git a/docs_roll/docs/User Guides/Advanced Features/remote_batch_transfer.md b/docs_roll/docs/User Guides/Advanced Features/remote_batch_transfer.md index 6a6d32cf3..19c5c433d 100644 --- a/docs_roll/docs/User Guides/Advanced Features/remote_batch_transfer.md +++ b/docs_roll/docs/User Guides/Advanced Features/remote_batch_transfer.md @@ -20,6 +20,7 @@ In RL training pipelines (especially VLM and Agentic scenarios), `DataProto` bat - **Transfer Backend**: A pluggable storage backend responsible for `put`, `get`, and `delete` operations. Currently supported backends: - `None` (Dummy): No remote storage; data stays local (default). - `TransferQueue`: Uses the [TransferQueue](https://github.com/kvcache-ai/TransferQueue) library for high-performance distributed key-value transfer. + - `Mooncake`: Uses Mooncake as an optional structured `DataProto` transfer backend for large tensor, non-tensor, and multimodal rollout payloads. ### How It Works @@ -41,10 +42,29 @@ transfer_backend: num_data_storage_units: 16 ``` +Mooncake can be enabled as an optional backend: + +```yaml +transfer_backend: + backend_name: Mooncake + backend_config: + client_scope: node + local_hostname: 192.168.0.1 + metadata_server: P2PHANDSHAKE + global_segment_size: 8589934592 + local_buffer_size: 8589934592 + protocol: rdma + rdma_devices: erdma_0 + master_server_addr: 192.168.0.1:50051 +``` + - `backend_name`: The name of the transfer backend to use. - `null` (default): Disables remote transfer; all data stays local. This is the default behavior when `transfer_backend` is not configured. - `TransferQueue`: Uses the TransferQueue library for high-performance data transfer. -- `backend_config`: Backend-specific configuration dictionary. For TransferQueue, this corresponds to the TransferQueue initialization config. + - `Mooncake`: Uses Mooncake structured `DataProto` transfer. This backend supports tensor batch fields, `non_tensor_batch`, and `meta_info`, and is intended for large multi-node rollout payloads. +- `backend_config`: Backend-specific configuration dictionary. + - For TransferQueue, this corresponds to the TransferQueue initialization config. + - For Mooncake, this config can be passed explicitly as shown above, or loaded from Mooncake environment configuration. - `backend.SimpleStorage.num_data_storage_units`: The number of storage units to shard data across. Can be configured based on the number of CPU cores and cluster nodes. `msgpack` serialization has a maximum 4 GB limit per object, so larger data transfers require more storage units to shard `non_tensor_batch` into smaller pieces. ### Agentic Pipeline Optimization @@ -65,6 +85,7 @@ Manually calling `to_remote` inside environment workers is incompatible with fil | Backend | Status | Notes | |---------|--------|-------| | TransferQueue | End-to-end tested | Production-ready. Tested across RLVR, VLM, and Agentic pipelines. | +| Mooncake | Experimental | Optional structured `DataProto` backend for tensor, non-tensor, metadata, and multimodal rollout payloads. | | RayMemoryStore | Illustration only | Not tested. Provided as a reference implementation for the `ColumnRemoteBatch` pattern. | ### TODO diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Advanced Features/remote_batch_transfer.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Advanced Features/remote_batch_transfer.md index eede46713..8d90c260b 100644 --- a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Advanced Features/remote_batch_transfer.md +++ b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Advanced Features/remote_batch_transfer.md @@ -20,6 +20,7 @@ ROLL 框架支持 **RemoteBatch**，一种惰性数据传输机制，将数据 - **传输后端（Transfer Backend）**：负责 `put`、`get` 和 `delete` 操作的可插拔存储后端。目前支持的后端： - `None`（Dummy）：无远程存储，数据保留在本地（默认）。 - `TransferQueue`：使用 [TransferQueue](https://github.com/kvcache-ai/TransferQueue) 库进行高性能分布式键值传输。 + - `Mooncake`：作为可选的结构化 `DataProto` 传输后端，用于大规模 tensor、non-tensor 和多模态 rollout payload。 ### 工作原理 @@ -41,10 +42,29 @@ transfer_backend: num_data_storage_units: 16 ``` +Mooncake 可以作为可选 backend 启用： + +```yaml +transfer_backend: + backend_name: Mooncake + backend_config: + client_scope: node + local_hostname: 192.168.0.1 + metadata_server: P2PHANDSHAKE + global_segment_size: 8589934592 + local_buffer_size: 8589934592 + protocol: rdma + rdma_devices: erdma_0 + master_server_addr: 192.168.0.1:50051 +``` + - `backend_name`：要使用的传输后端名称。 - `null`（默认）：禁用远程传输，所有数据保留在本地。未配置 `transfer_backend` 时的默认行为。 - `TransferQueue`：使用 TransferQueue 库进行高性能数据传输。 -- `backend_config`：后端特定的配置字典。对于 TransferQueue，对应 TransferQueue 的初始化配置。 + - `Mooncake`：使用 Mooncake 结构化 `DataProto` 传输，支持 tensor batch、`non_tensor_batch` 和 `meta_info`，适用于大规模多节点 rollout payload。 +- `backend_config`：后端特定的配置字典。 + - 对于 TransferQueue，对应 TransferQueue 的初始化配置。 + - 对于 Mooncake，可以像上例一样显式配置，也可以从 Mooncake 环境配置中加载。 - `backend.SimpleStorage.num_data_storage_units`：数据分片的存储单元数量。可以根据 CPU 核数和集群节点数进行配置。`msgpack` 序列化单个对象有最大 4GB 的限制，因此传输大数据时需要更多的 storage unit 来将 `non_tensor_batch` 分片成更小的块。 ### Agentic Pipeline 优化 @@ -65,6 +85,7 @@ output_queue.put(batch) | 后端 | 状态 | 说明 | |------|------|------| | TransferQueue | 端到端已测试 | 生产可用。已在 RLVR、VLM 和 Agentic Pipeline 中测试通过。 | +| Mooncake | 实验性 | 可选的结构化 `DataProto` backend，支持 tensor、non-tensor、metadata 和多模态 rollout payload。 | | RayMemoryStore | 仅作示例 | 未经测试。仅作为 `ColumnRemoteBatch` 模式的参考实现提供。 | ### TODO diff --git a/roll/distributed/scheduler/transfer_backend.py b/roll/distributed/scheduler/transfer_backend.py index ca2b137ea..03f9d6754 100644 --- a/roll/distributed/scheduler/transfer_backend.py +++ b/roll/distributed/scheduler/transfer_backend.py @@ -7,6 +7,7 @@ import torch import numpy as np import sys +from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy if sys.version_info < (3, 13): import transfer_queue as tq @@ -22,6 +23,10 @@ logger = get_logger() +MOONCAKE_CLIENT_SCOPE_NODE = "node" +MOONCAKE_CLIENT_SCOPE_PROCESS = "process" +MOONCAKE_NODE_ACTOR_NAME_PREFIX = "MooncakeNodeTransfer" + # Global reference to keep SharedStorage actor alive _shared_storage = None @@ -34,6 +39,28 @@ def _check_transfer_queue_available(): ) +def _check_mooncake_available(): + try: + import mooncake.store # noqa: F401 + import mooncake.structured_object_store # noqa: F401 + except ImportError as exc: + raise ImportError("Mooncake transfer backend requires the mooncake Python package.") from exc + + +def _mooncake_client_scope(config: dict[str, Any] | None) -> str: + return (config or {}).get("client_scope", MOONCAKE_CLIENT_SCOPE_NODE) + + +def _prepare_mooncake_backend_config(config: TransferBackendArguments) -> None: + if config.backend_name != "Mooncake": + return + if config.backend_config is None: + config.backend_config = {} + if _mooncake_client_scope(config.backend_config) != MOONCAKE_CLIENT_SCOPE_NODE: + return + config.backend_config.setdefault("node_actor_session_id", uuid.uuid4().hex) + + def init_transfer_backend(config: TransferBackendArguments | None): global _shared_storage @@ -43,6 +70,7 @@ def init_transfer_backend(config: TransferBackendArguments | None): if config is None: config = TransferBackendArguments() + _prepare_mooncake_backend_config(config) ray.get(_shared_storage.put.remote(key="transfer_backend_config", data=config)) backend_name = config.backend_name @@ -53,6 +81,9 @@ def init_transfer_backend(config: TransferBackendArguments | None): _check_transfer_queue_available() init_transfer_queue_server(backend_config) logger.info(f"Initialized TransferQueue transfer backend: {config}") + elif backend_name == "Mooncake": + _check_mooncake_available() + logger.info(f"Initialized Mooncake transfer backend: {config}") else: raise ValueError(f"Unsupported transfer backend: {backend_name}") @@ -81,6 +112,11 @@ def init_client(): _client = DummyClient() elif config.backend_name == "TransferQueue": _client = TransferQueueClient() + elif config.backend_name == "Mooncake": + if _mooncake_client_scope(config.backend_config) == MOONCAKE_CLIENT_SCOPE_NODE: + _client = MooncakeNodeClientProxy(config.backend_config) + else: + _client = MooncakeClient(config.backend_config) else: raise ValueError(f"Unsupported transfer backend: {config.backend_name}") logger.info(f"Initialized transfer client: {_client.__class__.__name__}") @@ -119,7 +155,13 @@ def create_tensordict(fields: dict[str, torch.Tensor | np.ndarray]) -> TensorDic class DummyClient: - def put(self, partition, row_ids: list[str], fields: dict[str, torch.Tensor | np.ndarray], batch_size: int): + def put( + self, + partition, + row_ids: list[str], + fields: dict[str, torch.Tensor | np.ndarray], + batch_size: int, + ): return None def get(self, partition, keys: list[str], fields: list[Any]): @@ -154,7 +196,13 @@ def __init__(self): get_if_exists=True, ).remote() - def put(self, partition, row_ids: list[str], fields: dict[str, torch.Tensor | np.ndarray], batch_size: int): + def put( + self, + partition, + row_ids: list[str], + fields: dict[str, torch.Tensor | np.ndarray], + batch_size: int, + ): # TODO move RayMemoryStoreClient to another file from roll.distributed.scheduler.remote_protocol import ColumnRemoteBatch @@ -182,6 +230,219 @@ def delete(self, partition, keys: list[str], fields: list[Any]): pass +@ray.remote +class MooncakeNodeTransferActor: + def __init__(self, config: dict[str, Any] | None = None): + self.client = MooncakeClient(config) + + def put( + self, + partition, + row_ids: list[str], + fields: dict[str, torch.Tensor | np.ndarray], + batch_size: int, + ): + return self.client.put(partition, row_ids, fields, batch_size) + + def get(self, partition, keys: list[str], fields: list[Any]): + return self.client.get(partition, keys, fields) + + def delete(self, partition, keys: list[str], fields: list[Any]): + return self.client.delete(partition, keys, fields) + + +class MooncakeNodeClientProxy: + def __init__(self, config: dict[str, Any] | None = None): + self.config = dict(config or {}) + self.actor = self._get_node_actor() + + def put( + self, + partition, + row_ids: list[str], + fields: dict[str, torch.Tensor | np.ndarray], + batch_size: int, + ): + return ray.get(self.actor.put.remote(partition, row_ids, fields, batch_size)) + + def get(self, partition, keys: list[str], fields: list[Any]): + return ray.get(self.actor.get.remote(partition, keys, fields)) + + def delete(self, partition, keys: list[str], fields: list[Any]): + return ray.get(self.actor.delete.remote(partition, keys, fields)) + + def _get_node_actor(self): + node_id = ray.get_runtime_context().get_node_id() + session_id = self.config.get("node_actor_session_id", "default") + actor_name = f"{MOONCAKE_NODE_ACTOR_NAME_PREFIX}-{session_id}-{node_id[:16]}" + return MooncakeNodeTransferActor.options( + name=actor_name, + get_if_exists=True, + namespace=RAY_NAMESPACE, + max_concurrency=1, + num_cpus=0, + scheduling_strategy=NodeAffinitySchedulingStrategy(node_id=node_id, soft=False), + ).remote(self.config) + + +class MooncakeClient: + def __init__(self, config: dict[str, Any] | None = None): + _check_mooncake_available() + from mooncake.store import MooncakeDistributedStore + from mooncake.structured_object_store import BundleTransferPolicy, MooncakeBundleTransfer + + config = config or {} + store = MooncakeDistributedStore() + ret = self._setup_store(store, config) + if ret != 0: + raise RuntimeError(f"Mooncake store setup failed, return code={ret}") + + self.backend = MooncakeBundleTransfer(store, key_prefix=config.get("key_prefix", "roll")) + policy_config = config.get("transfer_policy") + self.transfer_policy = BundleTransferPolicy(**policy_config) if policy_config else None + + def put( + self, + partition, + row_ids: list[str], + fields: dict[str, torch.Tensor | np.ndarray], + batch_size: int, + ): + from roll.distributed.scheduler.protocol import DataProto + from roll.distributed.scheduler.remote_protocol import ColumnRemoteBatch + + batch_fields, non_tensor_fields = self._split_fields(fields) + meta_info = {"roll_row_ids": row_ids} + if batch_fields: + data = DataProto.from_dict(tensors=batch_fields, non_tensors=non_tensor_fields, meta_info=meta_info) + else: + data = DataProto( + batch=TensorDict({}, batch_size=[batch_size]), + non_tensor_batch=non_tensor_fields, + meta_info=meta_info, + ) + assert len(data) == batch_size + + ref = self.backend.put_dataproto(data, partition=partition, policy=self.transfer_policy) + field_refs = { + key: {"ref": ref, "kind": "batch" if key in batch_fields else "non_tensor"} + for key in fields.keys() + } + return ColumnRemoteBatch( + partition=partition, + device=data.batch.device if batch_fields else None, + fields=field_refs, + is_nested=False, + cache=create_tensordict(fields), + batch_size=batch_size, + ) + + def get(self, partition, keys: list[str], fields: list[Any]): + if not fields: + return TensorDict({}, batch_size=[0]) + + grouped: dict[str, dict[str, Any]] = {} + for key, field in zip(keys, fields): + ref = field["ref"] + group = grouped.setdefault(self._ref_key(ref), {"ref": ref, "batch": [], "non_tensor": []}) + group[field["kind"]].append(key) + + data_dict = {} + for group in grouped.values(): + data = self.backend.get_dataproto( + group["ref"], + batch_fields=group["batch"], + non_tensor_fields=group["non_tensor"], + ) + data_dict.update(data.get("batch", {})) + data_dict.update(data.get("non_tensor_batch", {})) + return create_tensordict({key: data_dict[key] for key in keys}) + + def delete(self, partition, keys: list[str], fields: list[Any]): + deleted = set() + for field in fields: + ref = field["ref"] + ref_key = self._ref_key(ref) + if ref_key in deleted: + continue + self.backend.cleanup_dataproto(ref) + deleted.add(ref_key) + + def _setup_store(self, store, config: dict[str, Any]) -> int: + setup_args = config.get("setup_args") + setup_kwargs = config.get("setup_kwargs") + if setup_args is not None: + return store.setup(*setup_args) + if setup_kwargs is not None: + return store.setup(**setup_kwargs) + if config.get("master_server_addr") or config.get("master_server_address"): + return store.setup( + self._require_config(config, "local_hostname"), + self._require_config(config, "metadata_server"), + config.get("global_segment_size", 3355443200), + config.get("local_buffer_size", 1073741824), + config.get("protocol", "tcp"), + config.get("rdma_devices") or config.get("device_name", ""), + config.get("master_server_addr") or config.get("master_server_address"), + ) + return self._setup_store_from_env(store, config) + + def _setup_store_from_env(self, store, config: dict[str, Any]) -> int: + from mooncake.mooncake_config import MooncakeConfig + + mooncake_config = MooncakeConfig.load_from_env() + return store.setup( + config.get("local_hostname", mooncake_config.local_hostname), + config.get("metadata_server", mooncake_config.metadata_server), + config.get("global_segment_size", mooncake_config.global_segment_size), + config.get("local_buffer_size", mooncake_config.local_buffer_size), + config.get("protocol", mooncake_config.protocol), + config.get("rdma_devices") or config.get("device_name", mooncake_config.device_name or ""), + config.get("master_server_addr") + or config.get("master_server_address", mooncake_config.master_server_address), + ) + + @staticmethod + def _split_fields( + fields: dict[str, torch.Tensor | np.ndarray], + ) -> tuple[dict[str, torch.Tensor], dict[str, np.ndarray]]: + tensors = {key: value for key, value in fields.items() if isinstance(value, torch.Tensor)} + non_tensors = {key: value for key, value in fields.items() if isinstance(value, np.ndarray)} + MooncakeClient._validate_fields(tensors, non_tensors, fields) + return tensors, non_tensors + + @staticmethod + def _validate_fields( + batch_fields: dict[str, Any], + non_tensor_fields: dict[str, Any], + all_fields: dict[str, Any] | None = None, + ) -> None: + unsupported = {} + for key, value in batch_fields.items(): + if not isinstance(value, torch.Tensor): + unsupported[key] = type(value) + for key, value in non_tensor_fields.items(): + if not isinstance(value, np.ndarray): + unsupported[key] = type(value) + if all_fields is not None: + known = set(batch_fields) | set(non_tensor_fields) + unsupported.update( + {key: type(value) for key, value in all_fields.items() if key not in known} + ) + if unsupported: + raise TypeError(f"Unsupported Mooncake fields: {unsupported}") + + @staticmethod + def _ref_key(ref) -> str: + return getattr(ref, "object_id", repr(ref)) + + @staticmethod + def _require_config(config: dict[str, Any], key: str) -> Any: + value = config.get(key) + if value is None or value == "": + raise ValueError(f"Mooncake backend_config requires {key!r} when master_server_addr is set") + return value + def init_transfer_queue_server(config): # Must create enough storage units or may encounter: # EncodeError: Can't encode Ext objects with data longer than 2**32 - 1. @@ -195,7 +456,13 @@ def __init__(self): _check_transfer_queue_available() tq.init() - def put(self, partition, row_ids: list[str], fields: dict[str, torch.Tensor | np.ndarray], batch_size: int): + def put( + self, + partition, + row_ids: list[str], + fields: dict[str, torch.Tensor | np.ndarray], + batch_size: int, + ): # TODO move TransferQueueClient to another file from roll.distributed.scheduler.remote_protocol import RowRemoteBatch diff --git a/tests/distributed/scheduler/test_mooncake_transfer_backend.py b/tests/distributed/scheduler/test_mooncake_transfer_backend.py new file mode 100644 index 000000000..dda46b55d --- /dev/null +++ b/tests/distributed/scheduler/test_mooncake_transfer_backend.py @@ -0,0 +1,145 @@ +import os +import shutil +import socket +import subprocess +import time + +import numpy as np +import pytest +import torch + +from roll.configs.base_config import TransferBackendArguments +from roll.distributed.scheduler.transfer_backend import ( + MOONCAKE_CLIENT_SCOPE_NODE, + MOONCAKE_CLIENT_SCOPE_PROCESS, + MooncakeClient, + _mooncake_client_scope, + _prepare_mooncake_backend_config, +) + + +def _wait_for_port(host: str, port: int, timeout: float = 10.0) -> None: + deadline = time.time() + timeout + while time.time() < deadline: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.settimeout(0.2) + if sock.connect_ex((host, port)) == 0: + return + time.sleep(0.1) + raise TimeoutError(f"Timed out waiting for {host}:{port}") + + +def _mooncake_master_endpoint() -> tuple[str, int]: + master = os.environ.get("MOONCAKE_MASTER", "") + if not master: + pytest.skip("Set MOONCAKE_MASTER to run the Mooncake RDMA backend test") + host, port = master.rsplit(":", 1) + return host, int(port) + + +@pytest.fixture(scope="module") +def mooncake_master(): + host, port = _mooncake_master_endpoint() + if shutil.which("mooncake_master") is None: + pytest.skip("mooncake_master is not available in PATH") + + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.settimeout(0.2) + if sock.connect_ex((host, port)) == 0: + yield f"{host}:{port}" + return + + process = subprocess.Popen( + [ + "mooncake_master", + f"--rpc_address={host}", + f"--rpc_port={port}", + "--logtostderr=true", + ], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + try: + _wait_for_port(host, port) + yield f"{host}:{port}" + finally: + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=5) + + +def test_mooncake_client_scope_defaults_to_node(): + config = TransferBackendArguments(backend_name="Mooncake", backend_config={}) + + _prepare_mooncake_backend_config(config) + + assert _mooncake_client_scope(config.backend_config) == MOONCAKE_CLIENT_SCOPE_NODE + assert config.backend_config["node_actor_session_id"] + + +def test_mooncake_process_scope_keeps_config_small(): + config = TransferBackendArguments( + backend_name="Mooncake", + backend_config={"client_scope": MOONCAKE_CLIENT_SCOPE_PROCESS}, + ) + + _prepare_mooncake_backend_config(config) + + assert "node_actor_session_id" not in config.backend_config + + +def test_mooncake_client_splits_roll_fields(): + fields = { + "tokens": torch.tensor([[1], [2]]), + "prompt": np.array(["a", "b"], dtype=object), + } + + tensors, non_tensors = MooncakeClient._split_fields(fields) + + assert tensors == {"tokens": fields["tokens"]} + assert non_tensors == {"prompt": fields["prompt"]} + + +def test_mooncake_client_rejects_unsupported_fields(): + with pytest.raises(TypeError, match="Unsupported Mooncake fields"): + MooncakeClient._split_fields({"bad": ["a", "b"]}) + + +def test_mooncake_client_real_rdma_round_trip(mooncake_master): + protocol = os.environ.get("MOONCAKE_PROTOCOL", "") + if protocol != "rdma": + pytest.skip("Set MOONCAKE_PROTOCOL=rdma to run the Mooncake RDMA backend test") + + local_hostname = os.environ.get("MOONCAKE_LOCAL_HOSTNAME", "") + rdma_devices = os.environ.get("MOONCAKE_DEVICE_NAME", "") + if not local_hostname or not rdma_devices: + pytest.skip("Set MOONCAKE_LOCAL_HOSTNAME and MOONCAKE_DEVICE_NAME for RDMA testing") + + client = MooncakeClient( + { + "local_hostname": local_hostname, + "metadata_server": os.environ.get("MOONCAKE_METADATA_SERVER", "P2PHANDSHAKE"), + "global_segment_size": int(os.environ.get("MOONCAKE_GLOBAL_SEGMENT_SIZE", 1024 * 1024 * 1024)), + "local_buffer_size": int(os.environ.get("MOONCAKE_LOCAL_BUFFER_SIZE", 1024 * 1024 * 1024)), + "protocol": protocol, + "rdma_devices": rdma_devices, + "master_server_addr": mooncake_master, + "transfer_policy": {"copy_mode": "auto"}, + } + ) + fields = { + "tokens": torch.tensor([[1, 2], [3, 4]]), + "prompt": np.array(["a", "b"], dtype=object), + } + + remote = client.put("rollout", ["0", "1"], fields, batch_size=2) + materialized = client.get("rollout", ["tokens", "prompt"], [remote.fields["tokens"], remote.fields["prompt"]]) + + assert torch.equal(materialized["tokens"], fields["tokens"]) + assert list(materialized["prompt"]) == ["a", "b"] + + client.delete("rollout", list(remote.fields.keys()), list(remote.fields.values()))