refactor(mooncake): unify tensor and non-tensor data paths into a single packed-buffer batch transfer

Author: xupinjie

transfer_queue/storage/clients/mooncake_client.py

@@ -13,17 +13,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import pickle
 import time
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import Any
 
+import numpy as np
 import torch
+from tensordict import TensorDictBase
 from torch import Tensor
 
 from transfer_queue.storage.clients.base import StorageClientFactory, StorageKVClient
+from transfer_queue.utils import serial_utils
 from transfer_queue.utils.logging_utils import get_logger
-from transfer_queue.utils.tensor_utils import allocate_empty_tensors, get_nbytes, merge_contiguous_memory
+from transfer_queue.utils.tensor_utils import allocate_empty_tensors, get_nbytes
 
 logger = get_logger(__name__)
 
@@ -40,12 +42,33 @@
 RETRY_DELAY_SECONDS = 1.0
 
 
+def _detach_from_buffer(obj: Any) -> Any:
+    """Deep-copy all tensor/array leaves so the result owns its storage."""
+    # TODO: replace with a keep-alive scheme on the source buffer to skip the copy.
+    if isinstance(obj, torch.Tensor):
+        return obj.clone()
+    if isinstance(obj, np.ndarray):
+        return obj.copy()
+    if isinstance(obj, dict):
+        return {k: _detach_from_buffer(v) for k, v in obj.items()}
+    if isinstance(obj, list):
+        return [_detach_from_buffer(v) for v in obj]
+    if isinstance(obj, tuple):
+        return tuple(_detach_from_buffer(v) for v in obj)
+    if isinstance(obj, TensorDictBase):
+        return obj.apply(lambda t: t.clone())
+    return obj
+
+
 @StorageClientFactory.register("MooncakeStoreClient")
 class MooncakeStoreClient(StorageKVClient):
     """
     Storage client for MooncakeStore.
     """
 
+    _logged_first_put: bool = False
+    _logged_first_get: bool = False
+
     def __init__(self, config: dict[str, Any]):
         super().__init__(config)
         if not MOONCAKE_STORE_IMPORTED:
@@ -98,55 +121,70 @@ def __init__(self, config: dict[str, Any]):
         if ret != 0:
             raise RuntimeError(f"Mooncake store setup failed with error code: {ret}")
 
-    def put(self, keys: list[str], values: list[Any]) -> None:
+    def put(self, keys: list[str], values: list[Any]) -> list[dict]:
         """Stores multiple key-value pairs to MooncakeStore.
 
         Args:
             keys (List[str]): List of unique string identifiers.
             values (List[Any]): List of values to store (tensors, scalars, dicts, etc.).
-        """
 
+        Returns:
+            List[Dict]: Per-key ``{"packed_size": int}`` metadata, in the same order as ``keys``.
+        """
         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_keys.append(key)
-                tensor_values.append(value)
-            else:
-                non_tensor_keys.append(key)
-                non_tensor_values.append(value)
+        if not type(self)._logged_first_put:
+            logger.info("[TQ-MOONCAKE-REFACTOR] put() entered: unified pack-into data path")
+            type(self)._logged_first_put = True
 
+        custom_meta: list[dict] = []
         futures = []
         with ThreadPoolExecutor(max_workers=MAX_WORKER_THREADS) as executor:
-            for i in range(0, len(tensor_keys), BATCH_SIZE_LIMIT):
-                batch_keys = tensor_keys[i : i + BATCH_SIZE_LIMIT]
-                batch_tensors = tensor_values[i : i + BATCH_SIZE_LIMIT]
-                futures.append(executor.submit(self._put_tensors_thread_worker, batch_keys, batch_tensors))
-
-            for i in range(0, len(non_tensor_keys), BATCH_SIZE_LIMIT):
-                batch_keys = non_tensor_keys[i : i + BATCH_SIZE_LIMIT]
-                batch_values = non_tensor_values[i : i + BATCH_SIZE_LIMIT]
-                futures.append(executor.submit(self._put_bytes_thread_worker, batch_keys, batch_values))
+            for start in range(0, len(values), BATCH_SIZE_LIMIT):
+                batch_keys = keys[start : start + BATCH_SIZE_LIMIT]
+                batch_values = values[start : start + BATCH_SIZE_LIMIT]
+
+                # Encode every value (msgpack zero-copy; transparent pickle fallback),
+                # then pack all encoded frames into ONE contiguous buffer. The worker
+                # registers this single buffer once and uses per-slice (ptr, size) for
+                # batch_upsert_from.
+                batch_items = [serial_utils.encode(v) for v in batch_values]
+                batch_sizes = [serial_utils.calc_packed_size(items) for items in batch_items]
+                # TODO: switch to a MooncakeStore-allocated buffer once such an API exists.
+                big_buf = torch.empty(sum(batch_sizes), dtype=torch.uint8)
+                big_buf_mv = big_buf.numpy().data
+                base_ptr = big_buf.data_ptr()
+
+                batch_ptrs: list[int] = []
+                offset = 0
+                for items, size in zip(batch_items, batch_sizes, strict=True):
+                    serial_utils.pack_into(big_buf_mv[offset : offset + size], items)
+                    batch_ptrs.append(base_ptr + offset)
+                    offset += size
+
+                custom_meta.extend({"packed_size": s} for s in batch_sizes)
+                futures.append(
+                    executor.submit(self._put_batch_worker, batch_keys, batch_ptrs, batch_sizes, big_buf)
+                )
 
             for future in as_completed(futures):
                 future.result()
 
-        return None
+        return custom_meta
+
+    def _put_batch_worker(
+        self, batch_keys: list[str], batch_ptrs: list[int], batch_sizes: list[int], big_buf: Tensor
+    ) -> None:
+        """Worker thread for putting one packed-buffer batch to MooncakeStore.
 
-    def _put_tensors_thread_worker(self, batch_keys: list[str], batch_tensors: list[Tensor]) -> None:
-        """Worker thread for putting batch of tensors to MooncakeStore."""
+        ``big_buf`` is passed only to keep the underlying storage alive while
+        ``batch_ptrs`` (per-value slices into it) are in flight.
+        """
 
-        batch_ptrs, batch_sizes, _contiguous_tensors = self._preprocess_tensors_for_put(batch_tensors)
-        batch_ptr_reduced, batch_sizes_reduced = merge_contiguous_memory(batch_ptrs, batch_sizes)
-        self._register_all_buffers(batch_ptr_reduced, batch_sizes_reduced)
+        self._store.register_buffer(big_buf.data_ptr(), big_buf.nbytes)
 
         try:
             results = self._store.batch_upsert_from(batch_keys, batch_ptrs, batch_sizes, config=self.replica_config)
@@ -206,98 +244,65 @@ def _put_tensors_thread_worker(self, batch_keys: list[str], batch_tensors: list[
                 )
 
         finally:
-            self._unregister_all_buffers(batch_ptr_reduced)
-
-    def _put_bytes_thread_worker(self, batch_keys: list[str], batch_values: list[Any]):
-        """Worker thread for putting batch of non-tensors to MooncakeStore."""
-
-        serialized_values = [pickle.dumps(v, protocol=pickle.HIGHEST_PROTOCOL) for v in batch_values]
-
-        # FIXME: When MooncakeStore supports per-key status codes for upsert_batch and get_batch,
-        #        switch the bytes write/read paths from whole-batch retry to per-key selective retry,
-        #        matching the tensor-path behaviour.
-        ret = self._store.upsert_batch(batch_keys, serialized_values, self.replica_config)
-        if ret == 0:
-            return
+            self._store.unregister_buffer(big_buf.data_ptr())
 
-        logger.error(
-            f"upsert_batch failed for {len(batch_keys)} keys with error code: {ret}. "
-            f"Retrying up to {MAX_RETRIES} times..."
-        )
-
-        for attempt in range(1, MAX_RETRIES + 1):
-            ret = self._store.upsert_batch(batch_keys, serialized_values, self.replica_config)
-            if ret == 0:
-                logger.info("upsert_batch succeeded after retransmission.")
-                return
-
-            logger.error(
-                f"upsert_batch retry {attempt}/{MAX_RETRIES} failed for {len(batch_keys)} keys with error code: {ret}."
-            )
-            if attempt < MAX_RETRIES:
-                time.sleep(RETRY_DELAY_SECONDS)
-
-        raise RuntimeError(
-            f"upsert_batch failed for {len(batch_keys)} keys with error code: {ret} after retrying {MAX_RETRIES} times."
-        )
-
-    def get(
-        self,
-        keys: list[str],
-        shapes: list[Any] | None = None,
-        dtypes: list[Any] | None = None,
-        custom_backend_meta: list[str] | None = None,
-    ) -> list[Any]:
+    def get(self, keys: list[str], **kwargs) -> list[Any]:
         """Get multiple key-value pairs from MooncakeStore.
 
         Args:
             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.
+            **kwargs: Must contain ``custom_backend_meta`` — per-key dicts
+                carrying ``"packed_size": int``.
 
         Returns:
             Retrieved values in the same order as input keys.
         """
+        if not type(self)._logged_first_get:
+            logger.info("[TQ-MOONCAKE-REFACTOR] get() entered: unified unpack+detach data path")
+            type(self)._logged_first_get = True
+
+        custom_backend_meta = kwargs.get("custom_backend_meta")
+        if custom_backend_meta is None:
+            raise ValueError("MooncakeStoreClient.get requires custom_backend_meta with per-key packed_size.")
+        if len(custom_backend_meta) != len(keys):
+            raise ValueError(
+                f"Length of custom_backend_meta ({len(custom_backend_meta)}) must match keys ({len(keys)})"
+            )
 
-        if shapes is None or dtypes is None:
-            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")
-
-        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)
+        try:
+            packed_sizes = [m["packed_size"] for m in custom_backend_meta]
+        except (KeyError, TypeError) as e:
+            raise ValueError("custom_backend_meta entries must be dicts with 'packed_size'") from e
 
-        results = [None] * len(keys)
+        results: list[Any] = [None] * len(keys)
 
+        # TODO: when MooncakeStore exposes a pre-registered receive-buffer API
+        # (symmetric to YuanRong's get_buffers), drop the local alloc + register
+        # below and hand decoded views straight from MooncakeStore's memory.
         futures = []
         with ThreadPoolExecutor(max_workers=MAX_WORKER_THREADS) as executor:
-            for i in range(0, len(tensor_indices), BATCH_SIZE_LIMIT):
-                batch_indexes = tensor_indices[i : i + BATCH_SIZE_LIMIT]
-                batch_keys = [keys[i] for i in batch_indexes]
-                batch_shapes = [shapes[i] for i in batch_indexes]
-                batch_dtypes = [dtypes[i] for i in batch_indexes]
+            for start in range(0, len(keys), BATCH_SIZE_LIMIT):
+                end = min(start + BATCH_SIZE_LIMIT, len(keys))
+                batch_keys = keys[start:end]
+                batch_shapes = [(packed_sizes[j],) for j in range(start, end)]
+                batch_dtypes = [torch.uint8] * (end - start)
+                batch_indexes = list(range(start, end))
                 futures.append(
                     executor.submit(
                         self._get_tensors_thread_worker, batch_keys, batch_shapes, batch_dtypes, batch_indexes
                     )
                 )
 
-            for i in range(0, len(non_tensor_indices), BATCH_SIZE_LIMIT):
-                batch_indexes = non_tensor_indices[i : i + BATCH_SIZE_LIMIT]
-                batch_keys = [keys[i] for i in batch_indexes]
-                futures.append(executor.submit(self._get_bytes_thread_worker, batch_keys, batch_indexes))
-
             for future in as_completed(futures):
-                retrieved_values, batch_indexes = future.result()
-                for idx, val in zip(batch_indexes, retrieved_values, strict=True):
-                    results[idx] = val
+                packed_tensors, indexes = future.result()
+                for idx, packed in zip(indexes, packed_tensors, strict=True):
+                    results[idx] = packed
+
+        for idx, packed in enumerate(results):
+            if packed is None:
+                continue
+            items = serial_utils.unpack_from(packed.numpy().data)
+            results[idx] = _detach_from_buffer(serial_utils.decode(items))
 
         return results
 
@@ -374,57 +379,6 @@ def _get_tensors_thread_worker(
 
         return batch_buffer_tensors, indexes
 
-    def _get_bytes_thread_worker(self, batch_keys: list[str], indexes: list[int]) -> tuple[list[Any], list[int]]:
-        raw_results = self._store.get_batch(batch_keys)
-        if len(raw_results) != len(batch_keys):
-            raise RuntimeError(f"get_batch returned {len(raw_results)} items, expected {len(batch_keys)}")
-
-        # FIXME: Use MooncakeStore provided ret codes to detect transmission failures when supported
-        # Currently we rely on empty bytes (b'') to detect transmission failures because
-        # MooncakeStore does not currently return a separate status code per key.
-        failed_indices = [i for i, result in enumerate(raw_results) if result == b""]
-        if failed_indices:
-            current_failed_keys = [batch_keys[i] for i in failed_indices]
-            current_failed_indices = failed_indices
-
-            logger.error(f"get_batch failed for keys {current_failed_keys}. Retrying up to {MAX_RETRIES} times...")
-
-            for attempt in range(1, MAX_RETRIES + 1):
-                retry_results = self._store.get_batch(current_failed_keys)
-
-                next_failed_keys = []
-                next_failed_indices = []
-
-                for i, result in enumerate(retry_results):
-                    original_idx = current_failed_indices[i]
-                    if result == b"":
-                        next_failed_keys.append(current_failed_keys[i])
-                        next_failed_indices.append(original_idx)
-                    else:
-                        # Write the successfully retried value back to its original slot immediately.
-                        raw_results[original_idx] = result
-
-                if not next_failed_indices:
-                    logger.info("get_batch succeeded after retransmission.")
-                    break  # All retries in this attempt succeeded.
-
-                logger.error(f"get_batch retry {attempt}/{MAX_RETRIES} failed for {len(next_failed_keys)} keys.")
-
-                # Narrow down to still-failed items for the next retry attempt.
-                current_failed_keys = next_failed_keys
-                current_failed_indices = next_failed_indices
-
-                if attempt < MAX_RETRIES:
-                    time.sleep(RETRY_DELAY_SECONDS)
-            else:
-                # All retries exhausted.
-                raise RuntimeError(
-                    f"get_batch failed for keys {current_failed_keys} after retrying {MAX_RETRIES} times."
-                )
-
-        deserialized_results = [pickle.loads(result) if result != b"" else None for result in raw_results]
-        return deserialized_results, indexes
-
     def clear(self, keys: list[str], custom_backend_meta: list[Any] | None = None) -> None:
         """Deletes multiple keys from MooncakeStore.
 
@@ -443,23 +397,6 @@ def close(self):
             self._store.close()
             self._store = None
 
-    @staticmethod
-    def _preprocess_tensors_for_put(values: list[Tensor]) -> tuple[list[int], list[int], list[Tensor]]:
-        ptr_list: list[int] = []
-        size_list: list[int] = []
-        tensor_list: list[Tensor] = []  # hold reference for the contiguous tensor
-        for t in values:
-            # TODO: support gpu direct rdma and use different data paths.
-            #       For GPU, it's more reasonable to perform data copy since
-            #       The register overhead is much higher than CPU
-            if t.device.type == "cuda":
-                t = t.cpu()
-            t = t.contiguous()
-            tensor_list.append(t)
-            ptr_list.append(t.data_ptr())
-            size_list.append(t.nbytes)
-        return ptr_list, size_list, tensor_list
-
     def _register_all_buffers(self, ptrs, sizes):
         for ptr, size in zip(ptrs, sizes, strict=True):
             self._store.register_buffer(ptr, size)