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
2 changes: 2 additions & 0 deletions .pre-commit-config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -488,6 +488,7 @@ common-files: &common_files |
tests/integration/defs/accuracy/accuracy_core.py |
tests/integration/defs/accuracy/scripts/collect_evaluated_accuracies.py |
tests/integration/defs/accuracy/scripts/compute_theta_and_thresholds.py |
tests/integration/defs/accuracy/test_cli_flow.py |
tests/integration/defs/accuracy/test_disaggregated_serving.py |
tests/integration/defs/accuracy/test_llm_api_autodeploy.py |
tests/integration/defs/accuracy/test_llm_api_pytorch.py |
Expand Down Expand Up @@ -1258,6 +1259,7 @@ legacy-files: &legacy_files |
tests/integration/defs/accuracy/accuracy_core.py |
tests/integration/defs/accuracy/scripts/collect_evaluated_accuracies.py |
tests/integration/defs/accuracy/scripts/compute_theta_and_thresholds.py |
tests/integration/defs/accuracy/test_cli_flow.py |
tests/integration/defs/accuracy/test_disaggregated_serving.py |
tests/integration/defs/accuracy/test_llm_api_autodeploy.py |
tests/integration/defs/accuracy/test_llm_api_pytorch.py |
Expand Down
1 change: 1 addition & 0 deletions legacy-files.txt
Original file line number Diff line number Diff line change
Expand Up @@ -480,6 +480,7 @@ tests/integration/defs/accuracy/__init__.py
tests/integration/defs/accuracy/accuracy_core.py
tests/integration/defs/accuracy/scripts/collect_evaluated_accuracies.py
tests/integration/defs/accuracy/scripts/compute_theta_and_thresholds.py
tests/integration/defs/accuracy/test_cli_flow.py
tests/integration/defs/accuracy/test_disaggregated_serving.py
tests/integration/defs/accuracy/test_llm_api_autodeploy.py
tests/integration/defs/accuracy/test_llm_api_pytorch.py
Expand Down
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -537,6 +537,7 @@ exclude = [
"tests/integration/defs/accuracy/accuracy_core.py",
"tests/integration/defs/accuracy/scripts/collect_evaluated_accuracies.py",
"tests/integration/defs/accuracy/scripts/compute_theta_and_thresholds.py",
"tests/integration/defs/accuracy/test_cli_flow.py",
"tests/integration/defs/accuracy/test_disaggregated_serving.py",
"tests/integration/defs/accuracy/test_llm_api_autodeploy.py",
"tests/integration/defs/accuracy/test_llm_api_pytorch.py",
Expand Down
1 change: 1 addition & 0 deletions ruff-legacy.toml
Original file line number Diff line number Diff line change
Expand Up @@ -497,6 +497,7 @@ include = [
"tests/integration/defs/accuracy/accuracy_core.py",
"tests/integration/defs/accuracy/scripts/collect_evaluated_accuracies.py",
"tests/integration/defs/accuracy/scripts/compute_theta_and_thresholds.py",
"tests/integration/defs/accuracy/test_cli_flow.py",
"tests/integration/defs/accuracy/test_disaggregated_serving.py",
"tests/integration/defs/accuracy/test_llm_api_autodeploy.py",
"tests/integration/defs/accuracy/test_llm_api_pytorch.py",
Expand Down
34 changes: 19 additions & 15 deletions tensorrt_llm/llmapi/mpi_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -241,21 +241,25 @@ def __init__(self,
n_workers: int,
wait_shutdown: bool = False,
env_overrides: Optional[Dict[str, str]] = None):
"""Args:
n_workers: number of MPI workers to spawn.
wait_shutdown: when True, ``shutdown()`` blocks until the spawned
worker processes have actually exited. ``MPIPoolExecutor.shutdown``
returns at disconnect, but a worker's GPU memory is only released
when its process exits; callers that start new GPU work right
after ``shutdown()`` (e.g. CI test infrastructure handing a
pre-spawned pool to the next test) race that release and can OOM.
Off by default: production teardown does not need the barrier and
keeps its current latency.
env_overrides: extra environment variables to set in the WORKERS at
spawn, on top of the TRTLLM*/TLLM* variables forwarded from the
parent. The parent process environment is never touched — this
replaces the racy "set os.environ around the spawn, then restore"
pattern for callers that spawn pools from background threads.
"""Spawn a pool of MPI worker processes.

Args:
n_workers: number of MPI workers to spawn.
wait_shutdown: when True, ``shutdown()`` blocks until the spawned
worker processes have actually exited.
``MPIPoolExecutor.shutdown`` returns at disconnect, but a
worker's GPU memory is only released when its process exits;
callers that start new GPU work right after ``shutdown()``
(e.g. CI test infrastructure handing a pre-spawned pool to the
next test) race that release and can OOM. Off by default:
production teardown does not need the barrier and keeps its
current latency.
Comment thread
Wanli-Jiang marked this conversation as resolved.
env_overrides: extra environment variables to set in the WORKERS at
spawn, on top of the TRTLLM*/TLLM* variables forwarded from the
parent. The parent process environment is never touched — this
replaces the racy "set os.environ around the spawn, then
restore" pattern for callers that spawn pools from background
threads.
"""
self.n_workers = n_workers
self._wait_shutdown = wait_shutdown
Expand Down
Loading