[https://nvbugs/5922880][fix] Enable HMAC authentication in VisualGen ZMQ IPC channels (#12680)

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>

Author: yibinl-nvidia

tensorrt_llm/_torch/visual_gen/executor.py

@@ -83,11 +83,14 @@ def __init__(
         response_queue_addr: str,
         device_id: int,
         diffusion_args: "VisualGenArgs",
+        req_hmac_key: Optional[bytes] = None,
+        resp_hmac_key: Optional[bytes] = None,
     ):
         self.request_queue_addr = request_queue_addr
         self.response_queue_addr = response_queue_addr
         self.device_id = device_id
         self.diffusion_args = diffusion_args
+        self.resp_hmac_key = resp_hmac_key
 
         self.pipeline = None  # initialized in _load_pipeline
         self.requests_ipc = None
@@ -99,10 +102,10 @@ def __init__(
         if self.rank == 0:
             logger.info(f"Worker {device_id}: Connecting to request queue")
             self.requests_ipc = ZeroMqQueue(
-                (request_queue_addr, None),
+                (request_queue_addr, req_hmac_key),
                 is_server=False,
                 socket_type=zmq.PULL,
-                use_hmac_encryption=False,
+                use_hmac_encryption=True,
             )
             self.sender_thread = threading.Thread(target=self._sender_loop, daemon=True)
             self.sender_thread.start()
@@ -113,10 +116,10 @@ def _sender_loop(self):
         """Background thread for sending responses."""
         logger.info(f"Worker {self.device_id}: Connecting to response queue")
         responses_ipc = ZeroMqQueue(
-            (self.response_queue_addr, None),
+            (self.response_queue_addr, self.resp_hmac_key),
             is_server=False,
             socket_type=zmq.PUSH,
-            use_hmac_encryption=False,
+            use_hmac_encryption=True,
         )
 
         while True:
@@ -214,6 +217,8 @@ def run_diffusion_worker(
     response_queue_addr: str,
     diffusion_args: "VisualGenArgs",
     log_level: str = "info",
+    req_hmac_key: Optional[bytes] = None,
+    resp_hmac_key: Optional[bytes] = None,
 ):
     """Entry point for worker process."""
     try:
@@ -248,6 +253,8 @@ def run_diffusion_worker(
             response_queue_addr=response_queue_addr,
             device_id=device_id,
             diffusion_args=diffusion_args,
+            req_hmac_key=req_hmac_key,
+            resp_hmac_key=resp_hmac_key,
         )
         executor.serve_forever()
         if executor.pipeline is not None:

tensorrt_llm/visual_gen/visual_gen.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 import asyncio
 import atexit
+import os
 import queue
 import socket
 import threading
@@ -86,6 +87,10 @@ def __init__(
         self.req_addr_connect = f"tcp://{self.host_ip}:{req_port}"
         self.resp_addr_connect = f"tcp://{self.host_ip}:{resp_port}"
 
+        # Generate shared HMAC keys for IPC authentication
+        self.req_hmac_key = os.urandom(32)
+        self.resp_hmac_key = os.urandom(32)
+
         # IPC setup
         self.requests_ipc = None
         self.responses_ipc = None
@@ -122,6 +127,8 @@ def __init__(
                     "request_queue_addr": self.req_addr_connect,
                     "response_queue_addr": self.resp_addr_connect,
                     "diffusion_args": self.diffusion_args,
+                    "req_hmac_key": self.req_hmac_key,
+                    "resp_hmac_key": self.resp_hmac_key,
                     "log_level": logger.level,
                 },
             )
@@ -205,16 +212,16 @@ def _init_ipc(self) -> bool:
         try:
             logger.info("DiffusionClient: Initializing IPC")
             self.requests_ipc = ZeroMqQueue(
-                (self.request_queue_addr, None),
+                (self.request_queue_addr, self.req_hmac_key),
                 is_server=True,
                 socket_type=zmq.PUSH,
-                use_hmac_encryption=False,
+                use_hmac_encryption=True,
             )
             self.responses_ipc = ZeroMqQueue(
-                (self.response_queue_addr, None),
+                (self.response_queue_addr, self.resp_hmac_key),
                 is_server=True,
                 socket_type=zmq.PULL,
-                use_hmac_encryption=False,
+                use_hmac_encryption=True,
             )
             logger.info("DiffusionClient: IPC ready")
             return True