fix wait_for lost message

test passed
the race condition is resolved completely

Signed-off-by: Superjomn <[email protected]>

Author: Superjomn

tensorrt_llm/executor/ipc.py

@@ -244,8 +244,46 @@ async def get_async(self) -> Any:
         return await self._recv_data_async()
 
     async def get_async_noblock(self, timeout: float = 0.5) -> Any:
+        """Get data with timeout using polling to avoid message drops.
+
+        This method uses ZMQ's NOBLOCK flag with polling instead of asyncio.wait_for
+        to prevent cancelling recv operations which can cause message drops.
+
+        Args:
+            timeout: Timeout in seconds
+
+        Returns:
+            The received object
+
+        Raises:
+            asyncio.TimeoutError: If timeout is reached without receiving data
+        """
+        self.setup_lazily()
         self._check_thread_safety()
-        return await asyncio.wait_for(self.get_async(), timeout)
+
+        # Use polling loop instead of asyncio.wait_for to avoid cancelling recv
+        # which can cause message drops
+        deadline = asyncio.get_event_loop().time() + timeout
+        while True:
+            try:
+                # Try non-blocking receive
+                if self.socket_type == zmq.ROUTER:
+                    identity, data = await self.socket.recv_multipart(
+                        flags=zmq.NOBLOCK)
+                    self._last_identity = identity
+                    return self._parse_data(data)
+                else:
+                    if self.use_hmac_encryption:
+                        data = await self.socket.recv(flags=zmq.NOBLOCK)
+                        return self._parse_data(data)
+                    else:
+                        return await self.socket.recv_pyobj(flags=zmq.NOBLOCK)
+            except zmq.Again:
+                # No message available yet
+                if asyncio.get_event_loop().time() >= deadline:
+                    raise asyncio.TimeoutError()
+                # Short sleep to avoid busy-waiting
+                await asyncio.sleep(0.01)
 
     def close(self):
         if self.socket:

tensorrt_llm/executor/rpc/rpc_client.py

@@ -1,5 +1,6 @@
 import asyncio
 import concurrent.futures
+import os
 import threading
 import time
 import uuid
@@ -104,11 +105,20 @@ def __init__(self,
         '''
         self._address = address
         self._timeout = timeout
+
+        # Check if PAIR mode is enabled via environment variable
+        use_pair_mode = os.environ.get('TLLM_LLMAPI_ZMQ_PAIR', '0') != '0'
+        socket_type = zmq.PAIR if use_pair_mode else zmq.DEALER
+
+        if use_pair_mode:
+            logger_debug(
+                "[client] Using zmq.PAIR socket type for RPC communication")
+
         self._client_socket = ZeroMqQueue(address=(address, hmac_key),
                                           is_server=False,
                                           is_async=True,
                                           use_hmac_encryption=False,
-                                          socket_type=zmq.DEALER,
+                                          socket_type=socket_type,
                                           name="rpc_client")
         self._pending_futures = {}
         # map request_id to the queue for streaming responses

tensorrt_llm/executor/rpc/rpc_server.py

@@ -1,5 +1,6 @@
 import asyncio
 import inspect
+import os
 import threading
 import time
 import traceback
@@ -95,11 +96,20 @@ def bind(self, address: str = "tcp://*:5555") -> None:
             address (str): The ZMQ address to bind the client-facing socket.
         """
         self._address = address
+
+        # Check if PAIR mode is enabled via environment variable
+        use_pair_mode = os.environ.get('TLLM_LLMAPI_ZMQ_PAIR', '0') != '0'
+        socket_type = zmq.PAIR if use_pair_mode else zmq.ROUTER
+
+        if use_pair_mode:
+            logger_debug(
+                "[server] Using zmq.PAIR socket type for RPC communication")
+
         self._client_socket = ZeroMqQueue(address=(address, self._hmac_key),
                                           is_server=True,
                                           is_async=True,
                                           use_hmac_encryption=False,
-                                          socket_type=zmq.ROUTER,
+                                          socket_type=socket_type,
                                           name="rpc_server")
         logger.info(f"RPCServer is bound to {self._address}")