[client, storage] refactor: unify dynamic ZMQ socket decorator between simple_backend_manager and client (#66)

as title.

---------

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

Author: ji-huazhong

recipe/simple_use_case/single_controller_demo.py

@@ -53,13 +53,24 @@ def compute_loss(data1, _data2):
 
 def compute_reward(response_ids: torch.Tensor) -> TensorDict:
     """Simulate a reward model that scores each token position in the response.
+    Returns a TensorDict with a ``"rm_score"`` field whose shape matches
+    ``response_ids`` (i.e. one scalar per response token).
+    """
+    time.sleep(1)
+    reward = torch.randn_like(response_ids, dtype=torch.float32)
+
+    return TensorDict({"rm_score": reward}, batch_size=response_ids.size(0))
+
+
+def compute_advantage(rewards: torch.Tensor) -> TensorDict:
+    """Simulate the process of compute advantage
 
     Returns a TensorDict with an ``"advantage"`` field whose shape matches
-    ``response_ids`` (i.e. one scalar per response token).
+    ``rewards`` (i.e. one scalar per reward).
     """
     time.sleep(1)
-    advantage = torch.randn_like(response_ids, dtype=torch.float32)
-    return TensorDict({"advantage": advantage}, batch_size=response_ids.size(0))
+    advantage = torch.randn_like(rewards, dtype=torch.float32)
+    return TensorDict({"advantage": advantage}, batch_size=rewards.size(0)) 
 
 
 class TrainingWorker:
@@ -89,7 +100,7 @@ def infer_batch(self, kv_meta: KVBatchMeta) -> KVBatchMeta:
         """Simulate forward-only inference"""
         # 1. Pull data from storage
         data = tq.kv_batch_get_by_meta(meta=kv_meta)
-        logger.info(f"compute_log_prob: got data {data}")
+        logger.info(f"infer_batch: got data {data}")
 
         # 2. Model forward
         output = compute_log_prob(data["prompt_ids"], data["response_ids"])
@@ -494,6 +505,13 @@ def fit(self):
             meta = tq.kv_batch_put(keys=meta.keys, partition_id=meta.partition_id, fields=reward_output)
             logger.info(f"demo reward KVBatchMeta: {meta}")
 
+            # ========================= Compute advantage =========================
+            meta.fields = ["response_ids", "ref_log_prob", "old_log_prob","rm_score"]
+            advantage_data = tq.kv_batch_get_by_meta(meta=meta)
+            advantage_output = compute_advantage(advantage_data["rm_score"])
+            meta = tq.kv_batch_put(keys=meta.keys, partition_id=meta.partition_id, fields=advantage_output)
+            logger.info(f"demo advantage KVBatchMeta: {meta}")
+
             # ========================= Update actor =========================
             meta.fields = [
                 "input_ids",

transfer_queue/client.py

@@ -14,12 +14,9 @@
 # limitations under the License.
 
 import asyncio
-import logging
 import os
 import threading
-from functools import wraps
 from typing import Any, Callable, Optional
-from uuid import uuid4
 
 import torch
 import zmq
@@ -34,25 +31,25 @@
     TransferQueueStorageManagerFactory,
 )
 from transfer_queue.utils.common import limit_pytorch_auto_parallel_threads
+from transfer_queue.utils.logging_utils import get_logger
 from transfer_queue.utils.zmq_utils import (
     ZMQMessage,
     ZMQRequestType,
     ZMQServerInfo,
-    create_zmq_socket,
-    format_zmq_address,
+    with_zmq_socket,
 )
 
-logger = logging.getLogger(__name__)
-logger.setLevel(os.getenv("TQ_LOGGING_LEVEL", logging.WARNING))
-
-# Ensure logger has a handler
-if not logger.hasHandlers():
-    handler = logging.StreamHandler()
-    handler.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(name)s - %(message)s"))
-    logger.addHandler(handler)
+logger = get_logger(__name__)
 
 TQ_NUM_THREADS = int(os.environ.get("TQ_NUM_THREADS", 8))
 
+# Pre-bound decorator for controller socket operations.
+with_controller_socket = with_zmq_socket(
+    "request_handle_socket",
+    get_identity=lambda self: self.client_id,
+    get_peer=lambda self, target: self._controller,
+)
+
 
 class AsyncTransferQueueClient:
     """Asynchronous client for interacting with TransferQueue controller and storage systems.
@@ -99,63 +96,8 @@ def initialize_storage_manager(
             manager_type, controller_info=self._controller, config=config
         )
 
-    # TODO (TQStorage): Provide a general dynamic socket function for both Client & Storage @huazhong.
-    @staticmethod
-    def dynamic_socket(socket_name: str):
-        """Decorator to auto-manage ZMQ sockets for Controller/Storage servers.
-
-        Handles socket lifecycle: create -> connect -> inject -> close.
-
-        Args:
-            socket_name: Port name from server config to use for ZMQ connection (e.g., "data_req_port")
-
-        Decorated Function Requirements:
-            1. Must be an async class method (needs `self`)
-            2. `self` must have:
-               - `_controller`: Server registry
-               - `client_id`: Unique client ID for socket identity
-            3. Receives ZMQ socket via `socket` keyword argument (injected by decorator)
-        """
-
-        def decorator(func: Callable):
-            @wraps(func)
-            async def wrapper(self, *args, **kwargs):
-                server_info = self._controller
-                if not server_info:
-                    raise RuntimeError("No controller registered")
-
-                context = zmq.asyncio.Context()
-                address = format_zmq_address(server_info.ip, server_info.ports.get(socket_name))
-                identity = f"{self.client_id}_to_{server_info.id}_{uuid4().hex[:8]}".encode()
-                sock = create_zmq_socket(context, zmq.DEALER, identity=identity, ip=server_info.ip)
-
-                try:
-                    sock.connect(address)
-                    logger.debug(
-                        f"[{self.client_id}]: Connected to Controller {server_info.id} at {address} "
-                        f"with identity {identity.decode()}"
-                    )
-
-                    kwargs["socket"] = sock
-                    return await func(self, *args, **kwargs)
-                except Exception as e:
-                    logger.error(f"[{self.client_id}]: Error in socket operation with Controller {server_info.id}: {e}")
-                    raise
-                finally:
-                    try:
-                        if not sock.closed:
-                            sock.close(linger=-1)
-                    except Exception as e:
-                        logger.warning(f"[{self.client_id}]: Error closing socket to Controller {server_info.id}: {e}")
-
-                    context.term()
-
-            return wrapper
-
-        return decorator
-
     # ==================== Basic API ====================
-    @dynamic_socket(socket_name="request_handle_socket")
+    @with_controller_socket
     async def async_get_meta(
         self,
         data_fields: list[str],
@@ -245,7 +187,7 @@ async def async_get_meta(
                 f"{response_msg.body.get('message', 'Unknown error')}"
             )
 
-    @dynamic_socket(socket_name="request_handle_socket")
+    @with_controller_socket
     async def async_set_custom_meta(
         self,
         metadata: BatchMeta,
@@ -545,7 +487,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
 
-    @dynamic_socket(socket_name="request_handle_socket")
+    @with_controller_socket
     async def _clear_meta_in_controller(self, metadata: BatchMeta, socket=None):
         """Clear metadata in the controller.
 
@@ -571,7 +513,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.")
 
-    @dynamic_socket(socket_name="request_handle_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.
 
@@ -601,7 +543,7 @@ async def _get_partition_meta(self, partition_id: str, socket=None) -> BatchMeta
 
         return response_msg.body["metadata"]
 
-    @dynamic_socket(socket_name="request_handle_socket")
+    @with_controller_socket
     async def _clear_partition_in_controller(self, partition_id, socket=None):
         """Clear the whole partition in the controller.
 
@@ -628,7 +570,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 ====================
-    @dynamic_socket(socket_name="request_handle_socket")
+    @with_controller_socket
     async def async_get_consumption_status(
         self,
         task_name: str,
@@ -691,7 +633,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
 
-    @dynamic_socket(socket_name="request_handle_socket")
+    @with_controller_socket
     async def async_get_production_status(
         self,
         data_fields: list[str],
@@ -823,7 +765,7 @@ async def async_check_production_status(
             return False
         return torch.all(production_status == 1).item()
 
-    @dynamic_socket(socket_name="request_handle_socket")
+    @with_controller_socket
     async def async_reset_consumption(
         self,
         partition_id: str,
@@ -885,7 +827,7 @@ async def async_reset_consumption(
         except Exception as e:
             raise RuntimeError(f"[{self.client_id}]: Error in reset_consumption: {str(e)}") from e
 
-    @dynamic_socket(socket_name="request_handle_socket")
+    @with_controller_socket
     async def async_get_partition_list(
         self,
         socket: Optional[zmq.asyncio.Socket] = None,
@@ -931,7 +873,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 ====================
-    @dynamic_socket(socket_name="request_handle_socket")
+    @with_controller_socket
     async def async_kv_retrieve_meta(
         self,
         keys: list[str] | str,
@@ -997,7 +939,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
 
-    @dynamic_socket(socket_name="request_handle_socket")
+    @with_controller_socket
     async def async_kv_retrieve_keys(
         self,
         global_indexes: list[int] | int,
@@ -1060,7 +1002,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
 
-    @dynamic_socket(socket_name="request_handle_socket")
+    @with_controller_socket
     async def async_kv_list(
         self,
         partition_id: Optional[str] = None,

transfer_queue/controller.py

@@ -14,7 +14,6 @@
 # limitations under the License.
 
 import copy
-import logging
 import os
 import time
 from collections import defaultdict
@@ -37,6 +36,7 @@
 )
 from transfer_queue.sampler import BaseSampler, SequentialSampler
 from transfer_queue.utils.enum_utils import TransferQueueRole
+from transfer_queue.utils.logging_utils import get_logger
 from transfer_queue.utils.perf_utils import IntervalPerfMonitor
 from transfer_queue.utils.zmq_utils import (
     ZMQMessage,
@@ -48,14 +48,7 @@
     get_node_ip_address_raw,
 )
 
-logger = logging.getLogger(__name__)
-logger.setLevel(os.getenv("TQ_LOGGING_LEVEL", logging.WARNING))
-
-# Ensure logger has a handler (for Ray Actor subprocess)
-if not logger.hasHandlers():
-    handler = logging.StreamHandler()
-    handler.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(name)s - %(message)s"))
-    logger.addHandler(handler)
+logger = get_logger(__name__)
 
 TQ_CONTROLLER_GET_METADATA_TIMEOUT = int(os.environ.get("TQ_CONTROLLER_GET_METADATA_TIMEOUT", 1))
 TQ_CONTROLLER_GET_METADATA_CHECK_INTERVAL = int(os.environ.get("TQ_CONTROLLER_GET_METADATA_CHECK_INTERVAL", 5))

transfer_queue/dataloader/streaming_dataloader.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import logging
 import os
 from typing import Optional
 
@@ -22,15 +21,9 @@
 
 from transfer_queue.dataloader.streaming_dataset import StreamingDataset
 from transfer_queue.metadata import BatchMeta
+from transfer_queue.utils.logging_utils import get_logger
 
-logger = logging.getLogger(__name__)
-logger.setLevel(os.getenv("TQ_LOGGING_LEVEL", logging.WARNING))
-
-# Ensure logger has a handler
-if not logger.hasHandlers():
-    handler = logging.StreamHandler()
-    handler.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(name)s - %(message)s"))
-    logger.addHandler(handler)
+logger = get_logger(__name__)
 
 
 def _identity_collate_fn(data: tuple[TensorDict, BatchMeta]) -> tuple[TensorDict, BatchMeta]:

transfer_queue/dataloader/streaming_dataset.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import logging
 import os
 import time
 import uuid
@@ -25,19 +24,13 @@
 
 from transfer_queue.client import TransferQueueClient
 from transfer_queue.metadata import BatchMeta
+from transfer_queue.utils.logging_utils import get_logger
 
 TQ_STREAMING_DATASET_EMPTY_BATCH_SLEEP_INTERVAL = float(
     os.environ.get("TQ_STREAMING_DATASET_EMPTY_BATCH_SLEEP_INTERVAL", 1)
 )  # in seconds
 
-logger = logging.getLogger(__name__)
-logger.setLevel(os.getenv("TQ_LOGGING_LEVEL", logging.WARNING))
-
-# Ensure logger has a handler
-if not logger.hasHandlers():
-    handler = logging.StreamHandler()
-    handler.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(name)s - %(message)s"))
-    logger.addHandler(handler)
+logger = get_logger(__name__)
 
 
 class StreamingDataset(IterableDataset):

transfer_queue/interface.py

@@ -39,10 +39,11 @@
     cleanup_yuanrong_resources,
     initialize_yuanrong_backend,
 )
+from transfer_queue.utils.logging_utils import get_logger
 from transfer_queue.utils.zmq_utils import process_zmq_server_info
 
-logger = logging.getLogger(__name__)
-logger.setLevel(os.getenv("TQ_LOGGING_LEVEL", logging.WARNING))
+
+logger = get_logger(__name__)
 
 _TRANSFER_QUEUE_CLIENT: Any = None
 _TRANSFER_QUEUE_STORAGE: Any = None

transfer_queue/metadata.py

@@ -16,7 +16,6 @@
 import copy
 import dataclasses
 import itertools
-import logging
 import os
 from collections import defaultdict
 from dataclasses import dataclass
@@ -27,14 +26,9 @@
 import torch
 from tensordict import TensorDict
 
-logger = logging.getLogger(__name__)
-logger.setLevel(os.getenv("TQ_LOGGING_LEVEL", logging.WARNING))
+from transfer_queue.utils.logging_utils import get_logger
 
-# Ensure logger has a handler
-if not logger.hasHandlers():
-    handler = logging.StreamHandler()
-    handler.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(name)s - %(message)s"))
-    logger.addHandler(handler)
+logger = get_logger(__name__)
 
 
 # ---------------------------------------------------------------------------