update

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

Author: ji-huazhong

transfer_queue/client.py

@@ -17,7 +17,7 @@
 import logging
 import os
 import threading
-from typing import Any, Optional
+from typing import Any, Callable, Optional
 
 import torch
 import zmq
@@ -36,7 +36,7 @@
     ZMQMessage,
     ZMQRequestType,
     ZMQServerInfo,
-    dynamic_zmq_socket,
+    with_zmq_socket,
 )
 
 logger = logging.getLogger(__name__)
@@ -51,10 +51,10 @@
 TQ_NUM_THREADS = int(os.environ.get("TQ_NUM_THREADS", 8))
 
 # Pre-bound decorator for controller socket operations.
-_controller_socket = dynamic_zmq_socket(
+with_controller_socket = with_zmq_socket(
     "request_handle_socket",
-    owner_id_attr="client_id",
-    server_attr="_controller",
+    get_identity=lambda self: self.client_id,
+    get_peer=lambda self, target: self._controller,
 )
 
 
@@ -104,7 +104,7 @@ def initialize_storage_manager(
         )
 
     # ==================== Basic API ====================
-    @_controller_socket
+    @with_controller_socket
     async def async_get_meta(
         self,
         data_fields: list[str],
@@ -194,7 +194,7 @@ async def async_get_meta(
                 f"{response_msg.body.get('message', 'Unknown error')}"
             )
 
-    @_controller_socket
+    @with_controller_socket
     async def async_set_custom_meta(
         self,
         metadata: BatchMeta,
@@ -494,7 +494,7 @@ async def async_clear_samples(self, metadata: BatchMeta):
         except Exception as e:
             raise RuntimeError(f"Error in clear_samples operation: {str(e)}") from e
 
-    @_controller_socket
+    @with_controller_socket
     async def _clear_meta_in_controller(self, metadata: BatchMeta, socket=None):
         """Clear metadata in the controller.
 
@@ -520,7 +520,7 @@ async def _clear_meta_in_controller(self, metadata: BatchMeta, socket=None):
         if response_msg.request_type != ZMQRequestType.CLEAR_META_RESPONSE:
             raise RuntimeError("Failed to clear samples metadata in controller.")
 
-    @_controller_socket
+    @with_controller_socket
     async def _get_partition_meta(self, partition_id: str, socket=None) -> BatchMeta:
         """Get metadata required for the whole partition from controller.
 
@@ -550,7 +550,7 @@ async def _get_partition_meta(self, partition_id: str, socket=None) -> BatchMeta
 
         return response_msg.body["metadata"]
 
-    @_controller_socket
+    @with_controller_socket
     async def _clear_partition_in_controller(self, partition_id, socket=None):
         """Clear the whole partition in the controller.
 
@@ -577,7 +577,7 @@ async def _clear_partition_in_controller(self, partition_id, socket=None):
             raise RuntimeError(f"Failed to clear partition {partition_id} in controller.")
 
     # ==================== Status Query API ====================
-    @_controller_socket
+    @with_controller_socket
     async def async_get_consumption_status(
         self,
         task_name: str,
@@ -640,7 +640,7 @@ async def async_get_consumption_status(
         except Exception as e:
             raise RuntimeError(f"[{self.client_id}]: Error in get_consumption_status: {str(e)}") from e
 
-    @_controller_socket
+    @with_controller_socket
     async def async_get_production_status(
         self,
         data_fields: list[str],
@@ -772,7 +772,7 @@ async def async_check_production_status(
             return False
         return torch.all(production_status == 1).item()
 
-    @_controller_socket
+    @with_controller_socket
     async def async_reset_consumption(
         self,
         partition_id: str,
@@ -834,7 +834,7 @@ async def async_reset_consumption(
         except Exception as e:
             raise RuntimeError(f"[{self.client_id}]: Error in reset_consumption: {str(e)}") from e
 
-    @_controller_socket
+    @with_controller_socket
     async def async_get_partition_list(
         self,
         socket: Optional[zmq.asyncio.Socket] = None,
@@ -880,7 +880,7 @@ async def async_get_partition_list(
             raise RuntimeError(f"[{self.client_id}]: Error in get_partition_list: {str(e)}") from e
 
     # ==================== KV Interface API ====================
-    @_controller_socket
+    @with_controller_socket
     async def async_kv_retrieve_meta(
         self,
         keys: list[str] | str,
@@ -946,7 +946,7 @@ async def async_kv_retrieve_meta(
         except Exception as e:
             raise RuntimeError(f"[{self.client_id}]: Error in kv_retrieve_keys: {str(e)}") from e
 
-    @_controller_socket
+    @with_controller_socket
     async def async_kv_retrieve_keys(
         self,
         global_indexes: list[int] | int,
@@ -1009,7 +1009,7 @@ async def async_kv_retrieve_keys(
         except Exception as e:
             raise RuntimeError(f"[{self.client_id}]: Error in kv_retrieve_indexes: {str(e)}") from e
 
-    @_controller_socket
+    @with_controller_socket
     async def async_kv_list(
         self,
         partition_id: Optional[str] = None,

transfer_queue/storage/managers/simple_backend_manager.py

@@ -20,7 +20,7 @@
 from collections import defaultdict
 from collections.abc import Mapping
 from operator import itemgetter
-from typing import Any, NamedTuple
+from typing import Any, Callable, NamedTuple
 
 import torch
 import zmq
@@ -34,7 +34,7 @@
     ZMQMessage,
     ZMQRequestType,
     ZMQServerInfo,
-    dynamic_zmq_socket,
+    with_zmq_socket,
 )
 
 logger = logging.getLogger(__name__)
@@ -49,11 +49,11 @@
 TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT = int(os.environ.get("TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT", 200))  # seconds
 
 # Pre-bound decorator for storage-unit socket operations.
-_storage_unit_socket = dynamic_zmq_socket(
+with_storage_unit_socket = with_zmq_socket(
     "put_get_socket",
-    owner_id_attr="storage_manager_id",
-    server_attr="storage_unit_infos",
-    target_kwarg="target_storage_unit",
+    get_identity=lambda self: self.storage_manager_id,
+    get_peer=lambda self, target: self.storage_unit_infos[target],
+    resolve_target=lambda args, kwargs: kwargs.get("target_storage_unit"),
     timeout=TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT,
 )
 
@@ -220,7 +220,7 @@ def _select_by_positions(field_data, positions: list[int]):
             return field_data[positions]
 
     async def put_data(
-        self, data: TensorDict, metadata: BatchMeta, data_parser: Optional[Callable[[Any], Any]] = None
+        self, data: TensorDict, metadata: BatchMeta, data_parser: Callable[[Any], Any] | None = None
     ) -> None:
         """
         Send data to remote StorageUnit based on metadata.
@@ -281,13 +281,13 @@ async def put_data(
             field_schema,
         )
 
-    @_storage_unit_socket
+    @with_storage_unit_socket
     async def _put_to_single_storage_unit(
         self,
         global_indexes: list[int],
         storage_data: dict[str, Any],
         target_storage_unit: str,
-        data_parser: Optional[Callable[[Any], Any]] = None,
+        data_parser: Callable[[Any], Any] | None = None,
         socket: zmq.Socket = None,
     ):
         """
@@ -417,7 +417,7 @@ async def get_data(self, metadata: BatchMeta) -> TensorDict:
 
         return TensorDict(tensor_data, batch_size=len(metadata))
 
-    @_storage_unit_socket
+    @with_storage_unit_socket
     async def _get_from_single_storage_unit(
         self,
         global_indexes: list[int],
@@ -489,7 +489,7 @@ async def clear_data(self, metadata: BatchMeta) -> None:
             if isinstance(result, Exception):
                 logger.error(f"[{self.storage_manager_id}]: Error in clear operation task {i}: {result}")
 
-    @_storage_unit_socket
+    @with_storage_unit_socket
     async def _clear_single_storage_unit(self, global_indexes, target_storage_unit=None, socket=None):
         try:
             request_msg = ZMQMessage.create(

transfer_queue/utils/zmq_utils.py

@@ -17,7 +17,6 @@
 import os
 import socket
 import time
-from collections.abc import Mapping
 from dataclasses import dataclass
 from functools import wraps
 from typing import Any, Callable, Optional, TypeAlias
@@ -304,12 +303,12 @@ def create_zmq_socket(
     return socket
 
 
-def dynamic_zmq_socket(
+def with_zmq_socket(
     socket_name: str,
     *,
-    owner_id_attr: str,
-    server_attr: str,
-    target_kwarg: Optional[str] = None,
+    get_identity: Callable[[Any], str],
+    get_peer: Callable[[Any, Optional[str]], ZMQServerInfo],
+    resolve_target: Optional[Callable[[tuple, dict], Optional[str]]] = None,
     timeout: Optional[int] = None,
 ):
     """Create a reusable async decorator for request sockets.
@@ -320,50 +319,32 @@ def dynamic_zmq_socket(
 
     Args:
         socket_name: Socket port key in ``ZMQServerInfo.ports``.
-        owner_id_attr: Attribute name on ``self`` used in identity/log prefix
-            (e.g., ``client_id`` or ``storage_manager_id``).
-        server_attr: Attribute name on ``self`` that stores server info.
-            - ``ZMQServerInfo`` for single-target calls.
-            - ``Mapping[str, ZMQServerInfo]`` for multi-target calls.
-        target_kwarg: Optional kwarg name that provides target server id when
-            ``server_attr`` is a mapping.
+        get_identity: Callable that extracts owner identity from ``self``.
+            Example: ``lambda self: self.client_id``
+        get_peer: Callable that returns ``ZMQServerInfo`` for the target.
+            For single-target scenarios, ignore the target parameter.
+            Example: ``lambda self, target: self.server_info``
+            Example: ``lambda self, target: self.storage_unit_infos[target]``
+        resolve_target: Optional callable that extracts target identifier from
+            function arguments. Receives (args, kwargs) and returns target name.
+            Example: ``lambda args, kwargs: kwargs.get("target_storage_unit")``
         timeout: Optional timeout (seconds) for both send/recv operations.
     """
 
     def decorator(func: Callable):
         @wraps(func)
         async def wrapper(self, *args, **kwargs):
-            owner_id = getattr(self, owner_id_attr, None)
+            owner_id = get_identity(self)
             if owner_id is None:
-                raise RuntimeError(f"Missing owner id attribute: {owner_id_attr}")
-
-            server_obj = getattr(self, server_attr, None)
-            if server_obj is None:
-                raise RuntimeError(f"Missing server registry attribute: {server_attr}")
+                raise RuntimeError("get_identity returned None")
 
             target_name: Optional[str] = None
-            if target_kwarg is not None:
-                target_name = kwargs.get(target_kwarg)
-                if target_name is None:
-                    for arg in args:
-                        if isinstance(arg, str):
-                            target_name = arg
-                            break
-
-            if isinstance(server_obj, ZMQServerInfo):
-                if target_name is not None and target_name != server_obj.id:
-                    raise RuntimeError(
-                        f"Target mismatch: target '{target_name}' does not match registered server '{server_obj.id}'"
-                    )
-                server_info = server_obj
-            elif isinstance(server_obj, Mapping):
-                if target_name is None:
-                    raise RuntimeError(f"Missing target server identifier via '{target_kwarg}'")
-                server_info = server_obj.get(target_name)
-                if server_info is None:
-                    raise RuntimeError(f"Server '{target_name}' not found in registered servers")
-            else:
-                raise RuntimeError(f"Unsupported server registry type for '{server_attr}': {type(server_obj).__name__}")
+            if resolve_target is not None:
+                target_name = resolve_target(args, kwargs)
+
+            server_info = get_peer(self, target_name)
+            if server_info is None:
+                raise RuntimeError(f"get_peer returned None for target '{target_name}'")
 
             port = server_info.ports.get(socket_name)
             if port is None: