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
+  │   │   ├──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, 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

@@ -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_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

@@ -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

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.factory 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
         )
 

transfer_queue/interface.py

@@ -32,7 +32,7 @@
 from transfer_queue.metadata import KVBatchMeta
 from transfer_queue.sampler import *  # noqa: F401
 from transfer_queue.sampler import BaseSampler
-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.logging_utils import get_logger
 from transfer_queue.utils.yuanrong_utils import (