provide decorator for TransferQueueStorageManagerFactory

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

Author: 0oshowero0

transfer_queue/storage/managers/factory.py

@@ -14,7 +14,6 @@
 from typing import Any
 
 from transfer_queue.storage.managers.base import TransferQueueStorageManager
-from transfer_queue.storage.managers.simple_backend_manager import AsyncSimpleStorageManager
 
 
 class TransferQueueStorageManagerFactory:
@@ -23,13 +22,17 @@ class TransferQueueStorageManagerFactory:
     _registry: dict[str, type[TransferQueueStorageManager]] = {}
 
     @classmethod
-    def register(cls, manager_type: str, manager_cls: type[TransferQueueStorageManager]):
-        if not issubclass(manager_cls, TransferQueueStorageManager):
-            raise TypeError(
-                f"manager_cls {getattr(manager_cls, '__name__', repr(manager_cls))} must be "
-                f"a subclass of TransferQueueStorageManager"
-            )
-        cls._registry[manager_type] = manager_cls
+    def register(cls, manager_type: str):
+        def decorator(manager_cls: type[TransferQueueStorageManager]):
+            if not issubclass(manager_cls, TransferQueueStorageManager):
+                raise TypeError(
+                    f"manager_cls {getattr(manager_cls, '__name__', repr(manager_cls))} must be "
+                    f"a subclass of TransferQueueStorageManager"
+                )
+            cls._registry[manager_type] = manager_cls
+            return manager_cls
+
+        return decorator
 
     @classmethod
     def create(cls, manager_type: str, config: dict[str, Any]) -> TransferQueueStorageManager:
@@ -38,7 +41,3 @@ def create(cls, manager_type: str, config: dict[str, Any]) -> TransferQueueStora
                 f"Unknown manager_type: {manager_type}. Supported managers include: {list(cls._registry.keys())}"
             )
         return cls._registry[manager_type](config)
-
-
-# Register all the StorageManager
-TransferQueueStorageManagerFactory.register("AsyncSimpleStorageManager", AsyncSimpleStorageManager)

transfer_queue/storage/managers/simple_backend_manager.py

@@ -26,6 +26,7 @@
 
 from transfer_queue.metadata import BatchMeta
 from transfer_queue.storage.managers.base import TransferQueueStorageManager
+from transfer_queue.storage.managers.factory import TransferQueueStorageManagerFactory
 from transfer_queue.storage.simple_backend import StorageMetaGroup
 from transfer_queue.utils.utils import limit_pytorch_auto_parallel_threads
 from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, ZMQServerInfo, create_zmq_socket
@@ -34,6 +35,7 @@
 logger.setLevel(os.getenv("TQ_LOGGING_LEVEL", logging.WARNING))
 
 
+@TransferQueueStorageManagerFactory.register("AsyncSimpleStorageManager")
 class AsyncSimpleStorageManager(TransferQueueStorageManager):
     """Asynchronous storage manager that handles multiple storage units.