From 3621710eea4e11d80445f878d9cdf24002ebcb04 Mon Sep 17 00:00:00 2001 From: Cruz Zhao Date: Mon, 29 Jun 2026 19:14:00 +0800 Subject: [PATCH 1/2] Add Mooncake transfer backend support Wire Mooncake into the existing DataProto transfer backend path with a node-scoped client by default to reuse per-node store setup and registered buffer pools, while keeping process-local clients configurable. Co-Authored-By: Claude Opus 4.6 --- .../distributed/scheduler/transfer_backend.py | 188 ++++++++++++++++++ .../test_mooncake_transfer_backend.py | 111 +++++++++++ 2 files changed, 299 insertions(+) create mode 100644 tests/distributed/scheduler/test_mooncake_transfer_backend.py diff --git a/roll/distributed/scheduler/transfer_backend.py b/roll/distributed/scheduler/transfer_backend.py index ca2b137ea..6363a8e84 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.dataproto_transfer # noqa: F401 + import mooncake.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__}") @@ -182,6 +218,158 @@ 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.dataproto_transfer import MooncakeDataProtoTransferBackend + from mooncake.store import MooncakeDistributedStore + from roll.distributed.scheduler.protocol import DataProto + + config = config or {} + store = MooncakeDistributedStore() + setup_args = config.get("setup_args") + setup_kwargs = config.get("setup_kwargs") + if setup_args is not None: + ret = store.setup(*setup_args) + elif setup_kwargs is not None: + ret = store.setup(**setup_kwargs) + else: + ret = self._setup_store_from_env(store) + if ret != 0: + raise RuntimeError(f"Mooncake store setup failed, return code={ret}") + + self.backend = MooncakeDataProtoTransferBackend( + store, + key_prefix=config.get("key_prefix", "roll"), + data_cls=DataProto, + ) + self.shard_policy = self._create_shard_policy(config.get("shard_policy")) + + 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 + + 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)} + if len(tensors) + len(non_tensors) != len(fields): + unsupported = { + key: type(value) for key, value in fields.items() if key not in tensors and key not in non_tensors + } + raise TypeError(f"Unsupported Mooncake fields: {unsupported}") + + if tensors: + data = DataProto.from_dict(tensors=tensors, non_tensors=non_tensors, meta_info={}) + else: + data = DataProto(batch=None, non_tensor_batch=non_tensors, meta_info={}) + assert len(data) == batch_size + ref = self.backend.put_dataproto(data, partition=partition, shard_policy=self.shard_policy) + field_refs = { + key: {"ref": ref, "kind": "batch" if key in tensors else "non_tensor"} + for key in fields.keys() + } + return ColumnRemoteBatch( + partition=partition, + device=data.batch.device if tensors 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]) + ref = fields[0]["ref"] + if any(field["ref"].object_id != ref.object_id for field in fields): + raise ValueError("Mooncake backend cannot materialize fields from different refs in one get") + batch_fields = [key for key, field in zip(keys, fields) if field["kind"] == "batch"] + non_tensor_fields = [key for key, field in zip(keys, fields) if field["kind"] == "non_tensor"] + data = self.backend.materialize_dataproto( + ref, + batch_fields=batch_fields, + non_tensor_fields=non_tensor_fields, + include_meta_info=False, + ) + data_dict = {} + if data._batch is not None: + data_dict.update(data._batch.to_dict()) + data_dict.update(data._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"] + if ref.object_id in deleted: + continue + self.backend.remove_dataproto(ref) + deleted.add(ref.object_id) + + def _setup_store_from_env(self, store): + from mooncake.mooncake_config import MooncakeConfig + + config = MooncakeConfig.load_from_env() + return store.setup( + config.local_hostname, + config.metadata_server, + config.global_segment_size, + config.local_buffer_size, + config.protocol, + config.device_name or "", + config.master_server_address, + ) + + def _create_shard_policy(self, config: dict[str, Any] | None): + if not config: + return None + from mooncake.dataproto_transfer import DataProtoShardPolicy + + return DataProtoShardPolicy(**config) + + 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. 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..f1564103f --- /dev/null +++ b/tests/distributed/scheduler/test_mooncake_transfer_backend.py @@ -0,0 +1,111 @@ +import sys +import types +from types import SimpleNamespace + +import numpy as np +import torch + +from roll.configs.base_config import TransferBackendArguments +from roll.distributed.scheduler.protocol import DataProto +from roll.distributed.scheduler.transfer_backend import ( + MOONCAKE_CLIENT_SCOPE_NODE, + MOONCAKE_CLIENT_SCOPE_PROCESS, + MooncakeClient, + _mooncake_client_scope, + _prepare_mooncake_backend_config, +) + + +class FakeMooncakeStore: + def setup(self, *args, **kwargs): + return 0 + + +class FakeMooncakeBackend: + refs = {} + removed = [] + + def __init__(self, store, key_prefix="dataproto", data_cls=None): + self.data_cls = data_cls + + def put_dataproto(self, data, partition="default", shard_policy=None): + object_id = f"{partition}/ref" + ref = SimpleNamespace(object_id=object_id, row_count=len(data), manifest_key="manifest", manifest={}) + self.refs[object_id] = data + return ref + + def materialize_dataproto( + self, + ref, + batch_fields=None, + non_tensor_fields=None, + include_meta_info=True, + ): + data = self.refs[ref.object_id] + tensors = {} + if data._batch is not None: + selected = set(batch_fields or []) + tensors = {key: value for key, value in data._batch.to_dict().items() if key in selected} + non_tensors = { + key: value for key, value in data._non_tensor_batch.items() if key in set(non_tensor_fields or []) + } + return DataProto.from_dict(tensors=tensors, non_tensors=non_tensors) if tensors else DataProto( + batch=None, non_tensor_batch=non_tensors + ) + + def remove_dataproto(self, ref): + self.removed.append(ref.object_id) + + +class FakeShardPolicy: + def __init__(self, **kwargs): + self.kwargs = kwargs + + +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_round_trip(monkeypatch): + mooncake = types.ModuleType("mooncake") + store_mod = types.ModuleType("mooncake.store") + transfer_mod = types.ModuleType("mooncake.dataproto_transfer") + store_mod.MooncakeDistributedStore = FakeMooncakeStore + transfer_mod.MooncakeDataProtoTransferBackend = FakeMooncakeBackend + transfer_mod.DataProtoShardPolicy = FakeShardPolicy + monkeypatch.setitem(sys.modules, "mooncake", mooncake) + monkeypatch.setitem(sys.modules, "mooncake.store", store_mod) + monkeypatch.setitem(sys.modules, "mooncake.dataproto_transfer", transfer_mod) + + FakeMooncakeBackend.refs = {} + FakeMooncakeBackend.removed = [] + client = MooncakeClient({"setup_args": [], "shard_policy": {"enabled": True}}) + 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())) + assert FakeMooncakeBackend.removed == ["rollout/ref"] From 79116c17977b6e3b540e27b7b868f40142708a67 Mon Sep 17 00:00:00 2001 From: Cruz Zhao Date: Tue, 30 Jun 2026 04:17:07 +0800 Subject: [PATCH 2/2] Complete Mooncake DataProto transfer backend support Use Mooncake structured object transfer as the optional DataProto backend while keeping ROLL's existing transfer_backend.put API and RemoteBatch semantics. Add real RDMA-backed Mooncake tests without fake Mooncake modules. Co-Authored-By: Claude Opus 4.6 --- .../remote_batch_transfer.md | 23 +- .../remote_batch_transfer.md | 23 +- .../distributed/scheduler/transfer_backend.py | 217 ++++++++++++------ .../test_mooncake_transfer_backend.py | 160 ++++++++----- 4 files changed, 289 insertions(+), 134 deletions(-) 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 6363a8e84..03f9d6754 100644 --- a/roll/distributed/scheduler/transfer_backend.py +++ b/roll/distributed/scheduler/transfer_backend.py @@ -41,8 +41,8 @@ def _check_transfer_queue_available(): def _check_mooncake_available(): try: - import mooncake.dataproto_transfer # noqa: F401 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 @@ -155,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]): @@ -190,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 @@ -223,7 +235,13 @@ 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): + 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]): @@ -238,7 +256,13 @@ 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): + 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]): @@ -264,55 +288,49 @@ def _get_node_actor(self): class MooncakeClient: def __init__(self, config: dict[str, Any] | None = None): _check_mooncake_available() - from mooncake.dataproto_transfer import MooncakeDataProtoTransferBackend from mooncake.store import MooncakeDistributedStore - from roll.distributed.scheduler.protocol import DataProto + from mooncake.structured_object_store import BundleTransferPolicy, MooncakeBundleTransfer config = config or {} store = MooncakeDistributedStore() - setup_args = config.get("setup_args") - setup_kwargs = config.get("setup_kwargs") - if setup_args is not None: - ret = store.setup(*setup_args) - elif setup_kwargs is not None: - ret = store.setup(**setup_kwargs) - else: - ret = self._setup_store_from_env(store) + ret = self._setup_store(store, config) if ret != 0: raise RuntimeError(f"Mooncake store setup failed, return code={ret}") - self.backend = MooncakeDataProtoTransferBackend( - store, - key_prefix=config.get("key_prefix", "roll"), - data_cls=DataProto, - ) - self.shard_policy = self._create_shard_policy(config.get("shard_policy")) - - def put(self, partition, row_ids: list[str], fields: dict[str, torch.Tensor | np.ndarray], batch_size: int): + 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 - 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)} - if len(tensors) + len(non_tensors) != len(fields): - unsupported = { - key: type(value) for key, value in fields.items() if key not in tensors and key not in non_tensors - } - raise TypeError(f"Unsupported Mooncake fields: {unsupported}") - - if tensors: - data = DataProto.from_dict(tensors=tensors, non_tensors=non_tensors, meta_info={}) + 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=None, non_tensor_batch=non_tensors, meta_info={}) + 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, shard_policy=self.shard_policy) + + ref = self.backend.put_dataproto(data, partition=partition, policy=self.transfer_policy) field_refs = { - key: {"ref": ref, "kind": "batch" if key in tensors else "non_tensor"} + 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 tensors else None, + device=data.batch.device if batch_fields else None, fields=field_refs, is_nested=False, cache=create_tensordict(fields), @@ -322,53 +340,108 @@ def put(self, partition, row_ids: list[str], fields: dict[str, torch.Tensor | np def get(self, partition, keys: list[str], fields: list[Any]): if not fields: return TensorDict({}, batch_size=[0]) - ref = fields[0]["ref"] - if any(field["ref"].object_id != ref.object_id for field in fields): - raise ValueError("Mooncake backend cannot materialize fields from different refs in one get") - batch_fields = [key for key, field in zip(keys, fields) if field["kind"] == "batch"] - non_tensor_fields = [key for key, field in zip(keys, fields) if field["kind"] == "non_tensor"] - data = self.backend.materialize_dataproto( - ref, - batch_fields=batch_fields, - non_tensor_fields=non_tensor_fields, - include_meta_info=False, - ) + + 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 = {} - if data._batch is not None: - data_dict.update(data._batch.to_dict()) - data_dict.update(data._non_tensor_batch) + 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"] - if ref.object_id in deleted: + ref_key = self._ref_key(ref) + if ref_key in deleted: continue - self.backend.remove_dataproto(ref) - deleted.add(ref.object_id) + self.backend.cleanup_dataproto(ref) + deleted.add(ref_key) - def _setup_store_from_env(self, store): + 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 - config = MooncakeConfig.load_from_env() + mooncake_config = MooncakeConfig.load_from_env() return store.setup( - config.local_hostname, - config.metadata_server, - config.global_segment_size, - config.local_buffer_size, - config.protocol, - config.device_name or "", - config.master_server_address, + 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), ) - def _create_shard_policy(self, config: dict[str, Any] | None): - if not config: - return None - from mooncake.dataproto_transfer import DataProtoShardPolicy + @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}") - return DataProtoShardPolicy(**config) + @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: @@ -383,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 index f1564103f..dda46b55d 100644 --- a/tests/distributed/scheduler/test_mooncake_transfer_backend.py +++ b/tests/distributed/scheduler/test_mooncake_transfer_backend.py @@ -1,12 +1,14 @@ -import sys -import types -from types import SimpleNamespace +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.protocol import DataProto from roll.distributed.scheduler.transfer_backend import ( MOONCAKE_CLIENT_SCOPE_NODE, MOONCAKE_CLIENT_SCOPE_PROCESS, @@ -16,50 +18,58 @@ ) -class FakeMooncakeStore: - def setup(self, *args, **kwargs): - return 0 - - -class FakeMooncakeBackend: - refs = {} - removed = [] - - def __init__(self, store, key_prefix="dataproto", data_cls=None): - self.data_cls = data_cls - - def put_dataproto(self, data, partition="default", shard_policy=None): - object_id = f"{partition}/ref" - ref = SimpleNamespace(object_id=object_id, row_count=len(data), manifest_key="manifest", manifest={}) - self.refs[object_id] = data - return ref - - def materialize_dataproto( - self, - ref, - batch_fields=None, - non_tensor_fields=None, - include_meta_info=True, - ): - data = self.refs[ref.object_id] - tensors = {} - if data._batch is not None: - selected = set(batch_fields or []) - tensors = {key: value for key, value in data._batch.to_dict().items() if key in selected} - non_tensors = { - key: value for key, value in data._non_tensor_batch.items() if key in set(non_tensor_fields or []) - } - return DataProto.from_dict(tensors=tensors, non_tensors=non_tensors) if tensors else DataProto( - batch=None, non_tensor_batch=non_tensors - ) - - def remove_dataproto(self, ref): - self.removed.append(ref.object_id) - - -class FakeShardPolicy: - def __init__(self, **kwargs): - self.kwargs = kwargs +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(): @@ -82,20 +92,45 @@ def test_mooncake_process_scope_keeps_config_small(): assert "node_actor_session_id" not in config.backend_config -def test_mooncake_client_round_trip(monkeypatch): - mooncake = types.ModuleType("mooncake") - store_mod = types.ModuleType("mooncake.store") - transfer_mod = types.ModuleType("mooncake.dataproto_transfer") - store_mod.MooncakeDistributedStore = FakeMooncakeStore - transfer_mod.MooncakeDataProtoTransferBackend = FakeMooncakeBackend - transfer_mod.DataProtoShardPolicy = FakeShardPolicy - monkeypatch.setitem(sys.modules, "mooncake", mooncake) - monkeypatch.setitem(sys.modules, "mooncake.store", store_mod) - monkeypatch.setitem(sys.modules, "mooncake.dataproto_transfer", transfer_mod) - - FakeMooncakeBackend.refs = {} - FakeMooncakeBackend.removed = [] - client = MooncakeClient({"setup_args": [], "shard_policy": {"enabled": True}}) +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), @@ -108,4 +143,3 @@ def test_mooncake_client_round_trip(monkeypatch): assert list(materialized["prompt"]) == ["a", "b"] client.delete("rollout", list(remote.fields.keys()), list(remote.fields.values())) - assert FakeMooncakeBackend.removed == ["rollout/ref"]