use new mooncake API and enable zero-copy

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

try: use new mooncake API and enable zero-copy

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

update

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

update

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

fix

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

Author: 0oshowero0

transfer_queue/storage/clients/mooncake_client.py

@@ -23,13 +23,15 @@
 
 from transfer_queue.storage.clients.base import TransferQueueStorageKVClient
 from transfer_queue.storage.clients.factory import StorageClientFactory
+from transfer_queue.utils.tensor_utils import allocate_empty_tensors, get_nbytes, merge_continues_memory
 
 logger = logging.getLogger(__name__)
 logger.setLevel(os.getenv("TQ_LOGGING_LEVEL", logging.WARNING))
 
 MOONCAKE_STORE_IMPORTED: bool = True
 try:
-    from mooncake.store import MooncakeDistributedStore
+    from mooncake.store import MooncakeDistributedStore, ReplicateConfig
+
 except ImportError:
     MOONCAKE_STORE_IMPORTED = False
 
@@ -78,10 +80,9 @@ def __init__(self, config: dict[str, Any]):
         if not self.metadata_server.startswith("etcd://") and not self.metadata_server.endswith("/metadata"):
             self.metadata_server = self.metadata_server + "/metadata"
 
-        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.replica_config = ReplicateConfig()
+        # FIXME: hard_pin is not supported yet
+        # self.replica_config.with_hard_pin = True
 
         self._store = MooncakeDistributedStore()
         ret = self._store.setup(
@@ -116,12 +117,8 @@ def put(self, keys: list[str], values: list[Any]) -> Optional[list[Any]]:
 
         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)
+                tensor_values.append(value)
             else:
                 non_tensor_keys.append(key)
                 non_tensor_values.append(pickle.dumps(value))
@@ -139,38 +136,50 @@ def _batch_put_tensors(self, keys: list[str], tensors: list[Tensor]):
             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)
+            batch_ptrs, batch_sizes = self._preprocess_tensors_for_put(batch_tensors)
+            batch_ptr_reduced, batch_sizes_reduced = merge_continues_memory(batch_ptrs, batch_sizes)
+            self._register_all_buffers(batch_ptr_reduced, batch_sizes_reduced)
+
+            results = self._store.batch_upsert_from(batch_keys, batch_ptrs, batch_sizes, config=self.replica_config)
             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}"
                 )
 
+            self._unregister_all_buffers(batch_ptr_reduced)
+
     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)
+            ret = self._store.upsert_batch(batch_keys, batch_values, self.replica_config)
             if ret != 0:
                 raise RuntimeError(f"put_batch failed with error code: {ret}")
 
-    def get(self, keys: list[str], shapes=None, dtypes=None, custom_backend_meta=None) -> list[Any]:
+    def get(
+        self,
+        keys: list[str],
+        shapes: Optional[list[Any]] = None,
+        dtypes: Optional[list[Any]] = None,
+        custom_backend_meta: Optional[list[str]] = None,
+    ) -> list[Any]:
         """Get multiple key-value pairs from MooncakeStore.
 
         Args:
-            keys (List[str]): Keys to fetch.
-            shapes (List[List[int]]): Expected tensor shapes (use [] for scalars).
-            dtypes (List[Optional[torch.dtype]]): Expected dtypes; use None for non-tensor data.
-            custom_backend_meta (List[str], optional): ...
+            keys: Keys to fetch.
+            shapes: Expected tensor shapes (use [] for scalars).
+            dtypes: Expected dtypes; use None for non-tensor data.
+            custom_backend_meta: Optional custom backend metadata.
 
         Returns:
-            List[Any]: Retrieved values in the same order as input keys.
+            Retrieved values in the same order as input keys.
         """
 
         if shapes is None or dtypes is None:
-            raise ValueError("MooncakeStoreClient needs shapes and dtypes")
+            raise ValueError("MooncakeStoreClient needs shapes and dtypes for zero-copy transfer.")
         if not (len(keys) == len(shapes) == len(dtypes)):
             raise ValueError("Lengths of keys, shapes, dtypes must match")
 
@@ -210,14 +219,25 @@ def _batch_get_tensors(self, keys: list[str], shapes: list, dtypes: list) -> lis
             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)
+            batch_nbytes = get_nbytes(batch_dtypes, batch_shapes)
+            batch_buffer_tensors, batch_buffer_ptrs = allocate_empty_tensors(batch_dtypes, batch_shapes)
 
-            if len(batch_results) != len(batch_keys):
-                raise RuntimeError(f"batch_get_tensor returned {len(batch_results)} items, expected {len(batch_keys)}")
+            batch_ptrs = batch_buffer_ptrs
 
-            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]}'")
+            self._register_all_buffers(batch_ptrs, batch_nbytes)
+            ret_codes = self._store.batch_get_into(batch_keys, batch_ptrs, batch_nbytes)
+            self._unregister_all_buffers(batch_ptrs)
+
+            if len(ret_codes) != len(batch_keys):
+                raise RuntimeError(f"batch_get_into returned {len(ret_codes)} results, expected {len(batch_keys)}")
+
+            # Check result codes and validate tensors
+            # Note: Positive values indicate success (bytes read), negative values indicate error
+            for j, (tensor, shape, dtype, ret_code) in enumerate(
+                zip(batch_buffer_tensors, batch_shapes, batch_dtypes, ret_codes, strict=True)
+            ):
+                if ret_code < 0:
+                    raise RuntimeError(f"batch_get_into failed for key '{batch_keys[j]}' with error code: {ret_code}")
                 if tensor.shape != torch.Size(shape):
                     raise RuntimeError(
                         f"Shape mismatch for key '{batch_keys[j]}': expected {shape}, got {tensor.shape}"
@@ -243,26 +263,35 @@ def _batch_get_bytes(self, keys: list[str]) -> list[bytes]:
     def clear(self, keys: list[str], custom_backend_meta=None):
         """Deletes multiple keys from MooncakeStore.
 
-
         Args:
             keys (List[str]): List of keys to remove.
             custom_backend_meta (List[Any], optional): ...
         """
-        global_indexes_patterns = {key.split("@")[0] + "@.*" for key in keys}
-        for p in global_indexes_patterns:
-            ret = self._store.remove_by_regex(p, force=True)
-            if ret < 0:
-                logger.warning(f"remove failed for key '{p}' with error code: {ret}")
-
-        # FIXME: controller returned BatchMeta may have mismatched fields in some case, preventing
-        #        key-value based backends to accurately clear all existing keys..
-        # for key in keys:
-        #     ret = self._store.remove(key)
-        #     if not (ret == 0 or ret == -704):
-        #         logger.warning(f"remove failed for key '{key}' with error code: {ret}")
+        rets = self._store.batch_remove(keys, force=True)
+        for i, ret in enumerate(rets):
+            if not (ret == 0 or ret == -704):
+                logger.error(f"remove failed for key '{keys[i]}' with error code: {ret}")
 
     def close(self):
         """Closes MooncakeStore."""
         if self._store:
             self._store.close()
             self._store = None
+
+    @staticmethod
+    def _preprocess_tensors_for_put(values: list[Tensor]) -> tuple[list[Any], list[Any]]:
+        ptr_list = []
+        size_list = []
+        for t in values:
+            t = t.contiguous()
+            ptr_list.append(t.data_ptr())
+            size_list.append(t.nbytes)
+        return ptr_list, size_list
+
+    def _register_all_buffers(self, ptrs, sizes):
+        for ptr, size in zip(ptrs, sizes, strict=False):
+            self._store.register_buffer(ptr, size)
+
+    def _unregister_all_buffers(self, ptrs):
+        for ptr in ptrs:
+            self._store.unregister_buffer(ptr)

transfer_queue/utils/tensor_utils.py

@@ -0,0 +1,172 @@
+# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+# Copyright 2025 The TransferQueue Team
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 logging
+import operator
+import os
+from functools import reduce
+
+import torch
+from torch import Tensor
+
+logger = logging.getLogger(__name__)
+logger.setLevel(os.getenv("TQ_LOGGING_LEVEL", logging.WARNING))
+
+
+def allocate_empty_tensors(dtypes: list[torch.dtype], shapes: list[tuple]) -> tuple[list[Tensor], list[int]]:
+    """Allocate empty tensors, grouping same dtypes into shared memory blocks.
+
+    Instead of allocating each tensor separately, this function groups tensors
+    by their dtype and allocates one large contiguous memory block per dtype.
+    Each tensor is then created as a view into this shared memory.
+
+    Args:
+        dtypes: List of torch dtypes for each tensor.
+        shapes: List of shapes (tuples) for each tensor.
+
+    Returns:
+        A tuple containing:
+            - List of tensors sharing memory within their dtype groups.
+            - List of memory pointers (data_ptr) for each tensor.
+
+    Example:
+        >>> dtypes = [torch.float32, torch.float32, torch.int32, torch.float32]
+        >>> shapes = [(10,), (20,), (5,), (15,)]
+        >>> tensors, ptrs = allocate_empty_tensors(dtypes, shapes)
+        >>> # tensors[0], [1], [3] share the same dtype and memory block
+    """
+    assert len(dtypes) == len(shapes), "dtypes and shapes must have the same length"
+
+    if len(dtypes) == 0:
+        return [], []
+
+    # Group indices by dtype
+    dtype_groups: dict[torch.dtype, list[int]] = {}
+    for i, dtype in enumerate(dtypes):
+        if dtype not in dtype_groups:
+            dtype_groups[dtype] = []
+        dtype_groups[dtype].append(i)
+
+    tensor_list = [torch.empty(()) for _ in range(len(dtypes))]
+    ptr_list = [0] * len(dtypes)
+
+    # For each dtype group, allocate one big tensor and create views
+    for dtype, indices in dtype_groups.items():
+        # Calculate total number of elements needed for this dtype
+        total_elements = 0
+        shape_info = []  # Store (index, shape, num_elements, offset)
+
+        for idx in indices:
+            shape = shapes[idx]
+            num_elements = reduce(operator.mul, shape)
+            shape_info.append((idx, shape, num_elements, total_elements))
+            total_elements += num_elements
+
+        # Allocate one big contiguous memory block for this dtype
+        big_tensor = torch.empty(total_elements, dtype=dtype)
+
+        # Create views into the big tensor for each small tensor
+        for idx, shape, num_elements, offset in shape_info:
+            # Use as_strided to create a view with the correct shape
+            small_tensor = big_tensor.as_strided(size=shape, stride=compute_stride(shape), storage_offset=offset)
+            tensor_list[idx] = small_tensor
+            ptr_list[idx] = small_tensor.data_ptr()
+
+    return tensor_list, ptr_list
+
+
+def compute_stride(shape: tuple[int, ...]) -> tuple[int, ...]:
+    """Compute stride for a contiguous row-major (C-style) tensor.
+
+    Args:
+        shape: The shape of the tensor.
+
+    Returns:
+        Stride tuple for contiguous storage.
+
+    Example:
+        >>> compute_stride((2, 3, 4))
+        (12, 4, 1)
+    """
+    stride = []
+    cumulative = 1
+    # Iterate from last dimension to first
+    for dim in reversed(shape):
+        stride.append(cumulative)
+        cumulative *= dim
+    return tuple(reversed(stride))
+
+
+def get_nbytes(dtypes, shapes) -> list[int]:
+    assert len(dtypes) == len(shapes)
+    nbytes = []
+    for i in range(len(dtypes)):
+        elem_size = torch.tensor([], dtype=dtypes[i]).element_size()
+        numel = reduce(operator.mul, shapes[i])
+        nbytes.append(elem_size * numel)
+
+    return nbytes
+
+
+def merge_continues_memory(ptrs: list[int], sizes: list[int]) -> tuple[list[int], list[int]]:
+    """Merge continuous memory regions to reduce register_buffer overhead
+
+    Args:
+        ptrs: List of memory pointers (starting addresses).
+        sizes: List of memory region sizes corresponding to each pointer.
+
+    Returns:
+        A tuple of (merged_ptrs, merged_sizes) where continuous regions
+        have been merged into single regions.
+
+    Example:
+        >>> merge_continues_memory([0, 10, 30], [10, 20, 10])
+        ([0, 30], [30, 10])
+
+        >>> merge_continues_memory([0, 5, 20], [5, 5, 10])
+        ([0, 20], [10, 10])
+    """
+    if not ptrs or not sizes:
+        return [], []
+
+    if len(ptrs) != len(sizes):
+        raise ValueError("ptrs and sizes must have the same length")
+
+    # Create list of (ptr, size) pairs and sort by pointer address
+    regions = sorted(zip(ptrs, sizes, strict=False), key=lambda x: x[0])
+
+    merged_ptrs = []
+    merged_sizes = []
+
+    # Initialize with the first region
+    current_ptr, current_size = regions[0]
+
+    for ptr, size in regions[1:]:
+        # Check if current region is contiguous with the next one
+        # A region is contiguous if: ptr == current_ptr + current_size
+        if ptr == current_ptr + current_size:
+            # Merge: extend the current region
+            current_size += size
+        else:
+            # Not contiguous: save the current region and start a new one
+            merged_ptrs.append(current_ptr)
+            merged_sizes.append(current_size)
+            current_ptr, current_size = ptr, size
+
+    # Add the last region
+    merged_ptrs.append(current_ptr)
+    merged_sizes.append(current_size)
+
+    return merged_ptrs, merged_sizes