refactor: simplify StorageManager naming

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

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
-  │   │   ├──simple_backend_manager.py  # AsyncSimpleStorageManager
+  │   │   ├──base.py                    # StorageManager, KVStorageManager, StorageManagerFactory
+  │   │   ├──simple_storage_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, StorageKVClientFactory
   │   │   ├── 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.
 

scripts/put_benchmark.py

@@ -30,7 +30,7 @@
 
 from transfer_queue import TransferQueueClient
 from transfer_queue.controller import TransferQueueController
-from transfer_queue.storage.simple_backend import SimpleStorageUnit
+from transfer_queue.storage.simple_storage import SimpleStorageUnit
 from transfer_queue.utils.common import get_placement_group
 from transfer_queue.utils.zmq_utils import process_zmq_server_info
 

tests/test_async_simple_storage_manager.py

@@ -24,7 +24,7 @@
 
 from transfer_queue.metadata import BatchMeta
 from transfer_queue.storage import AsyncSimpleStorageManager
-from transfer_queue.utils.enum_utils import TransferQueueRole
+from transfer_queue.utils.enum_utils import Role
 from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, ZMQServerInfo
 
 
@@ -35,13 +35,13 @@ async def mock_async_storage_manager():
     # Mock storage unit infos
     storage_unit_infos = {
         "storage_0": ZMQServerInfo(
-            role=TransferQueueRole.STORAGE,
+            role=Role.STORAGE,
             id="storage_0",
             ip="127.0.0.1",
             ports={"put_get_socket": 12345},
         ),
         "storage_1": ZMQServerInfo(
-            role=TransferQueueRole.STORAGE,
+            role=Role.STORAGE,
             id="storage_1",
             ip="127.0.0.1",
             ports={"put_get_socket": 12346},
@@ -50,7 +50,7 @@ async def mock_async_storage_manager():
 
     # Mock controller info
     controller_info = ZMQServerInfo(
-        role=TransferQueueRole.CONTROLLER,
+        role=Role.CONTROLLER,
         id="controller_0",
         ip="127.0.0.1",
         ports={"handshake_socket": 12347, "data_status_update_socket": 12348},
@@ -61,9 +61,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"
-    ) as mock_connect:
+    with patch("transfer_queue.storage.managers.base.StorageManager._connect_to_controller") as mock_connect:
         # Mock the manager without actually connecting
         manager = AsyncSimpleStorageManager.__new__(AsyncSimpleStorageManager)
         manager.storage_manager_id = "test_storage_manager"
@@ -148,7 +146,7 @@ async def test_async_storage_manager_error_handling():
     # Mock storage unit infos
     storage_unit_infos = {
         "storage_0": ZMQServerInfo(
-            role=TransferQueueRole.STORAGE,
+            role=Role.STORAGE,
             id="storage_0",
             ip="127.0.0.1",
             ports={"put_get_socket": 12345},
@@ -157,7 +155,7 @@ async def test_async_storage_manager_error_handling():
 
     # Mock controller info
     controller_info = ZMQServerInfo(
-        role=TransferQueueRole.CONTROLLER,
+        role=Role.CONTROLLER,
         id="controller_0",
         ip="127.0.0.1",
         ports={"handshake_socket": 12346, "data_status_update_socket": 12347},
@@ -242,19 +240,19 @@ async def test_get_data_routes_from_hash():
     """get_data should route using global_idx % num_su (hash routing)."""
     storage_unit_infos = {
         "storage_0": ZMQServerInfo(
-            role=TransferQueueRole.STORAGE,
+            role=Role.STORAGE,
             id="storage_0",
             ip="127.0.0.1",
             ports={"put_get_socket": 19010},
         ),
         "storage_1": ZMQServerInfo(
-            role=TransferQueueRole.STORAGE,
+            role=Role.STORAGE,
             id="storage_1",
             ip="127.0.0.1",
             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
@@ -295,19 +293,19 @@ async def test_clear_data_routes_from_hash():
     """clear_data should route using global_idx % num_su (hash routing)."""
     storage_unit_infos = {
         "storage_0": ZMQServerInfo(
-            role=TransferQueueRole.STORAGE,
+            role=Role.STORAGE,
             id="storage_0",
             ip="127.0.0.1",
             ports={"put_get_socket": 19020},
         ),
         "storage_1": ZMQServerInfo(
-            role=TransferQueueRole.STORAGE,
+            role=Role.STORAGE,
             id="storage_1",
             ip="127.0.0.1",
             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
@@ -346,19 +344,19 @@ async def test_hash_routing_stable_across_batch_sizes():
     """
     storage_unit_infos = {
         "storage_0": ZMQServerInfo(
-            role=TransferQueueRole.STORAGE,
+            role=Role.STORAGE,
             id="storage_0",
             ip="127.0.0.1",
             ports={"put_get_socket": 19030},
         ),
         "storage_1": ZMQServerInfo(
-            role=TransferQueueRole.STORAGE,
+            role=Role.STORAGE,
             id="storage_1",
             ip="127.0.0.1",
             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
@@ -407,19 +405,19 @@ async def test_hash_routing_stable_reversed_order():
     """
     storage_unit_infos = {
         "storage_0": ZMQServerInfo(
-            role=TransferQueueRole.STORAGE,
+            role=Role.STORAGE,
             id="storage_0",
             ip="127.0.0.1",
             ports={"put_get_socket": 19040},
         ),
         "storage_1": ZMQServerInfo(
-            role=TransferQueueRole.STORAGE,
+            role=Role.STORAGE,
             id="storage_1",
             ip="127.0.0.1",
             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_client.py

@@ -24,7 +24,7 @@
 
 from transfer_queue import TransferQueueClient
 from transfer_queue.metadata import BatchMeta
-from transfer_queue.utils.enum_utils import TransferQueueRole
+from transfer_queue.utils.enum_utils import Role
 from transfer_queue.utils.zmq_utils import (
     ZMQMessage,
     ZMQRequestType,
@@ -59,7 +59,7 @@ def __init__(self, controller_id="controller_0"):
         self.request_port = self._bind_to_random_port(self.request_socket)
 
         self.zmq_server_info = ZMQServerInfo(
-            role=TransferQueueRole.CONTROLLER,
+            role=Role.CONTROLLER,
             id=controller_id,
             ip="127.0.0.1",
             ports={
@@ -300,7 +300,7 @@ def __init__(self, storage_id="storage_0"):
         self.data_port = self._bind_to_random_port(self.data_socket)
 
         self.zmq_server_info = ZMQServerInfo(
-            role=TransferQueueRole.STORAGE,
+            role=Role.STORAGE,
             id=storage_id,
             ip="127.0.0.1",
             ports={
@@ -409,7 +409,7 @@ def client_setup(mock_controller, mock_storage):
 
     # Mock the storage manager to avoid handshake issues but mock all data operations
     with patch(
-        "transfer_queue.storage.managers.simple_backend_manager.AsyncSimpleStorageManager._connect_to_controller"
+        "transfer_queue.storage.managers.simple_storage_manager.AsyncSimpleStorageManager._connect_to_controller"
     ):
         config = {
             "controller_info": mock_controller.zmq_server_info,
@@ -502,7 +502,7 @@ def test_single_controller_multiple_storages():
 
         # Mock the storage manager to avoid handshake issues but mock all data operations
         with patch(
-            "transfer_queue.storage.managers.simple_backend_manager.AsyncSimpleStorageManager._connect_to_controller"
+            "transfer_queue.storage.managers.simple_storage_manager.AsyncSimpleStorageManager._connect_to_controller"
         ):
             config = {
                 "controller_info": controller.zmq_server_info,

tests/test_kv_storage_manager.py

@@ -106,7 +106,7 @@ def test_generate_values(test_data):
         assert torch.equal(values[i], expected_values[i])
 
 
-@patch("transfer_queue.storage.managers.base.StorageClientFactory.create")
+@patch("transfer_queue.storage.managers.base.StorageKVClientFactory.create")
 @patch.object(KVStorageManager, "_connect_to_controller", lambda self: None)
 def test_merge_tensors_to_tensordict(mock_create, test_data):
     """Test whether _merge_kv_to_tensordict can correctly reconstruct the TensorDict."""
@@ -268,7 +268,7 @@ def test_data_for_put_data():
     }
 
 
-STORAGE_CLIENT_FACTORY_PATH = "transfer_queue.storage.managers.base.StorageClientFactory"
+STORAGE_CLIENT_FACTORY_PATH = "transfer_queue.storage.managers.base.StorageKVClientFactory"
 
 
 @patch.object(KVStorageManager, "_connect_to_controller", lambda self: None)

tests/test_metadata.py

@@ -678,7 +678,7 @@ class TestStorageUnitDataStrict:
 
     def test_put_data_length_mismatch_raises(self):
         """put_data must raise when global_indexes and field values have different lengths."""
-        from transfer_queue.storage.simple_backend import StorageUnitData
+        from transfer_queue.storage.simple_storage import StorageUnitData
 
         sud = StorageUnitData(storage_size=10)
         # 3 indexes but only 2 values — must raise, not silently drop

tests/test_ray_p2p.py

@@ -23,8 +23,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.base import KVStorageManager, StorageManagerFactory
 from transfer_queue.utils.zmq_utils import ZMQServerInfo
 
 TEST_CONFIGS: list[tuple[tuple[int, int], torch.dtype]] = [
@@ -45,18 +44,18 @@
 
 # Step 1: Mock Controller Role
 try:
-    from transfer_queue.role import TransferQueueRole
+    from transfer_queue.role import Role
 except ImportError:
     from enum import Enum
 
-    class TransferQueueRole(Enum):
+    class Role(Enum):
         CONTROLLER = "controller"
         STORAGE = "storage"
 
 
 def create_mock_controller():
     return ZMQServerInfo(
-        role=TransferQueueRole.CONTROLLER,
+        role=Role.CONTROLLER,
         id="controller_0",
         ip="127.0.0.1",
         ports={
@@ -71,9 +70,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

tests/test_simple_storage_unit.py

@@ -21,7 +21,7 @@
 import torch
 import zmq
 
-from transfer_queue.storage.simple_backend import SimpleStorageUnit
+from transfer_queue.storage.simple_storage import SimpleStorageUnit
 from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType
 
 
@@ -420,7 +420,7 @@ def test_storage_unit_data_direct():
 
 def test_storage_unit_data_capacity_uses_active_keys():
     """Capacity check must use _active_keys, not scan field_data."""
-    from transfer_queue.storage.simple_backend import StorageUnitData
+    from transfer_queue.storage.simple_storage import StorageUnitData
 
     storage = StorageUnitData(storage_size=3)
 

tests/test_storage_client_factory.py

@@ -19,7 +19,7 @@
 import pytest
 import torch
 
-from transfer_queue.storage.clients.factory import StorageClientFactory
+from transfer_queue.storage.clients.base import StorageKVClientFactory
 from transfer_queue.storage.clients.yuanrong_client import YuanrongStorageClient
 
 
@@ -29,12 +29,12 @@ def setUp(self):
 
     @pytest.mark.skipif(find_spec("datasystem") is None, reason="datasystem is not available")
     def test_create_client(self):
-        self.assertIn("YuanrongStorageClient", StorageClientFactory._registry)
-        self.assertIs(StorageClientFactory._registry["YuanrongStorageClient"], YuanrongStorageClient)
-        StorageClientFactory.create("YuanrongStorageClient", self.cfg)
+        self.assertIn("YuanrongStorageClient", StorageKVClientFactory._registry)
+        self.assertIs(StorageKVClientFactory._registry["YuanrongStorageClient"], YuanrongStorageClient)
+        StorageKVClientFactory.create("YuanrongStorageClient", self.cfg)
 
         with self.assertRaises(ValueError) as cm:
-            StorageClientFactory.create("abc", self.cfg)
+            StorageKVClientFactory.create("abc", self.cfg)
         self.assertIn("Unknown StorageClient", str(cm.exception))
 
     @pytest.mark.skipif(
@@ -47,7 +47,7 @@ def test_client_create_empty_tensorlist(self):
         for t in tensors:
             shapes.append(t.shape)
             dtypes.append(t.dtype)
-        client = StorageClientFactory.create("YuanrongStorageClient", self.cfg)
+        client = StorageKVClientFactory.create("YuanrongStorageClient", self.cfg)
 
         empty_tensors = client._create_empty_npu_tensorlist(shapes, dtypes)
         self.assertEqual(len(tensors), len(empty_tensors))

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
         )