[https://nvbugs/6208457][fix] Pass IPC HMAC key through file descriptor

tensorrt_llm/commands/serve.py

@@ -25,7 +25,8 @@
 from tensorrt_llm._utils import mpi_rank
 from tensorrt_llm.commands.utils import (collect_explicit_cli_keys,
                                          get_is_diffusion_only_model)
-from tensorrt_llm.executor.utils import LlmLauncherEnvs
+from tensorrt_llm.executor.utils import (LlmLauncherEnvs,
+                                         set_spawn_proxy_process_ipc_hmac_key)
 from tensorrt_llm.inputs.multimodal import MultimodalServerConfig
 from tensorrt_llm.llmapi import (BuildConfig, CapacitySchedulerPolicy,
                                  DynamicBatchConfig, KvCacheConfig,
@@ -1544,11 +1545,13 @@ def _launch_disaggregated_leader(sub_comm, instance_idx: int, config_file: str,
     # This mimics the behavior of trtllm-llmapi-launch
     # TODO: Make the port allocation atomic
     free_ipc_addr = find_free_ipc_addr()
+    ipc_hmac_key = secrets.token_hex(32)
+    set_spawn_proxy_process_ipc_hmac_key(ipc_hmac_key)
+    os.environ.pop(
+        LlmLauncherEnvs.TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY_FD.value, None)
     os.environ[LlmLauncherEnvs.TLLM_SPAWN_PROXY_PROCESS] = "1"
     os.environ[
         LlmLauncherEnvs.TLLM_SPAWN_PROXY_PROCESS_IPC_ADDR.value] = free_ipc_addr
-    os.environ[LlmLauncherEnvs.TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY.
-               value] = secrets.token_hex(32)
     os.environ[DisaggLauncherEnvs.TLLM_DISAGG_RUN_REMOTE_MPI_SESSION_CLIENT.
                value] = "1"
     os.environ[DisaggLauncherEnvs.TLLM_DISAGG_INSTANCE_IDX] = str(instance_idx)
@@ -1564,7 +1567,6 @@ def _launch_disaggregated_leader(sub_comm, instance_idx: int, config_file: str,
 
     assert LlmLauncherEnvs.TLLM_SPAWN_PROXY_PROCESS in non_mpi_env
     assert LlmLauncherEnvs.TLLM_SPAWN_PROXY_PROCESS_IPC_ADDR in non_mpi_env
-    assert LlmLauncherEnvs.TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY in non_mpi_env
     assert DisaggLauncherEnvs.TLLM_DISAGG_INSTANCE_IDX in non_mpi_env
     assert DisaggLauncherEnvs.TLLM_DISAGG_RUN_REMOTE_MPI_SESSION_CLIENT in non_mpi_env
 
@@ -1587,13 +1589,24 @@ def _launch_disaggregated_leader(sub_comm, instance_idx: int, config_file: str,
     signal.signal(signal.SIGTERM, _signal_handler_cleanup_child)
     signal.signal(signal.SIGINT, _signal_handler_cleanup_child)
 
+    read_fd = -1
+    write_fd = -1
     try:
+        read_fd, write_fd = os.pipe()
+        os.write(write_fd, ipc_hmac_key.encode("ascii"))
+        os.close(write_fd)
+        write_fd = -1
+        non_mpi_env[LlmLauncherEnvs.TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY_FD.
+                    value] = str(read_fd)
         _child_p_global = subprocess.Popen(
             command,
             env=non_mpi_env,
             stdout=sys.stdout,  # Redirect to parent's stdout
             stderr=sys.stderr,  # Redirect to parent's stderr
+            pass_fds=(read_fd, ),
             start_new_session=True)
+        os.close(read_fd)
+        read_fd = -1
 
         logger.info(
             f"Parent process (PID {os.getpid()}) launched child process (PID {_child_p_global.pid})."
@@ -1607,6 +1620,11 @@ def _launch_disaggregated_leader(sub_comm, instance_idx: int, config_file: str,
         launch_remote_mpi_session_server(sub_comm)
 
     finally:
+        if write_fd != -1:
+            os.close(write_fd)
+        if read_fd != -1:
+            os.close(read_fd)
+
         # Restore original signal handlers
         signal.signal(signal.SIGTERM, original_sigterm_handler)
         signal.signal(signal.SIGINT, original_sigint_handler)

tensorrt_llm/executor/ipc.py

@@ -40,15 +40,19 @@ def __init__(self,
                  use_hmac_encryption: bool = True):
         '''
         Parameters:
-            address (tuple[str, Optional[bytes]], optional): The address (tcp-ip_port, hmac_auth_key) for the IPC. Defaults to None. If hmac_auth_key is None and use_hmac_encryption is False, the queue will not use HMAC encryption.
+            address (tuple[str, Optional[bytes]], optional): The address (tcp-ip_port, hmac_auth_key) for the IPC. Defaults to None. Servers generate an HMAC key when none is provided; clients must receive one.
             socket_type (int): The type of socket to use. Defaults to zmq.PAIR.
             is_server (bool): Whether the current process is the server or the client.
             is_async (bool): Whether to use asyncio for the socket. Defaults to False.
             name (str, optional): The name of the queue. Defaults to None.
             use_hmac_encryption (bool): Whether to use HMAC encryption for pickled data. Defaults to True.
         '''
 
-        assert use_hmac_encryption, "HMAC encryption is always required. Turning off HMAC encryption risks security vulnerability of unauthorized data serialization and deserialization. "
+        if not use_hmac_encryption:
+            raise ValueError(
+                "HMAC encryption is always required. Turning off HMAC "
+                "encryption risks unauthorized data serialization and "
+                "deserialization.")
 
         self.socket_type = socket_type
         self.address_endpoint = address[

tensorrt_llm/executor/utils.py

@@ -23,24 +23,71 @@ class LlmLauncherEnvs(StrEnum):
     # Spawn a process for the LLM-API Proxy
     TLLM_SPAWN_PROXY_PROCESS = "TLLM_SPAWN_PROXY_PROCESS"
     TLLM_SPAWN_PROXY_PROCESS_IPC_ADDR = "TLLM_SPAWN_PROXY_PROCESS_IPC_ADDR"
-    TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY = "TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY"
+    TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY_FD = (
+        "TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY_FD")
 
     # Whether to use periodical responses handler in await_responses
     TLLM_EXECUTOR_PERIODICAL_RESP_IN_AWAIT = "TLLM_EXECUTOR_PERIODICAL_RESP_IN_AWAIT"
 
 
+_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY: bytes | None = None
+
+
+def _normalize_spawn_proxy_process_ipc_hmac_key(key: str | bytes) -> bytes:
+    if isinstance(key, bytes):
+        if len(key) == 32:
+            return key
+        key = key.decode("ascii")
+
+    key_bytes = bytes.fromhex(key)
+    if len(key_bytes) != 32:
+        raise ValueError("IPC HMAC key must be 32 bytes.")
+    return key_bytes
+
+
+def set_spawn_proxy_process_ipc_hmac_key(key: str | bytes) -> None:
+    global _SPAWN_PROXY_PROCESS_IPC_HMAC_KEY
+    _SPAWN_PROXY_PROCESS_IPC_HMAC_KEY = (
+        _normalize_spawn_proxy_process_ipc_hmac_key(key))
+
+
+def _read_spawn_proxy_process_ipc_hmac_key_fd(fd_value: str) -> bytes:
+    fd = int(fd_value)
+    chunks: list[bytes] = []
+    try:
+        os.set_blocking(fd, True)
+        while True:
+            chunk = os.read(fd, 4096)
+            if not chunk:
+                break
+            chunks.append(chunk)
+    finally:
+        os.close(fd)
+
+    return _normalize_spawn_proxy_process_ipc_hmac_key(b"".join(chunks))
+
+
 def get_spawn_proxy_process_ipc_addr_env() -> str | None:
     ''' Get the IPC address for the spawn proxy process dynamically. '''
     return os.getenv(LlmLauncherEnvs.TLLM_SPAWN_PROXY_PROCESS_IPC_ADDR)
 
 
 def get_spawn_proxy_process_ipc_hmac_key_env() -> bytes:
     ''' Get the HMAC key for the spawn proxy process dynamically. '''
-    key = os.getenv("TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY")
-    assert key is not None, (
-        f"{LlmLauncherEnvs.TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY} is not set. "
+    global _SPAWN_PROXY_PROCESS_IPC_HMAC_KEY
+    if _SPAWN_PROXY_PROCESS_IPC_HMAC_KEY is not None:
+        return _SPAWN_PROXY_PROCESS_IPC_HMAC_KEY
+
+    key_fd = os.environ.pop(
+        LlmLauncherEnvs.TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY_FD, None)
+    if key_fd is not None:
+        _SPAWN_PROXY_PROCESS_IPC_HMAC_KEY = (
+            _read_spawn_proxy_process_ipc_hmac_key_fd(key_fd))
+        return _SPAWN_PROXY_PROCESS_IPC_HMAC_KEY
+
+    raise RuntimeError(
+        f"{LlmLauncherEnvs.TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY_FD} is not set. "
         "HMAC encryption is required for IPC communication.")
-    return bytes.fromhex(key)
 
 
 def get_spawn_proxy_process_env() -> bool:

tensorrt_llm/executor/worker.py

@@ -69,7 +69,8 @@ def __init__(
         # Setup RPC server for stats (skip init_rpc_worker to keep IPC response queue)
         # Only set up if rpc_addr is provided (for stats RPC support)
         if rpc_addr is not None:
-            assert hmac_key, "hmac_key is required when rpc_addr is set"
+            if not hmac_key:
+                raise ValueError("hmac_key is required when rpc_addr is set")
             self.rpc_addr = rpc_addr
             self.hmac_key = hmac_key
             self.start_rpc_server()  # Reuse from RpcWorkerMixin

tensorrt_llm/llmapi/trtllm-llmapi-launch

@@ -40,7 +40,17 @@ function maybe_export_free_tcp_addr_for_spawn_proxy_process {
 export tllm_mpi_size=$(mpi_world_size)
 log_stderr "tllm_mpi_size: $tllm_mpi_size"
 
-export TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY=$(openssl rand -hex 32)
+unset TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY_FD
+ipc_hmac_key=$(openssl rand -hex 32)
+
+function run_with_ipc_hmac_key {
+    local fd
+    exec {fd}<<<"$ipc_hmac_key"
+    TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY_FD="$fd" "$@"
+    local status=$?
+    exec {fd}<&-
+    return $status
+}
 
 if [ -z "$mpi_rank" ] || [ "$mpi_rank" -eq 0 ]; then
 
@@ -74,12 +84,13 @@ if [ -z "$mpi_rank" ] || [ "$mpi_rank" -eq 0 ]; then
         set +e
 
         # Execute the task with cleaned environment
-        "${task_with_command[@]}"
+        run_with_ipc_hmac_key "${task_with_command[@]}"
         task_exit_code=$?
         log_stderr "Rank${mpi_rank} Task exit code: $task_exit_code"
 
         # Stop the MPI Comm server
-        python3 -m tensorrt_llm.llmapi.mgmn_leader_node --action stop
+        run_with_ipc_hmac_key python3 -m tensorrt_llm.llmapi.mgmn_leader_node \
+            --action stop
         mpi_exit_code=$?
         log_stderr "Rank${mpi_rank} MPI Comm server exit code: $mpi_exit_code"
 
@@ -100,7 +111,7 @@ if [ -z "$mpi_rank" ] || [ "$mpi_rank" -eq 0 ]; then
 
     log_stderr "Rank${mpi_rank} run mgmn leader node with mpi_world_size: $(mpi_world_size) ..."
     log_stderr "Rank0 host: $HOSTNAME"
-    python3 -m tensorrt_llm.llmapi.mgmn_leader_node
+    run_with_ipc_hmac_key python3 -m tensorrt_llm.llmapi.mgmn_leader_node
     mgmn_leader_node_exit_code=$?
     log_stderr "Rank${mpi_rank} MGMN leader node exit code: $mgmn_leader_node_exit_code"
 

tests/unittest/executor/test_ipc.py

@@ -257,6 +257,18 @@ def test_hmac_validation_error_client_no_key(self):
                 use_hmac_encryption=True,  # But encryption enabled
             )
 
+    def test_hmac_encryption_cannot_be_disabled(self):
+        """Test that HMAC cannot be disabled by optimized Python removing asserts."""
+        with pytest.raises(ValueError, match="HMAC encryption is always required"):
+            ZeroMqQueue(
+                address=None,
+                socket_type=zmq.PAIR,
+                is_server=True,
+                is_async=False,
+                name="test_hmac_disabled",
+                use_hmac_encryption=False,
+            )
+
     def test_put_noblock_retry(self):
         """Test put_noblock with retry mechanism."""
         server = ZeroMqQueue(

tests/unittest/executor/test_launcher_envs.py

@@ -0,0 +1,79 @@
+import os
+import threading
+
+import pytest
+
+from tensorrt_llm.executor import utils as executor_utils
+from tensorrt_llm.executor.utils import LlmLauncherEnvs
+
+
+def _reset_ipc_hmac_key_env(monkeypatch):
+    monkeypatch.setattr(executor_utils, "_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY", None)
+    monkeypatch.delenv(
+        LlmLauncherEnvs.TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY_FD.value, raising=False
+    )
+
+
+def _write_key_fd(key_hex: str) -> int:
+    read_fd, write_fd = os.pipe()
+    os.write(write_fd, key_hex.encode("ascii"))
+    os.close(write_fd)
+    return read_fd
+
+
+def test_get_spawn_proxy_process_ipc_hmac_key_from_fd(monkeypatch):
+    _reset_ipc_hmac_key_env(monkeypatch)
+    key_hex = "01" * 32
+    read_fd = _write_key_fd(key_hex)
+    monkeypatch.setenv(LlmLauncherEnvs.TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY_FD.value, str(read_fd))
+
+    key = executor_utils.get_spawn_proxy_process_ipc_hmac_key_env()
+
+    assert key == bytes.fromhex(key_hex)
+    assert LlmLauncherEnvs.TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY_FD not in os.environ
+    with pytest.raises(OSError):
+        os.fstat(read_fd)
+
+
+def test_get_spawn_proxy_process_ipc_hmac_key_caches_fd_key(monkeypatch):
+    _reset_ipc_hmac_key_env(monkeypatch)
+    key_hex = "02" * 32
+    read_fd = _write_key_fd(key_hex)
+    monkeypatch.setenv(LlmLauncherEnvs.TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY_FD.value, str(read_fd))
+
+    key = executor_utils.get_spawn_proxy_process_ipc_hmac_key_env()
+
+    assert key == bytes.fromhex(key_hex)
+    assert LlmLauncherEnvs.TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY_FD not in os.environ
+    assert executor_utils.get_spawn_proxy_process_ipc_hmac_key_env() == key
+
+
+def test_get_spawn_proxy_process_ipc_hmac_key_from_nonblocking_fd(monkeypatch):
+    _reset_ipc_hmac_key_env(monkeypatch)
+    key_hex = "03" * 32
+    read_fd, write_fd = os.pipe()
+    os.set_blocking(read_fd, False)
+
+    def write_key():
+        try:
+            os.write(write_fd, key_hex.encode("ascii"))
+        finally:
+            os.close(write_fd)
+
+    writer = threading.Timer(0.05, write_key)
+    writer.start()
+    monkeypatch.setenv(LlmLauncherEnvs.TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY_FD.value, str(read_fd))
+
+    try:
+        key = executor_utils.get_spawn_proxy_process_ipc_hmac_key_env()
+    finally:
+        writer.join()
+
+    assert key == bytes.fromhex(key_hex)
+
+
+def test_get_spawn_proxy_process_ipc_hmac_key_missing(monkeypatch):
+    _reset_ipc_hmac_key_env(monkeypatch)
+
+    with pytest.raises(RuntimeError, match="HMAC encryption is required"):
+        executor_utils.get_spawn_proxy_process_ipc_hmac_key_env()