[None][feat] Support externally provided MPI sessions with explicit ownership - #16053
Conversation
📝 WalkthroughWalkthroughThis PR adds MPI session ownership tracking ( ChangesMPI session ownership tracking
Estimated code review effort: 3 (Moderate) | ~25 minutes PP communicator reuse
Estimated code review effort: 2 (Simple) | ~10 minutes Sequence Diagram(s)sequenceDiagram
participant LLM as BaseLLM
participant Proxy as GenerationExecutorProxy/RpcProxy
participant Session as MpiSession
participant Validator as validate_session_world_size
LLM->>LLM: compute _owns_mpi_session from external session presence
LLM->>Proxy: pass mpi_session
Proxy->>Validator: validate_session_world_size(mpi_session, model_world_size)
Validator-->>Proxy: raise ValueError or pass
Proxy->>Proxy: set _owns_mpi_session (True/False)
Note over Proxy,Session: On shutdown()
alt owns session
Proxy->>Session: shutdown()
else externally owned
Proxy->>Proxy: wait on stored worker_futures
end
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tensorrt_llm/executor/executor.py (1)
664-667: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSetting a "private" attribute from outside the owning class.
executor._owns_mpi_session = Truereaches intoGenerationExecutorProxy's internal state from the factory method after construction (where the constructor itself would have computedFalsesince a non-Nonesession was passed in). Consider exposing this via an explicit constructor parameter (e.g.owns_mpi_session: bool = ...) threaded through_create_ipc_executor, which would be less fragile to future refactors ofGenerationExecutorProxy.__init__.🤖 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/executor.py` around lines 664 - 667, The factory is mutating GenerationExecutorProxy’s internal state by setting the private _owns_mpi_session attribute after construction, which is fragile. Update GenerationExecutorProxy.__init__ to accept an explicit owns_mpi_session parameter and thread it through _create_ipc_executor so the ownership decision is established during निर्माण instead of being patched afterward. Make sure the proxy constructor and its call sites use this parameter consistently, preserving the current behavior when the session is created locally.tensorrt_llm/executor/rpc_proxy.py (1)
280-297: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueBroad
except Exceptionflagged by static analysis; likely acceptable here.Ruff BLE001 flags for future in getattr(self, "worker_futures", []): try: future.result() except Exception as e: Catching broadly is reasonable here since arbitrary worker-task exceptions must not abort shutdown, mirroring the existing bare
except:pattern inproxy.py'sshutdown(). No action likely needed, but flagging per the coding guideline to "limit the except to the smallest set of errors possible" for completeness.🤖 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/rpc_proxy.py` around lines 280 - 297, The shutdown path in rpc_proxy.py is catching Exception broadly around worker_futures.result(), which triggers the Ruff BLE001 warning. Update the shutdown logic in the executor’s teardown flow to handle only the expected future/task failure cases, or if swallowing all worker-task errors is intentional, add an explicit local lint suppression at that catch site while keeping the warning log. Keep the behavior in shutdown and the mpi_session shutdown path unchanged otherwise.Sources: Coding guidelines, Linters/SAST tools
🤖 Prompt for all review comments with 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.
Inline comments:
In `@tensorrt_llm/executor/rpc_proxy.py`:
- Around line 308-326: In _create_mpi_session, initialize the instance state
before calling validate_session_world_size(): set self.mpi_session to None and
self._owns_mpi_session to False up front so a failed validation does not leave a
partially constructed RpcProxy that later crashes in shutdown() with an
AttributeError. Then keep the existing create/adopt branching intact, updating
self.mpi_session only after successful creation or validation.
---
Nitpick comments:
In `@tensorrt_llm/executor/executor.py`:
- Around line 664-667: The factory is mutating GenerationExecutorProxy’s
internal state by setting the private _owns_mpi_session attribute after
construction, which is fragile. Update GenerationExecutorProxy.__init__ to
accept an explicit owns_mpi_session parameter and thread it through
_create_ipc_executor so the ownership decision is established during निर्माण
instead of being patched afterward. Make sure the proxy constructor and its call
sites use this parameter consistently, preserving the current behavior when the
session is created locally.
In `@tensorrt_llm/executor/rpc_proxy.py`:
- Around line 280-297: The shutdown path in rpc_proxy.py is catching Exception
broadly around worker_futures.result(), which triggers the Ruff BLE001 warning.
Update the shutdown logic in the executor’s teardown flow to handle only the
expected future/task failure cases, or if swallowing all worker-task errors is
intentional, add an explicit local lint suppression at that catch site while
keeping the warning log. Keep the behavior in shutdown and the mpi_session
shutdown path unchanged otherwise.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 2cf75c6d-09b2-4033-bf65-a1ff946f98f2
📒 Files selected for processing (6)
tensorrt_llm/_torch/distributed/communicator.pytensorrt_llm/executor/executor.pytensorrt_llm/executor/proxy.pytensorrt_llm/executor/rpc_proxy.pytensorrt_llm/llmapi/llm.pytensorrt_llm/llmapi/mpi_session.py
|
/bot run --disable-fail-fast |
|
PR_Github #57984 [ run ] triggered by Bot. Commit: |
|
PR_Github #57984 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58147 [ run ] triggered by Bot. Commit: |
|
PR_Github #58147 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58417 [ run ] triggered by Bot. Commit: |
|
PR_Github #58417 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58462 [ run ] triggered by Bot. Commit: |
|
PR_Github #58462 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58545 [ run ] triggered by Bot. Commit: |
|
PR_Github #58545 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58568 [ run ] triggered by Bot. Commit: |
|
PR_Github #58568 [ run ] completed with state
|
…wnership Allow an LLM to adopt an externally created MpiPoolSession instead of always spawning its own, so callers (e.g. test infrastructure) can share one worker pool across several LLM instances: - proxy/rpc_proxy accept an external session and validate its size via a shared validate_session_world_size() in mpi_session.py (an external session must match the model world size exactly, else the wrong number of executor workers would start). - _owns_mpi_session distinguishes owned from borrowed sessions in llm.py, proxy.py and rpc_proxy.py: borrowed sessions are never shut down by the LLM or the executor, on normal or exceptional paths; shutdown_abort is equally gated. The rpc_proxy non-owning shutdown blocks on outstanding worker futures for parity with the owned path. - communicator.py: init_pp_comm reuses the existing PPCommNCCL when the world size is unchanged (reused worker processes would otherwise re-create NCCL comms), and warns when replacing a live comm of a different world size - the minimum needed for reused workers. Split out of NVIDIA#15777 per review (core MPI session lifecycle only; the HF weight-loader cache and the test-side automatic reuse are companion PRs). Signed-off-by: qgai <qgai@nvidia.com>
|
/bot run --disable-fail-fast |
f40a9ea to
3c53d96
Compare
|
PR_Github #58675 [ run ] triggered by Bot. Commit: |
|
PR_Github #58675 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58739 [ run ] triggered by Bot. Commit: |
|
PR_Github #58739 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #58770 [ run ] triggered by Bot. Commit: |
|
PR_Github #58770 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58782 [ run ] triggered by Bot. Commit: |
|
PR_Github #58782 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58792 [ run ] triggered by Bot. Commit: |
|
PR_Github #58792 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #58826 [ run ] triggered by Bot. Commit: |
|
PR_Github #58826 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58848 [ run ] triggered by Bot. Commit: |
|
PR_Github #58848 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58878 [ run ] triggered by Bot. Commit: |
|
PR_Github #58878 [ run ] completed with state |
6 participants
Summary
Core MPI session lifecycle only (split out of #15777 per review):
GenerationExecutorProxy/GenerationExecutorRpcProxyadopt a caller-providedMpiPoolSession; its size is validated against the model world size via a sharedvalidate_session_world_size()inmpi_session.py(an external session must match exactly, else the wrong number of executor workers would start)._owns_mpi_sessioninllm.py,proxy.py,rpc_proxy.py. Borrowed sessions are never shut down by the LLM or the executor — on normal or exceptional paths;shutdown_abortis equally gated. The rpc_proxy non-owning shutdown blocks on outstanding worker futures for parity with the owned path.init_pp_commreuses the existingPPCommNCCLwhen world size is unchanged, and warns when replacing a live comm of a different world size.No test-specific optimization and no weight-loader behavior in this PR.
Companion PRs
Each is independently mergeable; together they supersede #15777.
Test plan
Summary by CodeRabbit