diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 5dd282c814c2..b2d7154d17fd 100644 --- a/tensorrt_llm/commands/serve.py +++ b/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) diff --git a/tensorrt_llm/executor/utils.py b/tensorrt_llm/executor/utils.py index 12eaa9afefa9..421d943916fd 100644 --- a/tensorrt_llm/executor/utils.py +++ b/tensorrt_llm/executor/utils.py @@ -23,12 +23,49 @@ 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) @@ -36,11 +73,20 @@ def get_spawn_proxy_process_ipc_addr_env() -> str | None: 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: diff --git a/tensorrt_llm/llmapi/trtllm-llmapi-launch b/tensorrt_llm/llmapi/trtllm-llmapi-launch index c11a3bbeb171..d35b60046a26 100755 --- a/tensorrt_llm/llmapi/trtllm-llmapi-launch +++ b/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" diff --git a/tests/unittest/executor/test_launcher_envs.py b/tests/unittest/executor/test_launcher_envs.py new file mode 100644 index 000000000000..0d1f96ef1620 --- /dev/null +++ b/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()