[https://nvbugs/6523767][fix] Size MPI worker-identity barrier timeout to cover worker bootstrap - #16971
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (6)
WalkthroughMPI worker identity collection now uses a validated configurable timeout. Session prefetching derives shadow-build deadlines from that timeout, propagates acquisition failures, changes synchronous fallback behavior, and updates cleanup and reuse tests. ChangesMPI session timeout configuration
Session prefetch and reuse behavior
Estimated code review effort: 4 (Complex) | ~45 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 (1)
tests/unittest/llmapi/test_mpi_session.py (1)
282-299: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winStrengthen timeout coverage. Assert the default 300s path explicitly, and add a test that captures the timeout passed from
_collect_worker_identities()tofutures_wait.
- Added:
test_identity_timeout_covers_worker_bootstrap,test_identity_timeout_env_override- Modified/removed: none
- Test-list coverage:
tests/unittest/llmapi/test_mpi_session.pyis already listed asISOLATIONintests/integration/test_lists/test-db/l0_a100.yml- Verdict: sufficient
🤖 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 `@tests/unittest/llmapi/test_mpi_session.py` around lines 282 - 299, Strengthen timeout coverage by asserting the unset environment path in test_identity_timeout_covers_worker_bootstrap returns the explicit 300-second default, while retaining the bootstrap lower-bound check. Add a focused test around _collect_worker_identities that mocks or intercepts futures_wait and verifies it receives the timeout from _identity_barrier_timeout(), including the configured environment override.Source: Path instructions
🤖 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/llmapi/mpi_session.py`:
- Around line 241-244: Update the timeout parsing logic around the float
conversion to accept overrides only when math.isfinite(value) and value > 0.
Preserve rejection of NaN, zero, and negative values, and extend the timeout
tests to cover explicit infinity and overflow inputs such as "1e309".
---
Nitpick comments:
In `@tests/unittest/llmapi/test_mpi_session.py`:
- Around line 282-299: Strengthen timeout coverage by asserting the unset
environment path in test_identity_timeout_covers_worker_bootstrap returns the
explicit 300-second default, while retaining the bootstrap lower-bound check.
Add a focused test around _collect_worker_identities that mocks or intercepts
futures_wait and verifies it receives the timeout from
_identity_barrier_timeout(), including the configured environment override.
🪄 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: b0fceb6a-b813-41df-b9d2-50760cfde64e
📒 Files selected for processing (2)
tensorrt_llm/llmapi/mpi_session.pytests/unittest/llmapi/test_mpi_session.py
allisonlim-nv
left a comment
allisonlim-nv
left a comment
There was a problem hiding this comment.
LGTM; timeout size for MPI worker extended to 300s, env override, fails if unable to verify workers are healthy.
|
/bot run --disable-fail-fast |
|
PR_Github #62368 [ run ] triggered by Bot. Commit: |
There was a problem hiding this comment.
🧹 Nitpick comments (3)
tests/test_common/session_prefetcher.py (2)
74-91: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winHardcoded
_FALLBACK_IDENTITY_TIMEOUTcan silently drift from the real default.
_FALLBACK_IDENTITY_TIMEOUT = 300.0mirrorsmpi_session._DEFAULT_IDENTITY_TIMEOUTby literal value, not by reference (deliberately, per the docstring, to avoid an eagertensorrt_llmimport). If the lower-level default ever changes, this fallback — used exactly whenmpi_sessionisn't loaded yet — will silently disagree and compute an incorrect wait budget.Consider adding a lightweight cross-check test (e.g., import
tensorrt_llm.llmapi.mpi_session._DEFAULT_IDENTITY_TIMEOUTin a test and assert equality with this constant) so a future change to one default fails loudly instead of drifting.🤖 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 `@tests/test_common/session_prefetcher.py` around lines 74 - 91, Add a lightweight test covering `_FALLBACK_IDENTITY_TIMEOUT` and the lower-level `tensorrt_llm.llmapi.mpi_session._DEFAULT_IDENTITY_TIMEOUT`, asserting both defaults remain equal. Keep the lazy-import behavior in `_fallback_identity_timeout` unchanged while ensuring future default changes fail the test instead of silently diverging.
435-477: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMissing return-type annotations on reworked
_drain/take.
_drain(self, timeout: float | None = None):andtake(self, spec: int):are both part of this rework (adding the optional timeout param and the newTimeoutError-raising behavior) but neither declares a return type, contrary to the repo guideline requiring every function be annotated.✏️ Suggested annotations
- def _drain(self, timeout: float | None = None): + def _drain(self, timeout: float | None = None) -> "_Built | None":- def take(self, spec: int): + def take(self, spec: int) -> "object | None":If importing the real session type is undesirable at module load (per the "do not import TensorRT-LLM here" constraint), consider a
TYPE_CHECKING-guarded import to get a precise type instead ofobject.As per coding guidelines: "Annotate every function, use
Nonefor non-returning functions, avoidAnyand unnecessary type ignores... useLiteral,overload,TypeVar, orProtocolwhen appropriate."🤖 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 `@tests/test_common/session_prefetcher.py` around lines 435 - 477, Add explicit return-type annotations to the reworked _drain and take methods, using the precise session/prefetch result types already established by the module; if that type requires the real session import, place it behind TYPE_CHECKING. Preserve _drain’s completed-build-or-None result and take’s prefetched-session-or-None behavior, while retaining its TimeoutError path.Source: Coding guidelines
tests/test_common/session_reuse.py (1)
366-383: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winBroad
except Exceptionon the new prefetch fallback path.The new
except Exception as e: # prefetch is an optimization: fall back(flagged by Ruff BLE001) will also swallow programming errors inprefetcher.take()(e.g.AttributeError,TypeError) and silently degrade to a synchronous spawn instead of surfacing the bug. TheTimeoutErrorcase is already split out correctly above it; consider narrowing the remaining catch to the concrete failure modestake()can legitimately raise (or at least log at a level that makes such fallbacks visible/alertable).As per coding guidelines: "Catch specific exceptions instead of using broad or bare exception handling such as
except:."🤖 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 `@tests/test_common/session_reuse.py` around lines 366 - 383, Narrow the generic exception handler around prefetcher.take in the prefetcher fallback path to the concrete expected prefetch failure exceptions, while preserving separate TimeoutError propagation and synchronous fallback behavior. Do not swallow programming errors such as AttributeError or TypeError; let unexpected exceptions surface.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.
Nitpick comments:
In `@tests/test_common/session_prefetcher.py`:
- Around line 74-91: Add a lightweight test covering
`_FALLBACK_IDENTITY_TIMEOUT` and the lower-level
`tensorrt_llm.llmapi.mpi_session._DEFAULT_IDENTITY_TIMEOUT`, asserting both
defaults remain equal. Keep the lazy-import behavior in
`_fallback_identity_timeout` unchanged while ensuring future default changes
fail the test instead of silently diverging.
- Around line 435-477: Add explicit return-type annotations to the reworked
_drain and take methods, using the precise session/prefetch result types already
established by the module; if that type requires the real session import, place
it behind TYPE_CHECKING. Preserve _drain’s completed-build-or-None result and
take’s prefetched-session-or-None behavior, while retaining its TimeoutError
path.
In `@tests/test_common/session_reuse.py`:
- Around line 366-383: Narrow the generic exception handler around
prefetcher.take in the prefetcher fallback path to the concrete expected
prefetch failure exceptions, while preserving separate TimeoutError propagation
and synchronous fallback behavior. Do not swallow programming errors such as
AttributeError or TypeError; let unexpected exceptions surface.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: a5be9bb4-5193-4279-ac35-b5409c59fa24
📒 Files selected for processing (6)
tensorrt_llm/llmapi/mpi_session.pytests/test_common/session_prefetcher.pytests/test_common/session_reuse.pytests/unittest/llmapi/test_mpi_session.pytests/unittest/llmapi/test_session_prefetcher.pytests/unittest/llmapi/test_session_reuse.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tensorrt_llm/llmapi/mpi_session.py
|
/bot run --disable-fail-fast |
|
PR_Github #62411 [ run ] triggered by Bot. Commit: |
|
PR_Github #62368 [ run ] completed with state |
|
PR_Github #62411 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62489 [ run ] triggered by Bot. Commit: |
|
PR_Github #62489 [ run ] completed with state |
…otstrap The wait_shutdown identity barrier's hardcoded 60s deadline actually bounds the entire worker bootstrap, not the barrier: MPIPoolExecutor() returns immediately because mpi4py spawns from a background manager thread, so the first submit() + futures_wait pays process spawn plus 'import tensorrt_llm'. That bootstrap is ~50-65s on an idle node and up to ~117s on a contended one, so a healthy-but-slow node fails closed with 0/N identities at ~61s. Size the deadline against the bootstrap ceiling the repo already measures (session_prefetcher.py budgets 180s for the same work) and make it tunable via TRTLLM_MPI_IDENTITY_TIMEOUT, which the TRTLLM-prefix env filter in _start_mpi_pool forwards to workers automatically. Fail-closed teardown is unchanged, so genuinely wedged pools are still rejected rather than handed out with the contract silently downgraded. Fixing this library-side covers both affected call paths, including SessionReuseCache.acquire's zero-retry branch that every AutoDeploy node id takes. Signed-off-by: handongl <handongl@nvidia.com>
Signed-off-by: qgai <qgai@nvidia.com>
Signed-off-by: qgai <qgai@nvidia.com>
42a9149 to
97162f3
Compare
|
/bot run --help |
|
PR_Github #62971 Bot args parsing error: usage: /bot [-h] |
|
/bot run reuse-pipeline |
|
PR_Github #62983 Bot args parsing error: usage: /bot [-h] |
|
/bot reuse-pipeline |
|
PR_Github #62984 [ reuse-pipeline ] triggered by Bot. Commit: |
|
PR_Github #62984 [ reuse-pipeline ] completed with state |
6 participants
Summary
MpiPoolSession(wait_shutdown=True)collects worker identities by submitting a barrier task to a freshly builtMPIPoolExecutorand waiting 60s for all N results. Because mpi4py spawns lazily from its manager thread, that first submission actually has to cover the entire worker bootstrap — process spawn plusimport tensorrt_llm, which the repo elsewhere measures at ~50-65s idle and up to ~117s on a contended node. On a busy B200 post-merge machine the 60s deadline expired with 0/N identities collected, so the fail-closed path tore the pool down and raised, aborting the test._identity_barrier_timeout()with aTRTLLM_MPI_IDENTITY_TIMEOUToverride that rejects unparsable, non-positive, and NaN values in favor of the default. The fail-closed behavior and its error message are kept intentionally — the message now reports the effective deadline and points at the override — because handing out a pool that cannot honor thewait_shutdowncontract would trade a loud failure for a silent one. Unit tests assert the default exceeds the slowest measured bootstrap and cover the env-override parsing table.Test plan
Links
Dev Engineer Review
_DEFAULT_IDENTITY_TIMEOUT: float = 300.0and_identity_barrier_timeout() -> floatintensorrt_llm/llmapi/mpi_session.py._identity_barrier_timeout()readsTRTLLM_MPI_IDENTITY_TIMEOUTand validates it as a finite, positive float; invalid/unparseable/NaN/non-positive values fall back to the 300s default (with warning/fallback behavior preserved).MpiPoolSession._collect_worker_identitiesto use the computed timeout for identity-barrier futures instead of60.0.RuntimeErrormessaging to include the effective deadline and explicitly referenceTRTLLM_MPI_IDENTITY_TIMEOUTfor diagnosis of slow-but-healthy bootstrap.tests/test_common/session_prefetcher.pyto derive shadow build join/drain budget from MPI identity-barrier logic when available, falling back toTRTLLM_MPI_IDENTITY_TIMEOUT.SessionPrefetcherdraining coordination (_drain_lock,_build_timed_out) and changed_drainto accepttimeout: float | None = None.take()behavior to always route through_drain()without the previous wrong-size in-flight early-discard logic.realsession is successfully obtained; on misses, it avoids retry/degrade flows by constructing the synchronousreal_cls(..., wait_shutdown=True)session directly.tests/test_common/session_reuse.pynow avoids converting prefetcher failures into synchronous-pool fallback at the seam and removes the “spawn failed, retrying once” retry attempt; scheduling/restocking is ordered afterrealcreation and remains best-effort.QA Engineer Review
Test code changes (files under
tests/touched):tests/unittest/llmapi/test_mpi_session.pytest_identity_collection_uses_configured_timeouttest_identity_timeout_covers_worker_bootstraptest_identity_timeout_env_overridetest_prefetch_fallback_identity_timeout_matches_mpi_defaulttest_identity_collection_fails_closed_on_timeouttest_identity_collection_fails_closed_on_duplicate_pidstests/unittest/llmapi/test_session_prefetcher.pytest_shadow_wait_budget_tracks_identity_timeouttest_shadow_wait_budget_handles_partially_loaded_mpi_moduletest_take_wrong_size_in_flight_waits_before_misstest_take_timeout_is_terminal_until_build_exitstest_concurrent_take_timeout_waits_only_oncetest_factory_spawn_failure_propagates_without_retrytests/unittest/llmapi/test_session_reuse.pytest_shadow_timeout_does_not_start_sync_pooltest_prefetcher_programming_error_propagatestest_spawn_failure_propagates_without_retrytests/test_common/session_prefetcher.py,tests/test_common/session_reuse.pyIntegration mapping (
tests/integration/test_lists/) to the touched unit tests:tests/unittest/llmapi/test_mpi_session.py:tests/integration/test_lists/test-db/l0_dgx_h100.yml:unittest/llmapi/test_mpi_session.py::test_llmapi_launch_multiple_taskstests/integration/test_lists/test-db/l0_a100.yml:unittest/llmapi/test_mpi_session.py ISOLATIONtest_session_prefetcher.py,test_session_reuse.py) from the searches performed.Verdict: needs follow-up.