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

tensorrt_llm/commands/serve.py

@@ -25,7 +25,8 @@
 from tensorrt_llm._torch.visual_gen.config import VisualGenArgs
 from tensorrt_llm._utils import mpi_rank
 from tensorrt_llm.commands.utils import get_is_diffusion_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,
@@ -1404,11 +1405,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)
@@ -1424,7 +1427,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
 
@@ -1447,13 +1449,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})."
@@ -1467,6 +1480,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/utils.py

@@ -23,24 +23,70 @@ 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:
+        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 AssertionError(
+        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/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_launcher_envs.py

@@ -0,0 +1,54 @@
+import os
+
+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_missing(monkeypatch):
+    _reset_ipc_hmac_key_env(monkeypatch)
+
+    with pytest.raises(AssertionError, match="HMAC encryption is required"):
+        executor_utils.get_spawn_proxy_process_ipc_hmac_key_env()