Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 4 additions & 22 deletions tensorrt_llm/commands/serve.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Expand All @@ -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

Expand All @@ -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})."
Expand All @@ -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)
Expand Down
56 changes: 5 additions & 51 deletions tensorrt_llm/executor/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,70 +23,24 @@ 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"

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | 🏗️ Heavy lift

Keep the IPC HMAC key off the environment.

ZeroMqQueue uses this key to authenticate pickle traffic before pickle.loads() in tensorrt_llm/executor/ipc.py:33-103,354-389. Switching the contract back to TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY means the secret now propagates through process environments in both launcher paths instead of staying on a narrowly scoped channel like the previous inherited-FD mechanism. That weakens the only barrier guarding unauthorized IPC deserialization.

Also applies to: 39-43

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/executor/utils.py` at line 26, The new constant
TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY exposes the IPC HMAC secret via process
environment—restore the previous non-environment mechanism: remove usage of
TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY from any subprocess env setup and instead
pass the HMAC key over the narrowly-scoped inherited-FD / dedicated IPC channel
used by ZeroMqQueue (or via explicit file descriptor/pipe passed to the child),
update the code paths in ZeroMqQueue and the launcher/spawn code that currently
read/write this env var so they accept the key from the inherited-FD/IPC handle,
and delete the environment-constant definition so the secret is never placed
into os.environ or subprocess env.


# 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. '''
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:
Expand Down
19 changes: 4 additions & 15 deletions tensorrt_llm/llmapi/trtllm-llmapi-launch
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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"

Expand All @@ -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"

Expand Down
54 changes: 0 additions & 54 deletions tests/unittest/executor/test_launcher_envs.py

This file was deleted.

Loading