more

Signed-off-by: ji-huazhong <hzji210@gmail.com>

Author: ji-huazhong

recipe/simple_use_case/single_controller_demo.py

@@ -15,7 +15,6 @@
 
 import argparse
 import asyncio
-import logging
 import os
 import random
 import time

transfer_queue/client.py

@@ -22,14 +22,9 @@
 import zmq
 import zmq.asyncio
 from tensordict import TensorDict
-from torch import Tensor
 
-from transfer_queue.metadata import (
-    BatchMeta,
-)
-from transfer_queue.storage import (
-    StorageManagerFactory,
-)
+from transfer_queue.metadata import BatchMeta
+from transfer_queue.storage import StorageManagerFactory
 from transfer_queue.utils.common import limit_pytorch_auto_parallel_threads
 from transfer_queue.utils.logging_utils import get_logger
 from transfer_queue.utils.zmq_utils import (
@@ -576,7 +571,7 @@ async def async_get_consumption_status(
         task_name: str,
         partition_id: str,
         socket: zmq.asyncio.Socket | None = None,
-    ) -> tuple[Tensor | None, Tensor | None]:
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
         """Get consumption status for current partition in a specific task.
 
         Args:
@@ -639,7 +634,7 @@ async def async_get_production_status(
         data_fields: list[str],
         partition_id: str,
         socket: zmq.asyncio.Socket | None = None,
-    ) -> tuple[Tensor | None, Tensor | None]:
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
         """Get production status for specific data fields and partition.
 
         Args:
@@ -881,21 +876,20 @@ async def async_kv_retrieve_meta(
         create: bool = False,
         socket: zmq.asyncio.Socket | None = None,
     ) -> BatchMeta:
-        """Asynchronously retrieve BatchMeta from the controller using user-specified keys.
+        """Asynchronously retrieve BatchMeta by user-defined keys.
+
+        Retrieves metadata for given keys from a specified partition.
+        If keys do not exist and `create=True`, they will be automatically registered.
 
         Args:
-            keys: List of keys to retrieve from the controller
+            keys: List of keys to retrieve.
             partition_id: The ID of the logical partition to search for keys.
-            create: Whether to register new keys if not found.
-            socket: ZMQ socket (injected by decorator)
+            create:  If True, automatically create entries for missing keys.
+            socket: ZMQ socket injected by @with_controller_socket.
 
         Returns:
-            metadata: BatchMeta of the corresponding keys
-
-        Raises:
-            TypeError: If `keys` is not a list of string or a string
+            BatchMeta: Metadata for the requested keys.
         """
-
         if isinstance(keys, str):
             keys = [keys]
         elif isinstance(keys, list):
@@ -919,25 +913,23 @@ async def async_kv_retrieve_meta(
         )
 
         try:
-            assert socket is not None
+            assert socket is not None, "Socket must be initialized before use"
             await socket.send_multipart(request_msg.serialize())
             response_serialized = await socket.recv_multipart(copy=False)
             response_msg = ZMQMessage.deserialize(response_serialized)
             logger.debug(
-                f"[{self.client_id}]: Client get kv_retrieve_keys response: {response_msg} "
+                f"[{self.client_id}] Received KV_RETRIEVE_META response: {response_msg} "
                 f"from controller {self._controller.id}"
             )
 
             if response_msg.request_type == ZMQRequestType.KV_RETRIEVE_META_RESPONSE:
-                metadata = response_msg.body.get("metadata", BatchMeta.empty())
-                return metadata
-            else:
-                raise RuntimeError(
-                    f"[{self.client_id}]: Failed to retrieve keys from controller {self._controller.id}: "
-                    f"{response_msg.body.get('message', 'Unknown error')}"
-                )
+                return response_msg.body.get("metadata", BatchMeta.empty())
+
+            raise RuntimeError(
+                f"[{self.client_id}] Failed to retrieve metadata {response_msg.body.get('message', 'Unknown error')}"
+            )
         except Exception as e:
-            raise RuntimeError(f"[{self.client_id}]: Error in kv_retrieve_keys: {str(e)}") from e
+            raise RuntimeError(f"[{self.client_id}] Failed in async_kv_retrieve_meta: {e}") from e
 
     @with_controller_socket
     async def async_kv_retrieve_keys(
@@ -1356,7 +1348,7 @@ def get_consumption_status(
         self,
         task_name: str,
         partition_id: str,
-    ) -> tuple[Tensor | None, Tensor | None]:
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
         """Synchronously get consumption status for a specific task and partition.
 
         Args:
@@ -1384,7 +1376,7 @@ def get_production_status(
         self,
         data_fields: list[str],
         partition_id: str,
-    ) -> tuple[Tensor | None, Tensor | None]:
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
         """Synchronously get production status for specific data fields and partition.
 
         Args:
@@ -1501,20 +1493,22 @@ def kv_retrieve_meta(
         partition_id: str,
         create: bool = False,
     ) -> BatchMeta:
-        """Synchronously retrieve BatchMeta from the controller using user-specified keys.
+        """Synchronously retrieve BatchMeta by user-defined keys.
+
+        Retrieves metadata for given keys from a specified partition.
+        If keys do not exist and `create=True`, they will be automatically registered.
 
         Args:
-            keys: List of keys to retrieve from the controller
-            partition_id: The ID of the logical partition to search for keys.
-            create: Whether to register new keys if not found.
+            keys: List of keys to retrieve from the controller.
+            partition_id: Logical partition to query.
+            create: If True, automatically create entries for non-existent keys.
 
         Returns:
-            metadata: BatchMeta of the corresponding keys
+            BatchMeta: Metadata for the requested keys.
 
         Raises:
             TypeError: If `keys` is not a list of string or a string
         """
-
         return self._kv_retrieve_meta(keys=keys, partition_id=partition_id, create=create)
 
     def kv_retrieve_keys(

transfer_queue/controller.py

@@ -45,7 +45,7 @@
     create_zmq_socket,
     format_zmq_address,
     get_free_port,
-    get_node_ip_address_raw,
+    get_node_ip_address,
 )
 
 logger = get_logger(__name__)
@@ -1577,17 +1577,17 @@ def kv_retrieve_meta(
             metadata: BatchMeta of the requested keys
         """
 
-        logger.debug(f"[{self.controller_id}]: Retrieve keys {keys} in partition {partition_id}")
+        logger.debug(f"[{self.controller_id}] Retrieve keys {keys} in partition {partition_id}")
 
+        # Ensure partition exists
         partition = self._get_partition(partition_id)
-
         if partition is None:
             if not create:
-                logger.warning(f"Partition {partition_id} were not found in controller!")
+                logger.warning(f"Partition {partition_id} not found!")
                 return BatchMeta.empty()
-            else:
-                self.create_partition(partition_id)
-                partition = self._get_partition(partition_id)
+
+            self.create_partition(partition_id)
+            partition = self._get_partition(partition_id)
 
         assert partition is not None
         global_indexes = partition.kv_retrieve_indexes(keys)
@@ -1631,9 +1631,7 @@ def kv_retrieve_meta(
             if col_idx < len(col_mask) and col_mask[col_idx]:
                 data_fields.append(field_name)
 
-        metadata = self.generate_batch_meta(partition_id, verified_global_indexes, data_fields, mode="force_fetch")
-
-        return metadata
+        return self.generate_batch_meta(partition_id, verified_global_indexes, data_fields, mode="force_fetch")
 
     def kv_retrieve_keys(
         self,
@@ -1674,7 +1672,7 @@ def kv_retrieve_keys(
     def _init_zmq_socket(self):
         """Initialize ZMQ sockets for communication."""
         self.zmq_context = zmq.Context()
-        self._node_ip = get_node_ip_address_raw()
+        self._node_ip = get_node_ip_address()
 
         while True:
             try:

transfer_queue/interface.py

@@ -474,37 +474,32 @@ def kv_batch_put(
     tags: list[dict[str, Any]] | None = None,
     data_parser: Callable[[Any], Any] | None = None,
 ) -> KVBatchMeta:
-    """Put multiple key-value pairs to TransferQueue in batch.
+    """Batch put multiple key-value pairs into the TransferQueue.
 
-    This method stores multiple key-value pairs in a single operation, which is more
-    efficient than calling kv_put multiple times.
+    This method stores multiple key-value entries in a single operation,
+    which is significantly more efficient than repeated calls to ``kv_put``.
 
     Args:
-        keys: List of user-specified keys for the data
-        partition_id: Logical partition to store the data in
-        fields: TensorDict containing data for all keys. Must have batch_size == len(keys).
-                If not provided, will only update the newly given tags to the keys.
-        tags: List of metadata tags, one for each key
-        data_parser: Optional callable to parse reference data (e.g., URLs) into real
-                     content. The input is a slice of the `fields` parameter passed to
-                     kv_put / kv_batch_put, in plain dict format (not TensorDict),
-                     mapping field_name -> batched values. For a regular tensor column
-                     the value is a batched tensor; for nested tensors (jagged or
-                     strided) and NonTensorStack columns the values are extracted into
-                     a list. It must modify values in-place based on the original keys;
-                     do not add or remove keys. The number of elements per column must
-                     also remain unchanged. Do not change the inner order of values
-                     within each column. Only supported by SimpleStorage.
+        keys: List of user-defined unique keys for the data entries.
+        partition_id: Logical partition where the data will be stored.
+        fields: TensorDict containing batched data for all keys. Must have ``batch_size == len(keys)``.
+            If not provided, only the associated tags will be updated.
+        tags: List of metadata dictionaries, one per key.  Length must match the number of keys.
+        data_parser: Optional callable to parse raw reference data (e.g., URLs) into real content
+            before storage. The input is a plain dict (not TensorDict) mapping field names to
+            batched values. The parser  **must modify data in-place** without adding/removing
+            keys or changing element counts/order. Only supported by ``SimpleStorage`` backend.
 
     Returns:
-        KVBatchMeta: Metadata containing the keys, tags, partition_id, and fields.
-                     The `fields` attribute includes all fields stored for these samples,
-                     including any new fields written by this put operation.
+        KVBatchMeta: Metadata object containing stored keys, tags, partition ID,
+            and field information. The ``fields`` attribute includes all
+            persisted fields for the written samples.
 
     Raises:
-        ValueError: If neither `fields` nor `tags` is provided
-        ValueError: If length of `keys` doesn't match length of `tags` or the batch_size of `fields` TensorDict
-        RuntimeError: If retrieved BatchMeta size doesn't match length of `keys`
+        ValueError: When both ``fields`` and ``tags`` are empty.
+        ValueError: When ``fields`` batch size mismatches key count.
+        ValueError: When ``tags`` length mismatches key count.
+        RuntimeError: When retrieved metadata size mismatches input key count.
 
     Example:
         >>> import transfer_queue as tq
@@ -517,49 +512,37 @@ def kv_batch_put(
         ... }, batch_size=3)
         >>> tags = [{"score": 0.9}, {"score": 0.85}, {"score": 0.95}]
         >>> meta = tq.kv_batch_put(keys=keys, partition_id="train", fields=fields, tags=tags)
-        >>> print(meta.fields)  # ['input_ids', 'attention_mask']
+        >>> print(meta.fields)
     """
+    num_keys = len(keys)
 
     if fields is None and tags is None:
         raise ValueError("Please provide at least one parameter of fields or tag.")
 
-    if fields is not None and fields.batch_size[0] != len(keys):
-        raise ValueError(
-            f"`keys` with length {len(keys)} does not match the `fields` TensorDict with "
-            f"batch_size {fields.batch_size[0]}"
-        )
+    if fields is not None and fields.batch_size[0] != num_keys:
+        raise ValueError(f"Length of `keys` ({num_keys}) does not match `fields` batch size ({fields.batch_size[0]}).")
 
     tq_client = _maybe_create_tq_client()
-
-    # 1. translate user-specified key to BatchMeta
     batch_meta = tq_client.kv_retrieve_meta(keys=keys, partition_id=partition_id, create=True)
 
-    if batch_meta.size != len(keys):
-        raise RuntimeError(
-            f"Retrieved BatchMeta size {batch_meta.size} does not match with input `keys` size {len(keys)}!"
-        )
+    if batch_meta.size != num_keys:
+        raise RuntimeError(f"Retrieved BatchMeta size {batch_meta.size} does not match input `keys` size {num_keys}.")
 
-    # 2. register the user-specified tags to BatchMeta
     if tags is not None:
-        if len(tags) != len(keys):
-            raise ValueError(f"keys with length {len(keys)} does not match length of tags {len(tags)}")
+        if len(tags) != num_keys:
+            raise ValueError(f"Length of `keys` ({num_keys}) does not match length of `tags` ({len(tags)}).")
         batch_meta.update_custom_meta(tags)
 
-    # 3. put data
     if fields is not None:
-        # After put, batch_meta.field_names will include the new fields written by user
         batch_meta = tq_client.put(fields, batch_meta, data_parser=data_parser)
-    else:
-        # Directly update custom_meta (tags) to controller
+    else:  # tags is not None
         tq_client.set_custom_meta(batch_meta)
 
-    fields_to_return = batch_meta.field_names
-
     return KVBatchMeta(
         keys=keys,
         tags=batch_meta.custom_meta,
         partition_id=partition_id,
-        fields=fields_to_return,
+        fields=batch_meta.field_names,
         extra_info=batch_meta.extra_info,
     )
 

transfer_queue/storage/clients/mooncake_client.py

@@ -63,9 +63,9 @@ def __init__(self, config: dict[str, Any]):
             self.device_name = ""
 
         if self.local_hostname is None or self.local_hostname == "":
-            from transfer_queue.utils.zmq_utils import get_node_ip_address_raw
+            from transfer_queue.utils.zmq_utils import get_node_ip_address
 
-            ip = get_node_ip_address_raw()
+            ip = get_node_ip_address()
             logger.info(f"Try to use Ray IP ({ip}) as local hostname for MooncakeStore.")
             self.local_hostname = ip
 

transfer_queue/storage/managers/base.py

@@ -17,7 +17,6 @@
 import itertools
 import os
 import time
-import warnings
 import weakref
 from abc import ABC, abstractmethod
 from concurrent.futures import ThreadPoolExecutor
@@ -388,32 +387,9 @@ def decorator(manager_cls: type[StorageManager]):
     @classmethod
     def create(cls, manager_type: str, controller_info: ZMQServerInfo, config: dict[str, Any]) -> StorageManager:
         """Create and return a StorageManager instance."""
-        if manager_type not in cls._registry:
-            if manager_type == "AsyncSimpleStorageManager":
-                warnings.warn(
-                    f"The manager_type {manager_type} will be deprecated in 0.1.7, please use SimpleStorage instead.",
-                    category=DeprecationWarning,
-                    stacklevel=2,
-                )
-                manager_type = "SimpleStorage"
-            elif manager_type == "MooncakeStorageManager":
-                warnings.warn(
-                    f"The manager_type {manager_type} will be deprecated in 0.1.7, please use MooncakeStore instead.",
-                    category=DeprecationWarning,
-                    stacklevel=2,
-                )
-                manager_type = "MooncakeStore"
-            elif manager_type == "YuanrongStorageManager":
-                warnings.warn(
-                    f"The manager_type {manager_type} will be deprecated in 0.1.7, please use Yuanrong instead.",
-                    category=DeprecationWarning,
-                    stacklevel=2,
-                )
-                manager_type = "Yuanrong"
-            else:
-                raise ValueError(
-                    f"Unknown manager_type: {manager_type}. Supported managers include: {list(cls._registry.keys())}"
-                )
+        assert manager_type in cls._registry, (
+            f"Unknown manager_type: {manager_type}. Supported managers include: {list(cls._registry.keys())}"
+        )
         return cls._registry[manager_type](controller_info, config)
 
 

transfer_queue/storage/simple_storage.py

@@ -34,7 +34,7 @@
     create_zmq_socket,
     format_zmq_address,
     get_free_port,
-    get_node_ip_address_raw,
+    get_node_ip_address,
 )
 
 logger = get_logger(__name__)
@@ -186,7 +186,7 @@ def _init_zmq_socket(self) -> None:
         - worker_socket (DEALER): Backend socket for worker communication.
         """
         self.zmq_context = zmq.Context()
-        self._node_ip = get_node_ip_address_raw()
+        self._node_ip = get_node_ip_address()
 
         # Frontend: ROUTER for receiving client requests
         self.put_get_socket = create_zmq_socket(self.zmq_context, zmq.ROUTER, self._node_ip)