refactor: simplify StorageManager naming

Rename internal storage manager classes by removing TransferQueue prefix:
- TransferQueueStorageManager → StorageManager
- TransferQueueStorageManagerFactory → StorageManagerFactory

These classes are internal components not exposed as public API,
so the shorter names improve readability without causing conflicts.

Author: ji-huazhong

README.md

@@ -50,7 +50,7 @@ TransferQueue offers **fine-grained, sub-sample-level** data management and **lo
 
 ### Control Plane: Panoramic Data Management
 
-In the control plane, `TransferQueueController` tracks the **production status** and **consumption status** of each training sample as metadata. Once all required data fields are ready (i.e., written to the `TransferQueueStorageManager`), the data sample can be consumed by downstream tasks.
+In the control plane, `TransferQueueController` tracks the **production status** and **consumption status** of each training sample as metadata. Once all required data fields are ready (i.e., written to the `StorageManager`), the data sample can be consumed by downstream tasks.
 
 We also track the consumption history for each computational task (e.g., `generate_sequences`, `compute_log_prob`, etc.). Therefore, even when different computational tasks require the same data field, they can consume the data independently without interfering with each other.
 
@@ -66,7 +66,7 @@ To make the data retrieval process more customizable, we provide a `Sampler` cla
 
 In the data plane, we utilize a pluggable design, enabling TransferQueue to integrate with different storage backends based on user requirements.
 
-Specifically, we provide a `TransferQueueStorageManager` abstraction class that defines the core APIs as follows:
+Specifically, we provide a `StorageManager` abstraction class that defines the core APIs as follows:
 
 - `async def put_data(self, data: TensorDict, metadata: BatchMeta) -> None`
 - `async def get_data(self, metadata: BatchMeta) -> TensorDict`
@@ -298,21 +298,19 @@ The data plane is organized as follows:
   │   │── simple_backend.py             # Default distributed storage backend (SimpleStorageUnit) by TQ 
   │   ├── managers/                     # Managers are upper level interfaces that encapsulate the interaction logic with TQ system.
   │   │   ├── __init__.py
-  │   │   ├──base.py                    # TransferQueueStorageManager, KVStorageManager
+  │   │   ├──base.py                    # StorageManager, KVStorageManager, StorageManagerFactory
   │   │   ├──simple_backend_manager.py  # AsyncSimpleStorageManager
   │   │   ├──yuanrong_manager.py        # YuanrongStorageManager
-  │   │   ├──mooncake_manager.py        # MooncakeStorageManager
-  │   │   └──factory.py                 # TransferQueueStorageManagerFactory
+  │   │   └──mooncake_manager.py        # MooncakeStorageManager
   │   └── clients/                      # Clients are lower level interfaces that directly manipulate the target storage backend.
   │   │   ├── __init__.py
-  │   │   ├── base.py                   # TransferQueueStorageKVClient
+  │   │   ├── base.py                   # StorageKVClient, StorageClientFactory
   │   │   ├── yuanrong_client.py        # YuanrongStorageClient
   │   │   ├── mooncake_client.py        # MooncakeStorageClient
-  │   │   ├── ray_storage_client.py     # RayStorageClient
-  │   │   └── factory.py                # TransferQueueStorageClientFactory
+  │   │   └── ray_storage_client.py     # RayStorageClient
 ```
 
-To integrate TransferQueue with a custom storage backend, start by implementing a subclass that inherits from `TransferQueueStorageManager`. This subclass acts as an adapter between the TransferQueue system and the target storage backend. For KV-based storage backends, you can simply inherit from `KVStorageManager`, which can serve as the general manager for all KV-based backends.
+To integrate TransferQueue with a custom storage backend, start by implementing a subclass that inherits from `StorageManager`. This subclass acts as an adapter between the TransferQueue system and the target storage backend. For KV-based storage backends, you can simply inherit from `KVStorageManager`, which can serve as the general manager for all KV-based backends.
 
 Distributed storage backends often come with their own native clients serving as the interface of the storage system. In such cases, a low-level adapter for this client can be written, following the examples provided in the `storage/clients` directory.
 

tests/test_async_simple_storage_manager.py

@@ -62,7 +62,7 @@ async def mock_async_storage_manager():
 
     # Mock the handshake process entirely to avoid ZMQ complexity
     with patch(
-        "transfer_queue.storage.managers.base.TransferQueueStorageManager._connect_to_controller"
+        "transfer_queue.storage.managers.base.StorageManager._connect_to_controller"
     ) as mock_connect:
         # Mock the manager without actually connecting
         manager = AsyncSimpleStorageManager.__new__(AsyncSimpleStorageManager)
@@ -254,7 +254,7 @@ async def test_get_data_routes_from_hash():
             ports={"put_get_socket": 19011},
         ),
     }
-    with patch("transfer_queue.storage.managers.base.TransferQueueStorageManager._connect_to_controller"):
+    with patch("transfer_queue.storage.managers.base.StorageManager._connect_to_controller"):
         manager = AsyncSimpleStorageManager.__new__(AsyncSimpleStorageManager)
         manager.storage_manager_id = "test_get"
         manager.storage_unit_infos = storage_unit_infos
@@ -307,7 +307,7 @@ async def test_clear_data_routes_from_hash():
             ports={"put_get_socket": 19021},
         ),
     }
-    with patch("transfer_queue.storage.managers.base.TransferQueueStorageManager._connect_to_controller"):
+    with patch("transfer_queue.storage.managers.base.StorageManager._connect_to_controller"):
         manager = AsyncSimpleStorageManager.__new__(AsyncSimpleStorageManager)
         manager.storage_manager_id = "test_clear"
         manager.storage_unit_infos = storage_unit_infos
@@ -358,7 +358,7 @@ async def test_hash_routing_stable_across_batch_sizes():
             ports={"put_get_socket": 19031},
         ),
     }
-    with patch("transfer_queue.storage.managers.base.TransferQueueStorageManager._connect_to_controller"):
+    with patch("transfer_queue.storage.managers.base.StorageManager._connect_to_controller"):
         manager = AsyncSimpleStorageManager.__new__(AsyncSimpleStorageManager)
         manager.storage_manager_id = "test_hash_batch"
         manager.storage_unit_infos = storage_unit_infos
@@ -419,7 +419,7 @@ async def test_hash_routing_stable_reversed_order():
             ports={"put_get_socket": 19041},
         ),
     }
-    with patch("transfer_queue.storage.managers.base.TransferQueueStorageManager._connect_to_controller"):
+    with patch("transfer_queue.storage.managers.base.StorageManager._connect_to_controller"):
         manager = AsyncSimpleStorageManager.__new__(AsyncSimpleStorageManager)
         manager.storage_manager_id = "test_hash_order"
         manager.storage_unit_infos = storage_unit_infos

tests/test_ray_p2p.py

@@ -24,7 +24,7 @@
 from transfer_queue.client import TransferQueueClient
 from transfer_queue.metadata import BatchMeta
 from transfer_queue.storage.managers.base import KVStorageManager
-from transfer_queue.storage.managers.factory import TransferQueueStorageManagerFactory
+from transfer_queue.storage.managers.factory import StorageManagerFactory
 from transfer_queue.utils.zmq_utils import ZMQServerInfo
 
 TEST_CONFIGS: list[tuple[tuple[int, int], torch.dtype]] = [
@@ -71,9 +71,9 @@ def create_mock_controller():
 def ensure_mock_storage_manager_registered():
     """Ensure MockKVStorageManager is registered in current process."""
 
-    if "KV_MOCK" not in TransferQueueStorageManagerFactory._registry:
+    if "KV_MOCK" not in StorageManagerFactory._registry:
 
-        @TransferQueueStorageManagerFactory.register("KV_MOCK")
+        @StorageManagerFactory.register("KV_MOCK")
         class MockKVStorageManager(KVStorageManager):
             def _connect_to_controller(self):
                 pass

transfer_queue/client.py

@@ -28,7 +28,7 @@
     BatchMeta,
 )
 from transfer_queue.storage import (
-    TransferQueueStorageManagerFactory,
+    StorageManagerFactory,
 )
 from transfer_queue.utils.common import limit_pytorch_auto_parallel_threads
 from transfer_queue.utils.logging_utils import get_logger
@@ -92,7 +92,7 @@ def initialize_storage_manager(
                     - zmq_info: ZMQ server information about the storage units
 
         """
-        self.storage_manager = TransferQueueStorageManagerFactory.create(
+        self.storage_manager = StorageManagerFactory.create(
             manager_type, controller_info=self._controller, config=config
         )
 

transfer_queue/storage/__init__.py

@@ -17,17 +17,17 @@
     AsyncSimpleStorageManager,
     MooncakeStorageManager,
     RayStorageManager,
-    TransferQueueStorageManager,
-    TransferQueueStorageManagerFactory,
+    StorageManager,
+    StorageManagerFactory,
     YuanrongStorageManager,
 )
 from .simple_backend import SimpleStorageUnit, StorageUnitData
 
 __all__ = [
     "SimpleStorageUnit",
     "StorageUnitData",
-    "TransferQueueStorageManager",
-    "TransferQueueStorageManagerFactory",
+    "StorageManager",
+    "StorageManagerFactory",
     "AsyncSimpleStorageManager",
     "MooncakeStorageManager",
     "YuanrongStorageManager",

transfer_queue/storage/clients/__init__.py

@@ -14,14 +14,14 @@
 # limitations under the License.
 
 # This module is currently empty but reserved for future client implementations
-from .base import TransferQueueStorageKVClient
+from .base import StorageKVClient
 from .factory import StorageClientFactory
 from .mooncake_client import MooncakeStoreClient
 from .ray_storage_client import RayStorageClient
 from .yuanrong_client import YuanrongStorageClient
 
 __all__ = [
-    "TransferQueueStorageKVClient",
+    "StorageKVClient",
     "StorageClientFactory",
     "RayStorageClient",
     "MooncakeStoreClient",

transfer_queue/storage/clients/base.py

@@ -17,7 +17,7 @@
 from typing import Any, Optional
 
 
-class TransferQueueStorageKVClient(ABC):
+class StorageKVClient(ABC):
     """
     Abstract base class for storage client.
     Subclasses must implement the core methods: put, get, and clear.
@@ -68,3 +68,44 @@ def get(self, keys: list[str], shapes=None, dtypes=None, custom_backend_meta=Non
     def clear(self, keys: list[str], custom_backend_meta=None) -> None:
         """Clear key-value pairs in the storage backend."""
         raise NotImplementedError("Subclasses must implement clear")
+
+
+class StorageClientFactory:
+    """
+    Factory class for creating storage client instances.
+    Uses a decorator-based registration mechanism to map client names to classes.
+    """
+
+    # Class variable: maps client names to their corresponding classes
+    _registry: dict[str, type[StorageKVClient]] = {}
+
+    @classmethod
+    def register(cls, client_type: str):
+        """
+        Decorator to register a concrete client class with the factory.
+        Args:
+            client_type (str): The name used to identify the client
+        Returns:
+            Callable: The decorator function that returns the original class
+        """
+
+        def decorator(client_class: type[StorageKVClient]) -> type[StorageKVClient]:
+            cls._registry[client_type] = client_class
+            return client_class
+
+        return decorator
+
+    @classmethod
+    def create(cls, client_type: str, config: dict) -> StorageKVClient:
+        """
+        Create and return an instance of the storage client by name.
+        Args:
+            client_type (str): The registered name of the client
+        Returns:
+            StorageClientFactory: An instance of the requested client
+        Raises:
+            ValueError: If no client is registered with the given name
+        """
+        if client_type not in cls._registry:
+            raise ValueError(f"Unknown StorageClient: {client_type}")
+        return cls._registry[client_type](config)

transfer_queue/storage/clients/factory.py

@@ -20,8 +20,7 @@
 import torch
 from torch import Tensor
 
-from transfer_queue.storage.clients.base import TransferQueueStorageKVClient
-from transfer_queue.storage.clients.factory import StorageClientFactory
+from transfer_queue.storage.clients.base import StorageKVClient, StorageClientFactory
 from transfer_queue.utils.logging_utils import get_logger
 from transfer_queue.utils.tensor_utils import allocate_empty_tensors, get_nbytes, merge_contiguous_memory
 
@@ -39,7 +38,7 @@
 
 
 @StorageClientFactory.register("MooncakeStoreClient")
-class MooncakeStoreClient(TransferQueueStorageKVClient):
+class MooncakeStoreClient(StorageKVClient):
     """
     Storage client for MooncakeStore.
     """

transfer_queue/storage/clients/mooncake_client.py

@@ -19,7 +19,7 @@
 import ray
 import torch
 
-from transfer_queue.storage.clients.base import TransferQueueStorageKVClient
+from transfer_queue.storage.clients.base import StorageKVClient
 from transfer_queue.storage.clients.factory import StorageClientFactory
 
 
@@ -47,7 +47,7 @@ def clear_obj_ref(self, keys: list[str]):
 
 
 @StorageClientFactory.register("RayStorageClient")
-class RayStorageClient(TransferQueueStorageKVClient):
+class RayStorageClient(StorageKVClient):
     """
     Storage client for Ray RDT.
     """