[Feat] Support Mooncake Store backend

transfer_queue/storage/__init__.py

@@ -13,7 +13,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .managers import AsyncSimpleStorageManager, TransferQueueStorageManager, TransferQueueStorageManagerFactory
+from .managers import (
+    AsyncSimpleStorageManager,
+    MooncakeStorageManager,
+    TransferQueueStorageManager,
+    TransferQueueStorageManagerFactory,
+    YuanrongStorageManager,
+)
 from .simple_backend import SimpleStorageUnit, StorageMetaGroup, StorageUnitData
 
 __all__ = [
@@ -23,4 +29,6 @@
     "TransferQueueStorageManager",
     "TransferQueueStorageManagerFactory",
     "AsyncSimpleStorageManager",
+    "MooncakeStorageManager",
+    "YuanrongStorageManager",
 ]

transfer_queue/storage/clients/__init__.py

@@ -16,10 +16,12 @@
 # This module is currently empty but reserved for future client implementations
 from .base import TransferQueueStorageKVClient
 from .factory import StorageClientFactory
+from .mooncake_client import MooncakeStorageClient
 from .yuanrong_client import YuanrongStorageClient
 
 __all__ = [
     "TransferQueueStorageKVClient",
     "StorageClientFactory",
+    "MooncakeStorageClient",
     "YuanrongStorageClient",
 ]

transfer_queue/storage/clients/mooncake_client.py

@@ -0,0 +1,188 @@
+import logging
+import os
+import pickle
+from typing import Any
+
+import torch
+from torch import Tensor
+
+from transfer_queue.storage.clients.base import TransferQueueStorageKVClient
+from transfer_queue.storage.clients.factory import StorageClientFactory
+
+logger = logging.getLogger(__name__)
+logger.setLevel(os.getenv("TQ_LOGGING_LEVEL", logging.WARNING))
+
+MOONCAKE_STORE_IMPORTED: bool = True
+try:
+    from mooncake.store import MooncakeDistributedStore
+except ImportError:
+    MOONCAKE_STORE_IMPORTED = False
+
+BATCH_SIZE_LIMIT: int = 500
+
+
+@StorageClientFactory.register("MooncakeStorageClient")
+class MooncakeStorageClient(TransferQueueStorageKVClient):
+    def __init__(self, config: dict[str, Any]):
+        if not MOONCAKE_STORE_IMPORTED:
+            raise ImportError("Mooncake Store not installed. Please install via: pip install mooncake-transfer-engine")
+
+        self.local_hostname = config.get("local_hostname", "localhost")
+        self.metadata_server = config.get("metadata_server")
+        self.global_segment_size = config.get("global_segment_size", 512 * 1024 * 1024)
+        self.local_buffer_size = config.get("local_buffer_size", 128 * 1024 * 1024)
+        self.protocol = config.get("protocol", "tcp")
+        self.device_name = config.get("device_name", "")
+        self.master_server_address = config.get("master_server_address")
+
+        if self.metadata_server is None:
+            raise ValueError("Missing 'metadata_server' in config")
+        if self.master_server_address is None:
+            raise ValueError("Missing 'master_server_address' in config")
+
+        self._store = MooncakeDistributedStore()
+        ret = self._store.setup(
+            self.local_hostname,
+            self.metadata_server,
+            self.global_segment_size,
+            self.local_buffer_size,
+            self.protocol,
+            self.device_name,
+            self.master_server_address,
+        )
+        if ret != 0:
+            raise RuntimeError(f"Mooncake store setup failed with error code: {ret}")
+
+    def put(self, keys: list[str], values: list[Any]):
+        if not isinstance(keys, list) or not isinstance(values, list):
+            raise ValueError("keys and values must be lists")
+        if len(keys) != len(values):
+            raise ValueError("Number of keys must match number of values")
+
+        tensor_keys = []
+        tensor_values = []
+        non_tensor_keys = []
+        non_tensor_values = []
+
+        for key, value in zip(keys, values, strict=True):
+            if isinstance(value, torch.Tensor):
+                tensor = value.contiguous()
+                # TODO: use gpu direct rdma instead
+                if tensor.device.type == "cuda":
+                    tensor = tensor.cpu()
+                tensor_keys.append(key)
+                tensor_values.append(tensor)
+            else:
+                non_tensor_keys.append(key)
+                non_tensor_values.append(pickle.dumps(value))
+
+        if tensor_keys:
+            self._batch_put_tensors(tensor_keys, tensor_values)
+
+        if non_tensor_keys:
+            self._batch_put_bytes(non_tensor_keys, non_tensor_values)
+
+    def _batch_put_tensors(self, keys: list[str], tensors: list[Tensor]):
+        for i in range(0, len(keys), BATCH_SIZE_LIMIT):
+            batch_keys = keys[i : i + BATCH_SIZE_LIMIT]
+            batch_tensors = tensors[i : i + BATCH_SIZE_LIMIT]
+
+            results = self._store.batch_put_tensor(batch_keys, batch_tensors)
+            if not all(r == 0 for r in results):
+                failed_indices = [j for j, r in enumerate(results) if r != 0]
+                error_codes = [results[j] for j in failed_indices]
+                raise RuntimeError(
+                    f"batch_put_tensor failed for indices {failed_indices} with error codes: {error_codes}"
+                )
+
+    def _batch_put_bytes(self, keys: list[str], values: list[bytes]):
+        for i in range(0, len(keys), BATCH_SIZE_LIMIT):
+            batch_keys = keys[i : i + BATCH_SIZE_LIMIT]
+            batch_values = values[i : i + BATCH_SIZE_LIMIT]
+
+            ret = self._store.put_batch(batch_keys, batch_values)
+            if ret != 0:
+                raise RuntimeError(f"put_batch failed with error code: {ret}")
+
+    def get(self, keys: list[str], shapes=None, dtypes=None) -> list[Any]:
+        if shapes is None or dtypes is None:
+            raise ValueError("MooncakeStorageClient needs shapes and dtypes")
+        if not (len(keys) == len(shapes) == len(dtypes)):
+            raise ValueError("Lengths of keys, shapes, dtypes must match")
+
+        tensor_indices = []
+        non_tensor_indices = []
+
+        for i, dtype in enumerate(dtypes):
+            if dtype is not None:
+                tensor_indices.append(i)
+            else:
+                non_tensor_indices.append(i)
+
+        results = [None] * len(keys)
+
+        if tensor_indices:
+            tensor_keys = [keys[i] for i in tensor_indices]
+            tensor_shapes = [shapes[i] for i in tensor_indices]
+            tensor_dtypes = [dtypes[i] for i in tensor_indices]
+            tensor_results = self._batch_get_tensors(tensor_keys, tensor_shapes, tensor_dtypes)
+            # TODO: optimize these for loops
+            for idx, tensor in zip(tensor_indices, tensor_results, strict=True):
+                results[idx] = tensor
+
+        if non_tensor_indices:
+            non_tensor_keys = [keys[i] for i in non_tensor_indices]
+            non_tensor_results = self._batch_get_bytes(non_tensor_keys)
+            for idx, data in zip(non_tensor_indices, non_tensor_results, strict=True):
+                results[idx] = pickle.loads(data)
+
+        return results
+
+    def _batch_get_tensors(self, keys: list[str], shapes: list, dtypes: list) -> list[Tensor]:
+        tensors = [None] * len(keys)
+
+        for i in range(0, len(keys), BATCH_SIZE_LIMIT):
+            batch_keys = keys[i : i + BATCH_SIZE_LIMIT]
+            batch_shapes = shapes[i : i + BATCH_SIZE_LIMIT]
+            batch_dtypes = dtypes[i : i + BATCH_SIZE_LIMIT]
+
+            batch_results = self._store.batch_get_tensor(batch_keys)
+
+            if len(batch_results) != len(batch_keys):
+                raise RuntimeError(f"batch_get_tensor returned {len(batch_results)} items, expected {len(batch_keys)}")
+
+            for j, (tensor, shape, dtype) in enumerate(zip(batch_results, batch_shapes, batch_dtypes, strict=True)):
+                if tensor is None:
+                    raise RuntimeError(f"batch_get_tensor returned None for key '{batch_keys[j]}'")
+                if tensor.shape != torch.Size(shape):
+                    raise RuntimeError(
+                        f"Shape mismatch for key '{batch_keys[j]}': expected {shape}, got {tensor.shape}"
+                    )
+                if tensor.dtype != dtype:
+                    raise RuntimeError(
+                        f"Dtype mismatch for key '{batch_keys[j]}': expected {dtype}, got {tensor.dtype}"
+                    )
+                tensors[i + j] = tensor
+
+        return tensors
+
+    def _batch_get_bytes(self, keys: list[str]) -> list[bytes]:
+        results = []
+        for i in range(0, len(keys), BATCH_SIZE_LIMIT):
+            batch_keys = keys[i : i + BATCH_SIZE_LIMIT]
+            batch_results = self._store.get_batch(batch_keys)
+            if len(batch_results) != len(batch_keys):
+                raise RuntimeError(f"get_batch returned {len(batch_results)} items, expected {len(batch_keys)}")
+            results.extend(batch_results)
+        return results
+
+    def clear(self, keys: list[str]):
+        for key in keys:
+            ret = self._store.remove(key)
+            if ret != 0:
+                logger.warning(f"remove failed for key '{key}' with error code: {ret}")
+
+    def close(self):
+        if self._store:
+            self._store.close()
+            self._store = None

transfer_queue/storage/managers/__init__.py

@@ -15,6 +15,7 @@
 
 from .base import TransferQueueStorageManager
 from .factory import TransferQueueStorageManagerFactory
+from .mooncake_manager import MooncakeStorageManager
 from .simple_backend_manager import AsyncSimpleStorageManager
 from .yuanrong_manager import YuanrongStorageManager
 
@@ -23,4 +24,5 @@
     "TransferQueueStorageManagerFactory",
     "AsyncSimpleStorageManager",
     "YuanrongStorageManager",
+    "MooncakeStorageManager",
 ]

transfer_queue/storage/managers/base.py

@@ -12,7 +12,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import asyncio
 import itertools
 import logging
 import os
@@ -432,9 +432,16 @@ async def put_data(self, data: TensorDict, metadata: BatchMeta) -> None:
         if not metadata.field_names:
             logger.warning("Attempted to put data, but metadata contains no fields.")
             return
+
+        # For each field, extract dtype and shape for each sample
+        num_samples = len(metadata.global_indexes)
+        if num_samples == 0:
+            return
+
         keys = self._generate_keys(data.keys(), metadata.global_indexes)
         values = self._generate_values(data)
-        self.storage_client.put(keys=keys, values=values)
+        loop = asyncio.get_event_loop()
+        await loop.run_in_executor(None, self.storage_client.put, keys, values)
 
         per_field_dtypes = {}
         per_field_shapes = {}
@@ -444,9 +451,9 @@ async def put_data(self, data: TensorDict, metadata: BatchMeta) -> None:
             per_field_dtypes[global_idx] = {}
             per_field_shapes[global_idx] = {}
 
-        # For each field, extract dtype and shape for each sample
         for field_name, field_data in data.items():
-            for i, data_item in enumerate(field_data):
+            for i in range(num_samples):
+                data_item = field_data[i]
                 global_idx = metadata.global_indexes[i]
                 per_field_dtypes[global_idx][field_name] = (
                     getattr(data_item, "dtype", None) if isinstance(data_item, Tensor) else None
@@ -484,5 +491,5 @@ async def clear_data(self, metadata: BatchMeta) -> None:
         if not metadata.field_names:
             logger.warning("Attempted to clear data, but metadata contains no fields.")
             return
-        keys = self._generate_keys(metadata)
+        keys = self._generate_keys(metadata.field_names, metadata.global_indexes)
         self.storage_client.clear(keys=keys)

transfer_queue/storage/managers/mooncake_manager.py

@@ -0,0 +1,31 @@
+import logging
+import os
+from typing import Any
+
+from transfer_queue.storage.managers.base import KVStorageManager
+from transfer_queue.storage.managers.factory import TransferQueueStorageManagerFactory
+
+logger = logging.getLogger(__name__)
+logger.setLevel(os.getenv("TQ_LOGGING_LEVEL", logging.WARNING))
+
+
+@TransferQueueStorageManagerFactory.register("MooncakeStorageManager")
+class MooncakeStorageManager(KVStorageManager):
+    def __init__(self, config: dict[str, Any]):
+        # Required: Address of the HTTP metadata server (e.g., "localhost:8080")
+        metadata_server = config.get("metadata_server", None)
+        # Required: Address of the master server RPC endpoint (e.g., "localhost:8081")
+        master_server_address = config.get("master_server_address", None)
+        # Optional: Name of the storage client, defaults to "MooncakeStorageClient" if not provided
+        client_name = config.get("client_name", None)
+
+        if metadata_server is None or not isinstance(metadata_server, str):
+            raise ValueError("Missing or invalid 'metadata_server' in config")
+        if master_server_address is None or not isinstance(master_server_address, str):
+            raise ValueError("Missing or invalid 'master_server_address' in config")
+        if client_name is None:
+            logger.info("Missing 'client_name' in config, using default value('MooncakeStorageClient')")
+            config["client_name"] = "MooncakeStorageClient"
+        elif client_name != "MooncakeStorageClient":
+            raise ValueError(f"Invalid 'client_name': {client_name} in config. Expecting 'MooncakeStorageClient'")
+        super().__init__(config)