add ipc TLLM_LLMAPI_ZMQ_DEBUG

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

Author: Superjomn

tensorrt_llm/executor/ipc.py

@@ -68,6 +68,11 @@ def __init__(self,
 
         self._setup_lock = threading.Lock()
 
+        # Thread safety debugging
+        self._zmq_thread_id = None
+        self._zmq_debug_enabled = os.environ.get('TLLM_LLMAPI_ZMQ_DEBUG',
+                                                 '0') != '0'
+
         # Check HMAC key condition
         if self.use_hmac_encryption and not self.is_server and self.hmac_key is None:
             raise ValueError(
@@ -114,12 +119,34 @@ def setup_lazily(self):
             self.poller = zmq.Poller()
             self.poller.register(self.socket, zmq.POLLIN)
 
+    def _check_thread_safety(self):
+        """Check if the current thread is the same as the thread that first used the socket."""
+        if not self._zmq_debug_enabled:
+            return
+
+        current_thread_id = threading.get_ident()
+
+        if self._zmq_thread_id is None:
+            # First call - capture the thread ID
+            self._zmq_thread_id = current_thread_id
+            logger_debug(
+                f"ZMQ socket [{self.name}] initialized on thread {current_thread_id}",
+                "cyan")
+        elif self._zmq_thread_id != current_thread_id:
+            # Thread mismatch - raise error
+            raise RuntimeError(
+                f"ZMQ thread safety violation detected in [{self.name}]: "
+                f"Socket created on thread {self._zmq_thread_id}, "
+                f"but accessed from thread {current_thread_id}. "
+                f"ZMQ sockets are not thread-safe!")
+
     def poll(self, timeout: int) -> bool:
         """
         Parameters:
             timeout (int): Timeout in seconds
         """
         self.setup_lazily()
+        self._check_thread_safety()
 
         events = dict(self.poller.poll(timeout=timeout * 1000))
         if self.socket in events and events[self.socket] == zmq.POLLIN:
@@ -129,6 +156,7 @@ def poll(self, timeout: int) -> bool:
 
     def put(self, obj: Any):
         self.setup_lazily()
+        self._check_thread_safety()
         with nvtx_range_debug("send", color="blue", category="IPC"):
             if self.use_hmac_encryption or self.socket_type == zmq.ROUTER:
                 # Need manual serialization for encryption or ROUTER multipart
@@ -156,6 +184,7 @@ def put_noblock(self,
         assert retry >= 0 and retry <= 10, "Retry must be between 0 and 10, adjust the wait_time if needed"
 
         self.setup_lazily()
+        self._check_thread_safety()
         with nvtx_range_debug("send", color="blue", category="IPC"):
 
             data = self._prepare_data(obj)
@@ -170,6 +199,7 @@ def put_noblock(self,
 
     async def put_async(self, obj: Any):
         self.setup_lazily()
+        self._check_thread_safety()
         try:
             if self.use_hmac_encryption or self.socket_type == zmq.ROUTER:
                 # Need manual serialization for encryption or ROUTER multipart
@@ -190,6 +220,7 @@ async def put_async(self, obj: Any):
 
     async def put_async_noblock(self, obj: Any):
         self.setup_lazily()
+        self._check_thread_safety()
         try:
             if self.use_hmac_encryption:
                 data = pickle.dumps(obj)  # nosec B301
@@ -204,13 +235,16 @@ async def put_async_noblock(self, obj: Any):
 
     def get(self) -> Any:
         self.setup_lazily()
+        self._check_thread_safety()
         return self._recv_data()
 
     async def get_async(self) -> Any:
         self.setup_lazily()
+        self._check_thread_safety()
         return await self._recv_data_async()
 
     async def get_async_noblock(self, timeout: float = 0.5) -> Any:
+        self._check_thread_safety()
         return await asyncio.wait_for(self.get_async(), timeout)
 
     def close(self):
@@ -319,6 +353,7 @@ def notify_with_retry(self, message, max_retries=5, timeout=1):
             raise ValueError(
                 "notify_with_retry is only supported for DEALER socket for now")
 
+        self._check_thread_safety()
         retry_count = 0
 
         while retry_count < max_retries: