diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index ef99deacb0c0..5608872da2dd 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -24,8 +24,7 @@ from tensorrt_llm._tensorrt_engine import LLM from tensorrt_llm._utils import mpi_rank from tensorrt_llm.commands.utils import get_is_diffusion_model -from tensorrt_llm.executor.utils import (LlmLauncherEnvs, - set_spawn_proxy_process_ipc_hmac_key) +from tensorrt_llm.executor.utils import LlmLauncherEnvs from tensorrt_llm.inputs.multimodal import MultimodalServerConfig from tensorrt_llm.llmapi import (BuildConfig, CapacitySchedulerPolicy, DynamicBatchConfig, KvCacheConfig, @@ -1404,13 +1403,11 @@ 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) @@ -1426,6 +1423,7 @@ 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 @@ -1448,24 +1446,13 @@ 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})." @@ -1479,11 +1466,6 @@ 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 421d943916fd..12eaa9afefa9 100644 --- a/tensorrt_llm/executor/utils.py +++ b/tensorrt_llm/executor/utils.py @@ -23,49 +23,12 @@ 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_FD = ( - "TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY_FD") + TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY = "TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY" # 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) @@ -73,20 +36,11 @@ 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. ''' - 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. " + 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. " "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 d35b60046a26..c11a3bbeb171 100755 --- a/tensorrt_llm/llmapi/trtllm-llmapi-launch +++ b/tensorrt_llm/llmapi/trtllm-llmapi-launch @@ -40,17 +40,7 @@ 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" -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 -} +export TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY=$(openssl rand -hex 32) if [ -z "$mpi_rank" ] || [ "$mpi_rank" -eq 0 ]; then @@ -84,13 +74,12 @@ if [ -z "$mpi_rank" ] || [ "$mpi_rank" -eq 0 ]; then set +e # Execute the task with cleaned environment - run_with_ipc_hmac_key "${task_with_command[@]}" + "${task_with_command[@]}" task_exit_code=$? log_stderr "Rank${mpi_rank} Task exit code: $task_exit_code" # Stop the MPI Comm server - run_with_ipc_hmac_key python3 -m tensorrt_llm.llmapi.mgmn_leader_node \ - --action stop + 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" @@ -111,7 +100,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" - run_with_ipc_hmac_key python3 -m tensorrt_llm.llmapi.mgmn_leader_node + 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 deleted file mode 100644 index 0d1f96ef1620..000000000000 --- a/tests/unittest/executor/test_launcher_envs.py +++ /dev/null @@ -1,54 +0,0 @@ -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()