remove unnecessary poller

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

Author: ji-huazhong

transfer_queue/controller.py

@@ -27,7 +27,6 @@
 import numpy as np
 import ray
 import torch
-import zmq
 from omegaconf import DictConfig
 from torch import Tensor
 
@@ -1047,51 +1046,44 @@ def _start_daemon_threads(self):
         )
 
     def _wait_connection(self):
-        """Wait for storage instances to complete handshake with retransmission support."""
+        """Wait for storage instances to complete handshake."""
         handshake_socket = self._transport.get_socket("handshake_socket")
-        poller = zmq.Poller()
-        poller.register(handshake_socket, zmq.POLLIN)
 
         logger.debug(f"Controller {self.controller_id} started waiting for storage connections...")
 
         while True:
-            socks = dict(poller.poll(1000))
-
-            if handshake_socket in socks:
-                try:
-                    messages = handshake_socket.recv_multipart(copy=False)
-                    identity = messages.pop(0)
-                    serialized_msg = messages
-                    request_msg = ZMQMessage.deserialize(serialized_msg)
+            try:
+                messages = handshake_socket.recv_multipart(copy=False)
+                identity = messages.pop(0)
+                serialized_msg = messages
+                request_msg = ZMQMessage.deserialize(serialized_msg)
 
-                    if request_msg.request_type == ZMQRequestType.HANDSHAKE:
-                        storage_manager_id = request_msg.sender_id
+                if request_msg.request_type == ZMQRequestType.HANDSHAKE:
+                    storage_manager_id = request_msg.sender_id
 
-                        # Always send ACK for HANDSHAKE
-                        response_msg = ZMQMessage.create(
-                            request_type=ZMQRequestType.HANDSHAKE_ACK,
-                            sender_id=self.controller_id,
-                            body={},
-                        ).serialize()
-                        handshake_socket.send_multipart([identity, *response_msg])
-
-                        # Track new connections
-                        if storage_manager_id not in self._connected_storage_managers:
-                            self._connected_storage_managers.add(storage_manager_id)
-                            storage_manager_type = request_msg.body.get("storage_manager_type", "Unknown")
-                            logger.debug(
-                                f"[{self.controller_id}]: received handshake from "
-                                f"storage manager {storage_manager_id} (type: {storage_manager_type}). "
-                                f"Total connected: {len(self._connected_storage_managers)}"
-                            )
-                        else:
-                            logger.debug(
-                                f"[{self.controller_id}]: received duplicate handshake from "
-                                f"storage manager {storage_manager_id}. Resending ACK."
-                            )
+                    response_msg = ZMQMessage.create(
+                        request_type=ZMQRequestType.HANDSHAKE_ACK,
+                        sender_id=self.controller_id,
+                        body={},
+                    ).serialize()
+                    handshake_socket.send_multipart([identity, *response_msg])
+
+                    if storage_manager_id not in self._connected_storage_managers:
+                        self._connected_storage_managers.add(storage_manager_id)
+                        storage_manager_type = request_msg.body.get("storage_manager_type", "Unknown")
+                        logger.debug(
+                            f"[{self.controller_id}]: received handshake from "
+                            f"storage manager {storage_manager_id} (type: {storage_manager_type}). "
+                            f"Total connected: {len(self._connected_storage_managers)}"
+                        )
+                    else:
+                        logger.debug(
+                            f"[{self.controller_id}]: received duplicate handshake from "
+                            f"storage manager {storage_manager_id}. Resending ACK."
+                        )
 
-                except Exception as e:
-                    logger.error(f"[{self.controller_id}]: error processing handshake: {e}")
+            except Exception as e:
+                logger.error(f"[{self.controller_id}]: error processing handshake: {e}")
 
     def _process_request(self):
         """Main request processing loop - adapted for partition-based operations."""
@@ -1708,7 +1700,6 @@ def get_metadata(
 
         elif mode == "force_fetch":
             batch_global_indexes = self.index_manager.get_indexes_for_partition(partition_id)
-            consumed_indexes = []
 
         # Package into metadata
         metadata = self.generate_batch_meta(partition_id, batch_global_indexes, data_fields, mode)

transfer_queue/utils/zmq_utils.py

@@ -104,9 +104,7 @@ class ZMQRequestType(ExplicitEnum):
 
 
 class ZMQServerInfo:
-    """
-    TransferQueue server info class.
-    """
+    """TransferQueue server info class."""
 
     def __init__(self, role: Role, id: str, ip: str, ports: dict[str, int]):
         self.role = role
@@ -133,9 +131,7 @@ def __str__(self) -> str:
 
 @dataclass
 class ZMQMessage:
-    """
-    ZMQMessage class for TransferQueue communication.
-    """
+    """ZMQMessage class for TransferQueue communication."""
 
     request_type: ZMQRequestType
     sender_id: str
@@ -192,10 +188,7 @@ def deserialize(cls, frames: list) -> "ZMQMessage":
 
 
 class ZMQServerTransport:
-    """Unified ZMQ transport abstraction for Controller / StorageUnit.
-    Encapsulates socket creation, binding, inproc endpoint, daemon thread,
-    ZMQ proxy lifecycle, and unified resource cleanup.
-    """
+    """Unified management of ZMQ Router Sockets, port binding, daemon threads, and message I/O."""
 
     def __init__(self, node_ip: str, ctx: zmq.Context | None = None):
         self.node_ip = node_ip
@@ -222,14 +215,17 @@ def create_router_socket(self, name: str) -> None:
                 logger.warning(f"ZMQ bind {name} failed, retrying...")
 
     def get_socket(self, name: str) -> zmq.Socket:
+        """Get ZMQ socket by name."""
         return self.sockets[name]
 
     def start_daemon_thread(self, target, name: str) -> None:
+        """Start a daemon thread with the given target functions."""
         t = threading.Thread(target=target, name=name, daemon=True)
         t.start()
         self.threads.append(t)
 
     def build_server_info(self, role: Role, server_id: str) -> ZMQServerInfo:
+        """Build ZMQServerInfo."""
         return ZMQServerInfo(
             role=role,
             id=server_id,